fix a bug of `LOAD_NAME`, `STORE_NAME` and `DELETE_NAME`

include/pocketpy/common.h

@@ -20,7 +20,7 @@
 #include <variant>
 #include <type_traits>
 
-#define PK_VERSION				"1.0.8"
+#define PK_VERSION				"1.0.9"
 
 #include "config.h"
 

include/pocketpy/frame.h

@@ -22,10 +22,7 @@ struct FastLocals{
     FastLocals(const CodeObject* co, PyObject** a): varnames_inv(&co->varnames_inv), a(a) {}
     FastLocals(const FastLocals& other): varnames_inv(other.varnames_inv), a(other.a) {}
 
-    PyObject* try_get(StrName name);
-    bool contains(StrName name);
-    void erase(StrName name);
-    bool try_set(StrName name, PyObject* value);
+    PyObject** try_get_name(StrName name);
     NameDict_ to_namedict();
 };
 

include/pocketpy/vm.h

@@ -351,6 +351,7 @@ public:
     void IndexError(const Str& msg){ _error("IndexError", msg); }
     void ValueError(const Str& msg){ _error("ValueError", msg); }
     void NameError(StrName name){ _error("NameError", fmt("name ", name.escape() + " is not defined")); }
+    void UnboundLocalError(StrName name){ _error("UnboundLocalError", fmt("local variable ", name.escape() + " referenced before assignment")); }
     void KeyError(PyObject* obj){ _error("KeyError", PK_OBJ_GET(Str, py_repr(obj))); }
     void BinaryOptError(const char* op) { TypeError(fmt("unsupported operand type(s) for ", op)); }
 

src/ceval.cpp

@@ -111,11 +111,15 @@ __NEXT_STEP:;
         if(_0 == PY_NULL) vm->NameError(co->varnames[byte.arg]);
         PUSH(_0);
     } DISPATCH();
-    TARGET(LOAD_NAME)
+    TARGET(LOAD_NAME) {
         heap._auto_collect();
         _name = StrName(byte.arg);
-        _0 = frame->_locals.try_get(_name);
-        if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
+        PyObject** slot = frame->_locals.try_get_name(_name);
+        if(slot != nullptr) {
+            if(*slot == PY_NULL) vm->UnboundLocalError(_name);
+            PUSH(*slot);
+            DISPATCH();
+        }
         _0 = frame->f_closure_try_get(_name);
         if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
         _0 = frame->f_globals().try_get(_name);
@@ -123,7 +127,7 @@ __NEXT_STEP:;
         _0 = vm->builtins->attr().try_get(_name);
         if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
         vm->NameError(_name);
-        DISPATCH();
+    } DISPATCH();
     TARGET(LOAD_NONLOCAL) {
         heap._auto_collect();
         _name = StrName(byte.arg);
@@ -164,16 +168,17 @@ __NEXT_STEP:;
     TARGET(STORE_FAST)
         frame->_locals[byte.arg] = POPX();
         DISPATCH();
-    TARGET(STORE_NAME)
+    TARGET(STORE_NAME){
         _name = StrName(byte.arg);
         _0 = POPX();
         if(frame->_callable != nullptr){
-            bool ok = frame->_locals.try_set(_name, _0);
-            if(!ok) vm->NameError(_name);
+            PyObject** slot = frame->_locals.try_get_name(_name);
+            if(slot == nullptr) vm->UnboundLocalError(_name);
+            *slot = _0;
         }else{
             frame->f_globals().set(_name, _0);
         }
-        DISPATCH();
+    } DISPATCH();
     TARGET(STORE_GLOBAL)
         frame->f_globals().set(StrName(byte.arg), POPX());
         DISPATCH();
@@ -202,8 +207,9 @@ __NEXT_STEP:;
     TARGET(DELETE_NAME)
         _name = StrName(byte.arg);
         if(frame->_callable != nullptr){
-            if(!frame->_locals.contains(_name)) vm->NameError(_name);
-            frame->_locals.erase(_name);
+            PyObject** slot = frame->_locals.try_get_name(_name);
+            if(slot == nullptr) vm->UnboundLocalError(_name);
+            *slot = PY_NULL;
         }else{
             if(!frame->f_globals().contains(_name)) vm->NameError(_name);
             frame->f_globals().erase(_name);

src/frame.cpp

@@ -1,28 +1,10 @@
 #include "pocketpy/frame.h"
 
 namespace pkpy{
-
-    PyObject* FastLocals::try_get(StrName name){
+    PyObject** FastLocals::try_get_name(StrName name){
         int index = varnames_inv->try_get(name);
         if(index == -1) return nullptr;
-        return a[index];
-    }
-
-    bool FastLocals::contains(StrName name){
-        return varnames_inv->contains(name);
-    }
-
-    void FastLocals::erase(StrName name){
-        int index = varnames_inv->try_get(name);
-        if(index == -1) FATAL_ERROR();
-        a[index] = nullptr;
-    }
-
-    bool FastLocals::try_set(StrName name, PyObject* value){
-        int index = varnames_inv->try_get(name);
-        if(index == -1) return false;
-        a[index] = value;
-        return true;
+        return &a[index];
     }
 
     NameDict_ FastLocals::to_namedict(){

src/vm.cpp

@@ -450,7 +450,7 @@ PyObject* VM::new_module(StrName name) {
 static std::string _opcode_argstr(VM* vm, Bytecode byte, const CodeObject* co){
     std::string argStr = byte.arg == -1 ? "" : std::to_string(byte.arg);
     switch(byte.op){
-        case OP_LOAD_CONST:
+        case OP_LOAD_CONST: case OP_FORMAT_STRING:
             if(vm != nullptr){
                 argStr += fmt(" (", CAST(Str, vm->py_repr(co->consts[byte.arg])), ")");
             }

tests/99_bugs.py

@@ -56,4 +56,19 @@ def f():
 f()
 
 # class A: a=b=1
-# class A: a, b = 1, 2
+# class A: a, b = 1, 2
+
+bmi = 0.0
+
+def test(a):
+    if a:
+        bmi = 1.4
+    return f'{bmi:.2f}'
+
+assert test(1) == '1.40'
+
+try:
+    assert test(0) == '0.00'
+    exit(1)
+except UnboundLocalError:
+    pass