[AST] Allow limiting the scope of common AST traversals (getParents, RAV).

[sammccall]

Summary

The goal is to allow analyses such as clang-tidy checks to run on a
subset of the AST, e.g. "only on main-file decls" for interactive tools.

Today, these become "problematically global" by running RecursiveASTVisitors
rooted at the TUDecl, or by navigating up via ASTContext::getParent().

The scope is restricted using a set of top-level-decls that RecursiveASTVisitors
should be rooted at. This also applies to the visitor that populates the
parent map, and so the top-level-decls are considered to have no parents.

This patch makes the traversal scope a mutable property of ASTContext.
The more obvious way to do this is to pass the top-level decls to
relevant functions directly, but this has some problems:

• it's error-prone: accidentally mixing restricted and unrestricted scopes is a performance trap. Interleaving multiple analyses is common (many clang-tidy checks run matchers or RAVs from matcher callbacks)
• it doesn't map well to the actual use cases, where we really do want *all* traversals to be restricted.
• it involves a lot of plumbing in parts of the code that don't care about traversals.

This approach was tried out in D54259 and D54261, I wanted to like it
but it feels pretty awful in practice.

Caveats: to get scope-limiting behavior of RecursiveASTVisitors, callers
have to call the new TraverseAST(Ctx) function instead of TraverseDecl(TU).
I think this is an improvement to the API regardless.

[ioeric]

The new API and the refactoring look good to me. Just a nit and a question.

lib/AST/ASTContext.cpp

10292

is this comment outdated?

lib/ASTMatchers/ASTMatchFinder.cpp

676

