make range faster

include/pocketpy/iter.h

@@ -6,7 +6,7 @@
 
 namespace pkpy{
 
-struct RangeIter{
+struct RangeIter{       // step > 0
     Range r;
     i64 current;
     RangeIter(Range r) : r(r), current(r.start) {}
@@ -14,6 +14,14 @@ struct RangeIter{
     static void _register(VM* vm, PyVar mod, PyVar type);
 };
 
+struct RangeIterR{      // step < 0
+    Range r;
+    i64 current;
+    RangeIterR(Range r) : r(r), current(r.start) {}
+
+    static void _register(VM* vm, PyVar mod, PyVar type);
+};
+
 struct ArrayIter{
     PyVar ref;
     PyVar* begin;

src/iter.cpp

@@ -6,12 +6,19 @@ namespace pkpy{
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0){ return _0; });
         vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0) -> unsigned{
             RangeIter& self = PK_OBJ_GET(RangeIter, _0);
-            if(self.r.step > 0){
-                if(self.current >= self.r.stop) return 0;
-            }else{
-                if(self.current <= self.r.stop) return 0;
-            }
-            vm->s_data.push(VAR(self.current));
+            if(self.current >= self.r.stop) return 0;
+            vm->s_data.emplace(VM::tp_int, self.current);
+            self.current += self.r.step;
+            return 1;
+        });
+    }
+
+    void RangeIterR::_register(VM* vm, PyVar mod, PyVar type){
+        vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0){ return _0; });
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0) -> unsigned{
+            RangeIterR& self = PK_OBJ_GET(RangeIterR, _0);
+            if(self.current <= self.r.stop) return 0;
+            vm->s_data.emplace(VM::tp_int, self.current);
             self.current += self.r.step;
             return 1;
         });

src/pocketpy.cpp

@@ -363,7 +363,14 @@ void __init_builtins(VM* _vm) {
         return VAR(r);
     });
 
-    _vm->bind__iter__(VM::tp_range, [](VM* vm, PyVar obj) { return vm->new_user_object<RangeIter>(PK_OBJ_GET(Range, obj)); });
+    _vm->bind__iter__(VM::tp_range, [](VM* vm, PyVar obj) {
+        const Range& r = PK_OBJ_GET(Range, obj);
+        if(r.step > 0){
+            return vm->new_user_object<RangeIter>(r);
+        }else{
+            return vm->new_user_object<RangeIterR>(r);
+        }
+    });
     
     // tp_nonetype
     _vm->bind__repr__(_vm->_tp(_vm->None), [](VM* vm, PyVar _0) -> Str {
@@ -1491,6 +1498,7 @@ void __init_builtins(VM* _vm) {
     });
 
     _vm->register_user_class<RangeIter>(_vm->builtins, "_range_iter");
+    _vm->register_user_class<RangeIterR>(_vm->builtins, "_range_iter_r");
     _vm->register_user_class<ArrayIter>(_vm->builtins, "_array_iter");
     _vm->register_user_class<StringIter>(_vm->builtins, "_string_iter");
     _vm->register_user_class<Generator>(_vm->builtins, "generator");

tests/05_list.py

@@ -134,3 +134,7 @@ a, b = [1, 2]
 assert a == 1 and b == 2
 assert [1, 2].__getitem__(0) == 1
 
+assert list(range(1, 5)) == [1, 2, 3, 4]
+assert list(range(1, 5, 2)) == [1, 3]
+assert list(range(5, 1, -1)) == [5, 4, 3, 2]
+assert list(range(5, 1, -2)) == [5, 3]