fix bugs

docs/C-API/stack.md

@@ -150,4 +150,20 @@ PK_EXPORT bool pkpy_is_none(pkpy_vm*, int i);
 
     ```
     [value] -> [repr(value)]
-    ```
+    ```
+
++ `bool pkpy_py_str(pkpy_vm*)`
+
+    Get the str of the value at the top of the stack.
+
+    ```
+    [value] -> [str(value)]
+    ```
+
++ `bool pkpy_py_import(pkpy_vm*, const char* name)`
+
+    Import a module and push it onto the stack.
+
+    ```
+    [] -> [module]
+    ```

include/pocketpy/frame.h

@@ -75,6 +75,14 @@ struct ValueStack {
     ValueStack& operator=(ValueStack&&) = delete;
 };
 
+struct UnwindTarget{
+    UnwindTarget* next;
+    int iblock;
+    int offset;
+
+    UnwindTarget(int iblock, int offset): next(nullptr), iblock(iblock), offset(offset) {}
+};
+
 struct Frame {
     const Bytecode* _ip;
     // This is for unwinding only, use `actual_sp_base()` for value stack access
@@ -85,21 +93,24 @@ struct Frame {
     PyVar _callable;    // a function object or nullptr (global scope)
     FastLocals _locals;
 
+    // This list will be freed in __pop_frame
+    UnwindTarget* _uw_list;
+
     NameDict& f_globals() { return _module->attr(); }
     PyVar f_closure_try_get(StrName name);
     int ip() const { return _ip - co->codes.data(); }
 
     // function scope
     Frame(PyVar* p0, const CodeObject* co, PyVar _module, PyVar _callable, PyVar* _locals_base)
-            : _ip(co->codes.data()-1), _sp_base(p0), co(co), _module(_module), _callable(_callable), _locals(co, _locals_base) { }
+            : _ip(co->codes.data()-1), _sp_base(p0), co(co), _module(_module), _callable(_callable), _locals(co, _locals_base), _uw_list(nullptr) { }
 
     // exec/eval
     Frame(PyVar* p0, const CodeObject* co, PyVar _module, PyVar _callable, FastLocals _locals)
-            : _ip(co->codes.data()-1), _sp_base(p0), co(co), _module(_module), _callable(_callable), _locals(_locals) { }
+            : _ip(co->codes.data()-1), _sp_base(p0), co(co), _module(_module), _callable(_callable), _locals(_locals), _uw_list(nullptr) { }
 
     // global scope
     Frame(PyVar* p0, const CodeObject_& co, PyVar _module)
-            : _ip(co->codes.data()-1), _sp_base(p0), co(co.get()), _module(_module), _callable(nullptr), _locals(co.get(), p0) {}
+            : _ip(co->codes.data()-1), _sp_base(p0), co(co.get()), _module(_module), _callable(nullptr), _locals(co.get(), p0), _uw_list(nullptr) { }
 
     PyVar* actual_sp_base() const { return _locals.a; }
     ArgsView stack_view(ValueStack* _s) const { return ArgsView(actual_sp_base(), _s->_sp); }
@@ -116,6 +127,10 @@ struct Frame {
 
     int curr_lineno() const { return co->lines[ip()].lineno; }
 
+    void set_unwind_target(PyVar* _sp);
+    UnwindTarget* find_unwind_target(int iblock);
+    void free_unwind_target();
+
     void _gc_mark(VM* vm) const;
 };
 

include/pocketpy/opcodes.h

@@ -140,6 +140,7 @@ OPCODE(ADD_CLASS_ANNOTATION)
 OPCODE(WITH_ENTER)
 OPCODE(WITH_EXIT)
 /**************************/
+OPCODE(TRY_ENTER)
 OPCODE(EXCEPTION_MATCH)
 OPCODE(RAISE)
 OPCODE(RAISE_ASSERT)

include/pocketpy/pocketpy_c.h

@@ -77,6 +77,7 @@ PK_EXPORT bool pkpy_unpack_sequence(pkpy_vm*, int size);
 PK_EXPORT bool pkpy_get_unbound_method(pkpy_vm*, pkpy_CName name);
 PK_EXPORT bool pkpy_py_repr(pkpy_vm*);
 PK_EXPORT bool pkpy_py_str(pkpy_vm*);
+PK_EXPORT bool pkpy_py_import(pkpy_vm*, pkpy_CString name);
 
 /* Error Handling */
 PK_EXPORT bool pkpy_error(pkpy_vm*, const char* name, pkpy_CString msg);

src/ceval.cpp

@@ -986,6 +986,10 @@ __NEXT_STEP:
         POP();
         DISPATCH()
     /*****************************************/
+    case OP_TRY_ENTER: {
+        frame->set_unwind_target(s_data._sp);
+        DISPATCH()
+    }
     case OP_EXCEPTION_MATCH: {
         PyVar assumed_type = POPX();
         check_type(assumed_type, tp_type);

src/compiler.cpp

@@ -703,6 +703,7 @@ __EAT_DOTS_END:
 
     void Compiler::compile_try_except() {
         ctx()->enter_block(CodeBlockType::TRY_EXCEPT);
+        ctx()->emit_(OP_TRY_ENTER, BC_NOARG, prev().line);
         compile_block_body();
         small_vector_2<int, 6> patches;
         patches.push_back(

src/frame.cpp

@@ -24,15 +24,17 @@ namespace pkpy{
     }
 
     int Frame::prepare_jump_exception_handler(ValueStack* _s){
-        PyVar obj = _s->popx();
         // try to find a parent try block
         int i = co->lines[ip()].iblock;
         while(i >= 0){
             if(co->blocks[i].type == CodeBlockType::TRY_EXCEPT) break;
-            i = _exit_block(_s, i);
+            i = co->blocks[i].parent;
         }
-        _s->push(obj);
         if(i < 0) return -1;
+        PyVar obj = _s->popx();                     // pop exception object
+        UnwindTarget* uw = find_unwind_target(i);
+        _s->reset(actual_sp_base() + uw->offset);   // unwind the stack
+        _s->push(obj);                              // push it back
         return co->blocks[i].end;
     }
 
@@ -69,4 +71,32 @@ namespace pkpy{
         }
     }
 
+    void Frame::set_unwind_target(PyVar* _sp){
+        int iblock = co->lines[ip()].iblock;
+        UnwindTarget* existing = find_unwind_target(iblock);
+        if(existing){
+            existing->offset = _sp - actual_sp_base();
+        }else{
+            UnwindTarget* prev = _uw_list;
+            _uw_list = new UnwindTarget(iblock, _sp - actual_sp_base());
+            _uw_list->next = prev;
+        }
+    }
+
+    UnwindTarget* Frame::find_unwind_target(int iblock){
+        UnwindTarget* p;
+        for(p=_uw_list; p!=nullptr; p=p->next){
+            if(p->iblock == iblock) return p;
+        }
+        return nullptr;
+    }
+
+    void Frame::free_unwind_target(){
+        while(_uw_list != nullptr){
+            UnwindTarget* p = _uw_list;
+            _uw_list = p->next;
+            delete p;
+        }
+    }
+
 }   // namespace pkpy

src/iter.cpp

@@ -55,6 +55,17 @@ namespace pkpy{
         frame._locals.a = vm->s_data._sp;
         // restore the context
         for(PyVar obj: s_backup) vm->s_data.push(obj);
+        // relocate stack objects (their addresses become invalid)
+        for(PyVar* p=frame.actual_sp_base(); p!=vm->s_data.end(); p++){
+            if(p->type == VM::tp_stack_memory){
+                // TODO: refactor this
+                int count = p->as<StackMemory>().count;
+                if(count < 0){
+                    void* new_p = p + count;
+                    p[1]._1 = reinterpret_cast<i64>(new_p);
+                }
+            }
+        }
         s_backup.clear();
         vm->callstack.emplace(std::move(frame));
 

src/pocketpy_c.cpp

@@ -489,6 +489,16 @@ bool pkpy_py_str(pkpy_vm* vm_handle) {
     return true;
 }
 
+bool pkpy_py_import(pkpy_vm* vm_handle, pkpy_CString name) {
+    VM* vm = (VM*) vm_handle;
+    PK_ASSERT_NO_ERROR()
+    PK_PROTECTED(
+        PyVar module = vm->py_import(name);
+        vm->s_data.push(module);
+    )
+    return true;
+}
+
 /* Error Handling */
 bool pkpy_error(pkpy_vm* vm_handle, const char* name, pkpy_CString message) {
     VM* vm = (VM*) vm_handle;

src/vm.cpp

@@ -1644,7 +1644,9 @@ void NextBreakpoint::_step(VM* vm){
 #endif
 
 void VM::__pop_frame(){
-    s_data.reset(callstack.top()._sp_base);
+    Frame& frame = callstack.top();
+    s_data.reset(frame._sp_base);
+    frame.free_unwind_target();
     callstack.pop();
 
 #if PK_ENABLE_PROFILER

src2/main.cpp

@@ -65,7 +65,7 @@ int main(int argc, char** argv){
     pkpy_vm* vm = pkpy_new_vm(true);
 
     pkpy_push_function(vm, "input(prompt=None) -> str", f_input);
-    pkpy_eval(vm, "__import__('builtins')");
+    pkpy_py_import(vm, "builtins");
     pkpy_setattr(vm, pkpy_name("input"));
 
     if(argc == 1){

src2/pocketpy_c.c

@@ -203,6 +203,11 @@ bool pkpy_py_str(pkpy_vm* vm) {
     return returnValue;
 }
 
+bool pkpy_py_import(pkpy_vm* vm, pkpy_CString name) {
+    bool returnValue;
+    return returnValue;
+}
+
 bool pkpy_error(pkpy_vm* vm, const char* name, pkpy_CString msg) {
     bool returnValue;
     return returnValue;

tests/70_collections.py

@@ -210,9 +210,12 @@ for e in [d, deque('abc'), deque('ab'), deque(), list(d)]:
     assertEqual(d == e, type(d) == type(e) and list(d) == list(e))
     assertEqual(d != e, not (type(d) == type(e) and list(d) == list(e)))
 
-args = map(deque, ('', 'a', 'b', 'ab', 'ba', 'abc', 'xba', 'xabc', 'cba'))
-for x in args:
-    for y in args:
+def get_args():
+    args = map(deque, ('', 'a', 'b', 'ab', 'ba', 'abc', 'xba', 'xabc', 'cba'))
+    return args
+
+for x in get_args():
+    for y in get_args():
         assertEqual(x == y, list(x) == list(y))
         assertEqual(x != y, list(x) != list(y))
         # assertEqual(x <  y, list(x) <  list(y))   # not currently supported

tests/99_bugs.py

@@ -1,3 +1,11 @@
+# multi-loop generator
+out = []
+args = map(lambda x: x+1, [1, 2, 3])
+for x in args:
+    for y in args:
+        out.append((x, y))
+assert out == [(2, 3), (2, 4)]
+
 # multi loop bug
 out = []
 a = [1, 2]