WIP on adding refcounts. I think this will be quite a bit more complicated...

WIP on adding refcounts.  I think this will be quite a bit more complicated because we have to change the exception control flow to add decrefs

src/CMakeLists.txt

@@ -45,6 +45,7 @@ add_library(PYSTON_OBJECTS OBJECT ${OPTIONAL_SRCS}
 		codegen/irgen/future.cpp
 		codegen/irgen/hooks.cpp
 		codegen/irgen/irgenerator.cpp
+		codegen/irgen/refcounts.cpp
 		codegen/irgen/util.cpp
 		codegen/memmgr.cpp
 		codegen/opt/aa.cpp

src/asm_writing/rewriter.h

@@ -27,6 +27,7 @@
 
 #include "asm_writing/assembler.h"
 #include "asm_writing/icinfo.h"
+#include "core/types.h"
 #include "core/threading.h"
 
 namespace pyston {
@@ -200,17 +201,6 @@ class Rewriter;
 class RewriterVar;
 class RewriterAction;
 
-enum class RefType {
-    UNKNOWN
-#ifndef NDEBUG
-        // Set this to non-zero to make it possible for the debugger to
-        = 0
-#endif
-        ,
-    OWNED,
-    BORROWED,
-};
-
 // This might make more sense as an inner class of Rewriter, but
 // you can't forward-declare that :/
 class RewriterVar {

src/codegen/compvars.cpp

@@ -1122,8 +1122,10 @@ public:
         llvm::Value* boxed;
         if (llvm::ConstantInt* llvm_val = llvm::dyn_cast<llvm::ConstantInt>(unboxed)) {
             boxed = embedRelocatablePtr(emitter.getIntConstant(llvm_val->getSExtValue()), g.llvm_value_type_ptr);
+            emitter.setType(boxed, RefType::BORROWED);
         } else {
             boxed = emitter.getBuilder()->CreateCall(g.funcs.boxInt, unboxed);
+            emitter.setType(boxed, RefType::OWNED);
         }
         return new ConcreteCompilerVariable(other_type, boxed);
     }
@@ -2652,9 +2654,9 @@ public:
         auto slice = var->getValue();
 
         ConcreteCompilerVariable* cstart, *cstop, *cstep;
-        cstart = slice.start ? slice.start->makeConverted(emitter, slice.start->getBoxType()) : getNone();
-        cstop = slice.stop ? slice.stop->makeConverted(emitter, slice.stop->getBoxType()) : getNone();
-        cstep = slice.step ? slice.step->makeConverted(emitter, slice.step->getBoxType()) : getNone();
+        cstart = slice.start ? slice.start->makeConverted(emitter, slice.start->getBoxType()) : emitter.getNone();
+        cstop = slice.stop ? slice.stop->makeConverted(emitter, slice.stop->getBoxType()) : emitter.getNone();
+        cstep = slice.step ? slice.step->makeConverted(emitter, slice.step->getBoxType()) : emitter.getNone();
 
         std::vector<llvm::Value*> args;
         args.push_back(cstart->getValue());
@@ -2676,11 +2678,6 @@ UnboxedSlice extractSlice(CompilerVariable* slice) {
     return static_cast<UnboxedSliceType::VAR*>(slice)->getValue();
 }
 
-ConcreteCompilerVariable* getNone() {
-    llvm::Constant* none = embedRelocatablePtr(None, g.llvm_value_type_ptr, "cNone");
-    return new ConcreteCompilerVariable(typeFromClass(none_cls), none);
-}
-
 class UndefType : public ConcreteCompilerType {
 public:
     std::string debugName() override { return "undefType"; }

src/codegen/compvars.h

@@ -339,8 +339,6 @@ public:
 // Emit the test for whether one variable 'is' another one.
 ConcreteCompilerVariable* doIs(IREmitter& emitter, CompilerVariable* lhs, CompilerVariable* rhs, bool negate);
 
-ConcreteCompilerVariable* getNone();
-
 // These functions all return an INT variable, from either an unboxed representation (makeInt) or
 // a boxed representation (makeUnboxedInt)
 CompilerVariable* makeInt(int64_t);

src/codegen/irgen.cpp

@@ -1083,10 +1083,14 @@ CompiledFunction* doCompile(FunctionMetadata* md, SourceInfo* source, ParamNames
     else
         phis = computeRequiredPhis(*param_names, source->cfg, liveness, source->getScopeInfo());
 
-    IRGenState irstate(md, cf, source, std::move(phis), param_names, getGCBuilder(), dbg_funcinfo);
+    RefcountTracker refcounter;
+
+    IRGenState irstate(md, cf, source, std::move(phis), param_names, getGCBuilder(), dbg_funcinfo, &refcounter);
 
     emitBBs(&irstate, types, entry_descriptor, blocks);
 
+    RefcountTracker::addRefcounts(&irstate);
+
     // De-opt handling:
 
     delete types;

src/codegen/irgen.h

@@ -21,6 +21,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/ValueMap.h"
 
 #include "core/options.h"
 #include "core/types.h"
@@ -29,6 +30,7 @@ namespace pyston {
 
 class AST_expr;
 class AST_stmt;
+class CFGBlock;
 class GCBuilder;
 class IREmitter;
 
@@ -102,8 +104,11 @@ public:
     virtual void checkAndPropagateCapiException(const UnwindInfo& unw_info, llvm::Value* returned_val,
                                                 llvm::Value* exc_val, bool double_check = false) = 0;
 
-    virtual Box* getIntConstant(int64_t n) = 0;
-    virtual Box* getFloatConstant(double d) = 0;
+    virtual BORROWED(Box*) getIntConstant(int64_t n) = 0;
+    virtual BORROWED(Box*) getFloatConstant(double d) = 0;
+
+    virtual void setType(llvm::Value* v, RefType reftype) = 0;
+    virtual ConcreteCompilerVariable* getNone() = 0;
 };
 
 extern const std::string CREATED_CLOSURE_NAME;
@@ -183,6 +188,24 @@ public:
     void calculateModuleHash(const llvm::Module* M, EffortLevel effort);
     bool haveCacheFileForHash();
 };
+
+class IRGenState;
+
+class RefcountTracker {
+private:
+    struct RefcountState {
+        RefType reftype;
+
+        llvm::SmallVector<llvm::Instruction*, 2> ref_consumers;
+    };
+    llvm::ValueMap<llvm::Value*, RefcountState> vars;
+public:
+
+    void setType(llvm::Value* v, RefType reftype);
+    void refConsumed(llvm::Value* v, llvm::Instruction*);
+    static void addRefcounts(IRGenState* state);
+};
+
 }
 
 #endif

src/codegen/irgen/irgenerator.cpp

@@ -46,7 +46,7 @@ extern "C" void dumpLLVM(void* _v) {
 
 IRGenState::IRGenState(FunctionMetadata* md, CompiledFunction* cf, SourceInfo* source_info,
                        std::unique_ptr<PhiAnalysis> phis, ParamNames* param_names, GCBuilder* gc,
-                       llvm::MDNode* func_dbg_info)
+                       llvm::MDNode* func_dbg_info, RefcountTracker* refcount_tracker)
     : md(md),
       cf(cf),
       source_info(source_info),
@@ -54,6 +54,7 @@ IRGenState::IRGenState(FunctionMetadata* md, CompiledFunction* cf, SourceInfo* s
       param_names(param_names),
       gc(gc),
       func_dbg_info(func_dbg_info),
+      refcount_tracker(refcount_tracker),
       scratch_space(NULL),
       frame_info(NULL),
       frame_info_arg(NULL),
@@ -494,9 +495,23 @@ public:
         setCurrentBasicBlock(normal_dest);
     }
 
-    Box* getIntConstant(int64_t n) override { return irstate->getSourceInfo()->parent_module->getIntConstant(n); }
+    Box* getIntConstant(int64_t n) override {
+        return autoDecref(irstate->getSourceInfo()->parent_module->getIntConstant(n));
+    }
+
+    Box* getFloatConstant(double d) override {
+        return autoDecref(irstate->getSourceInfo()->parent_module->getFloatConstant(d));
+    }
 
-    Box* getFloatConstant(double d) override { return irstate->getSourceInfo()->parent_module->getFloatConstant(d); }
+    void setType(llvm::Value* v, RefType reftype) {
+        irstate->getRefcounts()->setType(v, reftype);
+    }
+
+    ConcreteCompilerVariable* getNone() {
+        llvm::Constant* none = embedRelocatablePtr(None, g.llvm_value_type_ptr, "cNone");
+        setType(none, RefType::BORROWED);
+        return new ConcreteCompilerVariable(typeFromClass(none_cls), none);
+    }
 };
 
 IREmitter* createIREmitter(IRGenState* irstate, llvm::BasicBlock*& curblock, IRGenerator* irgenerator) {
@@ -580,6 +595,10 @@ public:
           types(types),
           state(RUNNING) {}
 
+    virtual CFGBlock* getCFGBlock() {
+        return myblock;
+    }
+
 private:
     OpInfo getOpInfoForNode(AST* ast, const UnwindInfo& unw_info) {
         assert(ast);
@@ -782,7 +801,7 @@ private:
                 return v;
             }
             case AST_LangPrimitive::NONE: {
-                return getNone();
+                return emitter.getNone();
             }
             case AST_LangPrimitive::NONZERO: {
                 assert(node->args.size() == 1);
@@ -818,7 +837,7 @@ private:
                 builder->CreateStore(converted_traceback->getValue(),
                                      builder->CreateConstInBoundsGEP2_32(exc_info, 0, 2));
 
-                return getNone();
+                return emitter.getNone();
             }
             case AST_LangPrimitive::UNCACHE_EXC_INFO: {
                 assert(node->args.empty());
@@ -834,7 +853,7 @@ private:
                 builder->CreateStore(v, builder->CreateConstInBoundsGEP2_32(exc_info, 0, 1));
                 builder->CreateStore(v, builder->CreateConstInBoundsGEP2_32(exc_info, 0, 2));
 
-                return getNone();
+                return emitter.getNone();
             }
             case AST_LangPrimitive::PRINT_EXPR: {
                 assert(node->args.size() == 1);
@@ -844,7 +863,7 @@ private:
 
                 emitter.createCall(unw_info, g.funcs.printExprHelper, converted->getValue());
 
-                return getNone();
+                return emitter.getNone();
             }
             default:
                 RELEASE_ASSERT(0, "%d", node->opcode);
@@ -1051,7 +1070,7 @@ private:
 
     ConcreteCompilerVariable* _getGlobal(AST_Name* node, const UnwindInfo& unw_info) {
         if (node->id.s() == "None")
-            return getNone();
+            return emitter.getNone();
 
         bool do_patchpoint = ENABLE_ICGETGLOBALS;
         if (do_patchpoint) {
@@ -1295,7 +1314,7 @@ private:
         ConcreteCompilerVariable* convertedGenerator = generator->makeConverted(emitter, generator->getBoxType());
 
 
-        CompilerVariable* value = node->value ? evalExpr(node->value, unw_info) : getNone();
+        CompilerVariable* value = node->value ? evalExpr(node->value, unw_info) : emitter.getNone();
         ConcreteCompilerVariable* convertedValue = value->makeConverted(emitter, value->getBoxType());
 
         llvm::Value* rtn
@@ -1903,6 +1922,7 @@ private:
             dest = d->makeConverted(emitter, d->getBoxType());
         } else {
             llvm::Value* sys_stdout_val = emitter.createCall(unw_info, g.funcs.getSysStdout);
+            emitter.setType(sys_stdout_val, RefType::BORROWED);
             dest = new ConcreteCompilerVariable(UNKNOWN, sys_stdout_val);
             // TODO: speculate that sys.stdout is a file?
         }
@@ -1927,7 +1947,7 @@ private:
 
         CompilerVariable* val;
         if (node->value == NULL) {
-            val = getNone();
+            val = emitter.getNone();
         } else {
             val = evalExpr(node->value, unw_info);
         }
@@ -1939,12 +1959,14 @@ private:
 
         ConcreteCompilerVariable* rtn = val->makeConverted(emitter, opt_rtn_type);
 
+        assert(rtn->getValue());
+        auto ret_inst = emitter.getBuilder()->CreateRet(rtn->getValue());
+
+        irstate->getRefcounts()->refConsumed(rtn->getValue(), ret_inst);
+
         symbol_table.clear();
 
         endBlock(DEAD);
-
-        assert(rtn->getValue());
-        emitter.getBuilder()->CreateRet(rtn->getValue());
     }
 
     void doBranch(AST_Branch* node, const UnwindInfo& unw_info) {

src/codegen/irgen/irgenerator.h

@@ -40,6 +40,7 @@ class GCBuilder;
 struct PatchpointInfo;
 class ScopeInfo;
 class TypeAnalysis;
+class RefcountTracker;
 
 typedef std::unordered_map<InternedString, CompilerVariable*> SymbolTable;
 typedef std::map<InternedString, CompilerVariable*> SortedSymbolTable;
@@ -66,6 +67,7 @@ private:
     ParamNames* param_names;
     GCBuilder* gc;
     llvm::MDNode* func_dbg_info;
+    RefcountTracker* refcount_tracker;
 
     llvm::AllocaInst* scratch_space;
     llvm::Value* frame_info;
@@ -76,7 +78,7 @@ private:
 
 public:
     IRGenState(FunctionMetadata* md, CompiledFunction* cf, SourceInfo* source_info, std::unique_ptr<PhiAnalysis> phis,
-               ParamNames* param_names, GCBuilder* gc, llvm::MDNode* func_dbg_info);
+               ParamNames* param_names, GCBuilder* gc, llvm::MDNode* func_dbg_info, RefcountTracker* refcount_tracker);
     ~IRGenState();
 
     CompiledFunction* getCurFunction() { return cf; }
@@ -106,6 +108,8 @@ public:
 
     llvm::MDNode* getFuncDbgInfo() { return func_dbg_info; }
 
+    RefcountTracker* getRefcounts() { return refcount_tracker; }
+
     ParamNames* getParamNames() { return param_names; }
 
     void setFrameInfoArgument(llvm::Value* v) { frame_info_arg = v; }
@@ -162,6 +166,7 @@ public:
     virtual llvm::BasicBlock* getCXXExcDest(llvm::BasicBlock* final_dest) = 0;
     virtual llvm::BasicBlock* getCAPIExcDest(llvm::BasicBlock* from_block, llvm::BasicBlock* final_dest,
                                              AST_stmt* current_stmt) = 0;
+    virtual CFGBlock* getCFGBlock() = 0;
 };
 
 class IREmitter;

src/codegen/irgen/refcounts.cpp

+// Copyright (c) 2014-2015 Dropbox, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "codegen/compvars.h"
+
+#include <cstdio>
+#include <sstream>
+
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "codegen/codegen.h"
+#include "codegen/gcbuilder.h"
+#include "codegen/irgen.h"
+#include "codegen/irgen/irgenerator.h"
+#include "codegen/irgen/util.h"
+#include "codegen/patchpoints.h"
+#include "core/options.h"
+#include "core/types.h"
+#include "runtime/float.h"
+#include "runtime/int.h"
+#include "runtime/objmodel.h"
+#include "runtime/types.h"
+#include "runtime/util.h"
+
+namespace pyston {
+
+void RefcountTracker::setType(llvm::Value* v, RefType reftype) {
+    auto& var = this->vars[v];
+
+    assert(var.reftype == reftype || var.reftype == RefType::UNKNOWN);
+    var.reftype = reftype;
+}
+
+void RefcountTracker::refConsumed(llvm::Value* v, llvm::Instruction* inst) {
+    auto& var = this->vars[v];
+
+    assert(var.reftype != RefType::UNKNOWN);
+    var.ref_consumers.push_back(inst);
+}
+
+void RefcountTracker::addRefcounts(IRGenState* irstate) {
+    llvm::Function* f = irstate->getLLVMFunction();
+    RefcountTracker* rt = irstate->getRefcounts();
+
+    fprintf(stderr, "Before refcounts:\n");
+    fprintf(stderr, "\033[35m");
+    dumpPrettyIR(f);
+    fprintf(stderr, "\033[0m");
+
+#ifndef NDEBUG
+    int num_untracked = 0;
+    auto check_val_missed = [&](llvm::Value* v) {
+        if (rt->vars.count(v))
+            return;
+
+        auto t = v->getType();
+        auto p = llvm::dyn_cast<llvm::PointerType>(t);
+        if (!p) {
+            //t->dump();
+            //printf("Not a pointer\n");
+            return;
+        }
+        auto s = llvm::dyn_cast<llvm::StructType>(p->getElementType());
+        if (!s) {
+            //t->dump();
+            //printf("Not a pointer\n");
+            return;
+        }
+
+        // Take care of inheritance.  It's represented as an instance of the base type at the beginning of the
+        // derived type, not as the types concatenated.
+        while (s->elements().size() > 0 && llvm::isa<llvm::StructType>(s->elements()[0]))
+            s = llvm::cast<llvm::StructType>(s->elements()[0]);
+
+        bool ok_type = false;
+        if (s->elements().size() >= 2 && s->elements()[0] == g.i64 && s->elements()[1] == g.llvm_class_type_ptr) {
+            //printf("This looks likes a class\n");
+            ok_type = true;
+        }
+
+        if (!ok_type) {
+#ifndef NDEBUG
+            if (s->getName().startswith("struct.pyston::Box") || (s->getName().startswith("Py") || s->getName().endswith("Object")) || s->getName().startswith("class.pyston::Box")) {
+                v->dump();
+                if (s && s->elements().size() >= 2) {
+                    s->elements()[0]->dump();
+                    s->elements()[1]->dump();
+                }
+                fprintf(stderr, "This is named like a refcounted object though it doesn't look like one");
+                assert(0);
+            }
+#endif
+            return;
+        }
+
+        if (rt->vars.count(v) == 0) {
+            num_untracked++;
+            printf("missed a refcounted object: ");
+            v->dump();
+            //abort();
+        }
+    };
+
+    for (auto&& g : f->getParent()->getGlobalList()) {
+        //g.dump();
+        check_val_missed(&g);
+    }
+
+    for (auto&& a : f->args()) {
+        check_val_missed(&a);
+    }
+
+    for (auto&& BB : *f) {
+        for (auto&& inst : BB) {
+            check_val_missed(&inst);
+            for (auto&& u : inst.uses()) {
+                check_val_missed(u.get());
+            }
+        }
+    }
+    ASSERT(num_untracked == 0, "");
+#endif
+
+    assert(0 && "implement me");
+}
+
+} // namespace pyston

src/core/types.h

@@ -72,6 +72,17 @@ enum Rewritable {
     REWRITABLE,
 };
 
+enum class RefType {
+    UNKNOWN
+#ifndef NDEBUG
+        // Set this to non-zero to make it possible for the debugger to
+        = 0
+#endif
+        ,
+    OWNED,
+    BORROWED,
+};
+
 template <typename T> struct ExceptionSwitchable {
 public:
     T capi_val;

src/runtime/builtin_modules/sys.cpp

@@ -100,7 +100,7 @@ BoxedList* getSysPath() {
     return static_cast<BoxedList*>(_sys_path);
 }
 
-Box* getSysStdout() {
+BORROWED(Box*) getSysStdout() {
     auto stdout_str = getStaticString("stdout");
     Box* sys_stdout = sys_module->getattr(stdout_str);
     RELEASE_ASSERT(sys_stdout, "lost sys.stdout??");