use std::cell::Cell; use num_bigint::{BigInt, Sign}; use num_integer::Integer; use num_traits::{One, Signed, Zero}; use crate::function::{OptionalArg, PyFuncArgs}; use crate::pyobject::{ PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; use crate::vm::VirtualMachine; use super::objint::{PyInt, PyIntRef}; use super::objiter; use super::objslice::{PySlice, PySliceRef}; use super::objtype::{self, PyClassRef}; /// range(stop) -> range object /// range(start, stop[, step]) -> range object /// /// Return an object that produces a sequence of integers from start (inclusive) /// to stop (exclusive) by step. range(i, j) produces i, i+1, i+2, ..., j-1. /// start defaults to 0, and stop is omitted! range(4) produces 0, 1, 2, 3. /// These are exactly the valid indices for a list of 4 elements. /// When step is given, it specifies the increment (or decrement). #[pyclass] #[derive(Debug, Clone)] pub struct PyRange { pub start: PyIntRef, pub stop: PyIntRef, pub step: PyIntRef, } impl PyValue for PyRange { fn class(vm: &VirtualMachine) -> PyClassRef { vm.ctx.range_type() } } impl PyRange { #[inline] fn offset(&self, value: &BigInt) -> Option { let start = self.start.as_bigint(); let stop = self.stop.as_bigint(); let step = self.step.as_bigint(); match step.sign() { Sign::Plus if value >= start && value < stop => Some(value - start), Sign::Minus if value <= self.start.as_bigint() && value > stop => Some(start - value), _ => None, } } #[inline] pub fn index_of(&self, value: &BigInt) -> Option { let step = self.step.as_bigint(); match self.offset(value) { Some(ref offset) if offset.is_multiple_of(step) => Some((offset / step).abs()), Some(_) | None => None, } } #[inline] pub fn is_empty(&self) -> bool { let start = self.start.as_bigint(); let stop = self.stop.as_bigint(); let step = self.step.as_bigint(); (start <= stop && step.is_negative()) || (start >= stop && step.is_positive()) } #[inline] pub fn forward(&self) -> bool { self.start.as_bigint() < self.stop.as_bigint() } #[inline] pub fn get(&self, index: &BigInt) -> Option { let start = self.start.as_bigint(); let stop = self.stop.as_bigint(); let step = self.step.as_bigint(); let index = index.clone(); if self.is_empty() { return None; } let length: BigInt = if start < stop { (stop - start - 1) / step + 1 } else { (start - stop - 1) / (-step) + 1 }; let index = if index.is_negative() { let new_index: BigInt = &length + &index; if new_index.is_negative() { return None; } length + index } else { if length <= index { return None; } index }; let result = start + step * index; Some(result) } #[inline] fn length(&self) -> PyInt { let start = self.start.as_bigint(); let stop = self.stop.as_bigint(); let step = self.step.as_bigint(); match step.sign() { Sign::Plus if start < stop => PyInt::new((stop - start - 1usize) / step + 1), Sign::Minus if start > stop => PyInt::new((start - stop - 1usize) / (-step) + 1), Sign::Plus | Sign::Minus => PyInt::new(0), Sign::NoSign => unreachable!(), } } } pub fn get_value(obj: &PyObjectRef) -> PyRange { obj.payload::().unwrap().clone() } pub fn init(context: &PyContext) { PyRange::extend_class(context, &context.types.range_type); PyRangeIterator::extend_class(context, &context.types.rangeiterator_type); } type PyRangeRef = PyRef; #[pyimpl] impl PyRange { #[pymethod(name = "__new__")] fn new(cls: PyClassRef, stop: PyIntRef, vm: &VirtualMachine) -> PyResult { PyRange { start: PyInt::new(BigInt::zero()).into_ref(vm), stop: stop.clone(), step: PyInt::new(BigInt::one()).into_ref(vm), } .into_ref_with_type(vm, cls) } fn new_from( cls: PyClassRef, start: PyIntRef, stop: PyIntRef, step: OptionalArg, vm: &VirtualMachine, ) -> PyResult { let step = step.unwrap_or_else(|| PyInt::new(BigInt::one()).into_ref(vm)); if step.as_bigint().is_zero() { return Err(vm.new_value_error("range() arg 3 must not be zero".to_string())); } PyRange { start, stop, step }.into_ref_with_type(vm, cls) } #[pyproperty(name = "start")] fn start(&self, _vm: &VirtualMachine) -> PyIntRef { self.start.clone() } #[pyproperty(name = "stop")] fn stop(&self, _vm: &VirtualMachine) -> PyIntRef { self.stop.clone() } #[pyproperty(name = "step")] fn step(&self, _vm: &VirtualMachine) -> PyIntRef { self.step.clone() } #[pymethod(name = "__iter__")] fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRangeIterator { PyRangeIterator { position: Cell::new(0), range: zelf, } } #[pymethod(name = "__reversed__")] fn reversed(&self, vm: &VirtualMachine) -> PyRangeIterator { let start = self.start.as_bigint(); let stop = self.stop.as_bigint(); let step = self.step.as_bigint(); // compute the last element that is actually contained within the range // this is the new start let remainder = ((stop - start) % step).abs(); let new_start = if remainder.is_zero() { stop - step } else { stop - &remainder }; let new_stop: BigInt = match step.sign() { Sign::Plus => start - 1, Sign::Minus => start + 1, Sign::NoSign => unreachable!(), }; let reversed = PyRange { start: PyInt::new(new_start).into_ref(vm), stop: PyInt::new(new_stop).into_ref(vm), step: PyInt::new(-step).into_ref(vm), }; PyRangeIterator { position: Cell::new(0), range: reversed.into_ref(vm), } } #[pymethod(name = "__len__")] fn len(&self, _vm: &VirtualMachine) -> PyInt { self.length() } #[pymethod(name = "__repr__")] fn repr(&self, _vm: &VirtualMachine) -> String { if self.step.as_bigint().is_one() { format!("range({}, {})", self.start, self.stop) } else { format!("range({}, {}, {})", self.start, self.stop, self.step) } } #[pymethod(name = "__bool__")] fn bool(&self, _vm: &VirtualMachine) -> bool { !self.is_empty() } #[pymethod(name = "__contains__")] fn contains(&self, needle: PyObjectRef, _vm: &VirtualMachine) -> bool { if let Ok(int) = needle.downcast::() { match self.offset(int.as_bigint()) { Some(ref offset) => offset.is_multiple_of(self.step.as_bigint()), None => false, } } else { false } } #[pymethod(name = "__eq__")] fn eq(&self, rhs: PyObjectRef, vm: &VirtualMachine) -> bool { if objtype::isinstance(&rhs, &vm.ctx.range_type()) { let rhs = get_value(&rhs); if self.length().as_bigint() != rhs.length().as_bigint() { return false; } if self.length().as_bigint().is_zero() { return true; } if self.start.as_bigint() != rhs.start.as_bigint() { return false; } let step = self.step.as_bigint(); if step.is_one() || step == rhs.step.as_bigint() { return true; } false } else { false } } #[pymethod(name = "__lt__")] fn lt(&self, _rhs: PyObjectRef, vm: &VirtualMachine) -> PyResult { Ok(vm.ctx.not_implemented()) } #[pymethod(name = "__gt__")] fn gt(&self, _rhs: PyObjectRef, vm: &VirtualMachine) -> PyResult { Ok(vm.ctx.not_implemented()) } #[pymethod(name = "__ge__")] fn ge(&self, _rhs: PyObjectRef, vm: &VirtualMachine) -> PyResult { Ok(vm.ctx.not_implemented()) } #[pymethod(name = "__le__")] fn le(&self, _rhs: PyObjectRef, vm: &VirtualMachine) -> PyResult { Ok(vm.ctx.not_implemented()) } #[pymethod(name = "index")] fn index(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { if let Ok(int) = needle.downcast::() { match self.index_of(int.as_bigint()) { Some(idx) => Ok(PyInt::new(idx)), None => Err(vm.new_value_error(format!("{} is not in range", int))), } } else { Err(vm.new_value_error("sequence.index(x): x not in sequence".to_string())) } } #[pymethod(name = "count")] fn count(&self, item: PyObjectRef, _vm: &VirtualMachine) -> PyInt { if let Ok(int) = item.downcast::() { if self.index_of(int.as_bigint()).is_some() { PyInt::new(1) } else { PyInt::new(0) } } else { PyInt::new(0) } } #[pymethod(name = "__getitem__")] 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 _tmp_len = self.length(); let range_length = _tmp_len.as_bigint(); 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); Ok(PyRange { start: PyInt::new(start).into_ref(vm), stop: PyInt::new(stop).into_ref(vm), step: PyInt::new(step).into_ref(vm), } .into_ref(vm) .into_object()) } RangeIndex::Int(index) => match self.get(index.as_bigint()) { Some(value) => Ok(PyInt::new(value).into_ref(vm).into_object()), None => Err(vm.new_index_error("range object index out of range".to_string())), }, } } #[pymethod(name = "__new__")] fn range_new(args: PyFuncArgs, vm: &VirtualMachine) -> PyResult { let range = if args.args.len() <= 2 { let (cls, stop) = args.bind(vm)?; PyRange::new(cls, stop, vm) } else { let (cls, start, stop, step) = args.bind(vm)?; PyRange::new_from(cls, start, stop, step, vm) }?; Ok(range.into_object()) } } #[pyclass] #[derive(Debug)] pub struct PyRangeIterator { position: Cell, range: PyRangeRef, } impl PyValue for PyRangeIterator { fn class(vm: &VirtualMachine) -> PyClassRef { vm.ctx.rangeiterator_type() } } type PyRangeIteratorRef = PyRef; #[pyimpl] impl PyRangeIterator { #[pymethod(name = "__next__")] fn next(&self, vm: &VirtualMachine) -> PyResult { let position = BigInt::from(self.position.get()); if let Some(int) = self.range.get(&position) { self.position.set(self.position.get() + 1); Ok(int) } else { Err(objiter::new_stop_iteration(vm)) } } #[pymethod(name = "__iter__")] fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRangeIteratorRef { zelf } } pub enum RangeIndex { Int(PyIntRef), Slice(PySliceRef), } impl TryFromObject for RangeIndex { fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { match_class!(match obj { i @ PyInt => Ok(RangeIndex::Int(i)), s @ PySlice => Ok(RangeIndex::Slice(s)), obj => Err(vm.new_type_error(format!( "sequence indices be integers or slices, not {}", obj.class(), ))), }) } }