Represent range with BigInts

jmhain · jmhain · commit 9ac8e55695f2 · 2019-02-04T17:38:23.000-08:00
diff --git a/tests/snippets/builtin_range.py b/tests/snippets/builtin_range.py
@@ -0,0 +1 @@
+assert range(2**63+1)[2**63] == 9223372036854775808
diff --git a/vm/src/obj/objiter.rs b/vm/src/obj/objiter.rs
@@ -9,7 +9,7 @@ use super::super::vm::VirtualMachine;
 use super::objbool;
 // use super::objstr;
 use super::objtype; // Required for arg_check! to use isinstance
-use num_bigint::ToBigInt;
+use num_bigint::{BigInt, ToBigInt};
 
 /*
  * This helper function is called at multiple places. First, it is called
@@ -121,7 +121,7 @@ fn iter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
             }
 
             PyObjectPayload::Range { ref range } => {
-                if let Some(int) = range.get(*position as i64) {
+                if let Some(int) = range.get(BigInt::from(*position)) {
                     *position += 1;
                     Ok(vm.ctx.new_int(int.to_bigint().unwrap()))
                 } else {
diff --git a/vm/src/obj/objrange.rs b/vm/src/obj/objrange.rs
@@ -4,27 +4,31 @@ use super::super::pyobject::{
 use super::super::vm::VirtualMachine;
 use super::objint;
 use super::objtype;
-use num_bigint::ToBigInt;
-use num_traits::ToPrimitive;
+use num_bigint::{BigInt, ToBigInt};
+use num_traits::{One, Signed, ToPrimitive, Zero};
 
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Clone)]
 pub struct RangeType {
     // Unfortunately Rust's built in range type doesn't support things like indexing
     // or ranges where start > end so we need to roll our own.
-    pub start: i64,
-    pub end: i64,
-    pub step: i64,
+    pub start: BigInt,
+    pub end: BigInt,
+    pub step: BigInt,
 }
 
 impl RangeType {
     #[inline]
     pub fn len(&self) -> usize {
-        ((self.end - self.start) / self.step).abs() as usize
+        ((self.end.clone() - self.start.clone()) / self.step.clone())
+            .abs()
+            .to_usize()
+            .unwrap()
     }
 
     #[inline]
     pub fn is_empty(&self) -> bool {
-        (self.start <= self.end && self.step < 0) || (self.start >= self.end && self.step > 0)
+        (self.start <= self.end && self.step.is_negative())
+            || (self.start >= self.end && self.step.is_positive())
     }
 
     #[inline]
@@ -33,8 +37,8 @@ impl RangeType {
     }
 
     #[inline]
-    pub fn get(&self, index: i64) -> Option<i64> {
-        let result = self.start + self.step * index;
+    pub fn get(&self, index: BigInt) -> Option<BigInt> {
+        let result = self.start.clone() + self.step.clone() * index;
 
         if self.forward() && !self.is_empty() && result < self.end {
             Some(result)
@@ -70,24 +74,24 @@ fn range_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     );
 
     let start = if let Some(_) = second {
-        objint::get_value(first).to_i64().unwrap()
+        objint::get_value(first)
     } else {
-        0i64
+        BigInt::zero()
     };
 
     let end = if let Some(pyint) = second {
-        objint::get_value(pyint).to_i64().unwrap()
+        objint::get_value(pyint)
     } else {
-        objint::get_value(first).to_i64().unwrap()
+        objint::get_value(first)
     };
 
     let step = if let Some(pyint) = step {
-        objint::get_value(pyint).to_i64().unwrap()
+        objint::get_value(pyint)
     } else {
-        1i64
+        BigInt::one()
     };
 
-    if step == 0 {
+    if step.is_zero() {
         Err(vm.new_value_error("range with 0 step size".to_string()))
     } else {
         Ok(PyObject::new(
@@ -128,15 +132,15 @@ fn range_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         args,
         required = [(zelf, Some(vm.ctx.range_type())), (subscript, None)]
     );
-    let zrange = if let PyObjectPayload::Range { range } = zelf.borrow().payload {
+    let zrange = if let PyObjectPayload::Range { ref range } = zelf.borrow().payload {
         range.clone()
     } else {
         unreachable!()
     };
 
     match subscript.borrow().payload {
         PyObjectPayload::Integer { ref value } => {
-            if let Some(int) = zrange.get(value.to_i64().unwrap()) {
+            if let Some(int) = zrange.get(value.clone()) {
                 Ok(PyObject::new(
                     PyObjectPayload::Integer {
                         value: int.to_bigint().unwrap(),
@@ -150,17 +154,17 @@ fn range_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         PyObjectPayload::Slice { start, stop, step } => {
             let new_start = if let Some(int) = start {
                 if let Some(i) = zrange.get(int.into()) {
-                    i as i64
+                    i
                 } else {
-                    zrange.start
+                    zrange.start.clone()
                 }
             } else {
-                zrange.start
+                zrange.start.clone()
             };
 
             let new_end = if let Some(int) = stop {
                 if let Some(i) = zrange.get(int.into()) {
-                    i as i64
+                    i
                 } else {
                     zrange.end
                 }

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+assert range(263+1)[263] == 9223372036854775808`