From 6b57c684e4754455d90ccd7541492f03225fc64a Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sun, 17 Mar 2024 19:22:17 +0800
Subject: [PATCH] some optimize

---
 src/iter.cpp     | 12 ++++++++----
 src/pocketpy.cpp |  1 +
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/iter.cpp b/src/iter.cpp
index e62108f1..3f2d6be6 100644
--- a/src/iter.cpp
+++ b/src/iter.cpp
@@ -7,11 +7,15 @@ namespace pkpy{
         vm->bind_notimplemented_constructor<RangeIter>(type);
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* obj){ return obj; });
         vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* obj){
-            RangeIter& self = _CAST(RangeIter&, obj);
-            bool has_next = self.r.step > 0 ? self.current < self.r.stop : self.current > self.r.stop;
-            if(!has_next) return vm->StopIteration;
+            RangeIter& self = PK_OBJ_GET(RangeIter, obj);
+            if(self.r.step > 0){
+                if(self.current >= self.r.stop) return vm->StopIteration;
+            }else{
+                if(self.current <= self.r.stop) return vm->StopIteration;
+            }
+            PyObject* ret = VAR(self.current);
             self.current += self.r.step;
-            return VAR(self.current - self.r.step);
+            return ret;
         });
     }
 
diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp
index 53e61949..6c35555a 100644
--- a/src/pocketpy.cpp
+++ b/src/pocketpy.cpp
@@ -352,6 +352,7 @@ void init_builtins(VM* _vm) {
             case 3: r.start = CAST(i64, args[0]); r.stop = CAST(i64, args[1]); r.step = CAST(i64, args[2]); break;
             default: vm->TypeError("expected 1-3 arguments, got " + std::to_string(args.size()));
         }
+        if(r.step == 0) vm->ValueError("range() arg 3 must not be zero");
         return VAR(r);
     });