Make RecursiveASTVisitor to traverse certain statements using data recursion

to avoid a stack overflow with extreme cases.

Part of rdar://10941790.

llvm-svn: 152820

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -148,6 +148,12 @@ public:
   /// TypeLocs.
   bool shouldWalkTypesOfTypeLocs() const { return true; }
 
+  /// \brief Return whether \param S should be traversed using data recursion
+  /// to avoid a stack overflow with extreme cases.
+  bool shouldUseDataRecursionFor(Stmt *S) const {
+    return isa<BinaryOperator>(S) || isa<UnaryOperator>(S) || isa<CaseStmt>(S);
+  }
+
   /// \brief Recursively visit a statement or expression, by
   /// dispatching to Traverse*() based on the argument's dynamic type.
   ///
@@ -397,8 +403,102 @@ private:
   bool TraverseDeclContextHelper(DeclContext *DC);
   bool TraverseFunctionHelper(FunctionDecl *D);
   bool TraverseVarHelper(VarDecl *D);
+
+  bool Walk(Stmt *S);
+
+  struct EnqueueJob {
+    Stmt *S;
+    Stmt::child_iterator StmtIt;
+
+    EnqueueJob(Stmt *S) : S(S), StmtIt() {
+      if (Expr *E = dyn_cast_or_null<Expr>(S))
+        S = E->IgnoreParens();
+    }
+  };
+  bool dataTraverse(Stmt *S);
 };
 
+template<typename Derived>
+bool RecursiveASTVisitor<Derived>::dataTraverse(Stmt *S) {
+
+  SmallVector<EnqueueJob, 16> Queue;
+  Queue.push_back(S);
+
+  while (!Queue.empty()) {
+    EnqueueJob &job = Queue.back();
+    Stmt *CurrS = job.S;
+    if (!CurrS) {
+      Queue.pop_back();
+      continue;
+    }
+
+    if (getDerived().shouldUseDataRecursionFor(CurrS)) {
+      if (job.StmtIt == Stmt::child_iterator()) {
+        if (!Walk(CurrS)) return false;
+        job.StmtIt = CurrS->child_begin();
+      } else {
+        ++job.StmtIt;
+      }
+
+      if (job.StmtIt != CurrS->child_end())
+        Queue.push_back(*job.StmtIt);
+      else
+        Queue.pop_back();
+      continue;
+    }
+
+    Queue.pop_back();
+    TRY_TO(TraverseStmt(CurrS));
+  }
+
+  return true;
+}
+
+template<typename Derived>
+bool RecursiveASTVisitor<Derived>::Walk(Stmt *S) {
+
+#define DISPATCH_WALK(NAME, CLASS, VAR) \
+  return getDerived().WalkUpFrom##NAME(static_cast<CLASS*>(VAR));
+
+  if (BinaryOperator *BinOp = dyn_cast<BinaryOperator>(S)) {
+    switch (BinOp->getOpcode()) {
+#define OPERATOR(NAME) \
+    case BO_##NAME: DISPATCH_WALK(Bin##NAME, BinaryOperator, S);
+
+    BINOP_LIST()
+#undef OPERATOR
+
+#define OPERATOR(NAME)                                          \
+    case BO_##NAME##Assign:                          \
+    DISPATCH_WALK(Bin##NAME##Assign, CompoundAssignOperator, S);
+
+    CAO_LIST()
+#undef OPERATOR
+    }
+  } else if (UnaryOperator *UnOp = dyn_cast<UnaryOperator>(S)) {
+    switch (UnOp->getOpcode()) {
+#define OPERATOR(NAME)                                                  \
+    case UO_##NAME: DISPATCH_WALK(Unary##NAME, UnaryOperator, S);
+
+    UNARYOP_LIST()
+#undef OPERATOR
+    }
+  }
+
+  // Top switch stmt: dispatch to TraverseFooStmt for each concrete FooStmt.
+  switch (S->getStmtClass()) {
+  case Stmt::NoStmtClass: break;
+#define ABSTRACT_STMT(STMT)
+#define STMT(CLASS, PARENT) \
+  case Stmt::CLASS##Class: DISPATCH_WALK(CLASS, CLASS, S);
+#include "clang/AST/StmtNodes.inc"
+  }
+
+#undef DISPATCH_WALK
+
+  return true;
+}
+
 #define DISPATCH(NAME, CLASS, VAR) \
   return getDerived().Traverse##NAME(static_cast<CLASS*>(VAR))
 
@@ -407,6 +507,9 @@ bool RecursiveASTVisitor<Derived>::TraverseStmt(Stmt *S) {
   if (!S)
     return true;
 
+  if (getDerived().shouldUseDataRecursionFor(S))
+    return dataTraverse(S);
+
   // If we have a binary expr, dispatch to the subcode of the binop.  A smart
   // optimizer (e.g. LLVM) will fold this comparison into the switch stmt
   // below.