BST: remove pointer to FunctionDef and ClassDef and the code object pointer

Replace it with a index into a table of definitions and code objects

This also now frees much more code objects early because they get proper refcounted now

src/analysis/type_analysis.cpp

@@ -478,7 +478,7 @@ private:
     }
 
     void visit_makeclass(BST_MakeClass* mkclass) override {
-        BST_ClassDef* node = mkclass->class_def;
+        auto* node = bst_cast<BST_ClassDef>(getCodeConstants().getFuncOrClass(mkclass->index_class_def).first);
 
         for (int i = 0; i < node->num_decorator; ++i) {
             getType(node->decorator[i]);
@@ -504,7 +504,7 @@ private:
     }
 
     void visit_makefunction(BST_MakeFunction* mkfn) override {
-        BST_FunctionDef* node = mkfn->function_def;
+        auto* node = bst_cast<BST_FunctionDef>(getCodeConstants().getFuncOrClass(mkfn->index_func_def).first);
 
         for (int i = 0; i < node->num_defaults + node->num_decorator; ++i) {
             getType(node->elts[i]);

src/codegen/ast_interpreter.cpp

@@ -74,7 +74,7 @@ public:
     static Box* executeInner(ASTInterpreter& interpreter, CFGBlock* start_block, BST_stmt* start_at);
 
 private:
-    Value createFunction(BST_FunctionDef* node);
+    Value createFunction(BST_FunctionDef* node, BoxedCode* node_code);
     Value doBinOp(BST_stmt* node, Value left, Value right, int op, BinExpType exp_type);
     void doStore(int vreg, STOLEN(Value) value);
     void doStoreArg(BST_Name* name, STOLEN(Value) value);
@@ -1125,8 +1125,7 @@ Value ASTInterpreter::visit_return(BST_Return* node) {
     return s;
 }
 
-Value ASTInterpreter::createFunction(BST_FunctionDef* node) {
-    BoxedCode* code = node->code;
+Value ASTInterpreter::createFunction(BST_FunctionDef* node, BoxedCode* code) {
     assert(code);
 
     std::vector<Box*> defaults;
@@ -1206,14 +1205,15 @@ Value ASTInterpreter::createFunction(BST_FunctionDef* node) {
 }
 
 Value ASTInterpreter::visit_makeFunction(BST_MakeFunction* mkfn) {
-    BST_FunctionDef* node = mkfn->function_def;
+    auto func_entry = getCodeConstants().getFuncOrClass(mkfn->index_func_def);
+    auto node = bst_cast<BST_FunctionDef>(func_entry.first);
 
     std::vector<Value> decorators;
     decorators.reserve(node->num_decorator);
     for (int i = 0; i < node->num_decorator; ++i)
         decorators.push_back(getVReg(node->elts[i]));
 
-    Value func = createFunction(node);
+    Value func = createFunction(node, func_entry.second);
 
     for (int i = decorators.size() - 1; i >= 0; i--) {
         func.o = runtimeCall(autoDecref(decorators[i].o), ArgPassSpec(1), autoDecref(func.o), 0, 0, 0, 0);
@@ -1228,8 +1228,9 @@ Value ASTInterpreter::visit_makeFunction(BST_MakeFunction* mkfn) {
 
 Value ASTInterpreter::visit_makeClass(BST_MakeClass* mkclass) {
     abortJITing();
-    BST_ClassDef* node = mkclass->class_def;
 
+    auto class_entry = getCodeConstants().getFuncOrClass(mkclass->index_class_def);
+    auto node = bst_cast<BST_ClassDef>(class_entry.first);
 
     BoxedTuple* bases_tuple = (BoxedTuple*)getVReg(node->vreg_bases_tuple).o;
     assert(bases_tuple->cls == tuple_cls);
@@ -1241,7 +1242,7 @@ Value ASTInterpreter::visit_makeClass(BST_MakeClass* mkclass) {
         decorators.push_back(getVReg(node->decorator[i]).o);
     }
 
-    BoxedCode* code = node->code;
+    BoxedCode* code = class_entry.second;
     assert(code);
 
     const ScopingResults& scope_info = code->source->scoping;

src/codegen/irgen/irgenerator.cpp

@@ -1488,8 +1488,8 @@ private:
     }
 
     CompilerVariable* evalMakeClass(BST_MakeClass* mkclass, const UnwindInfo& unw_info) {
-        assert(mkclass->type == BST_TYPE::MakeClass && mkclass->class_def->type == BST_TYPE::ClassDef);
-        BST_ClassDef* node = mkclass->class_def;
+        auto class_entry = irstate->getCodeConstants().getFuncOrClass(mkclass->index_class_def);
+        BST_ClassDef* node = bst_cast<BST_ClassDef>(class_entry.first);
 
         CompilerVariable* _bases_tuple = evalVReg(node->vreg_bases_tuple);
         ConcreteCompilerVariable* bases_tuple = _bases_tuple->makeConverted(emitter, _bases_tuple->getBoxType());
@@ -1499,7 +1499,7 @@ private:
             decorators.push_back(evalVReg(node->decorator[i]));
         }
 
-        BoxedCode* code = node->code;
+        BoxedCode* code = class_entry.second;
         assert(code);
         const ScopingResults& scope_info = code->source->scoping;
 
@@ -1542,8 +1542,7 @@ private:
         return cls;
     }
 
-    CompilerVariable* _createFunction(BST_FunctionDef* node, const UnwindInfo& unw_info) {
-        BoxedCode* code = node->code;
+    CompilerVariable* _createFunction(BST_FunctionDef* node, BoxedCode* code, const UnwindInfo& unw_info) {
         assert(code);
 
         std::vector<ConcreteCompilerVariable*> defaults;
@@ -1572,13 +1571,15 @@ private:
     }
 
     CompilerVariable* evalMakeFunction(BST_MakeFunction* mkfn, const UnwindInfo& unw_info) {
-        BST_FunctionDef* node = mkfn->function_def;
+        auto func_entry = irstate->getCodeConstants().getFuncOrClass(mkfn->index_func_def);
+        BST_FunctionDef* node = bst_cast<BST_FunctionDef>(func_entry.first);
+
         std::vector<CompilerVariable*> decorators;
         for (int i = 0; i < node->num_decorator; ++i) {
             decorators.push_back(evalVReg(node->elts[i]));
         }
 
-        CompilerVariable* func = _createFunction(node, unw_info);
+        CompilerVariable* func = _createFunction(node, func_entry.second, unw_info);
 
         for (int i = decorators.size() - 1; i >= 0; i--) {
             func = decorators[i]->call(emitter, getOpInfoForNode(node, unw_info), ArgPassSpec(1), { func }, NULL);

src/core/bst.cpp

@@ -710,7 +710,7 @@ void BST_MakeFunction::accept(BSTVisitor* v) {
     if (skip)
         return;
 
-    function_def->accept(v);
+    bst_cast<BST_FunctionDef>(v->getCodeConstants().getFuncOrClass(index_func_def).first)->accept(v);
     v->visit_vreg(&vreg_dst, true);
 }
 
@@ -723,7 +723,7 @@ void BST_MakeClass::accept(BSTVisitor* v) {
     if (skip)
         return;
 
-    class_def->accept(v);
+    bst_cast<BST_ClassDef>(v->getCodeConstants().getFuncOrClass(index_class_def).first)->accept(v);
     v->visit_vreg(&vreg_dst, true);
 }
 

src/core/bst.h

@@ -404,8 +404,6 @@ public:
 
 class BST_ClassDef : public BST_stmt {
 public:
-    BoxedCode* code;
-
     InternedString name;
     int vreg_bases_tuple;
     const int num_decorator;
@@ -471,8 +469,6 @@ class BST_FunctionDef : public BST_stmt {
 public:
     InternedString name; // if the name is not set this is a lambda
 
-    BoxedCode* code;
-
     const int num_decorator;
     const int num_defaults;
 
@@ -563,18 +559,20 @@ public:
 
 class BST_MakeFunction : public BST_stmt_with_dest {
 public:
-    BST_FunctionDef* function_def;
+    const int index_func_def;
 
-    BST_MakeFunction(BST_FunctionDef* fd) : BST_stmt_with_dest(BST_TYPE::MakeFunction, fd->lineno), function_def(fd) {}
+    BST_MakeFunction(BST_FunctionDef* fd, int index_func_def)
+        : BST_stmt_with_dest(BST_TYPE::MakeFunction, fd->lineno), index_func_def(index_func_def) {}
 
     BSTNODE(MakeFunction)
 };
 
 class BST_MakeClass : public BST_stmt_with_dest {
 public:
-    BST_ClassDef* class_def;
+    const int index_class_def;
 
-    BST_MakeClass(BST_ClassDef* cd) : BST_stmt_with_dest(BST_TYPE::MakeClass, cd->lineno), class_def(cd) {}
+    BST_MakeClass(BST_ClassDef* cd, int index_class_def)
+        : BST_stmt_with_dest(BST_TYPE::MakeClass, cd->lineno), index_class_def(index_class_def) {}
 
     BSTNODE(MakeClass)
 };

src/core/cfg.cpp

-// Copyright (c) 2014-2016 Dropbox, Inc.
+// Copyright (c) 2014-2016 Dropbox, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -1294,13 +1294,8 @@ private:
         static BoxedString* gen_name = getStaticString("<generator>");
 
         BoxedCode* code = cfgizer->runRecursively(new_body, gen_name, node->lineno, genexp_args, node);
-        // XXX bad!  this should be tracked ex through co_consts
-        // cur_code->co_consts.push_back(def->code)
-        constants.push_back(code);
-
         BST_FunctionDef* func = BST_FunctionDef::create(0, 0);
-        func->code = code;
-        BST_MakeFunction* mkfunc = new BST_MakeFunction(func);
+        BST_MakeFunction* mkfunc = new BST_MakeFunction(func, code_constants.addFuncOrClass(func, code));
         TmpValue func_var_name = pushBackCreateDst(mkfunc);
 
         return makeCall(func_var_name, { first });
@@ -1357,13 +1352,8 @@ private:
         static BoxedString* comp_name = getStaticString("<comperehension>");
 
         BoxedCode* code = cfgizer->runRecursively(new_body, comp_name, node->lineno, args, node);
-        // XXX bad!  this should be tracked ex through co_consts
-        // cur_code->co_consts.push_back(def->code)
-        constants.push_back(code);
-
         BST_FunctionDef* func = BST_FunctionDef::create(0, 0);
-        func->code = code;
-        BST_MakeFunction* mkfunc = new BST_MakeFunction(func);
+        BST_MakeFunction* mkfunc = new BST_MakeFunction(func, code_constants.addFuncOrClass(func, code));
         TmpValue func_var_name = pushBackCreateDst(mkfunc);
 
         return makeCall(func_var_name, { first });
@@ -1415,12 +1405,9 @@ private:
         }
 
         auto name = getStaticString("<lambda>");
-        bdef->code = cfgizer->runRecursively({ stmt }, name, node->lineno, node->args, node);
-        // XXX bad!  this should be tracked ex through co_consts
-        // cur_code->co_consts.push_back(def->code)
-        constants.push_back(bdef->code);
+        auto* code = cfgizer->runRecursively({ stmt }, name, node->lineno, node->args, node);
+        auto mkfn = new BST_MakeFunction(bdef, code_constants.addFuncOrClass(bdef, code));
 
-        auto mkfn = new BST_MakeFunction(bdef);
         return pushBackCreateDst(mkfn);
     }
 
@@ -1813,12 +1800,8 @@ public:
         TmpValue bases_name = pushBackCreateDst(bases);
         unmapExpr(bases_name, &def->vreg_bases_tuple);
 
-        def->code = cfgizer->runRecursively(node->body, node->name.getBox(), node->lineno, NULL, node);
-        // XXX bad!  this should be tracked ex through co_consts
-        // cur_code->co_consts.push_back(def->code)
-        constants.push_back(def->code);
-
-        auto mkclass = new BST_MakeClass(def);
+        auto* code = cfgizer->runRecursively(node->body, node->name.getBox(), node->lineno, NULL, node);
+        auto mkclass = new BST_MakeClass(def, code_constants.addFuncOrClass(def, code));
         auto tmp = pushBackCreateDst(mkclass);
         pushAssign(TmpValue(scoping->mangleName(def->name), node->lineno), tmp);
 
@@ -1839,12 +1822,8 @@ public:
             unmapExpr(remapExpr(node->args->defaults[i]), &def->elts[node->decorator_list.size() + i]);
         }
 
-        def->code = cfgizer->runRecursively(node->body, node->name.getBox(), node->lineno, node->args, node);
-        // XXX bad!  this should be tracked ex through co_consts
-        // cur_code->co_consts.push_back(def->code)
-        constants.push_back(def->code);
-
-        auto mkfunc = new BST_MakeFunction(def);
+        auto* code = cfgizer->runRecursively(node->body, node->name.getBox(), node->lineno, node->args, node);
+        auto mkfunc = new BST_MakeFunction(def, code_constants.addFuncOrClass(def, code));
         auto tmp = pushBackCreateDst(mkfunc);
         pushAssign(TmpValue(scoping->mangleName(def->name), node->lineno), tmp);
 

src/runtime/types.cpp

@@ -4047,6 +4047,10 @@ BORROWED(BoxedFloat*) CodeConstants::getFloatConstant(double d) const {
 void CodeConstants::dealloc() const {
     decrefArray(owned_refs.data(), owned_refs.size());
     owned_refs.clear();
+    for (auto&& e : funcs_and_classes) {
+        Py_DECREF(e.second);
+    }
+    funcs_and_classes.clear();
 }
 
 #ifndef Py_REF_DEBUG

src/runtime/types.h

@@ -1079,6 +1079,8 @@ private:
     // all objects we need to decref when the code object dies
     mutable std::vector<Box*> owned_refs;
 
+    mutable std::vector<std::pair<BST_stmt*, BoxedCode*>> funcs_and_classes;
+
     // Note: DenseMap doesn't work here since we don't prevent the tombstone/empty
     // keys from reaching it.
     mutable std::unordered_map<int64_t, BoxedInt*> int_constants;
@@ -1105,6 +1107,14 @@ public:
     BORROWED(BoxedInt*) getIntConstant(int64_t n) const;
     BORROWED(BoxedFloat*) getFloatConstant(double d) const;
 
+    std::pair<BST_stmt*, BORROWED(BoxedCode*)> getFuncOrClass(int constant) const {
+        return funcs_and_classes[constant];
+    }
+    int addFuncOrClass(BST_stmt* stmt, STOLEN(BoxedCode*) code) {
+        funcs_and_classes.emplace_back(stmt, code);
+        return funcs_and_classes.size() - 1;
+    }
+
     void dealloc() const;
 };
 

test/unittests/analysis.cpp

@@ -25,6 +25,18 @@ protected:
     }
 };
 
+static BoxedCode* getCodeObjectOfFirstMakeFunction(BoxedCode* module_code) {
+    BoxedCode* code = NULL;
+    for (BST_stmt* stmt : module_code->source->cfg->blocks[0]->body) {
+        if (stmt->type !=  BST_TYPE::MakeFunction)
+            continue;
+        code = module_code->code_constants.getFuncOrClass(bst_cast<BST_MakeFunction>(stmt)->index_func_def).second;
+        break;
+    }
+    assert(code);
+    return code;
+}
+
 // this test use functions (VRegInfo::getVReg) which are only available in a debug build
 #ifndef NDEBUG
 TEST_F(AnalysisTest, augassign) {
@@ -52,16 +64,10 @@ TEST_F(AnalysisTest, augassign) {
     ParamNames param_names(args, *module->interned_strings.get());
 
     // Hack to get at the cfg:
-    CFG* cfg = NULL;
-    for (BST_stmt* stmt : module_code->source->cfg->blocks[0]->body) {
-        if (stmt->type !=  BST_TYPE::MakeFunction)
-            continue;
-        cfg = bst_cast<BST_MakeFunction>(stmt)->function_def->code->source->cfg;
-        break;
-    }
-    assert(cfg);
+    BoxedCode* code = getCodeObjectOfFirstMakeFunction(module_code);
+    CFG* cfg = code->source->cfg;
 
-    std::unique_ptr<LivenessAnalysis> liveness = computeLivenessInfo(cfg, module_code->code_constants);
+    std::unique_ptr<LivenessAnalysis> liveness = computeLivenessInfo(cfg, code->code_constants);
     auto&& vregs = cfg->getVRegInfo();
 
     //cfg->print();
@@ -96,14 +102,7 @@ void doOsrTest(bool is_osr, bool i_maybe_undefined) {
     auto module_code = computeAllCFGs(module, true, future_flags, boxString(fn), main_module);
 
     // Hack to get at the cfg:
-    BoxedCode* code = NULL;
-    for (BST_stmt* stmt : module_code->source->cfg->blocks[0]->body) {
-        if (stmt->type !=  BST_TYPE::MakeFunction)
-            continue;
-        code = bst_cast<BST_MakeFunction>(stmt)->function_def->code;
-        break;
-    }
-    assert(code);
+    BoxedCode* code = getCodeObjectOfFirstMakeFunction(module_code);
     CFG* cfg = code->source->cfg;
     std::unique_ptr<LivenessAnalysis> liveness = computeLivenessInfo(cfg, module_code->code_constants);