From 4f97c0f053bfac41b762ea304ab9aff047347410 Mon Sep 17 00:00:00 2001
From: Jeong-Min Lee <falsetru@gmail.com>
Date: Sat, 11 Jun 2022 15:47:35 +0900
Subject: [PATCH] Extract saturate function.

---
 vm/src/builtins/list.rs | 47 +++++------------------------------------
 vm/src/sliceable.rs     | 15 ++++++++++++-
 2 files changed, 19 insertions(+), 43 deletions(-)

diff --git a/vm/src/builtins/list.rs b/vm/src/builtins/list.rs
index 1ef2f3927..677ea6bb8 100644
--- a/vm/src/builtins/list.rs
+++ b/vm/src/builtins/list.rs
@@ -10,7 +10,7 @@ use crate::{
     protocol::{PyIterReturn, PyMappingMethods, PySequence, PySequenceMethods},
     recursion::ReprGuard,
     sequence::{MutObjectSequenceOp, SequenceExt, SequenceMutExt},
-    sliceable::{SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
+    sliceable::{pyint_saturate_index, SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
     types::{
         AsMapping, AsSequence, Comparable, Constructor, Hashable, Initializer, IterNext,
         IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
@@ -19,7 +19,6 @@ use crate::{
     vm::VirtualMachine,
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
 };
-use num_traits::Signed;
 use std::{fmt, ops::DerefMut};
 
 /// Built-in mutable sequence.
@@ -277,54 +276,18 @@ impl PyList {
         vm: &VirtualMachine,
     ) -> PyResult<usize> {
         let len = self.len();
-        let start: OptionalArg<PyIntRef> = match start {
-            OptionalArg::Missing => OptionalArg::Missing,
-            OptionalArg::Present(obj) => OptionalArg::Present(obj.try_into_value(vm)?),
-        };
-        let stop: OptionalArg<PyIntRef> = match stop {
-            OptionalArg::Missing => OptionalArg::Missing,
-            OptionalArg::Present(obj) => OptionalArg::Present(obj.try_into_value(vm)?),
-        };
         let start: usize = match start {
             OptionalArg::Missing => 0,
             OptionalArg::Present(obj) => {
-                let bigint = obj.as_bigint();
-                if bigint.is_negative() {
-                    let idx: Result<usize, _> = bigint.abs().try_into();
-                    match idx {
-                        Ok(abs) => {
-                            if abs >= len {
-                                0
-                            } else {
-                                len - abs
-                            }
-                        }
-                        Err(_) => 0,
-                    }
-                } else {
-                    bigint.try_into().unwrap_or(len)
-                }
+                let int: PyIntRef = obj.try_into_value(vm)?;
+                pyint_saturate_index(int, len)
             }
         };
         let stop: usize = match stop {
             OptionalArg::Missing => len,
             OptionalArg::Present(obj) => {
-                let bigint = obj.as_bigint();
-                if bigint.is_negative() {
-                    let idx: Result<usize, _> = bigint.abs().try_into();
-                    match idx {
-                        Ok(abs) => {
-                            if abs >= len {
-                                0
-                            } else {
-                                len - abs
-                            }
-                        }
-                        Err(_) => 0,
-                    }
-                } else {
-                    bigint.try_into().unwrap_or(len)
-                }
+                let int: PyIntRef = obj.try_into_value(vm)?;
+                pyint_saturate_index(int, len)
             }
         };
         let index = self.mut_index_range(vm, &needle, start..stop)?;
diff --git a/vm/src/sliceable.rs b/vm/src/sliceable.rs
index 6fd2c5449..d32a8d131 100644
--- a/vm/src/sliceable.rs
+++ b/vm/src/sliceable.rs
@@ -1,6 +1,6 @@
 // export through slicable module, not slice.
 use crate::{
-    builtins::{int::PyInt, slice::PySlice},
+    builtins::{int::PyInt, slice::PySlice, PyIntRef},
     AsObject, PyObject, PyResult, VirtualMachine,
 };
 use num_traits::{Signed, ToPrimitive};
@@ -320,6 +320,19 @@ pub fn saturate_index(p: isize, len: usize) -> usize {
     p as usize
 }
 
+// Saturate p in range [0, len] inclusive
+pub fn pyint_saturate_index(p: PyIntRef, len: usize) -> usize {
+    let bigint = p.as_bigint();
+    if bigint.is_negative() {
+        bigint
+            .abs()
+            .try_into()
+            .map_or(0, |abs| len.saturating_sub(abs))
+    } else {
+        bigint.try_into().unwrap_or(len)
+    }
+}
+
 /// A saturated slice with values ranging in [isize::MIN, isize::MAX]. Used for
 /// slicable sequences that require indices in the aforementioned range.
 ///