break change of `m__next__`

include/pocketpy/frame.h

@@ -110,6 +110,10 @@ struct Frame {
     int _exit_block(ValueStack*, int);
     void jump_abs_break(ValueStack*, int);
 
+    void loop_break(ValueStack* s_data, const CodeObject*){
+        jump_abs_break(s_data, co->_get_block_codei(_ip).end);
+    }
+
     int curr_lineno() const { return co->lines[_ip].lineno; }
 
     void _gc_mark() const {

include/pocketpy/vm.h

@@ -57,8 +57,7 @@ struct PyTypeInfo{
     i64 (*m__hash__)(VM* vm, PyObject*) = nullptr;
     i64 (*m__len__)(VM* vm, PyObject*) = nullptr;
     PyObject* (*m__iter__)(VM* vm, PyObject*) = nullptr;
-    PyObject* (*m__next__)(VM* vm, PyObject*) = nullptr;
-    unsigned int (*m__next__unpack)(VM* vm, PyObject*) = nullptr;
+    unsigned (*m__next__)(VM* vm, PyObject*) = nullptr;
     PyObject* (*m__neg__)(VM* vm, PyObject*) = nullptr;
     PyObject* (*m__invert__)(VM* vm, PyObject*) = nullptr;
 
@@ -253,7 +252,7 @@ public:
     void bind__repr__(Type type, PyObject* (*f)(VM*, PyObject*));
     void bind__str__(Type type, PyObject* (*f)(VM*, PyObject*));
     void bind__iter__(Type type, PyObject* (*f)(VM*, PyObject*));
-    void bind__next__(Type type, PyObject* (*f)(VM*, PyObject*));
+    void bind__next__(Type type, unsigned (*f)(VM*, PyObject*));
     void bind__neg__(Type type, PyObject* (*f)(VM*, PyObject*));
     void bind__invert__(Type type, PyObject* (*f)(VM*, PyObject*));
     void bind__hash__(Type type, i64 (*f)(VM* vm, PyObject*));
@@ -310,6 +309,7 @@ public:
 
     i64 normalized_index(i64 index, int size);
     PyObject* py_next(PyObject* obj);
+    PyObject* _pack_next_retval(unsigned);
     bool py_callable(PyObject* obj);
     
     /***** Error Reporter *****/

src/array2d.cpp

@@ -356,19 +356,10 @@ struct Array2dIter{
     void _gc_mark() const{ PK_OBJ_MARK(ref); }
 
     static void _register(VM* vm, PyObject* mod, PyObject* type){
-        PyTypeInfo& info = vm->_all_types[PK_OBJ_GET(Type, type)];
-        info.subclass_enabled = false;
+        vm->_all_types[PK_OBJ_GET(Type, type)].subclass_enabled = false;
         vm->bind_notimplemented_constructor<Array2dIter>(type);
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0) { return _0; });
-        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
-            Array2dIter& self = PK_OBJ_GET(Array2dIter, _0);
-            Array2d& a = PK_OBJ_GET(Array2d, self.ref);
-            if(self.i == a.numel) return vm->StopIteration;
-            std::div_t res = std::div(self.i, a.n_cols);
-            return VAR(Tuple(VAR(res.rem), VAR(res.quot), a.data[self.i++]));
-        });
-
-        info.m__next__unpack = [](VM* vm, PyObject* _0) -> unsigned int{
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0) -> unsigned{
             Array2dIter& self = PK_OBJ_GET(Array2dIter, _0);
             Array2d& a = PK_OBJ_GET(Array2d, self.ref);
             if(self.i == a.numel) return 0;
@@ -377,7 +368,7 @@ struct Array2dIter{
             vm->s_data.push(VAR(res.quot));
             vm->s_data.push(a.data[self.i++]);
             return 3;
-        };
+        });
     }
 };
 

src/ceval.cpp

