diff --git a/include/pocketpy/iter.h b/include/pocketpy/iter.h index e3c8d6e5..4288d037 100644 --- a/include/pocketpy/iter.h +++ b/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; diff --git a/src/iter.cpp b/src/iter.cpp index b509afb8..2a205152 100644 --- a/src/iter.cpp +++ b/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; }); diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp index 34607bcd..25108899 100644 --- a/src/pocketpy.cpp +++ b/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(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(r); + }else{ + return vm->new_user_object(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(_vm->builtins, "_range_iter"); + _vm->register_user_class(_vm->builtins, "_range_iter_r"); _vm->register_user_class(_vm->builtins, "_array_iter"); _vm->register_user_class(_vm->builtins, "_string_iter"); _vm->register_user_class(_vm->builtins, "generator"); diff --git a/tests/05_list.py b/tests/05_list.py index f39858c1..021b8324 100644 --- a/tests/05_list.py +++ b/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]