3 files changed, 37 insertions, 20 deletions

diff --git a/lib/StaticAnalyzer/Core/Environment.cpp b/lib/StaticAnalyzer/Core/Environment.cpp
index bab89c545c..7dcd12c5bd 100644
--- a/lib/StaticAnalyzer/Core/Environment.cpp
+++ b/lib/StaticAnalyzer/Core/Environment.cpp
@@ -241,6 +241,10 @@ EnvironmentManager::removeDeadBindings(Environment Env,
       // Mark all symbols in the block expr's value live.
       RSScaner.scan(X);
       continue;
+    } else {
+      SymExpr::symbol_iterator SI = X.symbol_begin(), SE = X.symbol_end();
+      for (; SI != SE; ++SI)
+        SymReaper.maybeDead(*SI);
     }
   }
   
diff --git a/lib/StaticAnalyzer/Core/ExprEngine.cpp b/lib/StaticAnalyzer/Core/ExprEngine.cpp
index 045591c907..42bfe144db 100644
--- a/lib/StaticAnalyzer/Core/ExprEngine.cpp
+++ b/lib/StaticAnalyzer/Core/ExprEngine.cpp
@@ -268,22 +268,39 @@ static bool shouldRemoveDeadBindings(AnalysisManager &AMgr,
 
 void ExprEngine::removeDead(ExplodedNode *Pred, ExplodedNodeSet &Out,
                             const Stmt *ReferenceStmt,
-                            const StackFrameContext *LC,
+                            const LocationContext *LC,
                             const Stmt *DiagnosticStmt,
                             ProgramPoint::Kind K) {
   assert((K == ProgramPoint::PreStmtPurgeDeadSymbolsKind ||
-          ReferenceStmt == 0)
+          ReferenceStmt == 0 || isa<ReturnStmt>(ReferenceStmt))
           && "PostStmt is not generally supported by the SymbolReaper yet");
+  assert(LC && "Must pass the current (or expiring) LocationContext");
+
+  if (!DiagnosticStmt) {
+    DiagnosticStmt = ReferenceStmt;
+    assert(DiagnosticStmt && "Required for clearing a LocationContext");
+  }
+
   NumRemoveDeadBindings++;
   CleanedState = Pred->getState();
-  SymbolReaper SymReaper(LC, ReferenceStmt, SymMgr, getStoreManager());
+
+  // LC is the location context being destroyed, but SymbolReaper wants a
+  // location context that is still live. (If this is the top-level stack
+  // frame, this will be null.)
+  if (!ReferenceStmt) {
+    assert(K == ProgramPoint::PostStmtPurgeDeadSymbolsKind &&
+           "Use PostStmtPurgeDeadSymbolsKind for clearing a LocationContext");
+    LC = LC->getParent();
+  }
+
+  const StackFrameContext *SFC = LC ? LC->getCurrentStackFrame() : 0;
+  SymbolReaper SymReaper(SFC, ReferenceStmt, SymMgr, getStoreManager());
 
   getCheckerManager().runCheckersForLiveSymbols(CleanedState, SymReaper);
 
   // Create a state in which dead bindings are removed from the environment
   // and the store. TODO: The function should just return new env and store,
   // not a new state.
-  const StackFrameContext *SFC = LC->getCurrentStackFrame();
   CleanedState = StateMgr.removeDeadBindings(CleanedState, SFC, SymReaper);
 
   // Process any special transfer function for dead symbols.
@@ -345,8 +362,7 @@ void ExprEngine::ProcessStmt(const CFGStmt S,
   EntryNode = Pred;
   ExplodedNodeSet CleanedStates;
   if (shouldRemoveDeadBindings(AMgr, S, Pred, EntryNode->getLocationContext())){
-    removeDead(EntryNode, CleanedStates, currStmt,
-               Pred->getStackFrame(), currStmt);
+    removeDead(EntryNode, CleanedStates, currStmt, Pred->getLocationContext());
   } else
     CleanedStates.Add(EntryNode);
 
diff --git a/lib/StaticAnalyzer/Core/ExprEngineCallAndReturn.cpp b/lib/StaticAnalyzer/Core/ExprEngineCallAndReturn.cpp
index 3ead0817f7..a98e8b4922 100644
--- a/lib/StaticAnalyzer/Core/ExprEngineCallAndReturn.cpp
+++ b/lib/StaticAnalyzer/Core/ExprEngineCallAndReturn.cpp
@@ -168,27 +168,24 @@ static SVal adjustReturnValue(SVal V, QualType ExpectedTy, QualType ActualTy,
 void ExprEngine::removeDeadOnEndOfFunction(NodeBuilderContext& BC,
                                            ExplodedNode *Pred,
                                            ExplodedNodeSet &Dst) {
-  NodeBuilder Bldr(Pred, Dst, BC);
-
   // Find the last statement in the function and the corresponding basic block.
   const Stmt *LastSt = 0;
   const CFGBlock *Blk = 0;
   llvm::tie(LastSt, Blk) = getLastStmt(Pred);
   if (!Blk || !LastSt) {
+    Dst.Add(Pred);
     return;
   }
-  
-  // If the last statement is return, everything it references should stay live.
-  if (isa<ReturnStmt>(LastSt))
-    return;
 
   // Here, we call the Symbol Reaper with 0 stack context telling it to clean up
   // everything on the stack. We use LastStmt as a diagnostic statement, with 
-  // which the PreStmtPurgeDead point will be associated.
-  currBldrCtx = &BC;
-  removeDead(Pred, Dst, 0, 0, LastSt,
+  // which the program point will be associated. However, we only want to use
+  // LastStmt as a reference for what to clean up if it's a ReturnStmt;
+  // otherwise, everything is dead.
+  SaveAndRestore<const NodeBuilderContext *> NodeContextRAII(currBldrCtx, &BC);
+  removeDead(Pred, Dst, dyn_cast<ReturnStmt>(LastSt),
+             Pred->getLocationContext(), LastSt,
              ProgramPoint::PostStmtPurgeDeadSymbolsKind);
-  currBldrCtx = 0;
 }
 
 static bool wasDifferentDeclUsedForInlining(CallEventRef<> Call,
@@ -290,11 +287,11 @@ void ExprEngine::processCallExit(ExplodedNode *CEBNode) {
 
     NodeBuilderContext Ctx(getCoreEngine(), Blk, BindedRetNode);
     currBldrCtx = &Ctx;
-    // Here, we call the Symbol Reaper with 0 statement and caller location
+    // Here, we call the Symbol Reaper with 0 statement and callee location
     // context, telling it to clean up everything in the callee's context
-    // (and it's children). We use LastStmt as a diagnostic statement, which
-    // which the PreStmtPurge Dead point will be associated.
-    removeDead(BindedRetNode, CleanedNodes, 0, callerCtx, LastSt,
+    // (and its children). We use LastSt as a diagnostic statement, which
+    // which the program point will be associated.
+    removeDead(BindedRetNode, CleanedNodes, 0, calleeCtx, LastSt,
                ProgramPoint::PostStmtPurgeDeadSymbolsKind);
     currBldrCtx = 0;
   } else {