add FOR_ITER_UNPACK

include/pocketpy/iter.h

@@ -32,13 +32,9 @@ struct ArrayIter{
 struct StringIter{
     PY_CLASS(StringIter, builtins, _string_iterator)
     PyObject* ref;
-    Str* str;
+    int i;      // byte index
-    int index;      // byte index
+    StringIter(PyObject* ref) : ref(ref), i(0) {}
-
-    StringIter(PyObject* ref) : ref(ref), str(&PK_OBJ_GET(Str, ref)), index(0) {}
-
     void _gc_mark() const{ PK_OBJ_MARK(ref); }
-
     static void _register(VM* vm, PyObject* mod, PyObject* type);
 };
 
@@ -61,4 +57,15 @@ struct Generator{
     static void _register(VM* vm, PyObject* mod, PyObject* type);
 };
 
+struct DictItemsIter{
+    PY_CLASS(DictItemsIter, builtins, _dict_items_iterator)
+    PyObject* ref;
+    int i;
+    DictItemsIter(PyObject* ref) : ref(ref) {
+        i = PK_OBJ_GET(Dict, ref)._head_idx;
+    }
+    void _gc_mark() const{ PK_OBJ_MARK(ref); }
+    static void _register(VM* vm, PyObject* mod, PyObject* type);
+};
+
 } // namespace pkpy

include/pocketpy/opcodes.h

@@ -118,6 +118,7 @@ OPCODE(FOR_ITER)
 OPCODE(FOR_ITER_STORE_FAST)
 OPCODE(FOR_ITER_STORE_GLOBAL)
 OPCODE(FOR_ITER_YIELD_VALUE)
+OPCODE(FOR_ITER_UNPACK)
 /**************************/
 OPCODE(IMPORT_PATH)
 OPCODE(POP_IMPORT_STAR)

include/pocketpy/vm.h

