Zxc200611/pocketpy
1
0
Fork 0
mirror of https://github.com/pocketpy/pocketpy synced 2025-10-22 20:40:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

add FOR_ITER_UNPACK

Browse Source
This commit is contained in:
blueloveTH 2024-04-28 20:04:13 +08:00
parent f706944803
commit f619cee961
8 changed files with 120 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
include/pocketpy/iter.h
View File

@ -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

1
include/pocketpy/opcodes.h
View File

@ -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)

2
include/pocketpy/vm.h
View File

@ -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);

14
src/array2d.cpp
View File

@ -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;
};
}
};

68
src/ceval.cpp
View File

@ -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();
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");
}
_op_unpack_sequence(byte.arg);
} DISPATCH();
TARGET(UNPACK_EX) {
auto _lock = heap.gc_scope_lock(); // lock the gc via RAII!!

7
src/expr.cpp
View File

@ -365,7 +365,12 @@ namespace pkpy{
// build tuple and unpack it is meaningless
ctx->revert_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

42
src/iter.cpp
View File

@ -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<Generator>(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<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);
}

9
src/pocketpy.cpp
View File

@ -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>(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(){