Getting pretty far

Makefile

@@ -1077,12 +1077,12 @@ lint_%: %.cpp plugins/clang_linter.so
 
 REFCOUNT_CHECKER_BUILD_PATH := $(CMAKE_DIR_DBG)/plugins/refcount_checker/llvm/bin/refcount_checker
 REFCOUNT_CHECKER_RUN_PATH := $(CMAKE_DIR_DBG)/llvm/bin/refcount_checker
-$(REFCOUNT_CHECKER_RUN_PATH): plugins/refcount_checker/refcount_checker.cpp $(CMAKE_SETUP_DBG)
-	$(NINJA) -C $(CMAKE_DIR_DBG) refcount_checker copy_stdlib $(NINJAFLAGS)
-	cp $(REFCOUNT_CHECKER_BUILD_PATH) $(REFCOUNT_CHECKER_RUN_PATH)
+$(REFCOUNT_CHECKER_RUN_PATH): refcount_checker
 
 .PHONY: refcount_checker
-refcount_checker: $(REFCOUNT_CHECKER_RUN_PATH)
+refcount_checker:
+	$(NINJA) -C $(CMAKE_DIR_DBG) refcount_checker copy_stdlib $(NINJAFLAGS)
+	cp $(REFCOUNT_CHECKER_BUILD_PATH) $(REFCOUNT_CHECKER_RUN_PATH)
 
 .PHONY: refcheck_%
 refcheck_%: %.cpp refcount_checker

plugins/refcount_checker/refcount_checker.cpp

@@ -21,6 +21,7 @@
 #include "clang/Frontend/ASTConsumers.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/FrontendActions.h"
+#include "clang/Lex/Lexer.h"
 #include "clang/Tooling/CommonOptionsParser.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/Support/CommandLine.h"
@@ -72,9 +73,18 @@ static void dump(DeclContext* ctx) {
     }
 }
 
