Fix range() to support negative index

This commit is contained in:
Jeong YunWon
2019-04-26 23:49:51 +09:00
parent 6e867edb6d
commit 16152a1d32
2 changed files with 14 additions and 9 deletions

View File

@@ -54,3 +54,5 @@ assert list(reversed(range(1,10,5))) == [6, 1]
# range retains the original int refs
i = 2**64
assert range(i).stop is i
assert range(10)[-1] == 9

View File

@@ -1,5 +1,4 @@
use std::cell::Cell;
use std::ops::Mul;
use num_bigint::{BigInt, Sign};
use num_integer::Integer;
@@ -65,15 +64,18 @@ impl PyRange {
}
#[inline]
pub fn get<'a, T>(&'a self, index: T) -> Option<BigInt>
where
&'a BigInt: Mul<T, Output = BigInt>,
{
pub fn get(&self, index: &BigInt) -> Option<BigInt> {
let start = self.start.as_bigint();
let stop = self.stop.as_bigint();
let step = self.step.as_bigint();
let result = start + step * index;
let index = if index < &BigInt::zero() {
stop + index
} else {
index.clone()
};
let result = start + step * &index;
if (self.forward() && !self.is_empty() && &result < stop)
|| (!self.forward() && !self.is_empty() && &result > stop)
@@ -274,7 +276,7 @@ impl PyRange {
}
RangeIndex::Slice(slice) => {
let new_start = if let Some(int) = slice.start_index(vm)? {
if let Some(i) = self.get(int) {
if let Some(i) = self.get(&int) {
PyInt::new(i).into_ref(vm)
} else {
self.start.clone()
@@ -284,7 +286,7 @@ impl PyRange {
};
let new_end = if let Some(int) = slice.stop_index(vm)? {
if let Some(i) = self.get(int) {
if let Some(i) = self.get(&int) {
PyInt::new(i).into_ref(vm)
} else {
self.stop.clone()
@@ -339,7 +341,8 @@ type PyRangeIteratorRef = PyRef<PyRangeIterator>;
impl PyRangeIteratorRef {
fn next(self, vm: &VirtualMachine) -> PyResult<BigInt> {
if let Some(int) = self.range.get(self.position.get()) {
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 {