1 files changed, 9 insertions, 4 deletions
diff --git a/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp b/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp
index fe6e1fd6bbd2..b8dbb545041e 100644
--- a/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp
+++ b/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp
@@ -200,18 +200,20 @@ void AnalysisConsumer::HandleDeclContext(ASTContext &C, DeclContext *dc) {
         }
         break;
       }
-        
+       
+      case Decl::ObjCCategoryImpl:
       case Decl::ObjCImplementation: {
-        ObjCImplementationDecl* ID = cast<ObjCImplementationDecl>(*I);
+        ObjCImplDecl* ID = cast<ObjCImplDecl>(*I);
         HandleCode(ID);
         
-        for (ObjCImplementationDecl::method_iterator MI = ID->meth_begin(), 
+        for (ObjCContainerDecl::method_iterator MI = ID->meth_begin(), 
              ME = ID->meth_end(); MI != ME; ++MI) {
           checkerMgr->runCheckersOnASTDecl(*MI, *Mgr, BR);
 
           if ((*MI)->isThisDeclarationADefinition()) {
             if (!Opts.AnalyzeSpecificFunction.empty() &&
-                Opts.AnalyzeSpecificFunction != (*MI)->getSelector().getAsString())
+                Opts.AnalyzeSpecificFunction != 
+                  (*MI)->getSelector().getAsString())
               break;
             DisplayFunction(*MI);
             HandleCode(*MI);
@@ -232,6 +234,9 @@ void AnalysisConsumer::HandleTranslationUnit(ASTContext &C) {
   checkerMgr->runCheckersOnASTDecl(TU, *Mgr, BR);
   HandleDeclContext(C, TU);
 
+  // After all decls handled, run checkers on the entire TranslationUnit.
+  checkerMgr->runCheckersOnEndOfTranslationUnit(TU, *Mgr, BR);
+
   // Explicitly destroy the PathDiagnosticClient.  This will flush its output.
   // FIXME: This should be replaced with something that doesn't rely on
   // side-effects in PathDiagnosticClient's destructor. This is required when