@@ -760,45 +760,45 @@ __NEXT_STEP:;
         DISPATCH();
     TARGET(FOR_ITER){
         PyObject* _0 = py_next(TOP());
-        if(_0 != StopIteration){
-            PUSH(_0);
-        }else{
-            frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
-        }
+        if(_0 == StopIteration) frame->loop_break(&s_data, co);
+        else PUSH(_0);
     } DISPATCH();
     TARGET(FOR_ITER_STORE_FAST){
         PyObject* _0 = py_next(TOP());
-        if(_0 != StopIteration){
-            frame->_locals[byte.arg] = _0;
+        if(_0 == StopIteration){
+            frame->loop_break(&s_data, co);
         }else{
-            frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
+            frame->_locals[byte.arg] = _0;
         }
     } DISPATCH()
     TARGET(FOR_ITER_STORE_GLOBAL){
         PyObject* _0 = py_next(TOP());
-        if(_0 != StopIteration){
-            frame->f_globals().set(StrName(byte.arg), _0);
+        if(_0 == StopIteration){
+            frame->loop_break(&s_data, co);
         }else{
-            frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
+            frame->f_globals().set(StrName(byte.arg), _0);
         }
     } DISPATCH()
     TARGET(FOR_ITER_YIELD_VALUE){
         PyObject* _0 = py_next(TOP());
-        if(_0 != StopIteration){
+        if(_0 == StopIteration){
+            frame->loop_break(&s_data, co);
+        }else{
             PUSH(_0);
             return PY_OP_YIELD;
-        }else{
-            frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
         }
     } DISPATCH()
     TARGET(FOR_ITER_UNPACK){
         PyObject* _0 = TOP();
         const PyTypeInfo* _ti = _inst_type_info(_0);
-        if(_ti->m__next__unpack){
-            unsigned int n = _ti->m__next__unpack(this, _0);
+        if(_ti->m__next__){
+            unsigned n = _ti->m__next__(this, _0);
             if(n == 0){
                 // StopIteration
-                frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
+                frame->loop_break(&s_data, co);
+            }else if(n == 1){
+                // UNPACK_SEQUENCE
+                _op_unpack_sequence(byte.arg);
             }else{
                 if(n != byte.arg){
                     ValueError(_S("expected ", (int)byte.arg, " values to unpack, got ", (int)n));
@@ -806,14 +806,13 @@ __NEXT_STEP:;
             }
         }else{
             // FOR_ITER
-            if(_ti->m__next__) _0 = _ti->m__next__(this, _0);
-            else _0 = call_method(_0, __next__);
+            _0 = call_method(_0, __next__);
             if(_0 != StopIteration){
                 PUSH(_0);
                 // UNPACK_SEQUENCE
                 _op_unpack_sequence(byte.arg);
             }else{
-                frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
+                frame->loop_break(&s_data, co);
             }
         }
     } DISPATCH()

src/collections.cpp

@@ -30,20 +30,20 @@ namespace pkpy
 
         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)
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM *vm, PyObject *obj) -> unsigned
                          {
             PyDequeIter& self = _CAST(PyDequeIter&, obj);
             if(self.is_reversed){
-                if(self.rcurrent == self.rend) return vm->StopIteration;
-                PyObject* ret = *self.rcurrent;
+                if(self.rcurrent == self.rend) return 0;
+                vm->s_data.push(*self.rcurrent);
                 ++self.rcurrent;
-                return ret;
+                return 1;
             }
             else{
-                if(self.current == self.end) return vm->StopIteration;
-                PyObject* ret = *self.current;
+                if(self.current == self.end) return 0;
+                vm->s_data.push(*self.current);
                 ++self.current;
-                return ret;
+                return 1;
             } });
     }
     struct PyDeque

src/iter.cpp

@@ -6,16 +6,16 @@ namespace pkpy{
         vm->_all_types[PK_OBJ_GET(Type, type)].subclass_enabled = false;
         vm->bind_notimplemented_constructor<RangeIter>(type);
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){ return _0; });
-        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0) -> unsigned{
             RangeIter& self = PK_OBJ_GET(RangeIter, _0);
             if(self.r.step > 0){
-                if(self.current >= self.r.stop) return vm->StopIteration;
+                if(self.current >= self.r.stop) return 0;
             }else{
-                if(self.current <= self.r.stop) return vm->StopIteration;
+                if(self.current <= self.r.stop) return 0;
             }
-            PyObject* ret = VAR(self.current);
+            vm->s_data.push(VAR(self.current));
             self.current += self.r.step;
-            return ret;
+            return 1;
         });
     }
 
@@ -23,10 +23,11 @@ namespace pkpy{
         vm->_all_types[PK_OBJ_GET(Type, type)].subclass_enabled = false;
         vm->bind_notimplemented_constructor<ArrayIter>(type);
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){ return _0; });
-        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0) -> unsigned{
             ArrayIter& self = _CAST(ArrayIter&, _0);
-            if(self.current == self.end) return vm->StopIteration;
-            return *self.current++;
+            if(self.current == self.end) return 0;
+            vm->s_data.push(*self.current++);
+            return 1;
         });
     }
 
