Merge pull request #38883 from ileitch/SR-13766

[Index] Apply RelationContainedBy role to references contained by VarDecl.

Author: bnbarham

include/swift/IDE/SourceEntityWalker.h

@@ -88,6 +88,14 @@ class SourceEntityWalker {
   /// false, the remaining traversal is terminated and returns failure.
   virtual bool walkToDeclPost(Decl *D) { return true; }
 
+  /// This method is called in AST order, unlike \c walkToDecl* which are called
+  /// in source order. It is guaranteed to be called in balance with its
+  /// counterpart \c endBalancedASTOrderDeclVisit.
+  virtual void beginBalancedASTOrderDeclVisit(Decl *D){};
+
+  /// This method is called after declaration visitation in AST order.
+  virtual void endBalancedASTOrderDeclVisit(Decl *D){};
+
   /// This method is called when first visiting a statement, before walking into
   /// its children.  If it returns false, the subtree is skipped.
   virtual bool walkToStmtPre(Stmt *S) { return true; }
@@ -112,6 +120,15 @@ class SourceEntityWalker {
   /// returns false, the remaining traversal is terminated and returns failure.
   virtual bool walkToPatternPost(Pattern *P) { return true; }
 
+  /// This method is called when first visiting a type representation, before
+  /// walking into its children. If it returns false, the subtree is skipped.
+  virtual bool walkToTypeReprPre(TypeRepr *T) { return true; }
+
+  /// This method is called after visiting the children of a type
+  /// representation. If it returns false, the remaining traversal is terminated
+  /// and returns failure.
+  virtual bool walkToTypeReprPost(TypeRepr *T) { return true; }
+
   /// This method is called when a ValueDecl is referenced in source. If it
   /// returns false, the remaining traversal is terminated and returns failure.
   ///

lib/IDE/SourceEntityWalker.cpp

@@ -55,10 +55,12 @@ class SemaAnnotator : public ASTWalker {
     return SEWalker.shouldWalkIntoGenericParams();
   }
   bool walkToDeclPre(Decl *D) override;
+  bool walkToDeclPreProper(Decl *D);
   std::pair<bool, Expr *> walkToExprPre(Expr *E) override;
   bool walkToTypeReprPre(TypeRepr *T) override;
 
   bool walkToDeclPost(Decl *D) override;
+  bool walkToDeclPostProper(Decl *D);
   Expr *walkToExprPost(Expr *E) override;
   bool walkToTypeReprPost(TypeRepr *T) override;
 
@@ -118,6 +120,21 @@ bool SemaAnnotator::walkToDeclPre(Decl *D) {
     return isa<PatternBindingDecl>(D);
   }
 
+  SEWalker.beginBalancedASTOrderDeclVisit(D);
+  bool Continue = walkToDeclPreProper(D);
+
+  if (!Continue) {
+    // To satisfy the contract of balanced calls to
+    // begin/endBalancedASTOrderDeclVisit, we must call
+    // endBalancedASTOrderDeclVisit here if walkToDeclPost isn't going to be
+    // called.
+    SEWalker.endBalancedASTOrderDeclVisit(D);
+  }
+
+  return Continue;
+}
+
+bool SemaAnnotator::walkToDeclPreProper(Decl *D) {
   if (!handleCustomAttributes(D)) {
     Cancelled = true;
     return false;
@@ -203,6 +220,12 @@ bool SemaAnnotator::walkToDeclPre(Decl *D) {
 }
 
 bool SemaAnnotator::walkToDeclPost(Decl *D) {
+  bool Continue = walkToDeclPostProper(D);
+  SEWalker.endBalancedASTOrderDeclVisit(D);
+  return Continue;
+}
+
+bool SemaAnnotator::walkToDeclPostProper(Decl *D) {
   if (isDone())
     return false;
 
@@ -580,6 +603,12 @@ bool SemaAnnotator::walkToTypeReprPre(TypeRepr *T) {
   if (isDone())
     return false;
 
+  bool Continue = SEWalker.walkToTypeReprPre(T);
+  if (!Continue) {
+    Cancelled = true;
+    return false;
+  }
+
   if (auto IdT = dyn_cast<ComponentIdentTypeRepr>(T)) {
     if (ValueDecl *VD = IdT->getBoundDecl()) {
       if (auto *ModD = dyn_cast<ModuleDecl>(VD)) {
@@ -591,11 +620,19 @@ bool SemaAnnotator::walkToTypeReprPre(TypeRepr *T) {
                            ReferenceMetaData(SemaReferenceKind::TypeRef, None));
     }
   }
+
   return true;
 }
 
 bool SemaAnnotator::walkToTypeReprPost(TypeRepr *T) {
-  return !isDone();
+  if (isDone()) {
+    return false;
+  }
+
+  bool Continue = SEWalker.walkToTypeReprPost(T);
+  if (!Continue)
+    Cancelled = true;
+  return Continue;
 }
 
 std::pair<bool, Pattern *> SemaAnnotator::walkToPatternPre(Pattern *P) {