@@ -58,6 +58,7 @@ struct PyTypeInfo{
     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;
     PyObject* (*m__neg__)(VM* vm, PyObject*) = nullptr;
     PyObject* (*m__invert__)(VM* vm, PyObject*) = nullptr;
 
@@ -415,6 +416,7 @@ public:
     PyObject* _run_top_frame();
     void post_init();
     PyObject* _py_generator(Frame&& frame, ArgsView buffer);
+    void _op_unpack_sequence(uint16_t arg);
     void _prepare_py_call(PyObject**, ArgsView, ArgsView, const FuncDecl_&);
     // new style binding api
     PyObject* bind(PyObject*, const char*, const char*, NativeFuncC, UserData userdata={}, BindType bt=BindType::DEFAULT);

src/array2d.cpp

@@ -366,7 +366,8 @@ struct Array2dIter{
     void _gc_mark() const{ PK_OBJ_MARK(ref); }
 
     static void _register(VM* vm, PyObject* mod, PyObject* type){
-        vm->_all_types[PK_OBJ_GET(Type, type)].subclass_enabled = false;
+        PyTypeInfo& info = vm->_all_types[PK_OBJ_GET(Type, type)];
+        info.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){
@@ -376,6 +377,17 @@ struct Array2dIter{
             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{
+            Array2dIter& self = PK_OBJ_GET(Array2dIter, _0);
+            Array2d& a = PK_OBJ_GET(Array2d, self.ref);
+            if(self.i == a.numel) return 0;
+            std::div_t res = std::div(self.i, a.n_cols);
+            vm->s_data.push(VAR(res.rem));
+            vm->s_data.push(VAR(res.quot));
+            vm->s_data.push(a.data[self.i++]);
+            return 3;
+        };
     }
 };
 

src/ceval.cpp

@@ -35,6 +35,28 @@ namespace pkpy{
         }
 
 
+void VM::_op_unpack_sequence(uint16_t arg){
+    PyObject* _0 = POPX();
+    if(is_type(_0, VM::tp_tuple)){
+        // fast path for tuple
+        Tuple& tuple = PK_OBJ_GET(Tuple, _0);
+        if(tuple.size() == arg){
+            for(PyObject* obj: tuple) PUSH(obj);
+        }else{
+            ValueError(_S("expected ", (int)arg, " values to unpack, got ", (int)tuple.size()));
+        }
+    }else{
+        auto _lock = heap.gc_scope_lock();  // lock the gc via RAII!!
+        _0 = py_iter(_0);
+        for(int i=0; i<arg; i++){
+            PyObject* _1 = py_next(_0);
+            if(_1 == StopIteration) ValueError("not enough values to unpack");
+            PUSH(_1);
+        }
+        if(py_next(_0) != StopIteration) ValueError("too many values to unpack");
+    }
+}
+
 bool VM::py_lt(PyObject* _0, PyObject* _1){
     BINARY_F_COMPARE(__lt__, "<", __gt__);
     return ret == True;
@@ -769,6 +791,32 @@ __NEXT_STEP:;
             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(n == 0){
+                // StopIteration
+                frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
+            }else{
+                if(n != byte.arg){
+                    ValueError(_S("expected ", (int)byte.arg, " values to unpack, got ", (int)n));
+                }
+            }
+        }else{
+            // FOR_ITER
+            if(_ti->m__next__) _0 = _ti->m__next__(this, _0);
+            else _0 = call_method(_0, __next__);
+            if(_0 != StopIteration){
+                PUSH(_0);
+            }else{
+                frame->jump_abs_break(&s_data, co->_get_block_codei(frame->_ip).end);
+            }
+            // UNPACK_SEQUENCE
+            _op_unpack_sequence(byte.arg);
+        }
+    } DISPATCH()
     /*****************************************/
     TARGET(IMPORT_PATH){
         PyObject* _0 = co->consts[byte.arg];
@@ -798,25 +846,7 @@ __NEXT_STEP:;
     } DISPATCH();
     /*****************************************/
     TARGET(UNPACK_SEQUENCE){
-        PyObject* _0 = POPX();
+        _op_unpack_sequence(byte.arg);
-        if(is_type(_0, VM::tp_tuple)){
-            // fast path for tuple
-            Tuple& tuple = PK_OBJ_GET(Tuple, _0);
-            if(tuple.size() == byte.arg){
-                for(PyObject* obj: tuple) PUSH(obj);
-            }else{
-                ValueError(_S("expected ", (int)byte.arg, " values to unpack, got ", (int)tuple.size()));
-            }
-        }else{
-            auto _lock = heap.gc_scope_lock();  // lock the gc via RAII!!
-            _0 = py_iter(_0);
-            for(int i=0; i<byte.arg; i++){
-                PyObject* _1 = py_next(_0);
-                if(_1 == StopIteration) ValueError("not enough values to unpack");
-                PUSH(_1);
-            }
-            if(py_next(_0) != StopIteration) ValueError("too many values to unpack");
-        }
     } DISPATCH();
     TARGET(UNPACK_EX) {
         auto _lock = heap.gc_scope_lock();  // lock the gc via RAII!!

src/expr.cpp

@@ -364,9 +364,14 @@ namespace pkpy{
             if(prev.op == OP_BUILD_TUPLE && prev.arg == items.size()){
                 // build tuple and unpack it is meaningless
                 ctx->revert_last_emit_();
+            }else{
+                if(prev.op == OP_FOR_ITER){
+                    prev.op = OP_FOR_ITER_UNPACK;
+                    prev.arg = items.size();
                 }else{
                     ctx->emit_(OP_UNPACK_SEQUENCE, items.size(), line);
                 }
+            }
         }else{
             // starred assignment target must be in a tuple
             if(items.size() == 1) return false;

src/iter.cpp

@@ -36,11 +36,12 @@ namespace pkpy{
         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){
             StringIter& self = _CAST(StringIter&, _0);
-            if(self.index == self.str->size) return vm->StopIteration;
+            Str& s = PK_OBJ_GET(Str, self.ref);
-            int start = self.index;
+            if(self.i == s.size) return vm->StopIteration;
-            int len = utf8len(self.str->data[self.index]);
+            int start = self.i;
-            self.index += len;
+            int len = utf8len(s.data[self.i]);
-            return VAR(self.str->substr(start, len));
+            self.i += len;
+            return VAR(s.substr(start, len));
         });
     }
 
@@ -80,13 +81,38 @@ namespace pkpy{
     void Generator::_register(VM* vm, PyObject* mod, PyObject* type){
         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* obj){ return obj; });
+        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* obj){
+        vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
-            Generator& self = _CAST(Generator&, obj);
+            Generator& self = _CAST(Generator&, _0);
             return self.next(vm);
         });
     }
 
+    void DictItemsIter::_register(VM *vm, PyObject *mod, PyObject *type){
+        PyTypeInfo& info = vm->_all_types[PK_OBJ_GET(Type, type)];
+        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{
+            DictItemsIter& self = _CAST(DictItemsIter&, _0);
+            Dict& d = PK_OBJ_GET(Dict, self.ref);
+            if(self.i == -1) return 0;
+            vm->s_data.push(d._items[self.i].first);
+            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){
     return VAR_T(Generator, std::move(frame), buffer);
 }

src/pocketpy.cpp

@@ -1381,13 +1381,7 @@ void init_builtins(VM* _vm) {
     });
 
     _vm->bind_method<0>(VM::tp_dict, "items", [](VM* vm, ArgsView args) {
-        const Dict& self = _CAST(Dict&, args[0]);
+        return vm->heap.gcnew<DictItemsIter>(DictItemsIter::_type(vm), args[0]);
-        Tuple items(self.size());
-        int j = 0;
-        self.apply([&](PyObject* k, PyObject* v){
-            items[j++] = VAR(Tuple(k, v));
-        });
-        return VAR(std::move(items));
     });
 
     _vm->bind_method<1>(VM::tp_dict, "update", [](VM* vm, ArgsView args) {
@@ -1503,6 +1497,7 @@ void init_builtins(VM* _vm) {
     ArrayIter::register_class(_vm, _vm->builtins);
     StringIter::register_class(_vm, _vm->builtins);
     Generator::register_class(_vm, _vm->builtins);
+    DictItemsIter::register_class(_vm, _vm->builtins);
 }
 
 void VM::post_init(){