Refactor get_item for sequences

Author: youknowone

vm/src/obj/objlist.rs

@@ -18,7 +18,7 @@ use crate::pyobject::{
     PyResult, PyValue, TryFromObject, TypeProtocol,
 };
 use crate::sequence::{self, SimpleSeq};
-use crate::sliceable::{get_item, PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
+use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
 use crate::slots::{Comparable, Hashable, PyComparisonOp, Unhashable};
 use crate::vm::{ReprGuard, VirtualMachine};
 
@@ -190,12 +190,11 @@ impl PyList {
 
     #[pymethod(name = "__getitem__")]
     fn getitem(zelf: PyRef<Self>, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        Ok(
-            match get_item(vm, zelf.as_object(), &zelf.borrow_value(), needle)? {
-                Either::A(obj) => obj,
-                Either::B(vec) => vm.ctx.new_list(vec),
-            },
-        )
+        let result = match zelf.borrow_value().get_item(vm, needle, "list")? {
+            Either::A(obj) => obj,
+            Either::B(vec) => vm.ctx.new_list(vec),
+        };
+        Ok(result)
     }
 
     #[pymethod(name = "__iter__")]

vm/src/obj/objstr.rs

@@ -22,10 +22,11 @@ use crate::exceptions::IntoPyException;
 use crate::format::{FormatSpec, FormatString, FromTemplate};
 use crate::function::{OptionalArg, OptionalOption, PyFuncArgs};
 use crate::pyobject::{
-    BorrowValue, IdProtocol, IntoPyObject, ItemProtocol, PyClassImpl, PyComparisonValue, PyContext,
-    PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TryIntoRef, TypeProtocol,
+    BorrowValue, Either, IdProtocol, IntoPyObject, ItemProtocol, PyClassImpl, PyComparisonValue,
+    PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TryIntoRef,
+    TypeProtocol,
 };
-use crate::sliceable::{PySliceableSequence, SequenceIndex};
+use crate::sliceable::PySliceableSequence;
 use crate::slots::{Comparable, Hashable, PyComparisonOp};
 use crate::VirtualMachine;
 use rustpython_common::hash;
@@ -228,26 +229,12 @@ impl PyStr {
     }
 
     #[pymethod(name = "__getitem__")]
-    fn getitem(&self, needle: SequenceIndex, vm: &VirtualMachine) -> PyResult {
-        match needle {
-            SequenceIndex::Int(pos) => {
-                let index: usize = if pos.is_negative() {
-                    (self.value.chars().count() as isize + pos) as usize
-                } else {
-                    pos.abs() as usize
-                };
-
-                if let Some(character) = self.value.chars().nth(index) {
-                    Ok(vm.ctx.new_str(character.to_string()))
-                } else {
-                    Err(vm.new_index_error("string index out of range".to_owned()))
-                }
-            }
-            SequenceIndex::Slice(slice) => {
-                let string = self.get_slice_items(vm, &slice)?;
-                Ok(vm.ctx.new_str(string))
-            }
-        }
+    fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        let s = match self.get_item(vm, needle, "string")? {
+            Either::A(ch) => ch.to_string(),
+            Either::B(s) => s,
+        };
+        Ok(vm.ctx.new_str(s))
     }
 
     pub(crate) fn hash(&self, vm: &VirtualMachine) -> hash::PyHash {
@@ -1172,8 +1159,13 @@ pub fn borrow_value(obj: &PyObjectRef) -> &str {
 }
 
 impl PySliceableSequence for PyStr {
+    type Item = char;
     type Sliced = String;
 
+    fn do_get(&self, index: usize) -> Self::Item {
+        self.value.chars().nth(index).unwrap()
+    }
+
     fn do_slice(&self, range: Range<usize>) -> Self::Sliced {
         self.value
             .chars()

vm/src/obj/objtuple.rs

@@ -9,7 +9,7 @@ use crate::pyobject::{
     PyComparisonValue, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
 };
 use crate::sequence::{self, SimpleSeq};
-use crate::sliceable::get_item;
+use crate::sliceable::PySliceableSequence;
 use crate::slots::{Comparable, Hashable, PyComparisonOp};
 use crate::vm::{ReprGuard, VirtualMachine};
 use rustpython_common::hash::PyHash;
@@ -184,12 +184,11 @@ impl PyTuple {
 
     #[pymethod(name = "__getitem__")]
     fn getitem(zelf: PyRef<Self>, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        Ok(
-            match get_item(vm, zelf.as_object(), &zelf.elements, needle)? {
-                Either::A(obj) => obj,
-                Either::B(vec) => vm.ctx.new_tuple(vec),
-            },
-        )
+        let result = match zelf.elements.as_ref().get_item(vm, needle, "tuple")? {
+            Either::A(obj) => obj,
+            Either::B(vec) => vm.ctx.new_tuple(vec),
+        };
+        Ok(result)
     }
 
     #[pymethod(name = "index")]

vm/src/sliceable.rs

@@ -240,8 +240,10 @@ impl<T: Clone> PySliceableSequenceMut for Vec<T> {
 }
 
 pub trait PySliceableSequence {
+    type Item;
     type Sliced;
 
+    fn do_get(&self, index: usize) -> Self::Item;
     fn do_slice(&self, range: Range<usize>) -> Self::Sliced;
     fn do_slice_reverse(&self, range: Range<usize>) -> Self::Sliced;
     fn do_stepped_slice(&self, range: Range<usize>, step: usize) -> Self::Sliced;
@@ -332,25 +334,53 @@ pub trait PySliceableSequence {
             }
         }
     }
+
+    fn get_item(
+        &self,
+        vm: &VirtualMachine,
+        needle: PyObjectRef,
+        owner_type: &'static str,
+    ) -> PyResult<Either<Self::Item, Self::Sliced>> {
+        let needle = SequenceIndex::try_from_object_for(vm, needle, owner_type)?;
+        match needle {
+            SequenceIndex::Int(value) => {
+                let pos_index = self.wrap_index(value).ok_or_else(|| {
+                    vm.new_index_error(format!("{} index out of range", owner_type))
+                })?;
+                Ok(Either::A(self.do_get(pos_index)))
+            }
+            SequenceIndex::Slice(slice) => Ok(Either::B(self.get_slice_items(vm, &slice)?)),
+        }
+    }
 }
 
 impl<T: Clone> PySliceableSequence for [T] {
+    type Item = T;
     type Sliced = Vec<T>;
 
+    #[inline]
+    fn do_get(&self, index: usize) -> Self::Item {
+        self[index].clone()
+    }
+
+    #[inline]
     fn do_slice(&self, range: Range<usize>) -> Self::Sliced {
         self[range].to_vec()
     }
 
+    #[inline]
     fn do_slice_reverse(&self, range: Range<usize>) -> Self::Sliced {
         let mut slice = self[range].to_vec();
         slice.reverse();
         slice
     }
 
+    #[inline]
     fn do_stepped_slice(&self, range: Range<usize>, step: usize) -> Self::Sliced {
         self[range].iter().step_by(step).cloned().collect()
     }
 
+    #[inline]
     fn do_stepped_slice_reverse(&self, range: Range<usize>, step: usize) -> Self::Sliced {
         self[range].iter().rev().step_by(step).cloned().collect()
     }
@@ -376,8 +406,12 @@ pub enum SequenceIndex {
     Slice(PySliceRef),
 }
 
-impl TryFromObject for SequenceIndex {
-    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
+impl SequenceIndex {
+    fn try_from_object_for(
+        vm: &VirtualMachine,
+        obj: PyObjectRef,
+        owner_type: &'static str,
+    ) -> PyResult<Self> {
         match_class!(match obj {
             i @ PyInt => i
                 .borrow_value()
@@ -387,13 +421,20 @@ impl TryFromObject for SequenceIndex {
                     .new_index_error("cannot fit 'int' into an index-sized integer".to_owned())),
             s @ PySlice => Ok(SequenceIndex::Slice(s)),
             obj => Err(vm.new_type_error(format!(
-                "sequence indices be integers or slices, not {}",
-                obj.class(),
+                "{} indices must be integers or slices, not {}",
+                owner_type,
+                obj.lease_class().name,
             ))),
         })
     }
 }
 
+impl TryFromObject for SequenceIndex {
+    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
+        Self::try_from_object_for(vm, obj, "sequence")
+    }
+}
+
 /// Get the index into a sequence like type. Get it from a python integer
 /// object, accounting for negative index, and out of bounds issues.
 // pub fn get_sequence_index(vm: &VirtualMachine, index: &PyIntRef, length: usize) -> PyResult<usize> {
@@ -455,31 +496,6 @@ pub(crate) fn saturate_index(p: isize, len: usize) -> usize {
 //     start..stop
 // }
 
-pub fn get_item(
-    vm: &VirtualMachine,
-    sequence: &PyObjectRef,
-    elements: &[PyObjectRef],
-    subscript: PyObjectRef,
-) -> PyResult<Either<PyObjectRef, Vec<PyObjectRef>>> {
-    if let Some(i) = subscript.payload::<PyInt>() {
-        let value = i.borrow_value().to_isize().ok_or_else(|| {
-            vm.new_index_error("cannot fit 'int' into an index-sized integer".to_owned())
-        })?;
-        let pos_index = elements
-            .wrap_index(value)
-            .ok_or_else(|| vm.new_index_error("Index out of bounds!".to_owned()))?;
-        Ok(Either::A(elements[pos_index].clone()))
-    } else {
-        let slice = subscript.payload::<PySlice>().ok_or_else(|| {
-            vm.new_type_error(format!(
-                "{} indices must be integers or slices",
-                sequence.lease_class().name
-            ))
-        })?;
-        Ok(Either::B(elements.get_slice_items(vm, slice)?))
-    }
-}
-
 //Check if given arg could be used with PySliceableSequence.saturate_range()
 // pub fn is_valid_slice_arg(
 //     arg: OptionalArg<PyObjectRef>,