4 files changed, 31 insertions, 28 deletions

diff --git a/clang/Analysis/LiveVariables.cpp b/clang/Analysis/LiveVariables.cpp
index 3118a30d1fa..eba1a77941f 100644
--- a/clang/Analysis/LiveVariables.cpp
+++ b/clang/Analysis/LiveVariables.cpp
@@ -35,6 +35,7 @@ class RegisterDecls : public CFGRecStmtDeclVisitor<RegisterDecls> {
 public:
   RegisterDecls(LiveVariables::AnalysisDataTy& ad) : AD(ad) {}  
   void VisitVarDecl(VarDecl* VD) { AD.Register(VD); }
+  CFG& getCFG() { return AD.getCFG(); }
 };
 } // end anonymous namespace
 
@@ -52,34 +53,25 @@ namespace {
 static const bool Alive = true;
 static const bool Dead = false;  
 
-class TransferFuncs : public CFGStmtVisitor<TransferFuncs> {
+class TransferFuncs : public CFGRecStmtVisitor<TransferFuncs> {
   LiveVariables::AnalysisDataTy& AD;
   LiveVariables::ValTy LiveState;
 public:
   TransferFuncs(LiveVariables::AnalysisDataTy& ad) : AD(ad) {}
 
   LiveVariables::ValTy& getVal() { return LiveState; }
+  CFG& getCFG() { return AD.getCFG(); }
   
   void VisitDeclRefExpr(DeclRefExpr* DR);
   void VisitBinaryOperator(BinaryOperator* B);
   void VisitAssign(BinaryOperator* B);
   void VisitDeclStmt(DeclStmt* DS);
   void VisitUnaryOperator(UnaryOperator* U);
-  void VisitStmt(Stmt* S);
-  void VisitExpr(Expr* E);
-  void BlockStmt_VisitExpr(Expr *E);    
   void Visit(Stmt *S);
   
   DeclRefExpr* FindDeclRef(Stmt *S);
 };
-
-void TransferFuncs::VisitExpr(Expr * E) {
-  if (AD.getCFG().isBlkExpr(E)) return;
-  else VisitStmt(E);
-}
       
-void TransferFuncs::VisitStmt(Stmt* S) { VisitChildren(S); }
-
 void TransferFuncs::Visit(Stmt *S) {
   if (AD.Observer)
     AD.Observer->ObserveStmt(S,AD,LiveState);
@@ -160,11 +152,6 @@ void TransferFuncs::VisitDeclStmt(DeclStmt* DS) {
     LiveState(D,AD) = Dead;
 }
   
-void TransferFuncs::BlockStmt_VisitExpr(Expr* E) {
-  assert (AD.getCFG().isBlkExpr(E));
-  VisitChildren(E);
-}
-  
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
diff --git a/clang/Analysis/UninitializedValues.cpp b/clang/Analysis/UninitializedValues.cpp
index 6f24872cd23..0a496595f63 100644
--- a/clang/Analysis/UninitializedValues.cpp
+++ b/clang/Analysis/UninitializedValues.cpp
@@ -34,6 +34,7 @@ public:
   RegisterDecls(UninitializedValues::AnalysisDataTy& ad) :  AD(ad) {}
   
   void VisitBlockVarDecl(BlockVarDecl* VD) { AD.Register(VD); }
+  CFG& getCFG() { return AD.getCFG(); }
 };
   
 } // end anonymous namespace
@@ -58,6 +59,7 @@ public:
   }
   
   UninitializedValues::ValTy& getVal() { return V; }
+  CFG& getCFG() { return AD.getCFG(); }
   
   bool VisitDeclRefExpr(DeclRefExpr* DR);
   bool VisitBinaryOperator(BinaryOperator* B);
diff --git a/clang/include/clang/Analysis/Visitors/CFGRecStmtVisitor.h b/clang/include/clang/Analysis/Visitors/CFGRecStmtVisitor.h
index 002c2f0546a..bfe998bf9d3 100644
--- a/clang/include/clang/Analysis/Visitors/CFGRecStmtVisitor.h
+++ b/clang/include/clang/Analysis/Visitors/CFGRecStmtVisitor.h
@@ -22,18 +22,10 @@ template <typename ImplClass>
 class CFGRecStmtVisitor : public CFGStmtVisitor<ImplClass,void> {
 public:  
 
-  void Visit(Stmt* S) {
-    static_cast< CFGStmtVisitor<ImplClass>* >(this)->Visit(S);
+  void VisitStmt(Stmt* S) {
     static_cast< ImplClass* >(this)->VisitChildren(S);
   }
-  
-  void BlockStmt_Visit(Stmt* S) {
-    assert (S);
-    
-    static_cast< CFGStmtVisitor<ImplClass>* >(this)->BlockStmt_Visit(S);
-    static_cast< ImplClass* >(this)->VisitChildren(S);
-  }
-  
+      
   // Defining operator() allows the visitor to be used as a C++ style functor.
   void operator()(Stmt* S) { static_cast<ImplClass*>(this)->BlockStmt_Visit(S);}
 };
diff --git a/clang/include/clang/Analysis/Visitors/CFGStmtVisitor.h b/clang/include/clang/Analysis/Visitors/CFGStmtVisitor.h
index afc98730937..8ece7d0acb4 100644
--- a/clang/include/clang/Analysis/Visitors/CFGStmtVisitor.h
+++ b/clang/include/clang/Analysis/Visitors/CFGStmtVisitor.h
@@ -33,14 +33,36 @@ static_cast<ImplClass*>(this)->BlockStmt_Visit ## CLASS(static_cast<CLASS*>(S));
 
 template <typename ImplClass, typename RetTy=void>
 class CFGStmtVisitor : public StmtVisitor<ImplClass,RetTy> {
-public:
+  Stmt* CurrentBlkExpr;
+
+  struct NullifyStmt {
+    Stmt*& S;  
+    
+    NullifyStmt(Stmt*& s) : S(s) {}
+    ~NullifyStmt() { S = NULL; }
+  };
+  
+public:  
+  CFGStmtVisitor() : CurrentBlkExpr(NULL) {}  
+  
+  RetTy Visit(Stmt* S) {
+    if (S == CurrentBlkExpr || 
+        !static_cast<ImplClass*>(this)->getCFG().isBlkExpr(S))
+      return StmtVisitor<ImplClass,RetTy>::Visit(S);
+    else
+      return RetTy();
+  }
+  
   /// BlockVisit_XXX - Visitor methods for visiting the "root" statements in
   /// CFGBlocks.  Root statements are the statements that appear explicitly in 
   /// the list of statements in a CFGBlock.  For substatements, or when there
   /// is no implementation provided for a BlockStmt_XXX method, we default
   /// to using StmtVisitor's Visit method.
   RetTy BlockStmt_Visit(Stmt* S) {
-    switch (S->getStmtClass()) {    
+    CurrentBlkExpr = S;
+    NullifyStmt cleanup(CurrentBlkExpr);
+    
+    switch (S->getStmtClass()) {
       DISPATCH_CASE(CallExpr)
       DISPATCH_CASE(StmtExpr)
       DISPATCH_CASE(ConditionalOperator)