From b990f1834b4ba1c8e67ace010b0f554d0e59c1ea Mon Sep 17 00:00:00 2001
From: Ted Kremenek <kremenek@apple.com>
Date: Thu, 25 Oct 2007 22:24:19 +0000
Subject: [PATCH] Modified StmtIterator to now include visiting the
 initialization expression for EnumConstantDecls.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@43366 91177308-0d34-0410-b5e6-96231b3b80d8
---
 AST/StmtIterator.cpp             | 29 +++++++++++++++++++++--------
 include/clang/AST/Decl.h         |  2 ++
 include/clang/AST/StmtIterator.h |  4 ++--
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/AST/StmtIterator.cpp b/AST/StmtIterator.cpp
index 2d198c029d..961ca50f24 100644
--- a/AST/StmtIterator.cpp
+++ b/AST/StmtIterator.cpp
@@ -17,11 +17,23 @@
 
 using namespace clang;
 
+static inline bool declHasExpr(ScopedDecl *decl) {
+  if (VarDecl* D = dyn_cast<VarDecl>(decl))
+    if (D->getInit())
+      return true;
+  
+  if (EnumConstantDecl* D = dyn_cast<EnumConstantDecl>(decl))
+    if (D->getInitExpr())
+      return true;
+  
+  return false;  
+}
+
 void StmtIteratorBase::NextDecl() {
   assert (FirstDecl && Ptr.D);
 
   do Ptr.D = Ptr.D->getNextDeclarator();
-  while (Ptr.D != NULL && !isa<VarDecl>(Ptr.D));
+  while (Ptr.D != NULL && !declHasExpr(Ptr.D));
   
   if (Ptr.D == NULL) FirstDecl = NULL;
 }
@@ -29,12 +41,8 @@ void StmtIteratorBase::NextDecl() {
 StmtIteratorBase::StmtIteratorBase(ScopedDecl* d) {
   assert (d);
   
-  while (d != NULL) {
-    if (VarDecl* V = dyn_cast<VarDecl>(d))
-      if (V->getInit()) break;
-    
+  while (d != NULL && !declHasExpr(d))
     d = d->getNextDeclarator();
-  }
   
   FirstDecl = d;
   Ptr.D = d;
@@ -61,6 +69,11 @@ void StmtIteratorBase::PrevDecl() {
   Ptr.D = lastVD;
 }
 
-Stmt*& StmtIteratorBase::GetInitializer() const {
-  return reinterpret_cast<Stmt*&>(cast<VarDecl>(Ptr.D)->Init);
+Stmt*& StmtIteratorBase::GetDeclExpr() const {
+  if (VarDecl* D = dyn_cast<VarDecl>(Ptr.D))
+    return reinterpret_cast<Stmt*&>(D->Init);
+  else {
+    EnumConstantDecl* D = cast<EnumConstantDecl>(Ptr.D);
+    return reinterpret_cast<Stmt*&>(D->Init);
+  }
 }
diff --git a/include/clang/AST/Decl.h b/include/clang/AST/Decl.h
index 5788dc40f5..b8214a7120 100644
--- a/include/clang/AST/Decl.h
+++ b/include/clang/AST/Decl.h
@@ -430,6 +430,8 @@ public:
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Decl *D) { return D->getKind() == EnumConstant; }
   static bool classof(const EnumConstantDecl *D) { return true; }
+  
+  friend class StmtIteratorBase;
 };
 
 
diff --git a/include/clang/AST/StmtIterator.h b/include/clang/AST/StmtIterator.h
index 3752793c9f..14c6aed3d4 100644
--- a/include/clang/AST/StmtIterator.h
+++ b/include/clang/AST/StmtIterator.h
@@ -28,7 +28,7 @@ protected:
   
   void NextDecl();
   void PrevDecl();
-  Stmt*& GetInitializer() const;
+  Stmt*& GetDeclExpr() const;
 
   StmtIteratorBase(Stmt** s) : FirstDecl(NULL) { Ptr.S = s; }
   StmtIteratorBase(ScopedDecl* d);
@@ -84,7 +84,7 @@ public:
   }
   
   REFERENCE operator*() const { 
-    return (REFERENCE) (FirstDecl ? GetInitializer() : *Ptr.S);
+    return (REFERENCE) (FirstDecl ? GetDeclExpr() : *Ptr.S);
   }
   
   REFERENCE operator->() const { return operator*(); }