...

src/ceval.h

@@ -104,7 +104,7 @@ __NEXT_STEP:;
         _name = StrName(byte.arg);
         _0 = frame->_locals.try_get(_name);
         if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
-        _0 = frame->_closure.try_get(_name);
+        _0 = frame->f_closure_try_get(_name);
         if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
         _0 = frame->f_globals().try_get(_name);
         if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
@@ -115,9 +115,7 @@ __NEXT_STEP:;
     TARGET(LOAD_NONLOCAL) {
         heap._auto_collect();
         _name = StrName(byte.arg);
-        // _0 = frame->_locals.try_get(_name);
-        // if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
-        _0 = frame->_closure.try_get(_name);
+        _0 = frame->f_closure_try_get(_name);
         if(_0 != nullptr) { PUSH(_0); DISPATCH(); }
         _0 = frame->f_globals().try_get(_name);
         if(_0 != nullptr) { PUSH(_0); DISPATCH(); }

src/frame.h

@@ -178,18 +178,24 @@ struct Frame {
 
     PyObject* _module;
     FastLocals _locals;
-    FastLocals _closure;
+    PyObject* _callable;
 
     NameDict& f_globals() noexcept { return _module->attr(); }
+    
+    PyObject* f_closure_try_get(StrName name){
+        if(_callable == nullptr) return nullptr;
+        Function& fn = OBJ_GET(Function, _callable);
+        return fn._closure.try_get(name);
+    }
 
-    Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject* co, PyObject* _module, FastLocals&& _locals, const FastLocals& _closure)
-            : _s(_s), _sp_base(_sp_base), co(co), _module(_module), _locals(std::move(_locals)), _closure(_closure) { }
+    Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject* co, PyObject* _module, FastLocals&& _locals, PyObject* _callable)
+            : _s(_s), _sp_base(_sp_base), co(co), _module(_module), _locals(std::move(_locals)), _callable(_callable) { }
 
-    Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject* co, PyObject* _module, const FastLocals& _locals, const FastLocals& _closure)
-            : _s(_s), _sp_base(_sp_base), co(co), _module(_module), _locals(_locals), _closure(_closure) { }
+    Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject* co, PyObject* _module, const FastLocals& _locals, PyObject* _callable)
+            : _s(_s), _sp_base(_sp_base), co(co), _module(_module), _locals(_locals), _callable(_callable) { }
 
     Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject_& co, PyObject* _module)
-            : _s(_s), _sp_base(_sp_base), co(co.get()), _module(_module), _locals(), _closure() { }
+            : _s(_s), _sp_base(_sp_base), co(co.get()), _module(_module), _locals(), _callable(nullptr) { }
 
     Frame(const Frame& other) = delete;
     Frame& operator=(const Frame& other) = delete;
@@ -253,7 +259,7 @@ struct Frame {
         // TODO: fix here
         OBJ_MARK(_module);
         _locals._gc_mark();
-        _closure._gc_mark();
+        if(_callable != nullptr) OBJ_MARK(_callable);
         co->_gc_mark();
     }
 };

src/vm.h

@@ -912,15 +912,12 @@ inline PyObject* VM::_py_call(PyObject** sp_base, PyObject* callable, ArgsView a
     }
     PyObject* _module = fn._module != nullptr ? fn._module : top_frame()->_module;
     
-    // TODO: callable may be garbage collected if it is a temporary object
-    // very unlikely to happen, but still possible
     s_data.reset(sp_base);
-    PyObject** curr_sp = s_data._sp;
     if(co->is_generator){
-        PyObject* ret = PyIter(Generator(this, Frame(&s_data, curr_sp, co, _module, std::move(locals), fn._closure)));
+        PyObject* ret = PyIter(Generator(this, Frame(&s_data, s_data._sp, co, _module, std::move(locals), callable)));
         return ret;
     }
-    callstack.emplace(&s_data, curr_sp, co, _module, std::move(locals), fn._closure);
+    callstack.emplace(&s_data, s_data._sp, co, _module, std::move(locals), callable);
     return nullptr;
 }