AST/ASTWalker: Refactor for recursive MemberTypeRepr representation

Author: AnthonyLatsis

include/swift/AST/ASTWalker.h

@@ -103,6 +103,37 @@ enum class MacroWalking {
   None
 };
 
+/// A scheme for walking a `MemberTypeRepr`.
+enum class MemberTypeReprWalkingScheme {
+  /// Walk in source order, such that each subsequent dot-separated component is
+  /// a child of the previous one. For example, walk `A.B<T.U>.C` like so
+  /// (top-down order):
+  ///
+  /// ```
+  /// A
+  /// ╰─B
+  ///   ├─T
+  ///   │ ╰─U
+  ///   ╰─C
+  /// ```
+  SourceOrderRecursive,
+
+  /// Walk in AST order (that is, according to how member type
+  /// representations are modeled in the AST, such that each previous
+  /// dot-separated component is a child of the subsequent one), base before
+  /// generic arguments. For example, walk `A.B<T.U>.C` like so
+  /// (top-down order):
+  ///
+  /// ```
+  /// C
+  /// ╰─B
+  ///   ├─A
+  ///   ╰─U
+  ///     ╰─T
+  /// ```
+  ASTOrderRecursive
+};
+
 /// An abstract class used to traverse an AST.
 class ASTWalker {
 public:
@@ -563,6 +594,11 @@ class ASTWalker {
     return Action::Continue();
   }
 
+  /// This method configures how to walk `MemberTypeRepr` nodes.
+  virtual MemberTypeReprWalkingScheme getMemberTypeReprWalkingScheme() const {
+    return MemberTypeReprWalkingScheme::ASTOrderRecursive;
+  }
+
   /// This method configures whether the walker should explore into the generic
   /// params in AbstractFunctionDecl and NominalTypeDecl.
   virtual bool shouldWalkIntoGenericParams() { return false; }

include/swift/AST/TypeDeclFinder.h

@@ -19,7 +19,7 @@
 namespace swift {
 
 class BoundGenericType;
-class IdentTypeRepr;
+class DeclRefTypeRepr;
 class NominalType;
 class TypeAliasType;
 
@@ -57,15 +57,13 @@ class SimpleTypeDeclFinder : public TypeDeclFinder {
     : Callback(callback) {}
 };
 
-/// Walks a \c TypeRepr to find all \c IdentTypeRepr nodes with bound
-/// type declarations.
-///
-/// Subclasses can either override #visitTypeDecl if they only care about
-/// types on their own, or #visitIdentTypeRepr if they want to keep
-/// the TypeRepr around.
-class TypeReprIdentFinder : public ASTWalker {
-  /// The function to call when a \c IdentTypeRepr is seen.
-  llvm::function_ref<bool(const IdentTypeRepr *)> Callback;
+/// Walks a `TypeRepr` and reports all `DeclRefTypeRepr` nodes with bound
+/// type declarations by invoking a given callback. These nodes are reported in
+/// depth- and base-first AST order. For example, nodes in `A<T>.B<U>` will be
+/// reported in the following order: `TAUB`.
+class DeclRefTypeReprFinder : public ASTWalker {
+  /// The function to call when a `DeclRefTypeRepr` is seen.
+  llvm::function_ref<bool(const DeclRefTypeRepr *)> Callback;
 
   MacroWalking getMacroWalkingBehavior() const override {
     return MacroWalking::Arguments;
@@ -74,11 +72,10 @@ class TypeReprIdentFinder : public ASTWalker {
   PostWalkAction walkToTypeReprPost(TypeRepr *TR) override;
 
 public:
-  explicit TypeReprIdentFinder(
-      llvm::function_ref<bool(const IdentTypeRepr *)> callback)
-    : Callback(callback) {}
+  explicit DeclRefTypeReprFinder(
+      llvm::function_ref<bool(const DeclRefTypeRepr *)> callback)
+      : Callback(callback) {}
 };
-
 }
 
 #endif

lib/AST/ASTWalker.cpp

@@ -1438,6 +1438,8 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 #define TYPEREPR(Id, Parent) bool visit##Id##TypeRepr(Id##TypeRepr *T);
 #include "swift/AST/TypeReprNodes.def"
 
+  bool visitDeclRefTypeRepr(DeclRefTypeRepr *T);
+
   using Action = ASTWalker::Action;
 
   using PreWalkAction = ASTWalker::PreWalkAction;
@@ -1624,9 +1626,24 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     return false;
   }
 
+private:
+  /// Walk a `MemberTypeRepr` in source order such that each subsequent
+  /// dot-separated component is a child of the previous one
+  [[nodiscard]] bool doItInSourceOrderRecursive(MemberTypeRepr *T);
+
+public:
   /// Returns true on failure.
   [[nodiscard]]
   bool doIt(TypeRepr *T) {
+    if (auto *MTR = dyn_cast<MemberTypeRepr>(T)) {
+      switch (Walker.getMemberTypeReprWalkingScheme()) {
+      case MemberTypeReprWalkingScheme::SourceOrderRecursive:
+        return doItInSourceOrderRecursive(MTR);
+      case MemberTypeReprWalkingScheme::ASTOrderRecursive:
+        break;
+      }
+    }
+
     return traverse(
         Walker.walkToTypeReprPre(T),
         [&]() { return visit(T); },
@@ -1706,6 +1723,89 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
 } // end anonymous namespace
 
+bool Traversal::doItInSourceOrderRecursive(MemberTypeRepr *T) {
+  // Qualified types are modeled resursively such that each previous
+  // dot-separated component is a child of the next one. To walk a member type
+  // representation according to
+  // `MemberTypeReprWalkingScheme::SourceOrderRecursive`:
+
+  // 1. Pre-walk the dot-separated components in source order. If asked to skip
+  //    the children of a given component:
+  //    1. Set the depth at which to start post-walking later.
+  //    2. Skip its generic arguments and subsequent components.
+  std::function<bool(TypeRepr *, std::optional<unsigned> &, unsigned)>
+      doItInSourceOrderPre = [&](TypeRepr *T,
+                                 std::optional<unsigned> &StartPostWalkDepth,
+                                 unsigned Depth) {
+        if (auto *MemberTR = dyn_cast<MemberTypeRepr>(T)) {
+          if (doItInSourceOrderPre(MemberTR->getBase(), StartPostWalkDepth,
+                                   Depth + 1)) {
+            return true;
+          }
+
+          if (StartPostWalkDepth.has_value()) {
+            return false;
+          }
+        }
+
+        switch (this->Walker.walkToTypeReprPre(T).Action) {
+        case PreWalkAction::Stop:
+          return true;
+        case PreWalkAction::SkipChildren:
+          StartPostWalkDepth = Depth;
+          return false;
+        case PreWalkAction::SkipNode:
+          StartPostWalkDepth = Depth + 1;
+          return false;
+        case PreWalkAction::Continue:
+          break;
+        }
+
+        if (auto *DeclRefTR = dyn_cast<DeclRefTypeRepr>(T)) {
+          for (auto *Arg : DeclRefTR->getGenericArgs()) {
+            if (doIt(Arg)) {
+              return true;
+            }
+          }
+        } else if (visit(T)) {
+          return true;
+        }
+
+        return false;
+      };
+
+  // 2. Post-walk the components in reverse order, respecting the depth at which
+  //    to start post-walking if set.
+  std::function<bool(TypeRepr *, std::optional<unsigned>, unsigned)>
+      doItInSourceOrderPost = [&](TypeRepr *T,
+                                  std::optional<unsigned> StartPostWalkDepth,
+                                  unsigned Depth) {
+        if (!StartPostWalkDepth.has_value() || Depth >= *StartPostWalkDepth) {
+          switch (this->Walker.walkToTypeReprPost(T).Action) {
+          case PostWalkAction::Continue:
+            break;
+          case PostWalkAction::Stop:
+            return true;
+          }
+        }
+
+        if (auto *MemberTR = dyn_cast<MemberTypeRepr>(T)) {
+          return doItInSourceOrderPost(MemberTR->getBase(), StartPostWalkDepth,
+                                       Depth + 1);
+        }
+
+        return false;
+      };
+
+  std::optional<unsigned> StartPostWalkDepth;
+
+  if (doItInSourceOrderPre(T, StartPostWalkDepth, 0)) {
+    return true;
+  }
+
+  return doItInSourceOrderPost(T, StartPostWalkDepth, 0);
+}
+
 #pragma mark Statement traversal
 Stmt *Traversal::visitBreakStmt(BreakStmt *BS) {
   return BS;
@@ -2130,27 +2230,32 @@ bool Traversal::visitAttributedTypeRepr(AttributedTypeRepr *T) {
   return doIt(T->getTypeRepr());
 }
 
-bool Traversal::visitSimpleIdentTypeRepr(SimpleIdentTypeRepr *T) {
-  return false;
-}
+bool Traversal::visitDeclRefTypeRepr(DeclRefTypeRepr *T) {
+  if (auto *memberTR = dyn_cast<MemberTypeRepr>(T)) {
+    if (doIt(memberTR->getBase())) {
+      return true;
+    }
+  }
 
-bool Traversal::visitGenericIdentTypeRepr(GenericIdentTypeRepr *T) {
-  for (auto genArg : T->getGenericArgs()) {
-    if (doIt(genArg))
+  for (auto *genericArg : T->getGenericArgs()) {
+    if (doIt(genericArg)) {
       return true;
+    }
   }
+
   return false;
 }
 
-bool Traversal::visitMemberTypeRepr(MemberTypeRepr *T) {
-  if (doIt(T->getRoot()))
-    return true;
+bool Traversal::visitSimpleIdentTypeRepr(SimpleIdentTypeRepr *T) {
+  return visitDeclRefTypeRepr(T);
+}
 
-  for (auto comp : T->getMemberComponents()) {
-    if (doIt(comp))
-      return true;
-  }
-  return false;
+bool Traversal::visitGenericIdentTypeRepr(GenericIdentTypeRepr *T) {
+  return visitDeclRefTypeRepr(T);
+}
+
+bool Traversal::visitMemberTypeRepr(MemberTypeRepr *T) {
+  return visitDeclRefTypeRepr(T);
 }
 
 bool Traversal::visitFunctionTypeRepr(FunctionTypeRepr *T) {

lib/AST/NameLookup.cpp

@@ -3361,15 +3361,19 @@ CollectedOpaqueReprs swift::collectOpaqueTypeReprs(TypeRepr *r, ASTContext &ctx,
         if (!composition->isTypeReprAny())
           Reprs.push_back(composition);
         return Action::SkipNode();
-      } else if (auto generic = dyn_cast<GenericIdentTypeRepr>(repr)) {
-        if (generic->isProtocolOrProtocolComposition(dc)){
-          Reprs.push_back(generic);
-          return Action::SkipNode();
+      } else if (isa<DeclRefTypeRepr>(repr)) {
+        // We only care about the type of an outermost member type
+        // representation. For example, in `A<T>.B.C<U>`, check `C` and generic
+        // arguments `U` and `T`, but not `A` or `B`.
+        if (auto *parentMemberTR =
+                dyn_cast_or_null<MemberTypeRepr>(Parent.getAsTypeRepr())) {
+          if (repr == parentMemberTR->getBase()) {
+            return Action::Continue();
+          }
         }
-        return Action::Continue();
-      } else if (auto declRef = dyn_cast<DeclRefTypeRepr>(repr)) {
-        if (declRef->isProtocolOrProtocolComposition(dc))
-          Reprs.push_back(declRef);
+
+        if (repr->isProtocolOrProtocolComposition(dc))
+          Reprs.push_back(repr);
       }
       return Action::Continue();
     }

lib/AST/TypeDeclFinder.cpp

@@ -50,9 +50,9 @@ SimpleTypeDeclFinder::visitTypeAliasType(TypeAliasType *ty) {
 }
 
 ASTWalker::PostWalkAction
-TypeReprIdentFinder::walkToTypeReprPost(TypeRepr *TR) {
-  auto ITR = dyn_cast<IdentTypeRepr>(TR);
-  if (!ITR || !ITR->getBoundDecl())
+DeclRefTypeReprFinder::walkToTypeReprPost(TypeRepr *TR) {
+  auto *declRefTR = dyn_cast<DeclRefTypeRepr>(TR);
+  if (!declRefTR || !declRefTR->getBoundDecl())
     return Action::Continue();
-  return Action::StopIf(!Callback(ITR));
+  return Action::StopIf(!Callback(declRefTR));
 }

lib/IDE/SourceEntityWalker.cpp

@@ -62,6 +62,10 @@ class SemaAnnotator : public ASTWalker {
     return SEWalker.getMacroWalkingBehavior();
   }
 
+  MemberTypeReprWalkingScheme getMemberTypeReprWalkingScheme() const override {
+    return MemberTypeReprWalkingScheme::SourceOrderRecursive;
+  }
+
   PreWalkAction walkToDeclPre(Decl *D) override;
   PreWalkAction walkToDeclPreProper(Decl *D);
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override;
@@ -627,15 +631,15 @@ ASTWalker::PreWalkAction SemaAnnotator::walkToTypeReprPre(TypeRepr *T) {
   if (!Continue)
     return Action::Stop();
 
-  if (auto IdT = dyn_cast<IdentTypeRepr>(T)) {
-    if (ValueDecl *VD = IdT->getBoundDecl()) {
+  if (auto *DeclRefT = dyn_cast<DeclRefTypeRepr>(T)) {
+    if (ValueDecl *VD = DeclRefT->getBoundDecl()) {
       if (auto *ModD = dyn_cast<ModuleDecl>(VD)) {
-        auto ident = IdT->getNameRef().getBaseIdentifier();
-        auto Continue = passReference(ModD, {ident, IdT->getLoc()});
+        auto ident = DeclRefT->getNameRef().getBaseIdentifier();
+        auto Continue = passReference(ModD, {ident, DeclRefT->getLoc()});
         return Action::StopIf(!Continue);
       }
       auto Continue = passReference(
-          VD, Type(), IdT->getNameLoc(),
+          VD, Type(), DeclRefT->getNameLoc(),
           ReferenceMetaData(SemaReferenceKind::TypeRef, llvm::None));
       return Action::StopIf(!Continue);
     }

lib/IDE/SyntaxModel.cpp

@@ -402,7 +402,13 @@ class ModelASTWalker : public ASTWalker {
   PostWalkResult<Stmt *> walkToStmtPost(Stmt *S) override;
   PreWalkAction walkToDeclPre(Decl *D) override;
   PostWalkAction walkToDeclPost(Decl *D) override;
+
+  MemberTypeReprWalkingScheme getMemberTypeReprWalkingScheme() const override {
+    return MemberTypeReprWalkingScheme::SourceOrderRecursive;
+  }
+
   PreWalkAction walkToTypeReprPre(TypeRepr *T) override;
+
   bool shouldWalkIntoGenericParams() override { return true; }
 
 private:
@@ -1122,12 +1128,12 @@ ASTWalker::PreWalkAction ModelASTWalker::walkToTypeReprPre(TypeRepr *T) {
     if (!handleAttrs(AttrT->getAttrs()))
       return Action::SkipNode();
 
-  } else if (auto IdT = dyn_cast<IdentTypeRepr>(T)) {
-    if (!passTokenNodesUntil(IdT->getStartLoc(),
-                             ExcludeNodeAtLocation).shouldContinue)
+  } else if (auto *DeclRefT = dyn_cast<DeclRefTypeRepr>(T)) {
+    if (!passTokenNodesUntil(DeclRefT->getLoc(), ExcludeNodeAtLocation)
+             .shouldContinue)
       return Action::SkipNode();
     if (TokenNodes.empty() ||
-        TokenNodes.front().Range.getStart() != IdT->getStartLoc())
+        TokenNodes.front().Range.getStart() != DeclRefT->getLoc())
       return Action::SkipNode();
     if (!passNode({SyntaxNodeKind::TypeId, TokenNodes.front().Range}))
       return Action::SkipNode();

lib/Parse/ParseType.cpp

@@ -557,10 +557,11 @@ ParserResult<TypeRepr> Parser::parseTypeScalar(
       }
 
       PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
-        if (auto ident = dyn_cast<IdentTypeRepr>(T)) {
-          if (auto decl = ident->getBoundDecl()) {
-            if (auto genericParam = dyn_cast<GenericTypeParamDecl>(decl))
-              ident->overwriteNameRef(genericParam->createNameRef());
+        // Only unqualified identifiers can reference generic parameters.
+        if (auto *simpleIdentTR = dyn_cast<SimpleIdentTypeRepr>(T)) {
+          if (auto *genericParam = dyn_cast_or_null<GenericTypeParamDecl>(
+                  simpleIdentTR->getBoundDecl())) {
+            simpleIdentTR->overwriteNameRef(genericParam->createNameRef());
           }
         }
         return Action::Continue();