[ASTMatchers] Avoid recursion in ancestor matching to save stack space.

[sammccall]

Summary

A recent change increased the stack size of memoizedMatchesAncestorOfRecursively
leading to stack overflows on real code involving large fold expressions.
It's not totally unreasonable to choke on very deep ASTs, but as common
infrastructure it's be nice if ASTMatchFinder is more robust.
(It already uses data recursion for the regular "downward" traversal.)

[hokein]

clang/lib/ASTMatchers/ASTMatchFinder.cpp

705

this API makes sense, I'm wondering whether we should do it further (more aligned with matchesChildOf pattern):

• make matchesParentOf public;
• introduce a new matchesParentOf interface in ASTMatcherFinder, which calls this function here;
• then we can remove the special parent case in MatchASTVisitor::matchesAncestorOf, and the AncestorMatchMode enum is also not needed anymore;

I'm also ok with the current way.

738

nit: I think we should explicitly describe (either in the comment or name) these keys are for the single-parent *chain*. It took me a while to understand what are these keys for?

740

maybe

Finish => Memorize
Result => IsMatched

809

nit: assert(Parents.size() > 1).

[hokein]

bump this in case you have missed this patch :) our another internal pipeline seems to need this fix.

[sammccall]

clang/lib/ASTMatchers/ASTMatchFinder.cpp

705

I'm happy with that change, but I'd rather not bundle it into this one.
The goal here is to fix weird production characteristics (rather than APIs or algorithm output) and those are slippery, so let's change one thing at a time.

740

Finish => Memorize

Memoize and Memorize are different words, and memoize doesn't act on results but rather on functions. So I think this rename is confusing.

Result => IsMatched

Done

clang/lib/ASTMatchers/ASTMatchFinder.cpp

