4 files changed, 56 insertions, 11 deletions

diff --git a/lib/Analysis/GRExprEngine.cpp b/lib/Analysis/GRExprEngine.cpp
index 2ed35fa469..2ce8edd1cc 100644
--- a/lib/Analysis/GRExprEngine.cpp
+++ b/lib/Analysis/GRExprEngine.cpp
@@ -330,8 +330,9 @@ GRExprEngine::~GRExprEngine() {
 // Utility methods.
 //===----------------------------------------------------------------------===//
 
-void GRExprEngine::setTransferFunctions(GRTransferFuncs* tf) {
+void GRExprEngine::setTransferFunctionsAndCheckers(GRTransferFuncs* tf) {
   StateMgr.TF = tf;
+  StateMgr.Checkers = &Checkers;
   tf->RegisterChecks(*this);
   tf->RegisterPrinters(getStateManager().Printers);
 }
diff --git a/lib/Analysis/MallocChecker.cpp b/lib/Analysis/MallocChecker.cpp
index 2ed070a170..fab73ee7b1 100644
--- a/lib/Analysis/MallocChecker.cpp
+++ b/lib/Analysis/MallocChecker.cpp
@@ -23,13 +23,13 @@ using namespace clang;
 namespace {
 
 class RefState {
-  enum Kind { Allocated, Released, Escaped } K;
+  enum Kind { AllocateUnchecked, AllocateFailed, Released, Escaped } K;
   const Stmt *S;
 
 public:
   RefState(Kind k, const Stmt *s) : K(k), S(s) {}
 
-  bool isAllocated() const { return K == Allocated; }
+  bool isAllocated() const { return K == AllocateUnchecked; }
   bool isReleased() const { return K == Released; }
   bool isEscaped() const { return K == Escaped; }
 
@@ -37,7 +37,12 @@ public:
     return K == X.K && S == X.S;
   }
 
-  static RefState getAllocated(const Stmt *s) { return RefState(Allocated, s); }
+  static RefState getAllocateUnchecked(const Stmt *s) { 
+    return RefState(AllocateUnchecked, s); 
+  }
+  static RefState getAllocateFailed() {
+    return RefState(AllocateFailed, 0);
+  }
   static RefState getReleased(const Stmt *s) { return RefState(Released, s); }
   static RefState getEscaped(const Stmt *s) { return RefState(Escaped, s); }
 
@@ -62,6 +67,8 @@ public:
   void EvalDeadSymbols(CheckerContext &C,const Stmt *S,SymbolReaper &SymReaper);
   void EvalEndPath(GREndPathNodeBuilder &B, void *tag, GRExprEngine &Eng);
   void PreVisitReturnStmt(CheckerContext &C, const ReturnStmt *S);
+  const GRState *EvalAssume(const GRState *state, SVal Cond, bool Assumption);
+
 private:
   void MallocMem(CheckerContext &C, const CallExpr *CE);
   const GRState *MallocMemAux(CheckerContext &C, const CallExpr *CE,
@@ -74,6 +81,8 @@ private:
 };
 } // end anonymous namespace
 
+typedef llvm::ImmutableMap<SymbolRef, RefState> RegionStateTy;
+
 namespace clang {
   template <>
   struct GRStateTrait<RegionState> 
@@ -144,7 +153,7 @@ const GRState *MallocChecker::MallocMemAux(CheckerContext &C,
   SymbolRef Sym = RetVal.getAsLocSymbol();
   assert(Sym);
   // Set the symbol's state to Allocated.
-  return state->set<RegionState>(Sym, RefState::getAllocated(CE));
+  return state->set<RegionState>(Sym, RefState::getAllocateUnchecked(CE));
 }
 
 void MallocChecker::FreeMem(CheckerContext &C, const CallExpr *CE) {
@@ -298,3 +307,18 @@ void MallocChecker::PreVisitReturnStmt(CheckerContext &C, const ReturnStmt *S) {
 
   C.addTransition(state);
 }
+
+const GRState *MallocChecker::EvalAssume(const GRState *state, SVal Cond, 
+                                         bool Assumption) {
+  // If a symblic region is assumed to NULL, set its state to AllocateFailed.
+  // FIXME: should also check symbols assumed to non-null.
+
+  RegionStateTy RS = state->get<RegionState>();
+
+  for (RegionStateTy::iterator I = RS.begin(), E = RS.end(); I != E; ++I) {
+    if (state->getSymVal(I.getKey()))
+      state = state->set<RegionState>(I.getKey(),RefState::getAllocateFailed());
+  }
+
+  return state;
+}
diff --git a/lib/Analysis/SimpleConstraintManager.cpp b/lib/Analysis/SimpleConstraintManager.cpp
index 015db76080..23c3b41758 100644
--- a/lib/Analysis/SimpleConstraintManager.cpp
+++ b/lib/Analysis/SimpleConstraintManager.cpp
@@ -15,6 +15,7 @@
 #include "SimpleConstraintManager.h"
 #include "clang/Analysis/PathSensitive/GRExprEngine.h"
 #include "clang/Analysis/PathSensitive/GRState.h"
+#include "clang/Analysis/PathSensitive/Checker.h"
 
 namespace clang {
 
@@ -72,8 +73,17 @@ const GRState *SimpleConstraintManager::Assume(const GRState *state, Loc Cond,
   // EvalAssume is used to call into the GRTransferFunction object to perform
   // any checker-specific update of the state based on this assumption being
   // true or false.
-  return state ? state->getTransferFuncs().EvalAssume(state, Cond, Assumption)
-               : NULL;
+
+  if (!state)
+    return 0;
+
+  std::vector<std::pair<void *, Checker*> >::iterator 
+    I = state->checker_begin(), E = state->checker_end();
+
+  for (; I != E; ++I) {
+    state = I->second->EvalAssume(state, Cond, Assumption);
+  }
+  return state->getTransferFuncs().EvalAssume(state, Cond, Assumption);
 }
 
 const GRState *SimpleConstraintManager::AssumeAux(const GRState *state,
@@ -128,8 +138,18 @@ const GRState *SimpleConstraintManager::Assume(const GRState *state,
   // EvalAssume is used to call into the GRTransferFunction object to perform
   // any checker-specific update of the state based on this assumption being
   // true or false.
-  return state ? state->getTransferFuncs().EvalAssume(state, Cond, Assumption)
-               : NULL;
+
+  if (!state)
+    return 0;
+
+  std::vector<std::pair<void *, Checker*> >::iterator 
+    I = state->checker_begin(), E = state->checker_end();
+
+  for (; I != E; ++I) {
+    state = I->second->EvalAssume(state, Cond, Assumption);
+  }
+
+  return state->getTransferFuncs().EvalAssume(state, Cond, Assumption);
 }
 
 const GRState *SimpleConstraintManager::AssumeAux(const GRState *state,
diff --git a/lib/Frontend/AnalysisConsumer.cpp b/lib/Frontend/AnalysisConsumer.cpp
index dd5c6d3013..6824d8f4eb 100644
--- a/lib/Frontend/AnalysisConsumer.cpp
+++ b/lib/Frontend/AnalysisConsumer.cpp
@@ -373,7 +373,7 @@ static void ActionGRExprEngine(AnalysisConsumer &C, AnalysisManager& mgr,
   if (C.Opts.EnableExperimentalChecks)
     RegisterExperimentalChecks(Eng);
   
-  Eng.setTransferFunctions(tf);  
+  Eng.setTransferFunctionsAndCheckers(tf);  
 
   // Set the graph auditor.
   llvm::OwningPtr<ExplodedNode::Auditor> Auditor;
@@ -506,7 +506,7 @@ static void ActionInlineCall(AnalysisConsumer &C, AnalysisManager &mgr,
   
   // Make a fake transfer function. The GRTransferFunc interface will be 
   // removed.
-  Eng.setTransferFunctions(new GRTransferFuncs());  
+  Eng.setTransferFunctionsAndCheckers(new GRTransferFuncs());  
 
   // Register call inliner as the last checker.
   RegisterCallInliner(Eng);