1 files changed, 17 insertions, 6 deletions
diff --git a/lib/StaticAnalyzer/Core/CoreEngine.cpp b/lib/StaticAnalyzer/Core/CoreEngine.cpp
index da608f6c7558..4e2866c56f0e 100644
--- a/lib/StaticAnalyzer/Core/CoreEngine.cpp
+++ b/lib/StaticAnalyzer/Core/CoreEngine.cpp
@@ -18,7 +18,6 @@
 #include "clang/AST/StmtCXX.h"
 #include "clang/StaticAnalyzer/Core/PathSensitive/AnalysisManager.h"
 #include "clang/StaticAnalyzer/Core/PathSensitive/ExprEngine.h"
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/Casting.h"
 
@@ -310,8 +309,19 @@ void CoreEngine::HandleBlockEdge(const BlockEdge &L, ExplodedNode *Pred) {
     assert (L.getLocationContext()->getCFG()->getExit().size() == 0
             && "EXIT block cannot contain Stmts.");
 
+    // Get return statement..
+    const ReturnStmt *RS = nullptr;
+    if (!L.getSrc()->empty()) {
+      if (Optional<CFGStmt> LastStmt = L.getSrc()->back().getAs<CFGStmt>()) {
+        if ((RS = dyn_cast<ReturnStmt>(LastStmt->getStmt()))) {
+          if (!RS->getRetValue())
+            RS = nullptr;
+        }
+      }
+    }
+
     // Process the final state transition.
-    SubEng.processEndOfFunction(BuilderCtx, Pred);
+    SubEng.processEndOfFunction(BuilderCtx, Pred, RS);
 
     // This path is done. Don't enqueue any more nodes.
     return;
@@ -590,13 +600,14 @@ void CoreEngine::enqueueStmtNode(ExplodedNode *N,
     WList->enqueue(Succ, Block, Idx+1);
 }
 
-ExplodedNode *CoreEngine::generateCallExitBeginNode(ExplodedNode *N) {
+ExplodedNode *CoreEngine::generateCallExitBeginNode(ExplodedNode *N,
+                                                    const ReturnStmt *RS) {
   // Create a CallExitBegin node and enqueue it.
   const StackFrameContext *LocCtx
                          = cast<StackFrameContext>(N->getLocationContext());
 
   // Use the callee location context.
-  CallExitBegin Loc(LocCtx);
+  CallExitBegin Loc(LocCtx, RS);
 
   bool isNew;
   ExplodedNode *Node = G.getNode(Loc, N->getState(), false, &isNew);
@@ -620,12 +631,12 @@ void CoreEngine::enqueue(ExplodedNodeSet &Set,
   }
 }
 
-void CoreEngine::enqueueEndOfFunction(ExplodedNodeSet &Set) {
+void CoreEngine::enqueueEndOfFunction(ExplodedNodeSet &Set, const ReturnStmt *RS) {
   for (ExplodedNodeSet::iterator I = Set.begin(), E = Set.end(); I != E; ++I) {
     ExplodedNode *N = *I;
     // If we are in an inlined call, generate CallExitBegin node.
     if (N->getLocationContext()->getParent()) {
-      N = generateCallExitBeginNode(N);
+      N = generateCallExitBeginNode(N, RS);
       if (N)
         WList->enqueue(N);
     } else {