A quick search for "Found node that is not in the parent map." in my fav. search engine found a few prior discussions about this. E.g. @klimek seemed to want to remove the assertion (http://clang-developers.42468.n3.nabble.com/Question-about-failed-assertion-ASTMatchFinder-cc-quot-Found-node-that-is-not-in-the-parent-map-quot-td4038612.html), while @thakis seemed to favor keeping the assertion (https://bugs.chromium.org/p/chromium/issues/detail?id=580749).

Maybe we could still assertion if the scope if TU?

[sammccall]

Address review comments.
Remove special case for TUDecl since no-parents is now handled.

[ioeric]

lgtm

lib/ASTMatchers/ASTMatchFinder.cpp

677

nit: a,b,c

683

none_of?

[JonasToth]

Hi Sam,

this patch "broke" ExprMutAnalyzer, at least it creates an assertion failure that seems harmless. Your thought on it would be great!

The assertion in: ASTMatchFinder.cpp +680 is triggered in the Analysis to figure out if something is mutated, because Parents is empty and the traversal scope is the TU itself, but the node is a CXXConstructoDecl. This does not happen for alle cases though, please take a look at this reproducer: https://godbolt.org/z/0yOq2I

Only for the Lambda-Case problems occur. If I #if 0 the assertion out the issue goes away, and the real world code that triggered this behaviour is not bothered, too.

I am really puzzled why this issue occurs, but as you implemented the restrictions it would like to hear your opinion on the issue. If you want to reproduce yourself I am of course happy to help, but
https://github.com/JonasToth/clang/blob/fix_crash/unittests/Analysis/ExprMutationAnalyzerTest.cpp#L1134 is probably the easiest way to do so (just add this to trunk and run check-clang-unit).

Thank you for the time!

[JonasToth]

See PR39949 as well, as it seems to trigger the same or a similar problem.
@ioeric and @klimek maybe could give an opinion, too?

[klimek]

In D54309#1326852, @JonasToth wrote:

See PR39949 as well, as it seems to trigger the same or a similar problem.
@ioeric and @klimek maybe could give an opinion, too?

Sounds like we might not correctly set the parent map for CXXConstructorDecl then?

[JonasToth]

Sounds like we might not correctly set the parent map for CXXConstructorDecl then?

I unfortunatly don't know where to start to look for. Could you give me a hint what to inspect to figure that out?

[JonasToth]

new years ping on the crash-issue.

[sammccall]

In D54309#1320628, @JonasToth wrote:

this patch "broke" ExprMutAnalyzer, at least it creates an assertion failure that seems harmless. Your thought on it would be great!

As mentioned on the bug, this assertion already existed and this patch tries to preserve it in as many cases as possible.
You may be uncovering some unrelated bug (the kind the assertion was meant to catch). Does the assertion go away by reverting this patch (and keeping the old version of the assert)?

I've tried to check myself, but I can't use either of your reproducers. I'll try the one from the bug next (https://bugs.llvm.org/show_bug.cgi?id=39949)

The assertion in: ASTMatchFinder.cpp +680 is triggered in the Analysis to figure out if something is mutated, because Parents is empty and the traversal scope is the TU itself, but the node is a CXXConstructoDecl. This does not happen for alle cases though, please take a look at this reproducer: https://godbolt.org/z/0yOq2I

Only for the Lambda-Case problems occur. If I #if 0 the assertion out the issue goes away, and the real world code that triggered this behaviour is not bothered, too.

Oops, took me a while to realize this is a test for a check that hasn't landed :-)
The fact that this only triggers for a lambda case does suggest to me it's unrelated to this patch - those are exactly the sort of cases where apparently AST traversal has been incomplete in the past, yielding a partial parent map. My understanding is this tends to be harmless where it actually fires, but points to an AST inconsistency that is likely to cause problems in other cases. Thus there's pushback against removing the assertion (which I initially attempted to do).

I am really puzzled why this issue occurs, but as you implemented the restrictions it would like to hear your opinion on the issue. If you want to reproduce yourself I am of course happy to help, but
https://github.com/JonasToth/clang/blob/fix_crash/unittests/Analysis/ExprMutationAnalyzerTest.cpp#L1134 is probably the easiest way to do so (just add this to trunk and run check-clang-unit).

Is this still current? The code fails to parse (on 1137 a&& should be T&&). When I fix that the Results11 = match(withEnclosingCompound(...)) yields an empty vector. (The test does subsequently crash, but not in an interesting way).

[sammccall]

@JonasToth As noted on https://bugs.llvm.org/show_bug.cgi?id=39949, this assertion is a pre-existing problem with the parent map + lambdas, unrelated to this patch.

I've analyzed where it comes from (details in the bug) but don't know what the right fix should be. @klimek, can you take a look at the bug?

[JonasToth]

@sammccall Thank you for analyzing and investigation! Keeping the assertion is of course good and I am not sure on how to proceed.

Given the close branch for 8.0 we might opt for an hotfix that resolves the issue for now.
I think the bug-report is the better place to continue working on it.

include/clang/AST/ASTContext.h

 public:
   IdentifierTable &Idents;
   SelectorTable &Selectors;
   Builtin::Context &BuiltinInfo;
   mutable DeclarationNameTable DeclarationNames;
   IntrusiveRefCntPtr<ExternalASTSource> ExternalSource;
   ASTMutationListener *Listener = nullptr;
 
-  /// Contains parents of a node.
-  using ParentVector = llvm::SmallVector<ast_type_traits::DynTypedNode, 2>;
-
-  /// Maps from a node to its parents. This is used for nodes that have
-  /// pointer identity only, which are more common and we can save space by
-  /// only storing a unique pointer to them.
-  using ParentMapPointers =
-      llvm::DenseMap<const void *,
-                     llvm::PointerUnion4<const Decl *, const Stmt *,
-                                         ast_type_traits::DynTypedNode *,
-                                         ParentVector *>>;
-
-  /// Parent map for nodes without pointer identity. We store a full
-  /// DynTypedNode for all keys.
-  using ParentMapOtherNodes =
-      llvm::DenseMap<ast_type_traits::DynTypedNode,
-                     llvm::PointerUnion4<const Decl *, const Stmt *,
-                                         ast_type_traits::DynTypedNode *,
-                                         ParentVector *>>;
-
   /// Container for either a single DynTypedNode or for an ArrayRef to
   /// DynTypedNode. For use with ParentMap.
   class DynTypedNodeList {
     using DynTypedNode = ast_type_traits::DynTypedNode;
 
     llvm::AlignedCharArrayUnion<ast_type_traits::DynTypedNode,
                                 ArrayRef<DynTypedNode>> Storage;
     bool IsSingleNode;
public:
     bool empty() const { return begin() == end(); }
 
     const DynTypedNode &operator[](size_t N) const {
       assert(N < size() && "Out of bounds!");
       return *(begin() + N);
     }
   };
 
-  /// Returns the parents of the given node.
+  // A traversal scope limits the parts of the AST visible to certain analyses.
+  // RecursiveASTVisitor::TraverseAST will only visit reachable nodes, and
+  // getParents() will only observe reachable parent edges.
+  //
+  // The scope is defined by a set of "top-level" declarations.
+  // Initially, it is the entire TU: {getTranslationUnitDecl()}.
+  // Changing the scope clears the parent cache, which is expensive to rebuild.
+  std::vector<Decl *> getTraversalScope() const { return TraversalScope; }
+  void setTraversalScope(const std::vector<Decl *> &);
+
+  /// Returns the parents of the given node (within the traversal scope).
   ///
   /// Note that this will lazily compute the parents of all nodes
   /// and store them for later retrieval. Thus, the first call is O(n)
   /// in the number of AST nodes.
   ///
   /// Caveats and FIXMEs:
   /// Calculating the parent map over all AST nodes will need to load the
   /// full AST. This can be undesirable in the case where the full AST is
using DeallocationFunctionsAndArguments =
       llvm::SmallVector<std::pair<void (*)(void *), void *>, 16>;
   DeallocationFunctionsAndArguments Deallocations;
 
   // FIXME: This currently contains the set of StoredDeclMaps used
   // by DeclContext objects.  This probably should not be in ASTContext,
   // but we include it here so that ASTContext can quickly deallocate them.
   llvm::PointerIntPair<StoredDeclsMap *, 1> LastSDM;
 
-  std::unique_ptr<ParentMapPointers> PointerParents;
-  std::unique_ptr<ParentMapOtherNodes> OtherParents;
+  std::vector<Decl *> TraversalScope;
+  class ParentMap;
+  std::unique_ptr<ParentMap> Parents;
 
   std::unique_ptr<VTableContextBase> VTContext;
 
   void ReleaseDeclContextMaps();
-  void ReleaseParentMapEntries();
 
 public:
   enum PragmaSectionFlag : unsigned {
     PSF_None = 0,
     PSF_Read = 0x1,
     PSF_Write = 0x2,
     PSF_Execute = 0x4,
     PSF_Implicit = 0x8,

include/clang/AST/RecursiveASTVisitor.h

public:
 
   /// Return whether this visitor should recurse into implicit
   /// code, e.g., implicit constructors and destructors.
   bool shouldVisitImplicitCode() const { return false; }
 
   /// Return whether this visitor should traverse post-order.
   bool shouldTraversePostOrder() const { return false; }
 
+  /// Recursively visits an entire AST, starting from the top-level Decls
+  /// in the AST traversal scope (by default, the TranslationUnitDecl).
+  /// \returns false if visitation was terminated early.
+  bool TraverseAST(ASTContext &AST) {
+    for (Decl *D : AST.getTraversalScope())
+      if (!getDerived().TraverseDecl(D))
+        return false;
+    return true;
+  }
+
   /// Recursively visit a statement or expression, by
   /// dispatching to Traverse*() based on the argument's dynamic type.
   ///
   /// \returns false if the visitation was terminated early, true
   /// otherwise (including when the argument is nullptr).
   bool TraverseStmt(Stmt *S, DataRecursionQueue *Queue = nullptr);
 
   /// Invoked before visiting a statement or expression via data recursion.

lib/AST/ASTContext.cpp

: FunctionProtoTypes(this_()), TemplateSpecializationTypes(this_()),
       XRayFilter(new XRayFunctionFilter(LangOpts.XRayAlwaysInstrumentFiles,
                                         LangOpts.XRayNeverInstrumentFiles,
                                         LangOpts.XRayAttrListFiles, SM)),
       PrintingPolicy(LOpts), Idents(idents), Selectors(sels),
       BuiltinInfo(builtins), DeclarationNames(*this), Comments(SM),
       CommentCommandTraits(BumpAlloc, LOpts.CommentOpts),
       CompCategories(this_()), LastSDM(nullptr, 0) {
   TUDecl = TranslationUnitDecl::Create(*this);
+  TraversalScope = {TUDecl};
 }
 
 ASTContext::~ASTContext() {
-  ReleaseParentMapEntries();
-
   // Release the DenseMaps associated with DeclContext objects.
   // FIXME: Is this the ideal solution?
   ReleaseDeclContextMaps();
 
   // Call all of the deallocation functions on all of their targets.
   for (auto &Pair : Deallocations)
     (Pair.first)(Pair.second);
 
ASTContext::~ASTContext() {
   for (std::pair<const MaterializeTemporaryExpr *, APValue *> &MTVPair :
        MaterializedTemporaryValues)
     MTVPair.second->~APValue();
 
   for (const auto &Value : ModuleInitializers)
     Value.second->~PerModuleInitializers();
 }
 
-void ASTContext::ReleaseParentMapEntries() {
-  if (!PointerParents) return;
-  for (const auto &Entry : *PointerParents) {
+class ASTContext::ParentMap {
+  /// Contains parents of a node.
+  using ParentVector = llvm::SmallVector<ast_type_traits::DynTypedNode, 2>;
+
+  /// Maps from a node to its parents. This is used for nodes that have
+  /// pointer identity only, which are more common and we can save space by
+  /// only storing a unique pointer to them.
+  using ParentMapPointers = llvm::DenseMap<
+      const void *,
+      llvm::PointerUnion4<const Decl *, const Stmt *,
+                          ast_type_traits::DynTypedNode *, ParentVector *>>;
+
+  /// Parent map for nodes without pointer identity. We store a full
+  /// DynTypedNode for all keys.
+  using ParentMapOtherNodes = llvm::DenseMap<
+      ast_type_traits::DynTypedNode,
+      llvm::PointerUnion4<const Decl *, const Stmt *,
+                          ast_type_traits::DynTypedNode *, ParentVector *>>;
+
+  ParentMapPointers PointerParents;
+  ParentMapOtherNodes OtherParents;
+  class ASTVisitor;
+
+  static ast_type_traits::DynTypedNode
+  getSingleDynTypedNodeFromParentMap(ParentMapPointers::mapped_type U) {
+    if (const auto *D = U.dyn_cast<const Decl *>())
+      return ast_type_traits::DynTypedNode::create(*D);
+    if (const auto *S = U.dyn_cast<const Stmt *>())
+      return ast_type_traits::DynTypedNode::create(*S);
+    return *U.get<ast_type_traits::DynTypedNode *>();
+  }
+
+  template <typename NodeTy, typename MapTy>
+  static ASTContext::DynTypedNodeList getDynNodeFromMap(const NodeTy &Node,
+                                                        const MapTy &Map) {
+    auto I = Map.find(Node);
+    if (I == Map.end()) {
+      return llvm::ArrayRef<ast_type_traits::DynTypedNode>();
+    }
+    if (const auto *V = I->second.template dyn_cast<ParentVector *>()) {
+      return llvm::makeArrayRef(*V);
+    }
+    return getSingleDynTypedNodeFromParentMap(I->second);
+  }
+
+public:
+  ParentMap(ASTContext &Ctx);
+  ~ParentMap() {
+    for (const auto &Entry : PointerParents) {
       if (Entry.second.is<ast_type_traits::DynTypedNode *>()) {
         delete Entry.second.get<ast_type_traits::DynTypedNode *>();
       } else if (Entry.second.is<ParentVector *>()) {
         delete Entry.second.get<ParentVector *>();
       }
     }
-  for (const auto &Entry : *OtherParents) {
+    for (const auto &Entry : OtherParents) {
       if (Entry.second.is<ast_type_traits::DynTypedNode *>()) {
         delete Entry.second.get<ast_type_traits::DynTypedNode *>();
       } else if (Entry.second.is<ParentVector *>()) {
         delete Entry.second.get<ParentVector *>();
       }
     }
   }
 
+  DynTypedNodeList getParents(const ast_type_traits::DynTypedNode &Node) {
+    if (Node.getNodeKind().hasPointerIdentity())
+      return getDynNodeFromMap(Node.getMemoizationData(), PointerParents);
+    return getDynNodeFromMap(Node, OtherParents);
+  }
+};
+
+void ASTContext::setTraversalScope(const std::vector<Decl *> &TopLevelDecls) {
+  TraversalScope = TopLevelDecls;
+  Parents.reset();
+}
+
 void ASTContext::AddDeallocation(void (*Callback)(void*), void *Data) {
   Deallocations.push_back({Callback, Data});
 }
 
 void
 ASTContext::setExternalSource(IntrusiveRefCntPtr<ExternalASTSource> Source) {
   ExternalSource = std::move(Source);
 }
bool ASTContext::AtomicUsesUnsupportedLibcall(const AtomicExpr *E) const {
   CharUnits sizeChars = getTypeSizeInChars(AtomicTy);
   uint64_t Size = sizeChars.getQuantity();
   CharUnits alignChars = getTypeAlignInChars(AtomicTy);
   unsigned Align = alignChars.getQuantity();
   unsigned MaxInlineWidthInBits = getTargetInfo().getMaxAtomicInlineWidth();
   return (Size != Align || toBits(sizeChars) > MaxInlineWidthInBits);
 }
 
-static ast_type_traits::DynTypedNode getSingleDynTypedNodeFromParentMap(
-    ASTContext::ParentMapPointers::mapped_type U) {
-  if (const auto *D = U.dyn_cast<const Decl *>())
-    return ast_type_traits::DynTypedNode::create(*D);
-  if (const auto *S = U.dyn_cast<const Stmt *>())
-    return ast_type_traits::DynTypedNode::create(*S);
-  return *U.get<ast_type_traits::DynTypedNode *>();
-}
-
-namespace {
-
 /// Template specializations to abstract away from pointers and TypeLocs.
 /// @{
 template <typename T>
-ast_type_traits::DynTypedNode createDynTypedNode(const T &Node) {
+static ast_type_traits::DynTypedNode createDynTypedNode(const T &Node) {
   return ast_type_traits::DynTypedNode::create(*Node);
 }
 template <>
 ast_type_traits::DynTypedNode createDynTypedNode(const TypeLoc &Node) {
   return ast_type_traits::DynTypedNode::create(Node);
 }
 template <>
 ast_type_traits::DynTypedNode
 createDynTypedNode(const NestedNameSpecifierLoc &Node) {
   return ast_type_traits::DynTypedNode::create(Node);
 }
 /// @}
 
-  /// A \c RecursiveASTVisitor that builds a map from nodes to their
-  /// parents as defined by the \c RecursiveASTVisitor.
-  ///
-  /// Note that the relationship described here is purely in terms of AST
-  /// traversal - there are other relationships (for example declaration context)
-  /// in the AST that are better modeled by special matchers.
-  ///
-  /// FIXME: Currently only builds up the map using \c Stmt and \c Decl nodes.
-  class ParentMapASTVisitor : public RecursiveASTVisitor<ParentMapASTVisitor> {
-  public:
-    /// Builds and returns the translation unit's parent map.
-    ///
-    ///  The caller takes ownership of the returned \c ParentMap.
-    static std::pair<ASTContext::ParentMapPointers *,
-                     ASTContext::ParentMapOtherNodes *>
-    buildMap(TranslationUnitDecl &TU) {
-      ParentMapASTVisitor Visitor(new ASTContext::ParentMapPointers,
-                                  new ASTContext::ParentMapOtherNodes);
-      Visitor.TraverseDecl(&TU);
-      return std::make_pair(Visitor.Parents, Visitor.OtherParents);
-    }
+class ASTContext::ParentMap::ASTVisitor
+    : public RecursiveASTVisitor<ASTVisitor> {
+  ASTVisitor(ParentMap &Map) : Map(Map) {}
 
-  private:
-    friend class RecursiveASTVisitor<ParentMapASTVisitor>;
+  friend class RecursiveASTVisitor<ASTVisitor>;
 
-    using VisitorBase = RecursiveASTVisitor<ParentMapASTVisitor>;
+  using VisitorBase = RecursiveASTVisitor<ASTVisitor>;
 
-    ParentMapASTVisitor(ASTContext::ParentMapPointers *Parents,
-                        ASTContext::ParentMapOtherNodes *OtherParents)
-        : Parents(Parents), OtherParents(OtherParents) {}
+  bool shouldVisitTemplateInstantiations() const { return true; }
 
-    bool shouldVisitTemplateInstantiations() const {
-      return true;
-    }
-
-    bool shouldVisitImplicitCode() const {
-      return true;
-    }
+  bool shouldVisitImplicitCode() const { return true; }
 
-    template <typename T, typename MapNodeTy, typename BaseTraverseFn,
-              typename MapTy>
-    bool TraverseNode(T Node, MapNodeTy MapNode,
-                      BaseTraverseFn BaseTraverse, MapTy *Parents) {
-      if (!Node)
-        return true;
-      if (ParentStack.size() > 0) {
-        // FIXME: Currently we add the same parent multiple times, but only
-        // when no memoization data is available for the type.
-        // For example when we visit all subexpressions of template
-        // instantiations; this is suboptimal, but benign: the only way to
-        // visit those is with hasAncestor / hasParent, and those do not create
-        // new matches.
-        // The plan is to enable DynTypedNode to be storable in a map or hash
-        // map. The main problem there is to implement hash functions /
-        // comparison operators for all types that DynTypedNode supports that
-        // do not have pointer identity.
-        auto &NodeOrVector = (*Parents)[MapNode];
-        if (NodeOrVector.isNull()) {
-          if (const auto *D = ParentStack.back().get<Decl>())
-            NodeOrVector = D;
-          else if (const auto *S = ParentStack.back().get<Stmt>())
-            NodeOrVector = S;
-          else
-            NodeOrVector =
-                new ast_type_traits::DynTypedNode(ParentStack.back());
-        } else {
-          if (!NodeOrVector.template is<ASTContext::ParentVector *>()) {
-            auto *Vector = new ASTContext::ParentVector(
-                1, getSingleDynTypedNodeFromParentMap(NodeOrVector));
-            delete NodeOrVector
-                    .template dyn_cast<ast_type_traits::DynTypedNode *>();
-            NodeOrVector = Vector;
-          }
+  bool TraverseNode(T Node, MapNodeTy MapNode, BaseTraverseFn BaseTraverse,
+                    MapTy *Parents) {
+          NodeOrVector = new ast_type_traits::DynTypedNode(ParentStack.back());
+        if (!NodeOrVector.template is<ParentVector *>()) {
+          auto *Vector = new ParentVector(
 
-          auto *Vector =
-              NodeOrVector.template get<ASTContext::ParentVector *>();
-          // Skip duplicates for types that have memoization data.
-          // We must check that the type has memoization data before calling
-          // std::find() because DynTypedNode::operator== can't compare all
-          // types.
-          bool Found = ParentStack.back().getMemoizationData() &&
-                       std::find(Vector->begin(), Vector->end(),
-                                 ParentStack.back()) != Vector->end();
-          if (!Found)
-            Vector->push_back(ParentStack.back());
-        }
-      }
-      ParentStack.push_back(createDynTypedNode(Node));
-      bool Result = BaseTraverse();
-      ParentStack.pop_back();
-      return Result;
-    }
+        auto *Vector = NodeOrVector.template get<ParentVector *>();
 
-    bool TraverseDecl(Decl *DeclNode) {
-      return TraverseNode(DeclNode, DeclNode,
-                          [&] { return VisitorBase::TraverseDecl(DeclNode); },
-                          Parents);
-    }
+    return TraverseNode(
+        DeclNode, DeclNode, [&] { return VisitorBase::TraverseDecl(DeclNode); },
+        &Map.PointerParents);
 
-    bool TraverseStmt(Stmt *StmtNode) {
-      return TraverseNode(StmtNode, StmtNode,
-                          [&] { return VisitorBase::TraverseStmt(StmtNode); },
-                          Parents);
-    }
+    return TraverseNode(
+        StmtNode, StmtNode, [&] { return VisitorBase::TraverseStmt(StmtNode); },
+        &Map.PointerParents);
 
-    bool TraverseTypeLoc(TypeLoc TypeLocNode) {
-      return TraverseNode(
-          TypeLocNode, ast_type_traits::DynTypedNode::create(TypeLocNode),
-          [&] { return VisitorBase::TraverseTypeLoc(TypeLocNode); },
-          OtherParents);
-    }
+        &Map.OtherParents);
 
-    bool TraverseNestedNameSpecifierLoc(NestedNameSpecifierLoc NNSLocNode) {
-      return TraverseNode(
-          NNSLocNode, ast_type_traits::DynTypedNode::create(NNSLocNode),
-          [&] {
-            return VisitorBase::TraverseNestedNameSpecifierLoc(NNSLocNode);
-          },
-          OtherParents);
-    }
+        [&] { return VisitorBase::TraverseNestedNameSpecifierLoc(NNSLocNode); },
+        &Map.OtherParents);
 
-    ASTContext::ParentMapPointers *Parents;
-    ASTContext::ParentMapOtherNodes *OtherParents;
-    llvm::SmallVector<ast_type_traits::DynTypedNode, 16> ParentStack;
-  };
+  ParentMap &Map;
 
-} // namespace
-
-template <typename NodeTy, typename MapTy>
-static ASTContext::DynTypedNodeList getDynNodeFromMap(const NodeTy &Node,
-                                                      const MapTy &Map) {
-  auto I = Map.find(Node);
-  if (I == Map.end()) {
-    return llvm::ArrayRef<ast_type_traits::DynTypedNode>();
-  }
-  if (const auto *V =
-          I->second.template dyn_cast<ASTContext::ParentVector *>()) {
-    return llvm::makeArrayRef(*V);
-  }
-  return getSingleDynTypedNodeFromParentMap(I->second);
+ASTContext::ParentMap::ParentMap(ASTContext &Ctx) {
+  ASTVisitor(*this).TraverseAST(Ctx);
 }
 
 ASTContext::DynTypedNodeList
 ASTContext::getParents(const ast_type_traits::DynTypedNode &Node) {
-  if (!PointerParents) {
+  if (!Parents)
     // We always need to run over the whole translation unit, as
     // hasAncestor can escape any subtree.
-    auto Maps = ParentMapASTVisitor::buildMap(*getTranslationUnitDecl());
-    PointerParents.reset(Maps.first);
-    OtherParents.reset(Maps.second);
-  }
-  if (Node.getNodeKind().hasPointerIdentity())
-    return getDynNodeFromMap(Node.getMemoizationData(), *PointerParents);
-  return getDynNodeFromMap(Node, *OtherParents);
+    Parents = llvm::make_unique<ParentMap>(*this);
+  return Parents->getParents(Node);
 }
 
 bool
 ASTContext::ObjCMethodsAreEqual(const ObjCMethodDecl *MethodDecl,
                                 const ObjCMethodDecl *MethodImpl) {
   // No point trying to match an unavailable/deprecated mothod.
   if (MethodDecl->hasAttr<UnavailableAttr>()
       || MethodDecl->hasAttr<DeprecatedAttr>())

lib/ASTMatchers/ASTMatchFinder.cpp

private:
   // 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 ast_type_traits::DynTypedNode &Node, const DynTypedMatcher &Matcher,
       BoundNodesTreeBuilder *Builder, AncestorMatchMode MatchMode) {
-    if (Node.get<TranslationUnitDecl>() ==
-        ActiveASTContext->getTranslationUnitDecl())
+    if (Node.get<TranslationUnitDecl>() != nullptr)
       return false;
 
     // For AST-nodes that don't have an identity, we can't memoize.
     if (!Builder->isComparable())
       return matchesAncestorOfRecursively(Node, Matcher, Builder, MatchMode);
 
     MatchKey Key;
     Key.MatcherID = Matcher.getID();
bool memoizedMatchesAncestorOfRecursively(
     return CachedResult.ResultOfMatch;
   }
 
   bool matchesAncestorOfRecursively(const ast_type_traits::DynTypedNode &Node,
                                     const DynTypedMatcher &Matcher,
                                     BoundNodesTreeBuilder *Builder,
                                     AncestorMatchMode MatchMode) {
     const auto &Parents = ActiveASTContext->getParents(Node);
-    assert(!Parents.empty() && "Found node that is not in the parent map.");
+    if (Parents.empty())
+      return false;
     if (Parents.size() == 1) {
       // Only one parent - do recursive memoization.
       const ast_type_traits::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, Matcher, Builder,
                                                     MatchMode);
         // Once we get back from the recursive call, the result will be the
         // same as the parent's result.
       }
     } else {
void MatchFinder::match(const clang::ast_type_traits::DynTypedNode &Node,
   Visitor.set_active_ast_context(&Context);
   Visitor.match(Node);
 }
 
 void MatchFinder::matchAST(ASTContext &Context) {
   internal::MatchASTVisitor Visitor(&Matchers, Options);
   Visitor.set_active_ast_context(&Context);
   Visitor.onStartOfTranslationUnit();
-  Visitor.TraverseDecl(Context.getTranslationUnitDecl());
+  Visitor.TraverseAST(Context);
   Visitor.onEndOfTranslationUnit();
 }
 
 void MatchFinder::registerTestCallbackAfterParsing(
     MatchFinder::ParsingDoneTestCallback *NewParsingDone) {
   ParsingDone = NewParsingDone;
 }
 
 StringRef MatchFinder::MatchCallback::getID() const { return "<unknown>"; }
 
 } // end namespace ast_matchers
 } // end namespace clang

unittests/AST/ASTContextParentMapTest.cpp

 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/ASTContext.h"
 #include "MatchVerifier.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/Tooling/Tooling.h"
 #include "gtest/gtest.h"
+#include "gmock/gmock.h"
+
+using testing::ElementsAre;
 
 namespace clang {
 namespace ast_matchers {
 
 TEST(GetParents, ReturnsParentForDecl) {
   MatchVerifier<Decl> Verifier;
   EXPECT_TRUE(
       Verifier.match("class C { void f(); };",
TEST(GetParents, ReturnsMultipleParentsInTemplateInstantiations) {
   EXPECT_TRUE(TemplateVerifier.match(
       "template<typename T> struct C { void f() {} };"
       "void g() { C<int> c; c.f(); }",
       compoundStmt(allOf(
           hasAncestor(cxxRecordDecl(isTemplateInstantiation())),
           hasAncestor(cxxRecordDecl(unless(isTemplateInstantiation())))))));
 }
 
+TEST(GetParents, RespectsTraversalScope) {
+  auto AST =
+      tooling::buildASTFromCode("struct foo { int bar; };", "foo.cpp",
+                                std::make_shared<PCHContainerOperations>());
+  auto &Ctx = AST->getASTContext();
+  auto &TU = *Ctx.getTranslationUnitDecl();
+  auto &Foo = *TU.lookup(&Ctx.Idents.get("foo")).front();
+  auto &Bar = *cast<DeclContext>(Foo).lookup(&Ctx.Idents.get("bar")).front();
+
+  using ast_type_traits::DynTypedNode;
+  // Initially, scope is the whole TU.
+  EXPECT_THAT(Ctx.getParents(Bar), ElementsAre(DynTypedNode::create(Foo)));
+  EXPECT_THAT(Ctx.getParents(Foo), ElementsAre(DynTypedNode::create(TU)));
+
+  // Restrict the scope, now some parents are gone.
+  Ctx.setTraversalScope({&Foo});
+  EXPECT_THAT(Ctx.getParents(Bar), ElementsAre(DynTypedNode::create(Foo)));
+  EXPECT_THAT(Ctx.getParents(Foo), ElementsAre());
+
+  // Reset the scope, we get back the original results.
+  Ctx.setTraversalScope({&TU});
+  EXPECT_THAT(Ctx.getParents(Bar), ElementsAre(DynTypedNode::create(Foo)));
+  EXPECT_THAT(Ctx.getParents(Foo), ElementsAre(DynTypedNode::create(TU)));
+}
+
 } // end namespace ast_matchers
 } // end namespace clang

unittests/Tooling/CMakeLists.txt

add_clang_unittest(ToolingTests
   RecursiveASTVisitorTests/InitListExprPreOrder.cpp
   RecursiveASTVisitorTests/InitListExprPreOrderNoQueue.cpp
   RecursiveASTVisitorTests/IntegerLiteral.cpp
   RecursiveASTVisitorTests/LambdaDefaultCapture.cpp
   RecursiveASTVisitorTests/LambdaExpr.cpp
   RecursiveASTVisitorTests/NestedNameSpecifiers.cpp
   RecursiveASTVisitorTests/ParenExpr.cpp
   RecursiveASTVisitorTests/TemplateArgumentLocTraverser.cpp
+  RecursiveASTVisitorTests/TraversalScope.cpp
   RecursiveASTVisitorTestDeclVisitor.cpp
   RecursiveASTVisitorTestPostOrderVisitor.cpp
   RecursiveASTVisitorTestTypeLocVisitor.cpp
   RefactoringActionRulesTest.cpp
   RefactoringCallbacksTest.cpp
   RefactoringTest.cpp
   ReplacementsYamlTest.cpp
   RewriterTest.cpp

unittests/Tooling/RecursiveASTVisitorTests/TraversalScope.cpp

+//===- unittest/Tooling/RecursiveASTVisitorTests/TraversalScope.cpp -------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestVisitor.h"
+
+using namespace clang;
+
+namespace {
+
+class Visitor : public ExpectedLocationVisitor<Visitor, clang::TestVisitor> {
+public:
+  Visitor(ASTContext *Context) { this->Context = Context; }
+
+  bool VisitNamedDecl(NamedDecl *D) {
+    if (!D->isImplicit())
+      Match(D->getName(), D->getLocation());
+    return true;
+  }
+};
+
+TEST(RecursiveASTVisitor, RespectsTraversalScope) {
+  auto AST = tooling::buildASTFromCode(
+      R"cpp(
+struct foo {
+  struct bar {
+    struct baz {};
+  };
+};
+      )cpp",
+      "foo.cpp", std::make_shared<PCHContainerOperations>());
+  auto &Ctx = AST->getASTContext();
+  auto &TU = *Ctx.getTranslationUnitDecl();
+  auto &Foo = *TU.lookup(&Ctx.Idents.get("foo")).front();
+  auto &Bar = *cast<DeclContext>(Foo).lookup(&Ctx.Idents.get("bar")).front();
+
+  Ctx.setTraversalScope({&Bar});
+
+  Visitor V(&Ctx);
+  V.DisallowMatch("foo", 2, 8);
+  V.ExpectMatch("bar", 3, 10);
+  V.ExpectMatch("baz", 4, 12);
+  V.TraverseAST(Ctx);
+}
+
+} // end anonymous namespace