1 files changed, 52 insertions, 8 deletions

diff --git a/lib/StaticAnalyzer/Checkers/DirectIvarAssignment.cpp b/lib/StaticAnalyzer/Checkers/DirectIvarAssignment.cpp
index 2b35a74ed4..f7535090bd 100644
--- a/lib/StaticAnalyzer/Checkers/DirectIvarAssignment.cpp
+++ b/lib/StaticAnalyzer/Checkers/DirectIvarAssignment.cpp
@@ -14,6 +14,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "ClangSACheckers.h"
+#include "clang/AST/Attr.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/StmtVisitor.h"
 #include "clang/StaticAnalyzer/Core/BugReporter/BugReporter.h"
@@ -26,6 +27,26 @@ using namespace ento;
 
 namespace {
 
+/// The default method filter, which is used to filter out the methods on which
+/// the check should not be performed.
+///
+/// Checks for the init, dealloc, and any other functions that might be allowed
+/// to perform direct instance variable assignment based on their name.
+struct MethodFilter {
+  virtual bool operator()(ObjCMethodDecl *M) {
+    if (M->getMethodFamily() == OMF_init ||
+        M->getMethodFamily() == OMF_dealloc ||
+        M->getMethodFamily() == OMF_copy ||
+        M->getMethodFamily() == OMF_mutableCopy ||
+        M->getSelector().getNameForSlot(0).find("init") != StringRef::npos ||
+        M->getSelector().getNameForSlot(0).find("Init") != StringRef::npos)
+      return true;
+    return false;
+  }
+};
+
+static MethodFilter DefaultMethodFilter;
+
 class DirectIvarAssignment :
   public Checker<check::ASTDecl<ObjCImplementationDecl> > {
 
@@ -59,6 +80,10 @@ class DirectIvarAssignment :
   };
 
 public:
+  MethodFilter *ShouldSkipMethod;
+
+  DirectIvarAssignment() : ShouldSkipMethod(&DefaultMethodFilter) {}
+
   void checkASTDecl(const ObjCImplementationDecl *D, AnalysisManager& Mgr,
                     BugReporter &BR) const;
 };
@@ -118,14 +143,7 @@ void DirectIvarAssignment::checkASTDecl(const ObjCImplementationDecl *D,
     ObjCMethodDecl *M = *I;
     AnalysisDeclContext *DCtx = Mgr.getAnalysisDeclContext(M);
 
-    // Skip the init, dealloc functions and any functions that might be doing
-    // initialization based on their name.
-    if (M->getMethodFamily() == OMF_init ||
-        M->getMethodFamily() == OMF_dealloc ||
-        M->getMethodFamily() == OMF_copy ||
-        M->getMethodFamily() == OMF_mutableCopy ||
-        M->getSelector().getNameForSlot(0).find("init") != StringRef::npos ||
-        M->getSelector().getNameForSlot(0).find("Init") != StringRef::npos)
+    if ((*ShouldSkipMethod)(M))
       continue;
 
     const Stmt *Body = M->getBody();
@@ -175,6 +193,32 @@ void DirectIvarAssignment::MethodCrawler::VisitBinaryOperator(
 }
 }
 
+// Register the checker that checks for direct accesses in all functions,
+// except for the initialization and copy routines.
 void ento::registerDirectIvarAssignment(CheckerManager &mgr) {
   mgr.registerChecker<DirectIvarAssignment>();
 }
+
+// Register the checker that checks for direct accesses in functions annotated
+// with __attribute__((annotate("objc_no_direct_instance_variable_assignmemt"))).
+namespace {
+struct InvalidatorMethodFilter : MethodFilter {
+  virtual bool operator()(ObjCMethodDecl *M) {
+    for (specific_attr_iterator<AnnotateAttr>
+         AI = M->specific_attr_begin<AnnotateAttr>(),
+         AE = M->specific_attr_end<AnnotateAttr>(); AI != AE; ++AI) {
+      const AnnotateAttr *Ann = *AI;
+      if (Ann->getAnnotation() == "objc_no_direct_instance_variable_assignmemt")
+        return false;
+    }
+    return true;
+  }
+};
+
+InvalidatorMethodFilter AttrFilter;
+}
+
+void ento::registerDirectIvarAssignmentForAnnotatedFunctions(
+    CheckerManager &mgr) {
+  mgr.registerChecker<DirectIvarAssignment>()->ShouldSkipMethod = &AttrFilter;
+}