From 16152a1d3290be1fe8cd2773010483bf39628c68 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Fri, 26 Apr 2019 23:49:51 +0900
Subject: [PATCH] Fix range() to support negative index

---
 tests/snippets/builtin_range.py |  2 ++
 vm/src/obj/objrange.rs          | 21 ++++++++++++---------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/tests/snippets/builtin_range.py b/tests/snippets/builtin_range.py
index 2ea71f7a7..1dd5e7028 100644
--- a/tests/snippets/builtin_range.py
+++ b/tests/snippets/builtin_range.py
@@ -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
diff --git a/vm/src/obj/objrange.rs b/vm/src/obj/objrange.rs
index 1a21a49ad..56b7c0396 100644
--- a/vm/src/obj/objrange.rs
+++ b/vm/src/obj/objrange.rs
@@ -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 {