From 53e4043e0e9709c9500cb0494fe0be0ad2cbb127 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sun, 15 Oct 2023 04:44:10 +0800
Subject: [PATCH] fix a bug of `_class` in `Function`

---
 include/pocketpy/codeobject.h | 7 +++++--
 src/ceval.cpp                 | 6 +++---
 src/pocketpy.cpp              | 2 +-
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/include/pocketpy/codeobject.h b/include/pocketpy/codeobject.h
index 4a131cc8..b2f9c168 100644
--- a/include/pocketpy/codeobject.h
+++ b/include/pocketpy/codeobject.h
@@ -190,8 +190,12 @@ struct NativeFunc {
 
 struct Function{
     FuncDecl_ decl;
-    PyObject* _module;
+    PyObject* _module;  // weak ref
+    PyObject* _class;   // weak ref
     NameDict_ _closure;
+
+    explicit Function(FuncDecl_ decl, PyObject* _module, PyObject* _class, NameDict_ _closure):
+        decl(decl), _module(_module), _class(_class), _closure(_closure) {}
 };
 
 template<>
@@ -203,7 +207,6 @@ struct Py_<Function> final: PyObject {
     }
     void _obj_gc_mark() override {
         _value.decl->_gc_mark();
-        if(_value._module != nullptr) PK_OBJ_MARK(_value._module);
         if(_value._closure != nullptr) gc_mark_namedict(*_value._closure);
     }
 
diff --git a/src/ceval.cpp b/src/ceval.cpp
index 95b60898..7672048f 100644
--- a/src/ceval.cpp
+++ b/src/ceval.cpp
@@ -89,10 +89,10 @@ __NEXT_STEP:;
         PyObject* obj;
         if(decl->nested){
             NameDict_ captured = frame->_locals.to_namedict();
-            obj = VAR(Function({decl, frame->_module, captured}));
+            obj = VAR(Function(decl, frame->_module, nullptr, captured));
             captured->set(decl->code->name, obj);
         }else{
-            obj = VAR(Function({decl, frame->_module}));
+            obj = VAR(Function(decl, frame->_module, nullptr, nullptr));
         }
         PUSH(obj);
     } DISPATCH();
@@ -689,7 +689,7 @@ __NEXT_STEP:;
         StrName _name(byte.arg);
         PyObject* _0 = POPX();
         if(is_non_tagged_type(_0, tp_function)){
-            _0->attr().set(__class__, TOP());
+            PK_OBJ_GET(Function, _0)._class = TOP();
         }
         TOP()->attr().set(_name, _0);
     } DISPATCH();
diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp
index 0a6d2950..25a679be 100644
--- a/src/pocketpy.cpp
+++ b/src/pocketpy.cpp
@@ -123,7 +123,7 @@ void init_builtins(VM* _vm) {
         }else if(args.size() == 0){
             FrameId frame = vm->top_frame();
             if(frame->_callable != nullptr){
-                class_arg = frame->_callable->attr().try_get_likely_found(__class__);
+                class_arg = PK_OBJ_GET(Function, frame->_callable)._class;
                 if(frame->_locals.size() > 0) self_arg = frame->_locals[0];
             }
             if(class_arg == nullptr || self_arg == nullptr){