[NFC] AST: Refactor for recursive MemberTypeRepr representation

Author: AnthonyLatsis

include/swift/AST/Attr.h

@@ -54,6 +54,7 @@ class ASTContext;
 struct PrintOptions;
 class CustomAttr;
 class Decl;
+class DeclRefTypeRepr;
 class AbstractFunctionDecl;
 class FuncDecl;
 class ClassDecl;
@@ -1766,7 +1767,7 @@ class CustomAttr final : public DeclAttribute {
   ///
   /// For an identifier type repr, return a pair of `nullptr` and the
   /// identifier.
-  std::pair<IdentTypeRepr *, IdentTypeRepr *> destructureMacroRef();
+  std::pair<IdentTypeRepr *, DeclRefTypeRepr *> destructureMacroRef();
 
   /// Whether the attribute has any arguments.
   bool hasArgs() const { return argList != nullptr; }

include/swift/AST/TypeRepr.h

@@ -356,15 +356,19 @@ class DeclRefTypeRepr : public TypeRepr {
   /// is a \c MemberTypeRepr, and the method returns its last member component.
   IdentTypeRepr *getLastComponent();
 
+  DeclNameLoc getNameLoc() const;
+  DeclNameRef getNameRef() const;
+
   /// Returns whether this instance has been bound to a type declaration. This
   /// happens during type resolution.
   bool isBound() const;
 
   /// Returns the type declaration this instance has been bound to.
   TypeDecl *getBoundDecl() const;
 
-  /// The identifier that describes the last component.
-  DeclNameRef getNameRef() const;
+  DeclContext *getDeclContext() const;
+
+  void setValue(TypeDecl *TD, DeclContext *DC);
 
   /// Returns whether this instance has a generic argument list. That is, either
   /// a valid angle bracket source range, or a positive number of generic
@@ -373,12 +377,19 @@ class DeclRefTypeRepr : public TypeRepr {
 
   ArrayRef<TypeRepr *> getGenericArgs() const;
 
+  SourceRange getAngleBrackets() const;
+
   static bool classof(const TypeRepr *T) {
     return T->getKind() == TypeReprKind::SimpleIdent ||
            T->getKind() == TypeReprKind::GenericIdent ||
            T->getKind() == TypeReprKind::Member;
   }
   static bool classof(const DeclRefTypeRepr *T) { return true; }
+
+protected:
+  void printImpl(ASTPrinter &Printer, const PrintOptions &Opts) const;
+
+  friend class TypeRepr;
 };
 
 /// An identifier type with an optional set of generic arguments.
@@ -435,9 +446,8 @@ class IdentTypeRepr : public DeclRefTypeRepr {
   static bool classof(const IdentTypeRepr *T) { return true; }
 
 protected:
-  void printImpl(ASTPrinter &Printer, const PrintOptions &Opts) const;
-
   SourceLoc getLocImpl() const { return Loc.getBaseNameLoc(); }
+
   friend class TypeRepr;
 };
 
@@ -583,7 +593,6 @@ class MemberTypeRepr final : public DeclRefTypeRepr,
   SourceLoc getEndLocImpl() const { return getLastComponent()->getEndLoc(); }
   SourceLoc getLocImpl() const { return getLastComponent()->getLoc(); }
 
-  void printImpl(ASTPrinter &Printer, const PrintOptions &Opts) const;
   friend class TypeRepr;
 };
 

lib/AST/Attr.cpp

@@ -2737,18 +2737,18 @@ CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
       CustomAttr(atLoc, range, type, initContext, argList, implicit);
 }
 
