1 files changed, 26 insertions, 0 deletions
diff --git a/lib/Analysis/MallocChecker.cpp b/lib/Analysis/MallocChecker.cpp
index fdd6a3d026..93e708332e 100644
--- a/lib/Analysis/MallocChecker.cpp
+++ b/lib/Analysis/MallocChecker.cpp
@@ -60,6 +60,7 @@ public:
   void PostVisitCallExpr(CheckerContext &C, const CallExpr *CE);
   void EvalDeadSymbols(CheckerContext &C,const Stmt *S,SymbolReaper &SymReaper);
   void EvalEndPath(GREndPathNodeBuilder &B, void *tag, GRExprEngine &Eng);
+  void PreVisitReturnStmt(CheckerContext &C, const ReturnStmt *S);
 private:
   void MallocMem(CheckerContext &C, const CallExpr *CE);
   void FreeMem(CheckerContext &C, const CallExpr *CE);
@@ -190,3 +191,28 @@ void MallocChecker::EvalEndPath(GREndPathNodeBuilder &B, void *tag,
     }
   }
 }
+
+void MallocChecker::PreVisitReturnStmt(CheckerContext &C, const ReturnStmt *S) {
+  const Expr *RetE = S->getRetValue();
+  if (!RetE)
+    return;
+
+  const GRState *state = C.getState();
+
+  SymbolRef Sym = state->getSVal(RetE).getAsSymbol();
+
+  if (!Sym)
+    return;
+
+  const RefState *RS = state->get<RegionState>(Sym);
+  if (!RS)
+    return;
+
+  // FIXME: check other cases.
+  if (RS->isAllocated())
+    state = state->set<RegionState>(Sym, RefState::getEscaped(S));
+
+  ExplodedNode *N = C.GenerateNode(S, state);
+  if (N)
+    C.addTransition(N);
+}