From b86e803fec84f707345eb57a7c8eedc8ee70f321 Mon Sep 17 00:00:00 2001 From: Oscar Shrimpton Date: Thu, 10 Oct 2019 15:02:09 +0100 Subject: [PATCH] Implement .indices(len) of slice (Fixes #1431) range.__getitem now also uses slice.indices() internally. CPython: https://github.com/python/cpython/blob/master/Objects/sliceobject.c#L373 --- vm/src/obj/objrange.rs | 78 ++++-------------------------- vm/src/obj/objslice.rs | 105 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 112 insertions(+), 71 deletions(-) diff --git a/vm/src/obj/objrange.rs b/vm/src/obj/objrange.rs index 2ece52cb49..4cfa3aec4d 100644 --- a/vm/src/obj/objrange.rs +++ b/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 (mut substart, mut substop, mut substep) = + slice.inner_indices(&self.length(), vm)?; + let range_step = self.step(vm); + let range_start = self.start(vm); - 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); + 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()) diff --git a/vm/src/obj/objslice.rs b/vm/src/obj/objslice.rs index 2f1e7d56f6..cdbea2e68f 100644 --- a/vm/src/obj/objslice.rs +++ b/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 = 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 = 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 = 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::() { @@ -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::() { + 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> {