-std::pair<IdentTypeRepr *, IdentTypeRepr *> CustomAttr::destructureMacroRef() {
+std::pair<IdentTypeRepr *, DeclRefTypeRepr *>
+CustomAttr::destructureMacroRef() {
   TypeRepr *typeRepr = getTypeRepr();
   if (!typeRepr)
     return {nullptr, nullptr};
   if (auto *identType = dyn_cast<IdentTypeRepr>(typeRepr))
     return {nullptr, identType};
-  if (auto *memType = dyn_cast<MemberTypeRepr>(typeRepr))
-    if (auto *base = dyn_cast<IdentTypeRepr>(memType->getRoot()))
-      if (memType->getMemberComponents().size() == 1)
-        if (auto first =
-                dyn_cast<IdentTypeRepr>(memType->getMemberComponents().front()))
-          return {base, first};
+  if (auto *memType = dyn_cast<MemberTypeRepr>(typeRepr)) {
+    if (auto *base = dyn_cast<SimpleIdentTypeRepr>(memType->getBase())) {
+      return {base, memType};
+    }
+  }
   return {nullptr, nullptr};
 }
 

lib/AST/Expr.cpp

@@ -2236,57 +2236,47 @@ TypeExpr *TypeExpr::createForMemberDecl(DeclNameLoc ParentNameLoc,
   assert(ParentNameLoc.isValid());
   assert(NameLoc.isValid());
 
-  // The base component is the parent type.
-  auto *ParentComp = new (C) SimpleIdentTypeRepr(ParentNameLoc,
-                                                 Parent->createNameRef());
-  ParentComp->setValue(Parent, nullptr);
+  // The base is the parent type.
+  auto *BaseTR =
+      new (C) SimpleIdentTypeRepr(ParentNameLoc, Parent->createNameRef());
+  BaseTR->setValue(Parent, nullptr);
 
-  // The member component is the member we just found.
-  auto *MemberComp =
-      new (C) SimpleIdentTypeRepr(NameLoc, Decl->createNameRef());
-  MemberComp->setValue(Decl, nullptr);
+  auto *MemberTR =
+      MemberTypeRepr::create(C, BaseTR, NameLoc, Decl->createNameRef());
+  MemberTR->setValue(Decl, nullptr);
 
-  auto *TR = MemberTypeRepr::create(C, ParentComp, {MemberComp});
-  return new (C) TypeExpr(TR);
+  return new (C) TypeExpr(MemberTR);
 }
 
 TypeExpr *TypeExpr::createForMemberDecl(TypeRepr *ParentTR, DeclNameLoc NameLoc,
                                         TypeDecl *Decl) {
   ASTContext &C = Decl->getASTContext();
 
-  // Add a new component for the member we just found.
-  auto *NewComp = new (C) SimpleIdentTypeRepr(NameLoc, Decl->createNameRef());
-  NewComp->setValue(Decl, nullptr);
+  auto *MemberTR =
+      MemberTypeRepr::create(C, ParentTR, NameLoc, Decl->createNameRef());
+  MemberTR->setValue(Decl, nullptr);
 
-  TypeRepr *TR = nullptr;
-  if (auto *DeclRefTR = dyn_cast<DeclRefTypeRepr>(ParentTR)) {
-    // Create a new list of components.
-    SmallVector<IdentTypeRepr *, 4> Components;
-    if (auto *MemberTR = dyn_cast<MemberTypeRepr>(ParentTR)) {
-      auto MemberComps = MemberTR->getMemberComponents();
-      Components.append(MemberComps.begin(), MemberComps.end());
-    }
-
-    Components.push_back(NewComp);
-    TR = MemberTypeRepr::create(C, DeclRefTR->getRoot(), Components);
-  } else {
-    TR = MemberTypeRepr::create(C, ParentTR, NewComp);
-  }
-
-  return new (C) TypeExpr(TR);
+  return new (C) TypeExpr(MemberTR);
 }
 
 TypeExpr *TypeExpr::createForSpecializedDecl(DeclRefTypeRepr *ParentTR,
                                              ArrayRef<TypeRepr *> Args,
                                              SourceRange AngleLocs,
                                              ASTContext &C) {
-  auto *lastComp = ParentTR->getLastComponent();
+  DeclRefTypeRepr *specializedTR = nullptr;
 
-  if (!isa<SimpleIdentTypeRepr>(lastComp) || !lastComp->getBoundDecl())
+  auto *boundDecl = ParentTR->getBoundDecl();
+  if (!boundDecl || ParentTR->hasGenericArgList()) {
     return nullptr;
+  }
 
-  if (isa<TypeAliasDecl>(lastComp->getBoundDecl())) {
-    if (auto *memberTR = dyn_cast<MemberTypeRepr>(ParentTR)) {
+  if (isa<IdentTypeRepr>(ParentTR)) {
+    specializedTR = GenericIdentTypeRepr::create(
+        C, ParentTR->getNameLoc(), ParentTR->getNameRef(), Args, AngleLocs);
+    specializedTR->setValue(boundDecl, ParentTR->getDeclContext());
+  } else {
+    auto *const memberTR = cast<MemberTypeRepr>(ParentTR);
+    if (isa<TypeAliasDecl>(boundDecl)) {
       // If any of our parent types are unbound, bail out and let
       // the constraint solver can infer generic parameters for them.
       //
@@ -2299,48 +2289,30 @@ TypeExpr *TypeExpr::createForSpecializedDecl(DeclRefTypeRepr *ParentTR,
       //
       // FIXME: Once we can model generic typealiases properly, rip
       // this out.
-      auto isUnboundGenericComponent = [](IdentTypeRepr *ITR) -> bool {
-        if (isa<SimpleIdentTypeRepr>(ITR)) {
-          auto *decl = dyn_cast_or_null<GenericTypeDecl>(ITR->getBoundDecl());
+      MemberTypeRepr *currTR = memberTR;
+      while (auto *declRefBaseTR =
+                 dyn_cast<DeclRefTypeRepr>(currTR->getBase())) {
+        if (!declRefBaseTR->hasGenericArgList()) {
+          auto *decl =
+              dyn_cast_or_null<GenericTypeDecl>(declRefBaseTR->getBoundDecl());
           if (decl && decl->isGeneric())
-            return true;
+            return nullptr;
         }
 
-        return false;
-      };
-
-      for (auto *comp : memberTR->getMemberComponents().drop_back()) {
-        if (isUnboundGenericComponent(comp))
-          return nullptr;
-      }
-
-      if (auto *identRoot = dyn_cast<IdentTypeRepr>(memberTR->getRoot())) {
-        if (isUnboundGenericComponent(identRoot))
-          return nullptr;
+        currTR = dyn_cast<MemberTypeRepr>(declRefBaseTR);
+        if (!currTR) {
+          break;
+        }
       }
     }
-  }
 
-  auto *genericComp = GenericIdentTypeRepr::create(
-      C, lastComp->getNameLoc(), lastComp->getNameRef(), Args, AngleLocs);
-  genericComp->setValue(lastComp->getBoundDecl(), lastComp->getDeclContext());
-
-  TypeRepr *TR = nullptr;
-  if (auto *memberTR = dyn_cast<MemberTypeRepr>(ParentTR)) {
-    auto oldMemberComps = memberTR->getMemberComponents().drop_back();
-
-    // Create a new list of member components, replacing the last one with the
-    // new specialized one.
-    SmallVector<IdentTypeRepr *, 2> newMemberComps;
-    newMemberComps.append(oldMemberComps.begin(), oldMemberComps.end());
-    newMemberComps.push_back(genericComp);
-
-    TR = MemberTypeRepr::create(C, memberTR->getRoot(), newMemberComps);
-  } else {
-    TR = genericComp;
+    specializedTR =
+        MemberTypeRepr::create(C, memberTR->getBase(), ParentTR->getNameLoc(),
+                               ParentTR->getNameRef(), Args, AngleLocs);
+    specializedTR->setValue(boundDecl, ParentTR->getDeclContext());
   }
 
-  return new (C) TypeExpr(TR);
+  return new (C) TypeExpr(specializedTR);
 }
 
 // Create an implicit TypeExpr, with location information even though it

lib/AST/TypeCheckRequests.cpp

@@ -1822,11 +1822,8 @@ SourceRange UnresolvedMacroReference::getGenericArgsRange() const {
     auto [_, macro] = attr->destructureMacroRef();
     if (!macro)
       return SourceRange();
-    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(macro);
-    if (!genericTypeRepr)
-      return SourceRange();
 
-    return genericTypeRepr->getAngleBrackets();
+    return macro->getAngleBrackets();
   }
 
   llvm_unreachable("Unhandled case");
@@ -1840,11 +1837,8 @@ ArrayRef<TypeRepr *> UnresolvedMacroReference::getGenericArgs() const {
     auto [_, macro] = attr->destructureMacroRef();
     if (!macro)
       return {};
-    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(macro);
-    if (!genericTypeRepr)
-      return {};
 
-    return genericTypeRepr->getGenericArgs();
+    return macro->getGenericArgs();
   }
 
   llvm_unreachable("Unhandled case");

lib/AST/TypeRepr.cpp

@@ -203,6 +203,10 @@ const TypeRepr *DeclRefTypeRepr::getRoot() const {
   return cast<MemberTypeRepr>(this)->getRoot();
 }
 
+DeclNameLoc DeclRefTypeRepr::getNameLoc() const {
+  return const_cast<DeclRefTypeRepr *>(this)->getLastComponent()->getNameLoc();
+}
+
 bool DeclRefTypeRepr::isBound() const {
   return const_cast<DeclRefTypeRepr *>(this)->getLastComponent()->isBound();
 }
@@ -213,6 +217,16 @@ TypeDecl *DeclRefTypeRepr::getBoundDecl() const {
       ->getBoundDecl();
 }
 
+DeclContext *DeclRefTypeRepr::getDeclContext() const {
+  return const_cast<DeclRefTypeRepr *>(this)
+      ->getLastComponent()
+      ->getDeclContext();
+}
+
+void DeclRefTypeRepr::setValue(TypeDecl *TD, DeclContext *DC) {
+  getLastComponent()->setValue(TD, DC);
+}
+
 bool DeclRefTypeRepr::hasGenericArgList() const {
   return isa<GenericIdentTypeRepr>(
       const_cast<DeclRefTypeRepr *>(this)->getLastComponent());
@@ -227,6 +241,15 @@ ArrayRef<TypeRepr *> DeclRefTypeRepr::getGenericArgs() const {
   return {};
 }
 
+SourceRange DeclRefTypeRepr::getAngleBrackets() const {
+  auto *lastComp = const_cast<DeclRefTypeRepr *>(this)->getLastComponent();
+  if (auto *genericTR = dyn_cast<GenericIdentTypeRepr>(lastComp)) {
+    return genericTR->getAngleBrackets();
+  }
+
+  return SourceRange();
+}
+
 DeclNameRef DeclRefTypeRepr::getNameRef() const {
   return const_cast<DeclRefTypeRepr *>(this)->getLastComponent()->getNameRef();
 }
@@ -341,19 +364,13 @@ void AttributedTypeRepr::printAttrs(ASTPrinter &Printer,
   }
 }
 
-static void printGenericArgs(ASTPrinter &Printer, const PrintOptions &Opts,
-                             ArrayRef<TypeRepr *> Args) {
-  if (Args.empty())
-    return;
-
-  Printer << "<";
-  interleave(Args, [&](TypeRepr *Arg) { printTypeRepr(Arg, Printer, Opts); },
-             [&] { Printer << ", "; });
-  Printer << ">";
-}
+void DeclRefTypeRepr::printImpl(ASTPrinter &Printer,
+                                const PrintOptions &Opts) const {
+  if (auto *memberTR = dyn_cast<MemberTypeRepr>(this)) {
+    printTypeRepr(memberTR->getBase(), Printer, Opts);
+    Printer << ".";
+  }
 
-void IdentTypeRepr::printImpl(ASTPrinter &Printer,
-                              const PrintOptions &Opts) const {
   if (auto *TD = dyn_cast_or_null<TypeDecl>(getBoundDecl())) {
     if (auto MD = dyn_cast<ModuleDecl>(TD))
       Printer.printModuleRef(MD, getNameRef().getBaseIdentifier());
@@ -363,16 +380,13 @@ void IdentTypeRepr::printImpl(ASTPrinter &Printer,
     Printer.printName(getNameRef().getBaseIdentifier());
   }
 
-  if (auto GenIdT = dyn_cast<GenericIdentTypeRepr>(this))
-    printGenericArgs(Printer, Opts, GenIdT->getGenericArgs());
-}
-
-void MemberTypeRepr::printImpl(ASTPrinter &Printer,
-                               const PrintOptions &Opts) const {
-  printTypeRepr(getRoot(), Printer, Opts);
-  for (auto C : getMemberComponents()) {
-    Printer << ".";
-    printTypeRepr(C, Printer, Opts);
+  if (hasGenericArgList()) {
+    Printer << "<";
+    interleave(
+        getGenericArgs(),
+        [&](TypeRepr *Arg) { printTypeRepr(Arg, Printer, Opts); },
+        [&] { Printer << ", "; });
+    Printer << ">";
   }
 }