Merge pull request RustPython#3810 from youknowone/sequence-index-op

youknowone · web-flow · commit a30a170ac45b · 2022-06-25T02:11:33.000+09:00
SequenceIndexOp for saturate_index and wrap_index
diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs
@@ -25,7 +25,8 @@ mod array {
             },
             sequence::{SequenceExt, SequenceMutExt},
             sliceable::{
-                SaturatedSlice, SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp,
+                SaturatedSlice, SequenceIndex, SequenceIndexOp, SliceableSequenceMutOp,
+                SliceableSequenceOp,
             },
             types::{
                 AsBuffer, AsMapping, Comparable, Constructor, IterNext, IterNextIterable, Iterable,
@@ -128,7 +129,7 @@ mod array {
                     match self {
                         $(ArrayContentType::$n(v) => {
                             let val = <$t>::try_into_from_object(vm, obj)?;
-                            v.insert(v.saturate_index(i), val);
+                            v.insert(i.saturated_at(v.len()), val);
                         })*
                     }
                     Ok(())
diff --git a/stdlib/src/mmap.rs b/stdlib/src/mmap.rs
@@ -14,7 +14,7 @@ mod mmap {
         protocol::{
             BufferDescriptor, BufferMethods, PyBuffer, PyMappingMethods, PySequenceMethods,
         },
-        sliceable::{saturate_index, wrap_index, SaturatedSlice, SequenceIndex},
+        sliceable::{SaturatedSlice, SequenceIndex, SequenceIndexOp},
         types::{AsBuffer, AsMapping, AsSequence, Constructor},
         AsObject, FromArgs, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
         TryFromBorrowedObject, VirtualMachine,
@@ -592,11 +592,11 @@ mod mmap {
             let size = self.len();
             let start = options
                 .start
-                .map(|start| saturate_index(start, size))
+                .map(|start| start.saturated_at(size))
                 .unwrap_or_else(|| self.pos());
             let end = options
                 .end
-                .map(|end| saturate_index(end, size))
+                .map(|end| end.saturated_at(size))
                 .unwrap_or(size);
             (start, end)
         }
@@ -918,7 +918,8 @@ mod mmap {
         }
 
         fn getitem_by_index(&self, i: isize, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
-            let i = wrap_index(i, self.len())
+            let i = i
+                .wrapped_at(self.len())
                 .ok_or_else(|| vm.new_index_error("mmap index out of range".to_owned()))?;
 
             let b = match self.check_valid(vm)?.deref().as_ref().unwrap() {
@@ -999,7 +1000,8 @@ mod mmap {
             value: PyObjectRef,
             vm: &VirtualMachine,
         ) -> PyResult<()> {
-            let i = wrap_index(i, zelf.len())
+            let i = i
+                .wrapped_at(zelf.len())
                 .ok_or_else(|| vm.new_index_error("mmap index out of range".to_owned()))?;
 
             let b = value_from_object(vm, &value)?;
diff --git 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::{pyint_saturate_index, SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
+    sliceable::{SequenceIndex, SequenceIndexOp, SliceableSequenceMutOp, SliceableSequenceOp},
     types::{
         AsMapping, AsSequence, Comparable, Constructor, Hashable, Initializer, IterNext,
         IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
@@ -276,7 +276,7 @@ impl PyList {
         let len = self.len();
         let saturate = |obj: PyObjectRef, len| -> PyResult<_> {
             obj.try_into_value(vm)
-                .map(|int: PyIntRef| pyint_saturate_index(int, len))
+                .map(|int: PyIntRef| int.as_bigint().saturated_at(len))
         };
         let start = start.map_or(Ok(0), |obj| saturate(obj, len))?;
         let stop = stop.map_or(Ok(len), |obj| saturate(obj, len))?;
diff --git a/vm/src/builtins/memory.rs b/vm/src/builtins/memory.rs
@@ -17,7 +17,7 @@ use crate::{
     protocol::{
         BufferDescriptor, BufferMethods, PyBuffer, PyMappingMethods, PySequenceMethods, VecBuffer,
     },
-    sliceable::wrap_index,
+    sliceable::SequenceIndexOp,
     types::{AsBuffer, AsMapping, AsSequence, Comparable, Constructor, Hashable, PyComparisonOp},
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
     TryFromBorrowedObject, TryFromObject, VirtualMachine,
@@ -243,7 +243,8 @@ impl PyMemoryView {
             ));
         }
         let (shape, stride, suboffset) = self.desc.dim_desc[0];
-        let index = wrap_index(i, shape)
+        let index = i
+            .wrapped_at(shape)
             .ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?;
         let index = index as isize * stride + suboffset;
         let pos = (index + self.start as isize) as usize;
@@ -292,7 +293,8 @@ impl PyMemoryView {
             return Err(vm.new_not_implemented_error("sub-views are not implemented".to_owned()));
         }
         let (shape, stride, suboffset) = self.desc.dim_desc[0];
-        let index = wrap_index(i, shape)
+        let index = i
+            .wrapped_at(shape)
             .ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?;
         let index = index as isize * stride + suboffset;
         let pos = (index + self.start as isize) as usize;
@@ -316,7 +318,7 @@ impl PyMemoryView {
         if zelf.is(&src) {
             return if !is_equiv_structure(&zelf.desc, &dest.desc) {
                 Err(vm.new_value_error(
-                    "memoryview assigment: lvalue and rvalue have different structures".to_owned(),
+                    "memoryview assignment: lvalue and rvalue have different structures".to_owned(),
                 ))
             } else {
                 // assign self[:] to self
@@ -336,7 +338,7 @@ impl PyMemoryView {
 
         if !is_equiv_structure(&src.desc, &dest.desc) {
             return Err(vm.new_value_error(
-                "memoryview assigment: lvalue and rvalue have different structures".to_owned(),
+                "memoryview assignment: lvalue and rvalue have different structures".to_owned(),
             ));
         }
 
diff --git a/vm/src/builtins/tuple.rs b/vm/src/builtins/tuple.rs
@@ -8,7 +8,7 @@ use crate::{
     protocol::{PyIterReturn, PyMappingMethods, PySequenceMethods},
     recursion::ReprGuard,
     sequence::SequenceExt,
-    sliceable::pyint_saturate_index,
+    sliceable::SequenceIndexOp,
     sliceable::{SequenceIndex, SliceableSequenceOp},
     types::{
         AsMapping, AsSequence, Comparable, Constructor, Hashable, IterNext, IterNextIterable,
@@ -288,7 +288,7 @@ impl PyTuple {
         let len = self.len();
         let saturate = |obj: PyObjectRef, len| -> PyResult<_> {
             obj.try_into_value(vm)
-                .map(|int: PyIntRef| pyint_saturate_index(int, len))
+                .map(|int: PyIntRef| int.as_bigint().saturated_at(len))
         };
         let start = start.map_or(Ok(0), |i| saturate(i, len))?;
         let stop = stop.map_or(Ok(len), |i| saturate(i, len))?;
diff --git a/vm/src/protocol/buffer.rs b/vm/src/protocol/buffer.rs
@@ -7,7 +7,7 @@ use crate::{
         lock::{MapImmutable, PyMutex, PyMutexGuard},
     },
     object::PyObjectPayload,
-    sliceable::wrap_index,
+    sliceable::SequenceIndexOp,
     types::{Constructor, Unconstructible},
     AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromBorrowedObject,
     VirtualMachine,
@@ -255,7 +255,7 @@ impl BufferDescriptor {
             .cloned()
             .zip_eq(self.dim_desc.iter().cloned())
         {
-            let i = wrap_index(i, shape).ok_or_else(|| {
+            let i = i.wrapped_at(shape).ok_or_else(|| {
                 vm.new_index_error(format!("index out of bounds on dimension {}", i))
             })?;
             pos += i as isize * stride + suboffset;
diff --git a/vm/src/sliceable.rs b/vm/src/sliceable.rs
@@ -1,8 +1,9 @@
-// export through slicable module, not slice.
+// export through sliceable module, not slice.
 use crate::{
-    builtins::{int::PyInt, slice::PySlice, PyIntRef},
+    builtins::{int::PyInt, slice::PySlice},
     AsObject, PyObject, PyResult, VirtualMachine,
 };
+use num_bigint::BigInt;
 use num_traits::{Signed, ToPrimitive};
 use std::ops::Range;
 
@@ -12,7 +13,7 @@ where
 {
     type Item: Clone;
     fn do_set(&mut self, index: usize, value: Self::Item);
-    fn do_delele(&mut self, index: usize);
+    fn do_delete(&mut self, index: usize);
     /// as CPython, length of range and items could be different, function must act like Vec::splice()
     fn do_set_range(&mut self, range: Range<usize>, items: &[Self::Item]);
     fn do_delete_range(&mut self, range: Range<usize>);
@@ -33,7 +34,7 @@ where
         let pos = self
             .as_ref()
             .wrap_index(index)
-            .ok_or_else(|| vm.new_index_error("assigment index out of range".to_owned()))?;
+            .ok_or_else(|| vm.new_index_error("assignment index out of range".to_owned()))?;
         self.do_set(pos, value);
         Ok(())
     }
@@ -89,8 +90,8 @@ where
         let pos = self
             .as_ref()
             .wrap_index(index)
-            .ok_or_else(|| vm.new_index_error("assigment index out of range".to_owned()))?;
-        self.do_delele(pos);
+            .ok_or_else(|| vm.new_index_error("assignment index out of range".to_owned()))?;
+        self.do_delete(pos);
         Ok(())
     }
 
@@ -118,7 +119,7 @@ impl<T: Clone> SliceableSequenceMutOp for Vec<T> {
         self[index] = value;
     }
 
-    fn do_delele(&mut self, index: usize) {
+    fn do_delete(&mut self, index: usize) {
         self.remove(index);
     }
 
@@ -175,11 +176,11 @@ pub trait SliceableSequenceOp {
     fn len(&self) -> usize;
 
     fn wrap_index(&self, p: isize) -> Option<usize> {
-        wrap_index(p, self.len())
+        p.wrapped_at(self.len())
     }
 
     fn saturate_index(&self, p: isize) -> usize {
-        saturate_index(p, self.len())
+        p.saturated_at(self.len())
     }
 
     fn getitem_by_slice(
@@ -275,7 +276,7 @@ impl SequenceIndex {
         } else if let Some(slice) = obj.payload::<PySlice>() {
             slice.to_saturated(vm).map(Self::Slice)
         } else if let Some(i) = vm.to_index_opt(obj.to_owned()) {
-            // TODO: __index__ for indice is no more supported?
+            // TODO: __index__ for indices is no more supported?
             i?.try_to_primitive(vm)
                 .map_err(|_| {
                     vm.new_index_error("cannot fit 'int' into an index-sized integer".to_owned())
@@ -291,50 +292,53 @@ impl SequenceIndex {
     }
 }
 
-// Use PySliceableSequence::wrap_index for implementors
-pub fn wrap_index(p: isize, len: usize) -> Option<usize> {
-    let neg = p.is_negative();
-    let p = p.wrapping_abs() as usize;
-    if neg {
-        len.checked_sub(p)
-    } else if p >= len {
-        None
-    } else {
-        Some(p)
-    }
+pub trait SequenceIndexOp {
+    // Saturate p in range [0, len] inclusive
+    fn saturated_at(&self, len: usize) -> usize;
+    // Use PySliceableSequence::wrap_index for implementors
+    fn wrapped_at(&self, len: usize) -> Option<usize>;
 }
 
-// Saturate p in range [0, len] inclusive
-pub fn saturate_index(p: isize, len: usize) -> usize {
-    let len = len.to_isize().unwrap_or(isize::MAX);
-    let mut p = p;
-    if p < 0 {
-        p += len;
+impl SequenceIndexOp for isize {
+    fn saturated_at(&self, len: usize) -> usize {
+        let len = len.to_isize().unwrap_or(Self::MAX);
+        let mut p = *self;
         if p < 0 {
-            p = 0;
+            p += len;
         }
+        p.clamp(0, len) as usize
     }
-    if p > len {
-        p = len;
+
+    fn wrapped_at(&self, len: usize) -> Option<usize> {
+        let neg = self.is_negative();
+        let p = self.unsigned_abs();
+        if neg {
+            len.checked_sub(p)
+        } else if p >= len {
+            None
+        } else {
+            Some(p)
+        }
     }
-    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)
+impl SequenceIndexOp for BigInt {
+    fn saturated_at(&self, len: usize) -> usize {
+        if self.is_negative() {
+            self.abs()
+                .try_into()
+                .map_or(0, |abs| len.saturating_sub(abs))
+        } else {
+            self.try_into().unwrap_or(len)
+        }
+    }
+    fn wrapped_at(&self, _len: usize) -> Option<usize> {
+        unimplemented!("please add one once we need it")
     }
 }
 
 /// A saturated slice with values ranging in [isize::MIN, isize::MAX]. Used for
-/// slicable sequences that require indices in the aforementioned range.
+/// sliceable sequences that require indices in the aforementioned range.
 ///
 /// Invokes `__index__` on the PySliceRef during construction so as to separate the
 /// transformation from PyObject into isize and the adjusting of the slice to a given
@@ -383,16 +387,16 @@ impl SaturatedSlice {
             let stop = if self.stop == -1 {
                 len
             } else {
-                saturate_index(self.stop.saturating_add(1), len)
+                self.stop.saturating_add(1).saturated_at(len)
             };
             let start = if self.start == -1 {
                 len
             } else {
-                saturate_index(self.start.saturating_add(1), len)
+                self.start.saturating_add(1).saturated_at(len)
             };
             stop..start
         } else {
-            saturate_index(self.start, len)..saturate_index(self.stop, len)
+            self.start.saturated_at(len)..self.stop.saturated_at(len)
         };
 
         let (range, slice_len) = if range.start >= range.end {
diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs
@@ -13,8 +13,7 @@ mod _collections {
         protocol::{PyIterReturn, PySequenceMethods},
         recursion::ReprGuard,
         sequence::MutObjectSequenceOp,
-        sliceable,
-        sliceable::pyint_saturate_index,
+        sliceable::SequenceIndexOp,
         types::{
             AsSequence, Comparable, Constructor, Hashable, Initializer, IterNext, IterNextIterable,
             Iterable, PyComparisonOp, Unhashable,
@@ -167,7 +166,7 @@ mod _collections {
             let len = self.len();
             let saturate = |obj: PyObjectRef, len| -> PyResult<_> {
                 obj.try_into_value(vm)
-                    .map(|int: PyIntRef| pyint_saturate_index(int, len))
+                    .map(|int: PyIntRef| int.as_bigint().saturated_at(len))
             };
             let start = start.map_or(Ok(0), |i| saturate(i, len))?;
             let stop = stop.map_or(Ok(len), |i| saturate(i, len))?;
@@ -280,15 +279,15 @@ mod _collections {
         #[pymethod(magic)]
         fn getitem(&self, idx: isize, vm: &VirtualMachine) -> PyResult {
             let deque = self.borrow_deque();
-            sliceable::wrap_index(idx, deque.len())
+            idx.wrapped_at(deque.len())
                 .and_then(|i| deque.get(i).cloned())
                 .ok_or_else(|| vm.new_index_error("deque index out of range".to_owned()))
         }
 
         #[pymethod(magic)]
         fn setitem(&self, idx: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
             let mut deque = self.borrow_deque_mut();
-            sliceable::wrap_index(idx, deque.len())
+            idx.wrapped_at(deque.len())
                 .and_then(|i| deque.get_mut(i))
                 .map(|x| *x = value)
                 .ok_or_else(|| vm.new_index_error("deque index out of range".to_owned()))
@@ -297,7 +296,7 @@ mod _collections {
         #[pymethod(magic)]
         fn delitem(&self, idx: isize, vm: &VirtualMachine) -> PyResult<()> {
             let mut deque = self.borrow_deque_mut();
-            sliceable::wrap_index(idx, deque.len())
+            idx.wrapped_at(deque.len())
                 .and_then(|i| deque.remove(i).map(drop))
                 .ok_or_else(|| vm.new_index_error("deque index out of range".to_owned()))
         }
