[WIP]custom builtin support

src/codegen/irgen/hooks.cpp

@@ -384,23 +384,29 @@ static void pickGlobalsAndLocals(Box*& globals, Box*& locals) {
 
     RELEASE_ASSERT(globals && (globals->cls == module_cls || globals->cls == dict_cls), "Unspported globals type: %s",
                    globals ? globals->cls->tp_name : "NULL");
-//
-//     if (globals) {
-//         // From CPython (they set it to be f->f_builtins):
-//         Box* globals_dict;
-//         if (globals->cls == module_cls)
-//             globals_dict = globals->getAttrWrapper();
-//         else
-//             globals_dict = globals;
-//
-//         auto requested_builtins = PyDict_GetItemString(globals_dict, "__builtins__");
-//         if (requested_builtins == NULL)
-//             PyDict_SetItemString(globals_dict, "__builtins__", PyEval_GetBuiltins());
-//         else
-//             RELEASE_ASSERT(requested_builtins == builtins_module
-//                                || requested_builtins == builtins_module->getAttrWrapper(),
-//                            "we don't support overriding __builtins__");
-//     }
+
+    if (globals) {
+        // From CPython (they set it to be f->f_builtins):
+        Box* globals_dict;
+        if (globals->cls == module_cls)
+            globals_dict = globals->getAttrWrapper();
+        else
+            globals_dict = globals;
+
+        auto requested_builtins = PyDict_GetItemString(globals_dict, "__builtins__");
+        if (requested_builtins == NULL)
+            PyDict_SetItemString(globals_dict, "__builtins__", PyEval_GetBuiltins());
+        else
+            // ((FrameInfo*)cur_thread_state.frame_info)->builtins = requested_builtins;
+            PyDict_SetItemString(globals_dict, "__builtins__", requested_builtins);
+            // ((FrameInfo*)cur_thread_state.frame_info)->builtins = globals_dict;
+        //     if (requested_builtins != builtins_module
+        //                        && requested_builtins != builtins_module->getAttrWrapper()) {
+        //         PyDict_SetItemString(globals_dict, "__builtins__", requested_builtins);
+        //         builtins_module->clearAttrsForDealloc();
+        //         builtins_module = (BoxedModule*)requested_builtins;
+        //     }
+    }
 }
 
 extern "C" PyObject* PyEval_EvalCode(PyCodeObject* co, PyObject* globals, PyObject* locals) noexcept {

src/core/types.h

@@ -1053,6 +1053,8 @@ struct FrameInfo {
     FrameInfo* back;
     // TODO does this need to be owned?  how does cpython do it?
     BORROWED(BoxedCode*) code;
+    // What type of this field should be?
+    BORROWED(Box*) builtins;
 
     BORROWED(Box*) updateBoxedLocals();
 

src/runtime/builtin_modules/sys.cpp

@@ -490,7 +490,8 @@ static PyObject* sys_displayhook(PyObject* self, PyObject* o) noexcept {
         return NULL;
     }
     */
-    PyObject* builtins = builtins_module;
+    // PyObject* builtins = builtins_module;
+    PyObject* builtins = ((FrameInfo*)cur_thread_state.frame_info)->builtins;
 
     /* Print value except if None */
     /* After printing, also assign to '_' */

src/runtime/frame.cpp

@@ -235,9 +235,53 @@ void frameInvalidateBack(BoxedFrame* frame) {
 }
 
 extern "C" void initFrame(FrameInfo* frame_info) {
-    frame_info->back = (FrameInfo*)(cur_thread_state.frame_info);
+    FrameInfo* back = (FrameInfo*)(cur_thread_state.frame_info);
+    frame_info->back = back;
     cur_thread_state.frame_info = frame_info;
+
+    Box* builtins;
+    static BoxedString* builtins_str = getStaticString("__builtins__");
+    if (back == NULL || back->globals != frame_info->globals) {
+        if (PyModule_Check(frame_info->globals))
+            builtins = frame_info->globals->getattr(builtins_str);
+        else
+            builtins = PyDict_GetItem(frame_info->globals, builtins_str);
+
+        if (builtins && builtins->cls == attrwrapper_cls)
+            builtins = unwrapAttrWrapper(builtins);
+        if (builtins == NULL) {
+            builtins = PyDict_New();
+            if (builtins == NULL || PyDict_SetItemString(builtins, "None", Py_None) < 0)
+                RELEASE_ASSERT(0, "error");
+        } else
+            Py_INCREF(builtins);
+    } else {
+        builtins = incref(back->builtins);
+    }
+    frame_info->builtins = builtins;
 }
+//         Box* global_dict = PyModule_GetDict(frame_info->globals);
+//         builtins = PyDict_GetItem(global_dict, builtins_str);
+//         if (builtins) {
+//             if (PyModule_Check(builtins)) {
+//                 builtins = PyModule_GetDict(builtins);
+//                 assert(!builtins || PyDict_Check(builtins) || builtins->cls != attrwrapper_cls);
+//             } else if (!PyDict_Check(builtins) && builtins->cls != attrwrapper_cls)
+//                 builtins = NULL;
+//         }
+//         if (builtins == NULL) {
+//             builtins = PyDict_New();
+//             if (builtins == NULL || PyDict_SetItemString(builtins, "None", Py_None) < 0)
+//                 return;
+//         } else
+//             Py_INCREF(builtins);
+//     } else {
+//         builtins = builtins_module;
+//         // builtins = back->builtins;
+//         // assert(builtins != NULL && PyDict_Check
+//     }
+//     frame_info->builtins = builtins;
+// }
 
 FrameInfo* const FrameInfo::NO_DEINIT = (FrameInfo*)-2; // not -1 to not match memset(-1)
 
@@ -303,6 +347,7 @@ extern "C" void deinitFrame(FrameInfo* frame_info) noexcept {
     Py_CLEAR(frame_info->boxedLocals);
 
     Py_CLEAR(frame_info->globals);
+    Py_CLEAR(frame_info->builtins);
 
     assert(!PyErr_Occurred());
     if (restore_exc)

src/runtime/objmodel.cpp

@@ -7569,31 +7569,80 @@ extern "C" Box* getGlobal(Box* globals, BoxedString* name) {
         static StatCounter stat_builtins("getglobal_builtins");
         stat_builtins.log();
 
-        Box* rtn;
-        if (rewriter.get()) {
-            RewriterVar* builtins = rewriter->loadConst((intptr_t)builtins_module, Location::any());
-            GetattrRewriteArgs rewrite_args(rewriter.get(), builtins, rewriter->getReturnDestination());
-            rewrite_args.obj_shape_guarded = true; // always builtin module
-            rtn = builtins_module->getattr(name, &rewrite_args);
+        Box* builtins = ((FrameInfo*)cur_thread_state.frame_info)->builtins;
 
-            if (!rewrite_args.isSuccessful())
-                rewriter.reset(NULL);
-            else if (rtn) {
-                auto r_rtn = rewrite_args.getReturn(ReturnConvention::HAS_RETURN);
-                rewriter->commitReturning(r_rtn);
+        if (builtins->cls == module_cls) {
+            Box* rtn;
+
+            if (rewriter.get() && builtins == builtins_module) {
+                RewriterVar* r_builtins = rewriter->loadConst((intptr_t)builtins, Location::any());
+                GetattrRewriteArgs rewrite_args(rewriter.get(), r_builtins, rewriter->getReturnDestination());
+                rewrite_args.obj_shape_guarded = true; // always builtin module
+                rtn = builtins->getattr(name, &rewrite_args);
+
+                if (!rewrite_args.isSuccessful())
+                    rewriter.reset(NULL);
+                else if(rtn) {
+                    auto r_rtn = rewrite_args.getReturn(ReturnConvention::HAS_RETURN);
+                    rewriter->commitReturning(r_rtn);
+                } else {
+                    rewrite_args.getReturn(); // just make the asserts happy
+                    rewriter.reset(NULL);
+                }
             } else {
-                rewrite_args.getReturn(); // just to make the asserts happy
-                rewriter.reset(NULL);
+                rtn = builtins->getattr(name);
+            }
+
+            if (rtn) {
+                assert(rtn->ob_refcnt > 0);
+                Py_INCREF(rtn);
+                return rtn;
             }
         } else {
-            rtn = builtins_module->getattr(name);
-        }
+            ASSERT(builtins->cls == dict_cls || builtins->cls == attrwrapper_cls, "%s", builtins->cls->tp_name);
+            BoxedDict* d = static_cast<BoxedDict*>(builtins);
 
-        if (rtn) {
-            assert(rtn->ob_refcnt > 0);
-            Py_INCREF(rtn);
-            return rtn;
+            rewriter.reset(NULL);
+            REWRITE_ABORTED("Rewriting not implemented for getGlobals with a dict builtins yet");
+
+            auto it = d->d.find(name);
+            if (it != d->d.end()) {
+                assert(it->second->ob_refcnt > 0);
+                Py_INCREF(it->second);
+                return it->second;
+            }
         }
+
+        // Box* rtn;
+        // if (rewriter.get()) {
+        //     RewriterVar* builtins = rewriter->loadConst((intptr_t)((FrameInfo*)cur_thread_state.frame_info)->builtins, Location::any());
+        //     // RewriterVar* builtins = rewriter->loadConst((intptr_t)builtins_module, Location::any());
+        //     GetattrRewriteArgs rewrite_args(rewriter.get(), builtins, rewriter->getReturnDestination());
+        //     rewrite_args.obj_shape_guarded = true; // always builtin module
+        //     rtn = ((FrameInfo*)cur_thread_state.frame_info)->builtins->getattr(name, &rewrite_args);
+        //     // rtn = builtins_module->getattr(name, &rewrite_args);
+        //
+        //     if (!rewrite_args.isSuccessful())
+        //         rewriter.reset(NULL);
+        //     else if (rtn) {
+        //         auto r_rtn = rewrite_args.getReturn(ReturnConvention::HAS_RETURN);
+        //         rewriter->commitReturning(r_rtn);
+        //     } else {
+        //         rewrite_args.getReturn(); // just to make the asserts happy
+        //         rewriter.reset(NULL);
+        //     }
+        // } else {
+        //     rtn = ((FrameInfo*)cur_thread_state.frame_info)->builtins->getattr(name);
+        //     // rtn = builtins_module->getattr(name);
+        // }
+        //
+        // if (rtn) {
+        //     assert(rtn->ob_refcnt > 0);
+        //     Py_INCREF(rtn);
+        //     return rtn;
+        // } else {
+        //     raiseExcHelper(NameError, "name '%s' is not defined", name->data());
+        // }
     }
 
     assert(name->data()[name->size()] == '\0');

src/runtime/types.cpp

@@ -4814,6 +4814,7 @@ BORROWED(BoxedModule*) createModule(BoxedString* name, const char* fn, const cha
 
     if (name->s() == "__main__")
         module->setattr(autoDecref(internStringMortal("__builtins__")), builtins_module, NULL);
+        // module->setattr(autoDecref(internStringMortal("__builtins__")), builtins_module, NULL);
     return module;
 }