Implement .indices(len) of slice (Fixes RustPython#1431)

range.__getitem now also uses slice.indices() internally.
CPython: https://github.com/python/cpython/blob/master/Objects/sliceobject.c#L373

Author: Oscar Shrimpton

vm/src/obj/objrange.rs

@@ -339,77 +339,19 @@ impl PyRange {
     fn getitem(&self, subscript: RangeIndex, vm: &VirtualMachine) -> PyResult {
         match subscript {
             RangeIndex::Slice(slice) => {
-                let range_start = self.start.as_bigint();
-                let range_step = self.step.as_bigint();
-                let range_length = &self.length();
-
-                let substep = if let Some(slice_step) = slice.step_index(vm)? {
-                    if slice_step.is_zero() {
-                        return Err(vm.new_value_error("slice step cannot be zero".to_string()));
-                    }
-                    slice_step
-                } else {
-                    BigInt::one()
-                };
-
-                let negative_step = substep.is_negative();
-                let lower_bound = if negative_step {
-                    -BigInt::one()
-                } else {
-                    BigInt::zero()
-                };
-                let upper_bound = if negative_step {
-                    &lower_bound + range_length
-                } else {
-                    range_length.clone()
-                };
-
-                let substart = if let Some(slice_start) = slice.start_index(vm)? {
-                    if slice_start.is_negative() {
-                        let tmp = slice_start + range_length;
-                        if tmp < lower_bound {
-                            lower_bound.clone()
-                        } else {
-                            tmp.clone()
-                        }
-                    } else if slice_start > upper_bound {
-                        upper_bound.clone()
-                    } else {
-                        slice_start.clone()
-                    }
-                } else if negative_step {
-                    upper_bound.clone()
-                } else {
-                    lower_bound.clone()
-                };
-
-                let substop = if let Some(slice_stop) = slice.stop_index(vm)? {
-                    if slice_stop.is_negative() {
-                        let tmp = slice_stop + range_length;
-                        if tmp < lower_bound {
-                            lower_bound.clone()
-                        } else {
-                            tmp.clone()
-                        }
-                    } else if slice_stop > upper_bound {
-                        upper_bound.clone()
-                    } else {
-                        slice_stop.clone()
-                    }
-                } else if negative_step {
-                    lower_bound.clone()
-                } else {
-                    upper_bound.clone()
-                };
-
-                let step = range_step * &substep;
-                let start = range_start + (&substart * range_step);
-                let stop = range_start + (&substop * range_step);
+                let (mut substart, mut substop, mut substep) =
+                    slice.inner_indices(&self.length(), vm)?;
+                let range_step = self.step(vm);
+                let range_start = self.start(vm);
+
+                substep *= range_step.as_bigint();
+                substart = (substart * range_step.as_bigint()) + range_start.as_bigint();
+                substop = (substop * range_step.as_bigint()) + range_start.as_bigint();
 
                 Ok(PyRange {
-                    start: PyInt::new(start).into_ref(vm),
-                    stop: PyInt::new(stop).into_ref(vm),
-                    step: PyInt::new(step).into_ref(vm),
+                    start: PyInt::new(substart).into_ref(vm),
+                    stop: PyInt::new(substop).into_ref(vm),
+                    step: PyInt::new(substep).into_ref(vm),
                 }
                 .into_ref(vm)
                 .into_object())

vm/src/obj/objslice.rs

@@ -1,12 +1,13 @@
-use num_bigint::BigInt;
-
 use super::objint::PyInt;
 use super::objtype::{class_has_attr, PyClassRef};
 use crate::function::{OptionalArg, PyFuncArgs};
 use crate::pyobject::{
-    IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
+    IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryIntoRef,
+    TypeProtocol,
 };
 use crate::vm::VirtualMachine;
+use num_bigint::{BigInt, ToBigInt};
+use num_traits::{One, Signed, Zero};
 
 #[pyclass]
 #[derive(Debug)]
@@ -157,6 +158,92 @@ impl PySlice {
         Ok(eq)
     }
 
+    pub(crate) fn inner_indices(
+        &self,
+        length: &BigInt,
+        vm: &VirtualMachine,
+    ) -> PyResult<(BigInt, BigInt, BigInt)> {
+        // Calculate step
+        let step: BigInt;
+        if vm.is_none(&self.step(vm)) {
+            step = One::one();
+        } else {
+            // Clone the value, not the reference.
+            let this_step: PyRef<PyInt> = self.step(vm).try_into_ref(vm)?;
+            step = this_step.as_bigint().clone();
+
+            if step.is_zero() {
+                return Err(vm.new_value_error("slice step cannot be zero.".to_owned()));
+            }
+        }
+
+        // For convenience
+        let backwards = step.is_negative();
+
+        // Each end of the array
+        let lower = if backwards {
+            -1_i8.to_bigint().unwrap()
+        } else {
+            Zero::zero()
+        };
+
+        let upper = if backwards {
+            lower.clone() + length
+        } else {
+            length.clone()
+        };
+
+        // Calculate start
+        let mut start: BigInt;
+        if vm.is_none(&self.start(vm)) {
+            // Default
+            start = if backwards {
+                upper.clone()
+            } else {
+                lower.clone()
+            };
+        } else {
+            let this_start: PyRef<PyInt> = self.start(vm).try_into_ref(vm)?;
+            start = this_start.as_bigint().clone();
+
+            if start < Zero::zero() {
+                // From end of array
+                start += length;
+
+                if start < lower {
+                    start = lower.clone();
+                }
+            } else if start > upper {
+                start = upper.clone();
+            }
+        }
+
+        // Calculate Stop
+        let mut stop: BigInt;
+        if vm.is_none(&self.stop(vm)) {
+            stop = if backwards {
+                lower.clone()
+            } else {
+                upper.clone()
+            };
+        } else {
+            let this_stop: PyRef<PyInt> = self.stop(vm).try_into_ref(vm)?;
+            stop = this_stop.as_bigint().clone();
+
+            if stop < Zero::zero() {
+                // From end of array
+                stop += length;
+                if stop < lower {
+                    stop = lower.clone();
+                }
+            } else if stop > upper {
+                stop = upper.clone();
+            }
+        }
+
+        Ok((start, stop, step))
+    }
+
     #[pymethod(name = "__eq__")]
     fn eq(&self, rhs: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         if let Some(rhs) = rhs.payload::<PySlice>() {
@@ -221,6 +308,18 @@ impl PySlice {
     fn hash(&self, vm: &VirtualMachine) -> PyResult<()> {
         Err(vm.new_type_error("unhashable type".to_string()))
     }
+
+    #[pymethod(name = "indices")]
+    fn indices(&self, length: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        if let Some(length) = length.payload::<PyInt>() {
+            let (start, stop, step) = self.inner_indices(length.as_bigint(), vm)?;
+            Ok(vm
+                .ctx
+                .new_tuple(vec![vm.new_int(start), vm.new_int(stop), vm.new_int(step)]))
+        } else {
+            Ok(vm.ctx.not_implemented())
+        }
+    }
 }
 
 fn to_index_value(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Option<BigInt>> {