[NFC] AST/NameLookup: Refactor for recursive MemberTypeRepr representation

Author: AnthonyLatsis

include/swift/AST/TypeRepr.h

@@ -362,6 +362,11 @@ class DeclRefTypeRepr : public TypeRepr {
   /// The identifier that describes the last component.
   DeclNameRef getNameRef() const;
 
+  /// Returns whether this instance has a generic argument list. That is, either
+  /// a valid angle bracket source range, or a positive number of generic
+  /// arguments.
+  bool hasGenericArgList() const;
+
   ArrayRef<TypeRepr *> getGenericArgs() const;
 
   static bool classof(const TypeRepr *T) {
@@ -518,11 +523,14 @@ class MemberTypeRepr final : public DeclRefTypeRepr,
       private llvm::TrailingObjects<MemberTypeRepr, IdentTypeRepr *> {
   friend TrailingObjects;
 
+  const ASTContext *C;
+
   /// The root component, which is not necessarily an identifier type.
   TypeRepr *Root;
 
-  MemberTypeRepr(TypeRepr *Root, ArrayRef<IdentTypeRepr *> MemberComponents)
-      : DeclRefTypeRepr(TypeReprKind::Member), Root(Root) {
+  MemberTypeRepr(TypeRepr *Root, ArrayRef<IdentTypeRepr *> MemberComponents,
+                 const ASTContext &C)
+      : DeclRefTypeRepr(TypeReprKind::Member), C(&C), Root(Root) {
     Bits.MemberTypeRepr.NumMemberComponents = MemberComponents.size();
     assert(MemberComponents.size() > 0 &&
            "MemberTypeRepr requires at least 1 member component");
@@ -531,12 +539,24 @@ class MemberTypeRepr final : public DeclRefTypeRepr,
   }
 
 public:
+  static MemberTypeRepr *create(const ASTContext &C, TypeRepr *Base,
+                                DeclNameLoc NameLoc, DeclNameRef Name);
+
+  static MemberTypeRepr *create(const ASTContext &C, TypeRepr *Base,
+                                DeclNameLoc NameLoc, DeclNameRef Name,
+                                ArrayRef<TypeRepr *> GenericArgs,
+                                SourceRange AngleBrackets);
+
   static TypeRepr *create(const ASTContext &Ctx, TypeRepr *Root,
                           ArrayRef<IdentTypeRepr *> MemberComponents);
 
   static DeclRefTypeRepr *create(const ASTContext &Ctx,
                                  ArrayRef<IdentTypeRepr *> Components);
 
+  /// Returns the qualifier or base type representation. For example, `A.B`
+  /// for `A.B.C`.
+  TypeRepr *getBase() const;
+
   /// Returns the root qualifier. For example, `A` for `A.B.C`.
   TypeRepr *getRoot() const { return Root; }
 

lib/AST/NameLookup.cpp

@@ -2745,24 +2745,21 @@ resolveTypeDeclsToNominal(Evaluator &evaluator,
 
       // Recognize Swift.AnyObject directly.
       if (typealias->getName().is("AnyObject")) {
-        // TypeRepr version: Builtin.AnyObject
-        if (auto typeRepr = typealias->getUnderlyingTypeRepr()) {
-          if (auto memberTR = dyn_cast<MemberTypeRepr>(typeRepr)) {
-            auto memberComps = memberTR->getMemberComponents();
-            if (memberComps.size() == 1 &&
-                memberComps.front()->getNameRef().isSimpleName("AnyObject") &&
-                memberTR->getRoot()->isSimpleUnqualifiedIdentifier("Builtin")) {
-              anyObject = true;
-            }
-          }
-        }
-
         // Type version: an empty class-bound existential.
         if (typealias->hasInterfaceType()) {
           if (auto type = typealias->getUnderlyingType())
             if (type->isAnyObject())
               anyObject = true;
         }
+        // TypeRepr version: Builtin.AnyObject
+        else if (auto *memberTR = dyn_cast_or_null<MemberTypeRepr>(
+                     typealias->getUnderlyingTypeRepr())) {
+          if (!memberTR->hasGenericArgList() &&
+              memberTR->getNameRef().isSimpleName("AnyObject") &&
+              memberTR->getBase()->isSimpleUnqualifiedIdentifier("Builtin")) {
+            anyObject = true;
+          }
+        }
       }
 
       continue;
@@ -2922,50 +2919,25 @@ static DirectlyReferencedTypeDecls
 directReferencesForDeclRefTypeRepr(Evaluator &evaluator, ASTContext &ctx,
                                    DeclRefTypeRepr *repr, DeclContext *dc,
                                    bool allowUsableFromInline) {
-  DirectlyReferencedTypeDecls current;
-
-  auto *rootComp = repr->getRoot();
-  if (auto *identBase = dyn_cast<IdentTypeRepr>(rootComp)) {
-    // If we already set a declaration, use it.
-    if (auto *typeDecl = identBase->getBoundDecl()) {
-      current = {1, typeDecl};
-    } else {
-      // For the base component, perform unqualified name lookup.
-      current = directReferencesForUnqualifiedTypeLookup(
-          identBase->getNameRef(), identBase->getLoc(), dc,
-          LookupOuterResults::Excluded, allowUsableFromInline);
-    }
-  } else {
-    current = directReferencesForTypeRepr(evaluator, ctx, rootComp, dc,
-                                          allowUsableFromInline);
+  // If we already set a declaration, use it.
+  if (auto *typeDecl = repr->getBoundDecl()) {
+    return {typeDecl};
   }
 
-  auto *memberTR = dyn_cast<MemberTypeRepr>(repr);
-  if (!memberTR)
-    return current;
-
-  // If we didn't find anything, fail now.
-  if (current.empty())
-    return current;
-
-  for (const auto &component : memberTR->getMemberComponents()) {
-    // If we already set a declaration, use it.
-    if (auto typeDecl = component->getBoundDecl()) {
-      current = {1, typeDecl};
-      continue;
-    }
+  if (auto *memberTR = dyn_cast<MemberTypeRepr>(repr)) {
+    DirectlyReferencedTypeDecls baseTypeDecls = directReferencesForTypeRepr(
+        evaluator, ctx, memberTR->getBase(), dc, allowUsableFromInline);
 
-    // For subsequent components, perform qualified name lookup.
-    current =
-        directReferencesForQualifiedTypeLookup(evaluator, ctx, current,
-                                               component->getNameRef(), dc,
-                                               component->getLoc(),
-                                               allowUsableFromInline);
-    if (current.empty())
-      return current;
+    // For a qualified identifier, perform qualified name lookup.
+    return directReferencesForQualifiedTypeLookup(
+        evaluator, ctx, baseTypeDecls, repr->getNameRef(), dc, repr->getLoc(),
+        allowUsableFromInline);
   }
 
-  return current;
+  // For an unqualified identifier, perform unqualified name lookup.
+  return directReferencesForUnqualifiedTypeLookup(
+      repr->getNameRef(), repr->getLoc(), dc, LookupOuterResults::Excluded,
+      allowUsableFromInline);
 }
 
 static DirectlyReferencedTypeDecls
@@ -3654,11 +3626,12 @@ CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
                                assocType->getDescriptiveKind(),
                                assocType->getName());
 
-            auto *baseComp = new (ctx) SimpleIdentTypeRepr(
+            auto *baseTR = new (ctx) SimpleIdentTypeRepr(
                 identTypeRepr->getNameLoc(), DeclNameRef(moduleName));
 
             auto *newTE = new (ctx) TypeExpr(
-                MemberTypeRepr::create(ctx, baseComp, {identTypeRepr}));
+                MemberTypeRepr::create(ctx, baseTR, identTypeRepr->getNameLoc(),
+                                       identTypeRepr->getNameRef()));
             attr->resetTypeInformation(newTE);
             return nominal;
           }

lib/AST/TypeRepr.cpp

@@ -209,6 +209,11 @@ TypeDecl *DeclRefTypeRepr::getBoundDecl() const {
       ->getBoundDecl();
 }
 
+bool DeclRefTypeRepr::hasGenericArgList() const {
+  return isa<GenericIdentTypeRepr>(
+      const_cast<DeclRefTypeRepr *>(this)->getLastComponent());
+}
+
 ArrayRef<TypeRepr *> DeclRefTypeRepr::getGenericArgs() const {
   auto *lastComp = const_cast<DeclRefTypeRepr *>(this)->getLastComponent();
   if (auto *genericTR = dyn_cast<GenericIdentTypeRepr>(lastComp)) {
@@ -454,14 +459,29 @@ GenericIdentTypeRepr *GenericIdentTypeRepr::create(const ASTContext &C,
   return new (mem) GenericIdentTypeRepr(Loc, Id, GenericArgs, AngleBrackets);
 }
 
+MemberTypeRepr *MemberTypeRepr::create(const ASTContext &C, TypeRepr *Base,
+                                       DeclNameLoc NameLoc, DeclNameRef Name) {
+  assert(Base);
+  return cast<MemberTypeRepr>(DeclRefTypeRepr::create(C, Base, NameLoc, Name));
+}
+
+MemberTypeRepr *MemberTypeRepr::create(const ASTContext &C, TypeRepr *Base,
+                                       DeclNameLoc NameLoc, DeclNameRef Name,
+                                       ArrayRef<TypeRepr *> GenericArgs,
+                                       SourceRange AngleBrackets) {
+  assert(Base);
+  return cast<MemberTypeRepr>(DeclRefTypeRepr::create(
+      C, Base, NameLoc, Name, GenericArgs, AngleBrackets));
+}
+
 TypeRepr *MemberTypeRepr::create(const ASTContext &C, TypeRepr *Root,
                                  ArrayRef<IdentTypeRepr *> MemberComponents) {
   if (MemberComponents.empty())
     return Root;
 
   auto size = totalSizeToAlloc<IdentTypeRepr *>(MemberComponents.size());
   auto mem = C.Allocate(size, alignof(MemberTypeRepr));
-  return new (mem) MemberTypeRepr(Root, MemberComponents);
+  return new (mem) MemberTypeRepr(Root, MemberComponents, C);
 }
 
 DeclRefTypeRepr *MemberTypeRepr::create(const ASTContext &Ctx,
@@ -470,6 +490,11 @@ DeclRefTypeRepr *MemberTypeRepr::create(const ASTContext &Ctx,
       create(Ctx, Components.front(), Components.drop_front()));
 }
 
+TypeRepr *MemberTypeRepr::getBase() const {
+  return MemberTypeRepr::create(*C, getRoot(),
+                                getMemberComponents().drop_back());
+}
+
 PackTypeRepr::PackTypeRepr(SourceLoc keywordLoc, SourceRange braceLocs,
                            ArrayRef<TypeRepr*> elements)
   : TypeRepr(TypeReprKind::Pack),