@@ -34,14 +35,15 @@ namespace pkpy{
         vm->_all_types[PK_OBJ_GET(Type, type)].subclass_enabled = false;
         vm->bind_notimplemented_constructor<StringIter>(type);
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){ return _0; });
-        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0) -> unsigned{
             StringIter& self = _CAST(StringIter&, _0);
             Str& s = PK_OBJ_GET(Str, self.ref);
-            if(self.i == s.size) return vm->StopIteration;
+            if(self.i == s.size) return 0;
             int start = self.i;
             int len = utf8len(s.data[self.i]);
             self.i += len;
-            return VAR(s.substr(start, len));
+            vm->s_data.push(VAR(s.substr(start, len)));
+            return 1;
         });
     }
 
@@ -82,9 +84,12 @@ namespace pkpy{
         vm->_all_types[PK_OBJ_GET(Type, type)].subclass_enabled = false;
         vm->bind_notimplemented_constructor<Generator>(type);
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){ return _0; });
-        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0) -> unsigned{
             Generator& self = _CAST(Generator&, _0);
-            return self.next(vm);
+            PyObject* retval = self.next(vm);
+            if(retval == vm->StopIteration) return 0;
+            vm->s_data.push(retval);
+            return 1;
         });
     }
 
@@ -93,16 +98,7 @@ namespace pkpy{
         info.subclass_enabled = false;
         vm->bind_notimplemented_constructor<DictItemsIter>(type);
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){ return _0; });
-        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
-            DictItemsIter& self = _CAST(DictItemsIter&, _0);
-            Dict& d = PK_OBJ_GET(Dict, self.ref);
-            if(self.i == -1) return vm->StopIteration;
-            PyObject* retval = VAR(Tuple(d._items[self.i].first, d._items[self.i].second));
-            self.i = d._nodes[self.i].next;
-            return retval;
-        });
-
-        info.m__next__unpack = [](VM* vm, PyObject* _0) -> unsigned int{
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0) -> unsigned{
             DictItemsIter& self = _CAST(DictItemsIter&, _0);
             Dict& d = PK_OBJ_GET(Dict, self.ref);
             if(self.i == -1) return 0;
@@ -110,7 +106,7 @@ namespace pkpy{
             vm->s_data.push(d._items[self.i].second);
             self.i = d._nodes[self.i].next;
             return 2;
-        };
+        });
     }
 
 PyObject* VM::_py_generator(Frame&& frame, ArgsView buffer){

src/vm.cpp

@@ -252,7 +252,10 @@ namespace pkpy{
 
     PyObject* VM::py_next(PyObject* obj){
         const PyTypeInfo* ti = _inst_type_info(obj);
-        if(ti->m__next__) return ti->m__next__(this, obj);
+        if(ti->m__next__){
+            unsigned n = ti->m__next__(this, obj);
+            return _pack_next_retval(n);
+        }
         return call_method(obj, __next__);
     }
 
@@ -1312,6 +1315,24 @@ void VM::bind__delitem__(Type type, void (*f)(VM*, PyObject*, PyObject*)){
     PK_OBJ_GET(NativeFunc, nf).set_userdata(f);
 }
 
+
+    PyObject* VM::_pack_next_retval(unsigned n){
+        if(n == 0) return StopIteration;
+        if(n == 1) return s_data.popx();
+        PyObject* retval = VAR(s_data.view(n).to_tuple());
+        s_data._sp -= n;
+        return retval;
+    }
+
+    void VM::bind__next__(Type type, unsigned (*f)(VM*, PyObject*)){                         \
+        _all_types[type].m__next__ = f;                                                   \
+        PyObject* nf = bind_method<0>(_t(type), __next__, [](VM* vm, ArgsView args){       \
+            int n = lambda_get_userdata<unsigned(*)(VM*, PyObject*)>(args.begin())(vm, args[0]);\
+            return vm->_pack_next_retval(n);                                               \
+        });                                                                             \
+        PK_OBJ_GET(NativeFunc, nf).set_userdata(f);                                        \
+    }
+
 #define BIND_UNARY_SPECIAL(name)                                                        \
     void VM::bind##name(Type type, PyObject* (*f)(VM*, PyObject*)){                         \
         _all_types[type].m##name = f;                                                   \
@@ -1324,7 +1345,6 @@ void VM::bind__delitem__(Type type, void (*f)(VM*, PyObject*, PyObject*)){
     BIND_UNARY_SPECIAL(__repr__)
     BIND_UNARY_SPECIAL(__str__)
     BIND_UNARY_SPECIAL(__iter__)
-    BIND_UNARY_SPECIAL(__next__)
     BIND_UNARY_SPECIAL(__neg__)
     BIND_UNARY_SPECIAL(__invert__)
 #undef BIND_UNARY_SPECIAL