2 files changed, 28 insertions, 14 deletions

diff --git a/clang/include/clang/StaticAnalyzer/Core/BugReporter/PathDiagnostic.h b/clang/include/clang/StaticAnalyzer/Core/BugReporter/PathDiagnostic.h
index 326635d90ca..c9d36c293c4 100644
--- a/clang/include/clang/StaticAnalyzer/Core/BugReporter/PathDiagnostic.h
+++ b/clang/include/clang/StaticAnalyzer/Core/BugReporter/PathDiagnostic.h
@@ -16,6 +16,7 @@
 
 #include "clang/Basic/Diagnostic.h"
 #include "llvm/ADT/FoldingSet.h"
+#include "llvm/ADT/PointerUnion.h"
 #include <deque>
 #include <iterator>
 #include <string>
@@ -23,6 +24,7 @@
 
 namespace clang {
 
+class AnalysisContext;
 class BinaryOperator;
 class CompoundStmt;
 class Decl;
@@ -88,6 +90,9 @@ public:
   PathDiagnosticRange() : isPoint(false) {}
 };
 
+typedef llvm::PointerUnion<const LocationContext*, AnalysisContext*>
+                                                   LocationOrAnalysisContext;
+
 class PathDiagnosticLocation {
 private:
   enum Kind { RangeK, SingleLocK, StmtK, DeclK } K;
@@ -98,8 +103,10 @@ private:
   FullSourceLoc Loc;
   PathDiagnosticRange Range;
 
-  FullSourceLoc genLocation(const LocationContext *LC=0) const;
-  PathDiagnosticRange genRange(const LocationContext *LC=0) const;
+  FullSourceLoc
+    genLocation(LocationOrAnalysisContext LAC = (AnalysisContext*)0) const;
+  PathDiagnosticRange
+    genRange(LocationOrAnalysisContext LAC = (AnalysisContext*)0) const;
 
 public:
   PathDiagnosticLocation()
@@ -119,9 +126,9 @@ public:
 
   PathDiagnosticLocation(const Stmt *s,
                          const SourceManager &sm,
-                         const LocationContext *lc)
+                         LocationOrAnalysisContext lac)
     : K(StmtK), S(s), D(0), SM(&sm),
-      Loc(genLocation(lc)), Range(genRange(lc)) {}
+      Loc(genLocation(lac)), Range(genRange(lac)) {}
 
 
   PathDiagnosticLocation(const Decl *d, const SourceManager &sm)
@@ -132,7 +139,7 @@ public:
   // Create a location for the beginning of the statement.
   static PathDiagnosticLocation createBeginStmt(const Stmt *S,
                                                 const SourceManager &SM,
-                                                const LocationContext *LC);
+                                                LocationOrAnalysisContext LAC);
 
   /// Create the location for the operator of the binary expression.
   /// Assumes the statement has a valid location.
diff --git a/clang/lib/StaticAnalyzer/Core/PathDiagnostic.cpp b/clang/lib/StaticAnalyzer/Core/PathDiagnostic.cpp
index 1d03b669970..b8a811a6e18 100644
--- a/clang/lib/StaticAnalyzer/Core/PathDiagnostic.cpp
+++ b/clang/lib/StaticAnalyzer/Core/PathDiagnostic.cpp
@@ -131,16 +131,23 @@ void PathDiagnosticClient::HandlePathDiagnostic(const PathDiagnostic *D) {
 //===----------------------------------------------------------------------===//
 
 static SourceLocation getValidSourceLocation(const Stmt* S,
-                                             const LocationContext *LC) {
+                                             LocationOrAnalysisContext LAC) {
   SourceLocation L = S->getLocStart();
+  assert(!LAC.isNull() && "A valid LocationContext or AnalysisContext should "
+                          "be passed to PathDiagnosticLocation upon creation.");
 
   // S might be a temporary statement that does not have a location in the
   // source code, so find an enclosing statement and use it's location.
   if (!L.isValid()) {
-    const ParentMap &PM = LC->getParentMap();
+
+    ParentMap *PM = 0;
+    if (LAC.is<const LocationContext*>())
+      PM = &LAC.get<const LocationContext*>()->getParentMap();
+    else
+      PM = &LAC.get<AnalysisContext*>()->getParentMap();
 
     while (!L.isValid()) {
-      S = PM.getParent(S);
+      S = PM->getParent(S);
       L = S->getLocStart();
     }
   }
@@ -151,8 +158,8 @@ static SourceLocation getValidSourceLocation(const Stmt* S,
 PathDiagnosticLocation
   PathDiagnosticLocation::createBeginStmt(const Stmt *S,
                                           const SourceManager &SM,
-                                          const LocationContext *LC) {
-  return PathDiagnosticLocation(getValidSourceLocation(S, LC),
+                                          LocationOrAnalysisContext LAC) {
+  return PathDiagnosticLocation(getValidSourceLocation(S, LAC),
                                 SM, SingleLocK);
 }
 
@@ -246,7 +253,7 @@ PathDiagnosticLocation PathDiagnosticLocation::createSingleLocation(
 }
 
 FullSourceLoc
-  PathDiagnosticLocation::genLocation(const LocationContext *LC) const {
+  PathDiagnosticLocation::genLocation(LocationOrAnalysisContext LAC) const {
   assert(isValid());
   // Note that we want a 'switch' here so that the compiler can warn us in
   // case we add more cases.
@@ -255,7 +262,7 @@ FullSourceLoc
     case RangeK:
       break;
     case StmtK:
-      return FullSourceLoc(getValidSourceLocation(S, LC),
+      return FullSourceLoc(getValidSourceLocation(S, LAC),
                            const_cast<SourceManager&>(*SM));
     case DeclK:
       return FullSourceLoc(D->getLocation(), const_cast<SourceManager&>(*SM));
@@ -265,7 +272,7 @@ FullSourceLoc
 }
 
 PathDiagnosticRange
-  PathDiagnosticLocation::genRange(const LocationContext *LC) const {
+  PathDiagnosticLocation::genRange(LocationOrAnalysisContext LAC) const {
   assert(isValid());
   // Note that we want a 'switch' here so that the compiler can warn us in
   // case we add more cases.
@@ -300,7 +307,7 @@ PathDiagnosticRange
         case Stmt::BinaryConditionalOperatorClass:
         case Stmt::ConditionalOperatorClass:
         case Stmt::ObjCForCollectionStmtClass: {
-          SourceLocation L = getValidSourceLocation(S, LC);
+          SourceLocation L = getValidSourceLocation(S, LAC);
           return SourceRange(L, L);
         }
       }