public:
   bool matchesAncestorOf(const DynTypedNode &Node, ASTContext &Ctx,
                          const DynTypedMatcher &Matcher,
                          BoundNodesTreeBuilder *Builder,
                          AncestorMatchMode MatchMode) override {
     // Reset the cache outside of the recursive call to make sure we
     // don't invalidate any iterators.
     if (ResultCache.size() > MaxMemoizationEntries)
       ResultCache.clear();
-    return memoizedMatchesAncestorOfRecursively(Node, Ctx, Matcher, Builder,
-                                                MatchMode);
+    if (MatchMode == AncestorMatchMode::AMM_ParentOnly)
+      return matchesParentOf(Node, Matcher, Builder);
+    return matchesAnyAncestorOf(Node, Ctx, Matcher, Builder);
   }
 
   // Matches all registered matchers on the given node and calls the
   // result callback for every node that matches.
   void match(const DynTypedNode &Node) {
     // FIXME: Improve this with a switch or a visitor pattern.
     if (auto *N = Node.get<Decl>()) {
       match(*N);
void matchDispatch(const CXXCtorInitializer *Node) {
     matchWithoutFilter(*Node, Matchers->CtorInit);
   }
   void matchDispatch(const TemplateArgumentLoc *Node) {
     matchWithoutFilter(*Node, Matchers->TemplateArgumentLoc);
   }
   void matchDispatch(const void *) { /* Do nothing. */ }
   /// @}
 
+  // Returns whether a direct parent of \p Node matches \p Matcher.
+  // Unlike matchesAnyAncestorOf there's no memoization: it doesn't save much.
+  bool matchesParentOf(const DynTypedNode &Node, const DynTypedMatcher &Matcher,
+                       BoundNodesTreeBuilder *Builder) {
+    for (const auto &Parent : ActiveASTContext->getParents(Node)) {
+      BoundNodesTreeBuilder BuilderCopy = *Builder;
+      if (Matcher.matches(Parent, this, &BuilderCopy)) {
+        *Builder = std::move(BuilderCopy);
+        return true;
+      }
+    }
+    return false;
+  }
+
   // Returns whether an ancestor of \p Node matches \p Matcher.
   //
-  // The order of matching ((which can lead to different nodes being bound in
+  // The order of matching (which can lead to different nodes being bound in
   // case there are multiple matches) is breadth first search.
   //
   // To allow memoization in the very common case of having deeply nested
   // expressions inside a template function, we first walk up the AST, memoizing
   // the result of the match along the way, as long as there is only a single
   // parent.
   //
   // Once there are multiple parents, the breadth first search order does not
   // allow simple memoization on the ancestors. Thus, we only memoize as long
   // as there is a single parent.
-  bool memoizedMatchesAncestorOfRecursively(const DynTypedNode &Node,
-                                            ASTContext &Ctx,
-                                            const DynTypedMatcher &Matcher,
-                                            BoundNodesTreeBuilder *Builder,
-                                            AncestorMatchMode MatchMode) {
-    // For AST-nodes that don't have an identity, we can't memoize.
-    // When doing a single-level match, we don't need to memoize because
-    // ParentMap (in ASTContext) already memoizes the result.
-    if (!Builder->isComparable() ||
-        MatchMode == AncestorMatchMode::AMM_ParentOnly)
-      return matchesAncestorOfRecursively(Node, Ctx, Matcher, Builder,
-                                          MatchMode);
+  //
+  // We avoid a recursive implementation to prevent excessive stack use on
+  // very deep ASTs (similarly to RecursiveASTVisitor's data recursion).
+  bool matchesAnyAncestorOf(DynTypedNode Node, ASTContext &Ctx,
+                            BoundNodesTreeBuilder *Builder) {
 
-    MatchKey Key;
-    Key.MatcherID = Matcher.getID();
-    Key.Node = Node;
-    Key.BoundNodes = *Builder;
-    Key.Traversal = Ctx.getParentMapContext().getTraversalKind();
-    Key.Type = MatchType::Ancestors;
+    // Memoization keys that can be updated with the result.
+    // These are the memoizable nodes in the chain of unique parents, which
+    // terminates when a node has multiple parents, or matches, or is the root.
+    std::vector<MatchKey> Keys;
+    // When returning, update the memoization cache.
+    auto Finish = [&](bool Matched) {
+      for (const auto &Key : Keys) {
+        MemoizedMatchResult &CachedResult = ResultCache[Key];
+        CachedResult.ResultOfMatch = Matched;
+        CachedResult.Nodes = *Builder;
+      }
+      return Matched;
+    };
 
-    // Note that we cannot use insert and reuse the iterator, as recursive
-    // calls to match might invalidate the result cache iterators.
-    MemoizationMap::iterator I = ResultCache.find(Key);
+    // Loop while there's a single parent and we want to attempt memoization.
+    DynTypedNodeList Parents{ArrayRef<DynTypedNode>()}; // after loop: size != 1
+    for (;;) {
+      // A cache key only makes sense if memoization is possible.
+      if (Builder->isComparable()) {
+        Keys.emplace_back();
+        Keys.back().MatcherID = Matcher.getID();
+        Keys.back().Node = Node;
+        Keys.back().BoundNodes = *Builder;
+        Keys.back().Traversal = Ctx.getParentMapContext().getTraversalKind();
+        Keys.back().Type = MatchType::Ancestors;
+
+        // Check the cache.
+        MemoizationMap::iterator I = ResultCache.find(Keys.back());
         if (I != ResultCache.end()) {
+          Keys.pop_back(); // Don't populate the cache for the matching node!
           *Builder = I->second.Nodes;
-      return I->second.ResultOfMatch;
+          return Finish(I->second.ResultOfMatch);
+        }
       }
 
-    MemoizedMatchResult Result;
-    Result.Nodes = *Builder;
-    Result.ResultOfMatch = matchesAncestorOfRecursively(
-        Node, Ctx, Matcher, &Result.Nodes, MatchMode);
-
-    MemoizedMatchResult &CachedResult = ResultCache[Key];
-    CachedResult = std::move(Result);
+      Parents = ActiveASTContext->getParents(Node);
+      // Either no parents or multiple parents: leave chain+memoize mode and
+      // enter bfs+forgetful mode.
+      if (Parents.size() != 1)
+        break;
 
-    *Builder = CachedResult.Nodes;
-    return CachedResult.ResultOfMatch;
+      // Check the next parent.
+      Node = *Parents.begin();
+      BoundNodesTreeBuilder BuilderCopy = *Builder;
+      if (Matcher.matches(Node, this, &BuilderCopy)) {
+        *Builder = std::move(BuilderCopy);
+        return Finish(true);
       }
+    }
+    // We reached the end of the chain.
 
-  bool matchesAncestorOfRecursively(const DynTypedNode &Node, ASTContext &Ctx,
-                                    const DynTypedMatcher &Matcher,
-                                    BoundNodesTreeBuilder *Builder,
-                                    AncestorMatchMode MatchMode) {
-    const auto &Parents = ActiveASTContext->getParents(Node);
     if (Parents.empty()) {
       // Nodes may have no parents if:
       //  a) the node is the TranslationUnitDecl
       //  b) we have a limited traversal scope that excludes the parent edges
       //  c) there is a bug in the AST, and the node is not reachable
       // Usually the traversal scope is the whole AST, which precludes b.
       // Bugs are common enough that it's worthwhile asserting when we can.
 #ifndef NDEBUG
       if (!Node.get<TranslationUnitDecl>() &&
           /* Traversal scope is full AST if any of the bounds are the TU */
           llvm::any_of(ActiveASTContext->getTraversalScope(), [](Decl *D) {
             return D->getKind() == Decl::TranslationUnit;
           })) {
         llvm::errs() << "Tried to match orphan node:\n";
         Node.dump(llvm::errs(), *ActiveASTContext);
         llvm_unreachable("Parent map should be complete!");
       }
 #endif
-      return false;
-    }
-    if (Parents.size() == 1) {
-      // Only one parent - do recursive memoization.
-      const DynTypedNode Parent = Parents[0];
-      BoundNodesTreeBuilder BuilderCopy = *Builder;
-      if (Matcher.matches(Parent, this, &BuilderCopy)) {
-        *Builder = std::move(BuilderCopy);
-        return true;
-      }
-      if (MatchMode != ASTMatchFinder::AMM_ParentOnly) {
-        return memoizedMatchesAncestorOfRecursively(Parent, Ctx, Matcher,
-                                                    Builder, MatchMode);
-        // Once we get back from the recursive call, the result will be the
-        // same as the parent's result.
-      }
     } else {
-      // Multiple parents - BFS over the rest of the nodes.
-      llvm::DenseSet<const void *> Visited;
+      assert(Parents.size() > 1);
+      // BFS starting from the parents not yet considered.
+      // Memoization of newly visited nodes is not possible (but we still update
+      // results for the elements in the chain we found above).
       std::deque<DynTypedNode> Queue(Parents.begin(), Parents.end());
+      llvm::DenseSet<const void *> Visited;
       while (!Queue.empty()) {
         BoundNodesTreeBuilder BuilderCopy = *Builder;
         if (Matcher.matches(Queue.front(), this, &BuilderCopy)) {
           *Builder = std::move(BuilderCopy);
-          return true;
+          return Finish(true);
         }
-        if (MatchMode != ASTMatchFinder::AMM_ParentOnly) {
-          for (const auto &Parent :
-               ActiveASTContext->getParents(Queue.front())) {
-            // Make sure we do not visit the same node twice.
-            // Otherwise, we'll visit the common ancestors as often as there
-            // are splits on the way down.
-            if (Visited.insert(Parent.getMemoizationData()).second)
-              Queue.push_back(Parent);
-          }
-        }
+        for (const auto &Parent : ActiveASTContext->getParents(Queue.front())) {
         Queue.pop_front();
       }
     }
-    return false;
+    return Finish(false);
   }
 
   // Implements a BoundNodesTree::Visitor that calls a MatchCallback with
   // the aggregated bound nodes for each match.
   class MatchVisitor : public BoundNodesTreeBuilder::Visitor {
   public:
     MatchVisitor(ASTContext* Context,
                  MatchFinder::MatchCallback* Callback)