1 files changed, 53 insertions, 16 deletions

diff --git a/lib/StaticAnalyzer/Core/CFRefCount.cpp b/lib/StaticAnalyzer/Core/CFRefCount.cpp
index d9b1ce825c..9fd7fc7d4c 100644
--- a/lib/StaticAnalyzer/Core/CFRefCount.cpp
+++ b/lib/StaticAnalyzer/Core/CFRefCount.cpp
@@ -2494,6 +2494,23 @@ static QualType GetReturnType(const Expr* RetE, ASTContext& Ctx) {
   return RetTy;
 }
 
+
+// HACK: Symbols that have ref-count state that are referenced directly
+//  (not as structure or array elements, or via bindings) by an argument
+//  should not have their ref-count state stripped after we have
+//  done an invalidation pass.
+//
+// FIXME: This is a global to currently share between CFRefCount and
+// RetainReleaseChecker.  Eventually all functionality in CFRefCount should
+// be migrated to RetainReleaseChecker, and we can make this a non-global.
+llvm::DenseSet<SymbolRef> WhitelistedSymbols;
+namespace {
+struct ResetWhiteList {
+  ResetWhiteList() {}
+  ~ResetWhiteList() { WhitelistedSymbols.clear(); } 
+};
+}
+
 void CFRefCount::evalSummary(ExplodedNodeSet& Dst,
                              ExprEngine& Eng,
                              StmtNodeBuilder& Builder,
@@ -2510,12 +2527,9 @@ void CFRefCount::evalSummary(ExplodedNodeSet& Dst,
   SymbolRef ErrorSym = 0;
 
   llvm::SmallVector<const MemRegion*, 10> RegionsToInvalidate;
-
-  // HACK: Symbols that have ref-count state that are referenced directly
-  //  (not as structure or array elements, or via bindings) by an argument
-  //  should not have their ref-count state stripped after we have
-  //  done an invalidation pass.
-  llvm::DenseSet<SymbolRef> WhitelistedSymbols;
+  
+  // Use RAII to make sure the whitelist is properly cleared.
+  ResetWhiteList resetWhiteList;
 
   // Invalidate all instance variables of the receiver of a message.
   // FIXME: We should be able to do better with inter-procedural analysis.
@@ -2624,21 +2638,13 @@ void CFRefCount::evalSummary(ExplodedNodeSet& Dst,
 
   // NOTE: Even if RegionsToInvalidate is empty, we must still invalidate
   //  global variables.
+  // NOTE: RetainReleaseChecker handles the actual invalidation of symbols.
   state = state->invalidateRegions(RegionsToInvalidate.data(),
                                    RegionsToInvalidate.data() +
                                    RegionsToInvalidate.size(),
                                    Ex, Count, &IS,
                                    /* invalidateGlobals = */ true);
 
-  for (StoreManager::InvalidatedSymbols::iterator I = IS.begin(),
-       E = IS.end(); I!=E; ++I) {
-    SymbolRef sym = *I;
-    if (WhitelistedSymbols.count(sym))
-      continue;
-    // Remove any existing reference-count binding.
-    state = state->remove<RefBindings>(*I);
-  }
-
   // Evaluate the effect on the message receiver.
   if (!ErrorRange.isValid() && Receiver) {
     SymbolRef Sym = Receiver.getSValAsScalarOrLoc(state).getAsLocSymbol();
@@ -3418,12 +3424,43 @@ void CFRefCount::ProcessNonLeakError(ExplodedNodeSet& Dst,
 
 namespace {
 class RetainReleaseChecker
-  : public Checker< check::PostStmt<BlockExpr> > {
+  : public Checker< check::PostStmt<BlockExpr>, check::RegionChanges > {
 public:
+    bool wantsRegionUpdate;
+    
+    RetainReleaseChecker() : wantsRegionUpdate(true) {}
+    
+    
     void checkPostStmt(const BlockExpr *BE, CheckerContext &C) const;
+    const GRState *checkRegionChanges(const GRState *state,
+                            const StoreManager::InvalidatedSymbols *invalidated,
+                                      const MemRegion * const *begin,
+                                      const MemRegion * const *end) const;
+                                          
+    bool wantsRegionChangeUpdate(const GRState *state) const {
+      return wantsRegionUpdate;
+    }
 };
 } // end anonymous namespace
 
+const GRState *
+RetainReleaseChecker::checkRegionChanges(const GRState *state,
+                            const StoreManager::InvalidatedSymbols *invalidated,
+                                         const MemRegion * const *begin,
+                                         const MemRegion * const *end) const {
+  if (!invalidated)
+    return state;
+
+  for (StoreManager::InvalidatedSymbols::const_iterator I=invalidated->begin(),
+       E = invalidated->end(); I!=E; ++I) {
+    SymbolRef sym = *I;
+    if (WhitelistedSymbols.count(sym))
+      continue;
+    // Remove any existing reference-count binding.
+    state = state->remove<RefBindings>(sym);
+  }
+  return state;
+}
 
 void RetainReleaseChecker::checkPostStmt(const BlockExpr *BE,
                                          CheckerContext &C) const {