-class RefcheckingVisitor : public RecursiveASTVisitor<RefcheckingVisitor> {
+class FunctionRefchecker {
 private:
+    CompilerInstance& CI;
     ASTContext* Context;
+    SourceManager* SM;
+
+    enum class AnnotationType {
+        NONE,
+        BORROWED,
+        STOLEN,
+    };
+    AnnotationType return_ann;
 
     enum RefType {
         UNKNOWN,
@@ -196,12 +206,26 @@ private:
             return false;
 
         auto pointed_to = t->getPointeeType();
-        while (auto pt = dyn_cast<ParenType>(pointed_to))
+        do {
+            if (auto pt = dyn_cast<ParenType>(pointed_to)) {
                 pointed_to = pt->getInnerType();
+                continue;
+            }
+
+            if (auto pt = dyn_cast<TypedefType>(pointed_to)) {
+                pointed_to = pt->desugar();
+                continue;
+            }
+        } while (0);
 
         if (isa<BuiltinType>(pointed_to) || isa<FunctionType>(pointed_to))
             return false;
 
+        if (isa<TemplateTypeParmType>(pointed_to)) {
+            // TODO Hmm not sure what to do about templates
+            return false;
+        }
+
         auto cxx_record_decl = pointed_to->getAsCXXRecordDecl();
         if (!cxx_record_decl)
             t->dump();
@@ -211,18 +235,41 @@ private:
     }
 
     RefState* handle(Expr* expr, BlockState& state) {
-        if (isa<StringLiteral>(expr) || isa<IntegerLiteral>(expr)) {
+        if (isa<StringLiteral>(expr) || isa<IntegerLiteral>(expr) || isa<CXXBoolLiteralExpr>(expr)) {
             return NULL;
         }
 
-        if (isa<UnresolvedLookupExpr>(expr) || isa<CXXUnresolvedConstructExpr>(expr)
+        if (isa<UnresolvedLookupExpr>(expr) || isa<UnresolvedMemberExpr>(expr) || isa<CXXUnresolvedConstructExpr>(expr)
             || isa<CXXDependentScopeMemberExpr>(expr) || isa<DependentScopeDeclRefExpr>(expr)
-            || isa<CXXConstructExpr>(expr) || isa<PredefinedExpr>(expr)) {
+            || isa<CXXConstructExpr>(expr) || isa<PredefinedExpr>(expr) || isa<PackExpansionExpr>(expr)) {
             // Not really sure about this:
             assert(!isRefcountedType(expr->getType()));
             return NULL;
         }
 
+        if (isa<CXXDefaultArgExpr>(expr)) {
+            // Not really sure about this:
+            assert(!isRefcountedType(expr->getType()));
+            return NULL;
+        }
+
+        if (auto exprwc = dyn_cast<ExprWithCleanups>(expr)) {
+            // TODO: probably will need to be checking things here
+            return handle(exprwc->getSubExpr(), state);
+        }
+
+        if (auto mattmp = dyn_cast<MaterializeTemporaryExpr>(expr)) {
+            // not sure about this
+            return handle(mattmp->GetTemporaryExpr(), state);
+        }
+
+        if (auto bindtmp = dyn_cast<CXXBindTemporaryExpr>(expr)) {
+            // not sure about this
+            return handle(bindtmp->getSubExpr(), state);
+        }
+
+
+
         if (auto unaryop = dyn_cast<UnaryOperator>(expr)) {
             handle(unaryop->getSubExpr(), state);
             ASSERT(!isRefcountedType(unaryop->getType()), "implement me");
@@ -320,7 +367,8 @@ private:
                 handle(arg, state);
             }
 
-            bool can_throw = ft && !ft->isNothrow(*Context, false);
+            bool can_throw = ft && !isUnresolvedExceptionSpec(ft->getExceptionSpecType())
+                             && !ft->isNothrow(*Context, false);
             if (can_throw)
                 checkClean(state);
 
@@ -331,6 +379,19 @@ private:
             return NULL;
         }
 
+        if (auto newexpr = dyn_cast<CXXNewExpr>(expr)) {
+            for (auto plc :
+                 iterator_range<CXXNewExpr::arg_iterator>(newexpr->placement_arg_begin(), newexpr->placement_arg_end()))
+                handle(plc, state);
+
+            if (newexpr->hasInitializer())
+                handle(newexpr->getInitializer(), state);
+
+            if (isRefcountedType(newexpr->getType()))
+                return state.createBorrowed();
+            return NULL;
+        }
+
         if (auto condop = dyn_cast<ConditionalOperator>(expr)) {
             handle(condop->getCond(), state);
 
@@ -374,6 +435,31 @@ private:
             return;
         }
 
+        if (auto forstmt = dyn_cast<ForStmt>(stmt)) {
+            handle(forstmt->getInit(), state);
+            handle(forstmt->getCond(), state);
+
+            if (forstmt->getConditionVariable())
+                assert(!isRefcountedType(forstmt->getConditionVariable()->getType()));
+
+            BlockState old_state(state);
+            handle(forstmt->getBody(), state);
+            checkSameAndMerge(state, old_state);
+            return;
+        }
+
+        if (auto whilestmt = dyn_cast<WhileStmt>(stmt)) {
+            handle(whilestmt->getCond(), state);
+
+            if (whilestmt->getConditionVariable())
+                assert(!isRefcountedType(whilestmt->getConditionVariable()->getType()));
+
+            BlockState old_state(state);
+            handle(whilestmt->getBody(), state);
+            checkSameAndMerge(state, old_state);
+            return;
+        }
+
         if (auto ifstmt = dyn_cast<IfStmt>(stmt)) {
             handle(ifstmt->getCond(), state);
 
@@ -387,11 +473,7 @@ private:
         }
 
         if (auto declstmt = dyn_cast<DeclStmt>(stmt)) {
-            assert(declstmt->isSingleDecl());
-
-            auto decl = declstmt->getSingleDecl();
-            assert(decl);
-
+            for (auto decl : declstmt->decls()) {
                 auto vardecl = cast<VarDecl>(decl);
                 assert(vardecl);
 
@@ -407,16 +489,26 @@ private:
                     if (is_refcounted)
                         state.doAssign(vardecl, assigning);
                 }
+            }
             return;
         }
 
         if (auto rtnstmt = dyn_cast<ReturnStmt>(stmt)) {
             auto rstate = handle(rtnstmt->getRetValue(), state);
             if (isRefcountedType(rtnstmt->getRetValue()->getType())) {
+                if (return_ann != AnnotationType::BORROWED) {
                     ASSERT(rstate->num_refs > 0, "Returning an object with 0 refs!");
-                // TODO: handle borrowed returns
                     rstate->num_refs--;
                 }
+            }
+            return;
+        }
+
+        if (auto asmstmt = dyn_cast<AsmStmt>(stmt)) {
+            for (auto input : asmstmt->inputs())
+                handle(input, state);
+            for (auto input : asmstmt->outputs())
+                handle(input, state);
             return;
         }
 
@@ -424,6 +516,18 @@ private:
         RELEASE_ASSERT(0, "unhandled statement type: %s\n", stmt->getStmtClassName());
     }
 
+    AnnotationType getAnnotationType(SourceLocation loc) {
+        // see clang::DiagnosticRenderer::emitMacroExpansions for more info:
+        if (!loc.isMacroID())
+            return AnnotationType::NONE;
+        StringRef MacroName = Lexer::getImmediateMacroName(loc, *SM, CI.getLangOpts());
+        if (MacroName == "BORROWED")
+            return AnnotationType::BORROWED;
+        if (MacroName == "STOLEN")
+            return AnnotationType::STOLEN;
+        return getAnnotationType(SM->getImmediateMacroCallerLoc(loc));
+    }
+
     void checkFunction(FunctionDecl* func) {
         /*
         errs() << func->hasTrivialBody() << '\n';
@@ -440,6 +544,8 @@ private:
         errs() << "dumping:\n";
         func->dump(errs());
 
+        return_ann = getAnnotationType(func->getReturnTypeSourceRange().getBegin());
+
         BlockState state;
         for (auto param : func->params()) {
             if (isRefcountedType(param->getType())) {
@@ -453,8 +559,24 @@ private:
         checkClean(state);
     }
 
+    explicit FunctionRefchecker(CompilerInstance& CI)
+        : CI(CI), Context(&CI.getASTContext()), SM(&CI.getSourceManager()) {}
+
+public:
+    static void checkFunction(FunctionDecl* func, CompilerInstance& CI) {
+        FunctionRefchecker(CI).checkFunction(func);
+    }
+};
+
+class RefcheckingVisitor : public RecursiveASTVisitor<RefcheckingVisitor> {
+private:
+    CompilerInstance& CI;
+    ASTContext* Context;
+    SourceManager* SM;
+
 public:
-    explicit RefcheckingVisitor(ASTContext* Context) : Context(Context) {}
+    explicit RefcheckingVisitor(CompilerInstance& CI)
+        : CI(CI), Context(&CI.getASTContext()), SM(&CI.getSourceManager()) {}
 
     virtual ~RefcheckingVisitor() {}
 
@@ -481,7 +603,7 @@ public:
         // if (func->getNameInfo().getAsString() != "firstlineno")
         // return true;
 
-        checkFunction(func);
+        FunctionRefchecker::checkFunction(func, CI);
 
         return true /* keep going */;
     }
@@ -492,7 +614,7 @@ private:
     RefcheckingVisitor visitor;
 
 public:
-    explicit RefcheckingASTConsumer(ASTContext* Context) : visitor(Context) {}
+    explicit RefcheckingASTConsumer(CompilerInstance& CI) : visitor(CI) {}
 
     virtual void HandleTranslationUnit(ASTContext& Context) {
         visitor.TraverseDecl(Context.getTranslationUnitDecl());
@@ -503,12 +625,37 @@ public:
 class RefcheckingFrontendAction : public ASTFrontendAction {
 public:
     virtual std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance& CI, StringRef fname) {
-        return std::unique_ptr<ASTConsumer>(new RefcheckingASTConsumer(&CI.getASTContext()));
+        return std::unique_ptr<ASTConsumer>(new RefcheckingASTConsumer(CI));
+    }
+};
+
+// A way to inject refchecker-only compilation flags.
+// Not currently used, but uncomment the line in getCompileCommands() to define the `REFCHECKER` directive.
+class MyCompilationDatabase : public CompilationDatabase {
+private:
+    CompilationDatabase& base;
+public:
+    MyCompilationDatabase(CompilationDatabase& base) : base(base) {
+    }
+
+    virtual vector<CompileCommand> getCompileCommands(StringRef FilePath) const {
+        auto rtn = base.getCompileCommands(FilePath);
+        assert(rtn.size() == 1);
+        // rtn[0].CommandLine.push_back("-DREFCHECKER");
+        return rtn;
+    }
+
+    virtual vector<std::string> getAllFiles() const {
+        assert(0);
+    }
+
+    virtual vector<CompileCommand> getAllCompileCommands() const {
+        assert(0);
     }
 };
 
 int main(int argc, const char** argv) {
     CommonOptionsParser OptionsParser(argc, argv, RefcheckingToolCategory);
-    ClangTool Tool(OptionsParser.getCompilations(), OptionsParser.getSourcePathList());
+    ClangTool Tool(MyCompilationDatabase(OptionsParser.getCompilations()), OptionsParser.getSourcePathList());
     return Tool.run(newFrontendActionFactory<RefcheckingFrontendAction>().get());
 }