3 files changed, 22 insertions, 24 deletions
diff --git a/lib/StaticAnalyzer/Core/Checker.cpp b/lib/StaticAnalyzer/Core/Checker.cpp
index 07e0aac2d4..f793c7fe29 100644
--- a/lib/StaticAnalyzer/Core/Checker.cpp
+++ b/lib/StaticAnalyzer/Core/Checker.cpp
@@ -11,7 +11,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "clang/StaticAnalyzer/Core/PathSensitive/ProgramState.h"
 #include "clang/StaticAnalyzer/Core/Checker.h"
 
 using namespace clang;
diff --git a/lib/StaticAnalyzer/Core/ExprEngine.cpp b/lib/StaticAnalyzer/Core/ExprEngine.cpp
index 5087ae1d33..2d047d2c46 100644
--- a/lib/StaticAnalyzer/Core/ExprEngine.cpp
+++ b/lib/StaticAnalyzer/Core/ExprEngine.cpp
@@ -238,8 +238,13 @@ static bool shouldRemoveDeadBindings(AnalysisManager &AMgr,
 
 void ExprEngine::ProcessStmt(const CFGStmt S,
                              ExplodedNode *Pred) {
+  // TODO: Use RAII to remove the unnecessary, tagged nodes.
+  //RegisterCreatedNodes registerCreatedNodes(getGraph());
+
   // Reclaim any unnecessary nodes in the ExplodedGraph.
   G.reclaimRecentlyAllocatedNodes();
+  // Recycle any unused states in the ProgramStateManager.
+  StateMgr.recycleUnusedStates();
   
   currentStmt = S.getStmt();
   PrettyStackTraceLoc CrashInfo(getContext().getSourceManager(),
@@ -1851,7 +1856,7 @@ struct DOTGraphTraits<ExplodedNode*> :
     }
 
     ProgramStateRef state = N->getState();
-    Out << "\\|StateID: " << (void*) state.getPtr()
+    Out << "\\|StateID: " << (void*) state
         << " NodeID: " << (void*) N << "\\|";
     state->printDOT(Out);
 
diff --git a/lib/StaticAnalyzer/Core/ProgramState.cpp b/lib/StaticAnalyzer/Core/ProgramState.cpp
index 459bf83ce5..b8b9ad94df 100644
--- a/lib/StaticAnalyzer/Core/ProgramState.cpp
+++ b/lib/StaticAnalyzer/Core/ProgramState.cpp
@@ -25,26 +25,6 @@ using namespace ento;
 // FIXME: Move this elsewhere.
 ConstraintManager::~ConstraintManager() {}
 
-namespace clang { namespace  ento {
-/// Increments the number of times this state is referenced.
-
-void ProgramStateRetain(const ProgramState *state) {
-  ++const_cast<ProgramState*>(state)->refCount;
-}
-
-/// Decrement the number of times this state is referenced.
-void ProgramStateRelease(const ProgramState *state) {
-  assert(state->refCount > 0);
-  ProgramState *s = const_cast<ProgramState*>(state);
-  if (--s->refCount == 0) {
-    ProgramStateManager &Mgr = s->getStateManager();
-    Mgr.StateSet.RemoveNode(s);
-    s->~ProgramState();    
-    Mgr.freeStates.push_back(s);
-  }
-}
-}}
-
 ProgramState::ProgramState(ProgramStateManager *mgr, const Environment& env,
                  StoreRef st, GenericDataMap gdm)
   : stateMgr(mgr),
@@ -348,10 +328,23 @@ ProgramStateRef ProgramStateManager::getInitialState(const LocationContext *Init
   return getPersistentState(State);
 }
 
+void ProgramStateManager::recycleUnusedStates() {
+  for (std::vector<ProgramState*>::iterator i = recentlyAllocatedStates.begin(),
+       e = recentlyAllocatedStates.end(); i != e; ++i) {
+    ProgramState *state = *i;
+    if (state->referencedByExplodedNode())
+      continue;
+    StateSet.RemoveNode(state);
+    freeStates.push_back(state);
+    state->~ProgramState();
+  }
+  recentlyAllocatedStates.clear();
+}
+
 ProgramStateRef ProgramStateManager::getPersistentStateWithGDM(
                                                      ProgramStateRef FromState,
                                                      ProgramStateRef GDMState) {
-  ProgramState NewState(*FromState);
+  ProgramState NewState = *FromState;
   NewState.GDM = GDMState->GDM;
   return getPersistentState(NewState);
 }
@@ -375,11 +368,12 @@ ProgramStateRef ProgramStateManager::getPersistentState(ProgramState &State) {
   }
   new (newState) ProgramState(State);
   StateSet.InsertNode(newState, InsertPos);
+  recentlyAllocatedStates.push_back(newState);
   return newState;
 }
 
 ProgramStateRef ProgramState::makeWithStore(const StoreRef &store) const {
-  ProgramState NewSt(*this);
+  ProgramState NewSt = *this;
   NewSt.setStore(store);
   return getStateManager().getPersistentState(NewSt);
 }