fix hang on __index__

youknowone · Mar 12, 2022 · Oct 15, 2021 · Oct 16, 2021 · Oct 16, 2021 · Oct 16, 2021
commit 7b8c4091b59192cf1551dab9bef97ba3a1244878
diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs
@@ -14,7 +14,7 @@ mod array {
         vm::{
             builtins::{
                 PyByteArray, PyBytes, PyBytesRef, PyDictRef, PyFloat, PyInt, PyIntRef, PyList,
-                PyListRef, PySlice, PyStr, PyStrRef, PyTupleRef, PyTypeRef,
+                PyListRef, PyStr, PyStrRef, PyTupleRef, PyTypeRef,
             },
             class_or_notimplemented,
             function::{
@@ -25,14 +25,16 @@ mod array {
                 PyMappingMethods,
             },
             sequence::{SequenceMutOp, SequenceOp},
-            sliceable::{SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
+            sliceable::{
+                SaturatedSlice, SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp,
+            },
             types::{
                 AsBuffer, AsMapping, Comparable, Constructor, IterNext, IterNextIterable, Iterable,
                 PyComparisonOp,
             },
-            utils::Either,
             IdProtocol, PyComparisonValue, PyObject, PyObjectRef, PyObjectView, PyObjectWrap,
-            PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine,
+            PyRef, PyResult, PyValue, TryFromBorrowedObject, TryFromObject, TypeProtocol,
+            VirtualMachine,
         },
     };
     use itertools::Itertools;
@@ -251,39 +253,53 @@ mod array {
                     }
                 }
 
-                fn getitem_by_idx(
+                fn get(
                     &self,
                     i: usize,
                     vm: &VirtualMachine
-                ) -> PyResult<Option<PyObjectRef>> {
+                ) -> Option<PyResult> {
                     match self {
                         $(ArrayContentType::$n(v) => {
-                            v.get(i).map(|x| x.into_pyresult(vm)).transpose()
+                            v.get(i).map(|x| x.into_pyresult(vm))
                         })*
                     }
                 }
 
-                fn getitem(
-                    &self,
-                    needle: PyObjectRef,
+                fn getitem_by_index(&self, i: isize, vm: &VirtualMachine) -> PyResult {
+                    match self {
+                        $(ArrayContentType::$n(v) => {
+                            v.get_item_by_index(vm, i).map(|x| x.into_pyresult(vm))?
+                        })*
+                    }
+                }
+
+                fn getitem_by_slice(&self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult {
+                    match self {
+                        $(ArrayContentType::$n(v) => {
+                            let r = v.get_item_by_slice(vm, slice)?;
+                            let array = PyArray::from(ArrayContentType::$n(r));
+                            array.into_pyresult(vm)
+                        })*
+                    }
+                }
+
+                fn setitem_by_index(
+                    &mut self,
+                    i: isize,
+                    value: PyObjectRef,
                     vm: &VirtualMachine
-                ) -> PyResult {
+                ) -> PyResult<()> {
                     match self {
                         $(ArrayContentType::$n(v) => {
-                            match v.get_item(vm, &needle)? {
-                                Either::A(item) => item.into_pyresult(vm),
-                                Either::B(items) => {
-                                    let array: PyArray = ArrayContentType::$n(items).into();
-                                    array.into_pyresult(vm)
-                                }
-                            }
+                            let value = <$t>::try_into_from_object(vm, value)?;
+                            v.set_item_by_index(vm, i, value)
                         })*
                     }
                 }
 
                 fn setitem_by_slice(
                     &mut self,
-                    slice: &PySlice,
+                    slice: SaturatedSlice,
                     items: &ArrayContentType,
                     vm: &VirtualMachine
                 ) -> PyResult<()> {
@@ -300,7 +316,7 @@ mod array {
 
                 fn setitem_by_slice_no_resize(
                     &mut self,
-                    slice: &PySlice,
+                    slice: SaturatedSlice,
                     items: &ArrayContentType,
                     vm: &VirtualMachine
                 ) -> PyResult<()> {
@@ -315,24 +331,18 @@ mod array {
                     }
                 }
 
-                fn setitem_by_idx(
-                    &mut self,
-                    i: isize,
-                    value: PyObjectRef,
-                    vm: &VirtualMachine
-                ) -> PyResult<()> {
+                fn delitem_by_index(&mut self, i: isize, vm: &VirtualMachine) -> PyResult<()> {
                     match self {
                         $(ArrayContentType::$n(v) => {
-                            let value = <$t>::try_into_from_object(vm, value)?;
-                            v.set_item_by_index(vm, i, value)
+                            v.del_item_by_index(vm, i)
                         })*
                     }
                 }
 
-                fn delitem(&mut self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
+                fn delitem_by_slice(&mut self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult<()> {
                     match self {
                         $(ArrayContentType::$n(v) => {
-                            v.del_item(vm, needle)
+                            v.del_item_by_slice(vm, slice)
                         })*
                     }
                 }
@@ -413,7 +423,7 @@ mod array {
                     &'a self,
                     vm: &'vm VirtualMachine
                 ) -> impl Iterator<Item = PyResult> + 'a {
-                    (0..self.len()).map(move |i| self.getitem_by_idx(i, vm).map(Option::unwrap))
+                    (0..self.len()).map(move |i| self.get(i, vm).unwrap())
                 }
 
                 fn cmp(&self, other: &ArrayContentType) -> Result<Option<Ordering>, ()> {
@@ -933,36 +943,42 @@ mod array {
             self.copy()
         }
 
+        fn _getitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult {
+            match SequenceIndex::try_from_borrowed_object(vm, needle)? {
+                SequenceIndex::Int(i) => self.read().getitem_by_index(i, vm),
+                SequenceIndex::Slice(slice) => self.read().getitem_by_slice(slice, vm),
+            }
+        }
+
         #[pymethod(magic)]
         fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-            self.read().getitem(needle, vm)
+            self._getitem(&needle, vm)
         }
 
-        #[pymethod(magic)]
-        fn setitem(
+        fn _setitem(
             zelf: PyRef<Self>,
-            needle: PyObjectRef,
-            obj: PyObjectRef,
+            needle: &PyObject,
+            value: PyObjectRef,
             vm: &VirtualMachine,
         ) -> PyResult<()> {
-            match SequenceIndex::try_borrow_from_object(vm, &needle)? {
-                SequenceIndex::Int(i) => zelf.write().setitem_by_idx(i, obj, vm),
+            match SequenceIndex::try_from_borrowed_object(vm, needle)? {
+                SequenceIndex::Int(i) => zelf.write().setitem_by_index(i, value, vm),
                 SequenceIndex::Slice(slice) => {
                     let cloned;
                     let guard;
-                    let items = if zelf.is(&obj) {
+                    let items = if zelf.is(&value) {
                         cloned = zelf.read().clone();
                         &cloned
                     } else {
-                        match obj.payload::<PyArray>() {
+                        match value.payload::<PyArray>() {
                             Some(array) => {
                                 guard = array.read();
                                 &*guard
                             }
                             None => {
                                 return Err(vm.new_type_error(format!(
                                     "can only assign array (not \"{}\") to array slice",
-                                    obj.class().name()
+                                    value.class()
                                 )));
                             }
                         }
@@ -977,8 +993,25 @@ mod array {
         }
 
         #[pymethod(magic)]
-        fn delitem(zelf: PyRef<Self>, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
-            zelf.try_resizable(vm)?.delitem(&needle, vm)
+        fn setitem(
+            zelf: PyRef<Self>,
+            needle: PyObjectRef,
+            value: PyObjectRef,
+            vm: &VirtualMachine,
+        ) -> PyResult<()> {
+            Self::_setitem(zelf, &needle, value, vm)
+        }
+
+        fn _delitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
+            match SequenceIndex::try_from_borrowed_object(vm, needle)? {
+                SequenceIndex::Int(i) => self.try_resizable(vm)?.delitem_by_index(i, vm),
+                SequenceIndex::Slice(slice) => self.try_resizable(vm)?.delitem_by_slice(slice, vm),
+            }
+        }
+
+        #[pymethod(magic)]
+        fn delitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+            self._delitem(&needle, vm)
         }
 
         #[pymethod(magic)]
@@ -1211,14 +1244,14 @@ mod array {
         const MAPPING_METHODS: PyMappingMethods = PyMappingMethods {
             length: Some(|mapping, _vm| Ok(Self::mapping_downcast(mapping).len())),
             subscript: Some(|mapping, needle, vm| {
-                Self::mapping_downcast(mapping).getitem(needle.to_owned(), vm)
+                Self::mapping_downcast(mapping)._getitem(needle, vm)
             }),
             ass_subscript: Some(|mapping, needle, value, vm| {
                 let zelf = Self::mapping_downcast(mapping);
                 if let Some(value) = value {
-                    Self::setitem(zelf.to_owned(), needle.to_owned(), value, vm)
+                    Self::_setitem(zelf.to_owned(), needle, value, vm)
                 } else {
-                    Self::delitem(zelf.to_owned(), needle.to_owned(), vm)
+                    zelf._delitem(needle, vm)
                 }
             }),
         };
@@ -1264,8 +1297,8 @@ mod array {
     impl IterNext for PyArrayIter {
         fn next(zelf: &PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
             let pos = zelf.position.fetch_add(1, atomic::Ordering::SeqCst);
-            let r = if let Some(item) = zelf.array.read().getitem_by_idx(pos, vm)? {
-                PyIterReturn::Return(item)
+            let r = if let Some(item) = zelf.array.read().get(pos, vm) {
+                PyIterReturn::Return(item?)
             } else {
                 PyIterReturn::StopIteration(None)
             };

diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs
@@ -30,7 +30,8 @@ use crate::{
     },
     utils::Either,
     IdProtocol, PyClassImpl, PyComparisonValue, PyContext, PyObject, PyObjectRef, PyObjectView,
-    PyObjectWrap, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine,
+    PyObjectWrap, PyRef, PyResult, PyValue, TryFromBorrowedObject, TryFromObject, TypeProtocol,
+    VirtualMachine,
 };
 use bstr::ByteSlice;
 use std::{borrow::Cow, mem::size_of};
@@ -159,29 +160,19 @@ impl PyByteArray {
         self.inner().contains(needle, vm)
     }
 
-    fn setitem_by_idx(&self, i: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+    fn _setitem_by_index(&self, i: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
         let value = value_from_object(vm, &value)?;
-        let mut elements = self.borrow_buf_mut();
-        if let Some(i) = elements.wrap_index(i) {
-            elements[i] = value;
-            Ok(())
-        } else {
-            Err(vm.new_index_error("index out of range".to_owned()))
-        }
+        self.borrow_buf_mut().set_item_by_index(vm, i, value)
     }
 
-    #[pymethod(magic)]
-    fn setitem(
+    fn _setitem(
         zelf: PyRef<Self>,
-        needle: PyObjectRef,
+        needle: &PyObject,
         value: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
-        match SequenceIndex::try_borrow_from_object(vm, &needle)? {
-            SequenceIndex::Int(i) => {
-                let value = value_from_object(vm, &value)?;
-                zelf.borrow_buf_mut().set_item_by_index(vm, i, value)
-            }
+        match SequenceIndex::try_from_borrowed_object(vm, needle)? {
+            SequenceIndex::Int(i) => zelf._setitem_by_index(i, value, vm),
             SequenceIndex::Slice(slice) => {
                 let items = if zelf.is(&value) {
                     zelf.borrow_buf().to_vec()
@@ -198,6 +189,16 @@ impl PyByteArray {
         }
     }
 
+    #[pymethod(magic)]
+    fn setitem(
+        zelf: PyRef<Self>,
+        needle: PyObjectRef,
+        value: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult<()> {
+        Self::_setitem(zelf, &needle, value, vm)
+    }
+
     #[pymethod(magic)]
     fn iadd(zelf: PyRef<Self>, other: ArgBytesLike, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
         zelf.try_resizable(vm)?
@@ -206,26 +207,35 @@ impl PyByteArray {
         Ok(zelf)
     }
 
-    #[pymethod(magic)]
-    fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        self.inner().getitem(&needle, vm)
+    fn _getitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult {
+        match SequenceIndex::try_from_borrowed_object(vm, needle)? {
+            SequenceIndex::Int(i) => self.borrow_buf().get_item_by_index(vm, i).map(|x| vec![x]),
+            SequenceIndex::Slice(slice) => self.borrow_buf().get_item_by_slice(vm, slice),
+        }
+        .map(|x| vm.ctx.new_bytes(x).into())
     }
 
-    fn delitem_by_idx(&self, i: isize, vm: &VirtualMachine) -> PyResult<()> {
-        self.try_resizable(vm)?.elements.del_item_by_index(vm, i)
+    #[pymethod(magic)]
+    fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        self._getitem(&needle, vm)
     }
 
-    #[pymethod(magic)]
-    pub fn delitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
-        match SequenceIndex::try_borrow_from_object(vm, &needle)? {
-            SequenceIndex::Int(i) => self.delitem_by_idx(i, vm),
+    pub fn _delitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
+        match SequenceIndex::try_from_borrowed_object(vm, needle)? {
+            SequenceIndex::Int(i) => self.borrow_buf_mut().del_item_by_index(vm, i),
             SequenceIndex::Slice(slice) => {
+                // TODO: delete 0 elements don't need resizable
                 let elements = &mut self.try_resizable(vm)?.elements;
                 elements.del_item_by_slice(vm, slice)
             }
         }
     }
 
+    #[pymethod(magic)]
+    pub fn delitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+        self._delitem(&needle, vm)
+    }
+
     #[pystaticmethod]
     fn maketrans(from: PyBytesInner, to: PyBytesInner, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
         PyBytesInner::maketrans(from, to, vm)
@@ -777,13 +787,18 @@ impl PyByteArray {
                 .mul(n as isize, vm)
                 .map(|x| x.into_object(vm))
         }),
-        item: Some(|seq, i, vm| Self::sequence_downcast(seq).inner().item(i, vm)),
+        item: Some(|seq, i, vm| {
+            Self::sequence_downcast(seq)
+                .borrow_buf()
+                .get_item_by_index(vm, i)
+                .map(|x| vm.ctx.new_bytes(vec![x]).into())
+        }),
         ass_item: Some(|seq, i, value, vm| {
             let zelf = Self::sequence_downcast(seq);
             if let Some(value) = value {
-                zelf.setitem_by_idx(i, value, vm)
+                zelf._setitem_by_index(i, value, vm)
             } else {
-                zelf.delitem_by_idx(i, vm)
+                zelf.borrow_buf_mut().del_item_by_index(vm, i)
             }
         }),
         contains: Some(|seq, other, vm| {