diff --git a/lib/StaticAnalyzer/Core/BugReporter.cpp b/lib/StaticAnalyzer/Core/BugReporter.cpp
index df3ebb8..b82d123 100644
--- a/lib/StaticAnalyzer/Core/BugReporter.cpp
+++ b/lib/StaticAnalyzer/Core/BugReporter.cpp
@@ -1261,45 +1261,6 @@
   return S;
 }
 
-PathDiagnosticPiece*
-BugReport::getEndPath(BugReporterContext &BRC,
-                      const ExplodedNode *EndPathNode) {
-
-  const ProgramPoint &PP = EndPathNode->getLocation();
-  PathDiagnosticLocation L;
-
-  if (const BlockEntrance *BE = dyn_cast<BlockEntrance>(&PP)) {
-    const CFGBlock *block = BE->getBlock();
-    if (block->getBlockID() == 0) {
-      L = PathDiagnosticLocation(
-          EndPathNode->getLocationContext()->getDecl()->getBodyRBrace(),
-          BRC.getSourceManager());
-    }
-  }
-
-  if (!L.isValid()) {
-    const Stmt *S = getStmt();
-
-    if (!S)
-      return NULL;
-
-    L = PathDiagnosticLocation(S, BRC.getSourceManager());
-  }
-
-  BugReport::ranges_iterator Beg, End;
-  llvm::tie(Beg, End) = getRanges();
-
-  // Only add the statement itself as a range if we didn't specify any
-  // special ranges for this report.
-  PathDiagnosticPiece *P = new PathDiagnosticEventPiece(L, getDescription(),
-                                                        Beg == End);
-
-  for (; Beg != End; ++Beg)
-    P->addRange(*Beg);
-
-  return P;
-}
-
 std::pair<BugReport::ranges_iterator, BugReport::ranges_iterator>
 BugReport::getRanges() {
     // If no custom ranges, add the range of the statement corresponding to
@@ -1657,15 +1618,28 @@
   // Start building the path diagnostic...
   PathDiagnosticBuilder PDB(*this, R, BackMap.get(), getPathDiagnosticClient());
 
-  if (PathDiagnosticPiece *Piece = R->getEndPath(PDB, N))
-    PD.push_back(Piece);
-  else
-    return;
-
   // Register additional node visitors.
   R->addVisitor(new NilReceiverBRVisitor());
   R->addVisitor(new ConditionBRVisitor());
 
+  // Generate the very last diagnostic piece - the piece is visible before 
+  // the trace is expanded.
+  PathDiagnosticPiece *LastPiece = 0;
+  for (BugReport::visitor_iterator I = R->visitor_begin(),
+                                   E = R->visitor_end(); I!=E; ++I) {
+    if (PathDiagnosticPiece *Piece = (*I)->getEndPath(PDB, N, *R)) {
+      assert (!LastPiece &&
+              "There can only be one final piece in a diagnostic.");
+      LastPiece = Piece;
+    }
+  }
+  if (!LastPiece)
+    LastPiece = BugReporterVisitor::getDefaultEndPath(PDB, N, *R);
+  if (LastPiece)
+    PD.push_back(LastPiece);
+  else
+    return;
+
   switch (PDB.getGenerationScheme()) {
     case PathDiagnosticClient::Extensive:
       GenerateExtensivePathDiagnostic(PD, PDB, N);
diff --git a/lib/StaticAnalyzer/Core/BugReporterVisitors.cpp b/lib/StaticAnalyzer/Core/BugReporterVisitors.cpp
index 9fbccf8..9b5a60c 100644
--- a/lib/StaticAnalyzer/Core/BugReporterVisitors.cpp
+++ b/lib/StaticAnalyzer/Core/BugReporterVisitors.cpp
@@ -72,6 +72,54 @@
 //===----------------------------------------------------------------------===//
 // Definitions for bug reporter visitors.
 //===----------------------------------------------------------------------===//
+
+PathDiagnosticPiece*
+BugReporterVisitor::getEndPath(BugReporterContext &BRC,
+                               const ExplodedNode *EndPathNode,
+                               BugReport &BR) {
+  return 0;
+}
+
+PathDiagnosticPiece*
+BugReporterVisitor::getDefaultEndPath(BugReporterContext &BRC,
+                                      const ExplodedNode *EndPathNode,
+                                      BugReport &BR) {
+  const ProgramPoint &PP = EndPathNode->getLocation();
+  PathDiagnosticLocation L;
+
+  if (const BlockEntrance *BE = dyn_cast<BlockEntrance>(&PP)) {
+    const CFGBlock *block = BE->getBlock();
+    if (block->getBlockID() == 0) {
+      L = PathDiagnosticLocation(
+          EndPathNode->getLocationContext()->getDecl()->getBodyRBrace(),
+          BRC.getSourceManager());
+    }
+  }
+
+  if (!L.isValid()) {
+    const Stmt *S = BR.getStmt();
+
+    if (!S)
+      return NULL;
+
+    L = PathDiagnosticLocation(S, BRC.getSourceManager());
+  }
+
+  BugReport::ranges_iterator Beg, End;
+  llvm::tie(Beg, End) = BR.getRanges();
+
+  // Only add the statement itself as a range if we didn't specify any
+  // special ranges for this report.
+  PathDiagnosticPiece *P = new PathDiagnosticEventPiece(L,
+      BR.getDescription(),
+      Beg == End);
+  for (; Beg != End; ++Beg)
+    P->addRange(*Beg);
+
+  return P;
+}
+
+
 void FindLastStoreBRVisitor ::Profile(llvm::FoldingSetNodeID &ID) const {
   static int tag = 0;
   ID.AddPointer(&tag);
diff --git a/lib/StaticAnalyzer/Core/CFRefCount.cpp b/lib/StaticAnalyzer/Core/CFRefCount.cpp
index 9e6d829..836cb15 100644
--- a/lib/StaticAnalyzer/Core/CFRefCount.cpp
+++ b/lib/StaticAnalyzer/Core/CFRefCount.cpp
@@ -1960,23 +1960,38 @@
   //===---------===//
 
   class CFRefReportVisitor : public BugReporterVisitor {
+  protected:
     SymbolRef Sym;
     const CFRefCount &TF;
+    
   public:
-
     CFRefReportVisitor(SymbolRef sym, const CFRefCount &tf)
        : Sym(sym), TF(tf) {}
 
-    void Profile(llvm::FoldingSetNodeID &ID) const {
+    virtual void Profile(llvm::FoldingSetNodeID &ID) const {
       static int x = 0;
       ID.AddPointer(&x);
       ID.AddPointer(Sym);
     }
 
-    PathDiagnosticPiece *VisitNode(const ExplodedNode *N,
-                                   const ExplodedNode *PrevN,
-                                   BugReporterContext &BRC,
-                                   BugReport &BR);
+    virtual PathDiagnosticPiece *VisitNode(const ExplodedNode *N,
+                                           const ExplodedNode *PrevN,
+                                           BugReporterContext &BRC,
+                                           BugReport &BR);
+
+    virtual PathDiagnosticPiece *getEndPath(BugReporterContext &BRC,
+                                            const ExplodedNode *N,
+                                            BugReport &BR);
+  };
+
+  class CFRefLeakReportVisitor : public CFRefReportVisitor {
+  public:
+    CFRefLeakReportVisitor(SymbolRef sym, const CFRefCount &tf)
+       : CFRefReportVisitor(sym, tf) {}
+
+    PathDiagnosticPiece *getEndPath(BugReporterContext &BRC,
+                                    const ExplodedNode *N,
+                                    BugReport &BR);
   };
 
   class CFRefReport : public BugReport {
@@ -1985,9 +2000,10 @@
     const CFRefCount &TF;
   public:
     CFRefReport(CFRefBug& D, const CFRefCount &tf,
-                ExplodedNode *n, SymbolRef sym)
+                ExplodedNode *n, SymbolRef sym, bool registerVisitor = true)
       : BugReport(D, D.getDescription(), n), Sym(sym), TF(tf) {
-      addVisitor(new CFRefReportVisitor(sym, tf));
+      if (registerVisitor)
+        addVisitor(new CFRefReportVisitor(sym, tf));
     }
 
     CFRefReport(CFRefBug& D, const CFRefCount &tf,
@@ -2011,23 +2027,18 @@
 
     SymbolRef getSymbol() const { return Sym; }
 
-    PathDiagnosticPiece *getEndPath(BugReporterContext &BRC,
-                                    const ExplodedNode *N);
-
     std::pair<const char**,const char**> getExtraDescriptiveText();
   };
 
   class CFRefLeakReport : public CFRefReport {
     SourceLocation AllocSite;
     const MemRegion* AllocBinding;
+
   public:
     CFRefLeakReport(CFRefBug& D, const CFRefCount &tf,
                     ExplodedNode *n, SymbolRef sym,
                     ExprEngine& Eng);
 
-    PathDiagnosticPiece *getEndPath(BugReporterContext &BRC,
-                                    const ExplodedNode *N);
-
     SourceLocation getLocation() const { return AllocSite; }
   };
 } // end anonymous namespace
@@ -2384,17 +2395,19 @@
 }
 
 PathDiagnosticPiece*
-CFRefReport::getEndPath(BugReporterContext &BRC,
-                        const ExplodedNode *EndN) {
+CFRefReportVisitor::getEndPath(BugReporterContext &BRC,
+                               const ExplodedNode *EndN,
+                               BugReport &BR) {
   // Tell the BugReporterContext to report cases when the tracked symbol is
   // assigned to different variables, etc.
   BRC.addNotableSymbol(Sym);
-  return BugReport::getEndPath(BRC, EndN);
+  return BugReporterVisitor::getDefaultEndPath(BRC, EndN, BR);
 }
 
 PathDiagnosticPiece*
-CFRefLeakReport::getEndPath(BugReporterContext &BRC,
-                            const ExplodedNode *EndN){
+CFRefLeakReportVisitor::getEndPath(BugReporterContext &BRC,
+                                   const ExplodedNode *EndN,
+                                   BugReport &BR) {
 
   // Tell the BugReporterContext to report cases when the tracked symbol is
   // assigned to different variables, etc.
@@ -2493,7 +2506,7 @@
 CFRefLeakReport::CFRefLeakReport(CFRefBug& D, const CFRefCount &tf,
                                  ExplodedNode *n,
                                  SymbolRef sym, ExprEngine& Eng)
-: CFRefReport(D, tf, n, sym) {
+: CFRefReport(D, tf, n, sym, false) {
 
   // Most bug reports are cached at the location where they occurred.
   // With leaks, we want to unique them by the location where they were
@@ -2527,7 +2540,7 @@
   if (AllocBinding)
     os << " and stored into '" << AllocBinding->getString() << '\'';
 
-  addVisitor(new CFRefReportVisitor(sym, tf));
+  addVisitor(new CFRefLeakReportVisitor(sym, tf));
 }
 
 //===----------------------------------------------------------------------===//