diff --git a/include/pocketpy/iter.h b/include/pocketpy/iter.h index 116d2a95..2e3635bc 100644 --- a/include/pocketpy/iter.h +++ b/include/pocketpy/iter.h @@ -32,13 +32,9 @@ struct ArrayIter{ struct StringIter{ PY_CLASS(StringIter, builtins, _string_iterator) PyObject* ref; - Str* str; - int index; // byte index - - StringIter(PyObject* ref) : ref(ref), str(&PK_OBJ_GET(Str, ref)), index(0) {} - + int i; // byte index + StringIter(PyObject* ref) : ref(ref), i(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 \ No newline at end of file diff --git a/include/pocketpy/opcodes.h b/include/pocketpy/opcodes.h index 6a13a478..5bdf2581 100644 --- a/include/pocketpy/opcodes.h +++ b/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) diff --git a/include/pocketpy/vm.h b/include/pocketpy/vm.h index f4a6691a..35a9653d 100644 --- a/include/pocketpy/vm.h +++ b/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); diff --git a/src/array2d.cpp b/src/array2d.cpp index fb931e1d..104b170e 100644 --- a/src/array2d.cpp +++ b/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(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; + }; } }; diff --git a/src/ceval.cpp b/src/ceval.cpp index 5239197f..17faf236 100644 --- a/src/ceval.cpp +++ b/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; ijump_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(); - 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; irevert_last_emit_(); }else{ - ctx->emit_(OP_UNPACK_SEQUENCE, items.size(), line); + 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 diff --git a/src/iter.cpp b/src/iter.cpp index 12c70bac..48ebd6df 100644 --- a/src/iter.cpp +++ b/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; - int start = self.index; - int len = utf8len(self.str->data[self.index]); - self.index += len; - return VAR(self.str->substr(start, len)); + Str& s = PK_OBJ_GET(Str, self.ref); + if(self.i == s.size) return vm->StopIteration; + int start = self.i; + int len = utf8len(s.data[self.i]); + 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(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){ - Generator& self = _CAST(Generator&, 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* _0){ + 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(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); } diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp index bc53e99e..d3ca8281 100644 --- a/src/pocketpy.cpp +++ b/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]); - Tuple items(self.size()); - int j = 0; - self.apply([&](PyObject* k, PyObject* v){ - items[j++] = VAR(Tuple(k, v)); - }); - return VAR(std::move(items)); + return vm->heap.gcnew(DictItemsIter::_type(vm), args[0]); }); _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(){