Share module extraction logic

Author: ApolloZhu

include/swift/AST/Attr.h

@@ -64,6 +64,7 @@ class ModuleDecl;
 class PatternBindingInitializer;
 class TrailingWhereClause;
 class TypeExpr;
+class IdentTypeRepr;
 
 class alignas(1 << AttrAlignInBits) AttributeBase
     : public ASTAllocated<AttributeBase> {
@@ -1769,6 +1770,15 @@ class CustomAttr final : public DeclAttribute {
   TypeRepr *getTypeRepr() const;
   Type getType() const;
 
+  /// Destructure an attribute's type repr for a macro reference.
+  ///
+  /// For a 1-level member type repr whose base and member are both identifier
+  /// types, e.g. `Foo.Bar`, return a pair of the base and the member.
+  ///
+  /// For an identifier type repr, return a pair of `nullptr` and the
+  /// identifier.
+  std::pair<IdentTypeRepr *, IdentTypeRepr *> destructureMacroRef();
+
   /// Whether the attribute has any arguments.
   bool hasArgs() const { return argList != nullptr; }
 

lib/AST/Attr.cpp

@@ -2587,6 +2587,21 @@ CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
       CustomAttr(atLoc, range, type, initContext, argList, implicit);
 }
 
+std::pair<IdentTypeRepr *, IdentTypeRepr *> 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->getBaseComponent()))
+      if (memType->getMemberComponents().size() == 1)
+        if (auto first =
+                dyn_cast<IdentTypeRepr>(memType->getMemberComponents().front()))
+          return {base, first};
+  return {nullptr, nullptr};
+}
+
 TypeRepr *CustomAttr::getTypeRepr() const { return typeExpr->getTypeRepr(); }
 Type CustomAttr::getType() const { return typeExpr->getInstanceType(); }
 

lib/AST/NameLookup.cpp

@@ -3587,32 +3587,13 @@ CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
   // Look for names at module scope, so we don't trigger name lookup for
   // nested scopes. At this point, we're looking to see whether there are
   // any suitable macros.
-
-  // Handle a module-qualified name.
-  if (auto *memTypeRepr =
-          dyn_cast_or_null<MemberTypeRepr>(attr->getTypeRepr())) {
-    auto baseTypeRepr = memTypeRepr->getBaseComponent();
-    auto memberReprs = memTypeRepr->getMemberComponents();
-    auto *moduleNameRepr = dyn_cast<IdentTypeRepr>(baseTypeRepr);
-
-    if (moduleNameRepr && (memberReprs.size() == 1)) {
-      if (auto *macroNameRepr = dyn_cast<IdentTypeRepr>(memberReprs.front())) {
-        auto moduleName = moduleNameRepr->getNameRef();
-        auto macroName = macroNameRepr->getNameRef();
-        auto macros = namelookup::lookupMacros(dc, moduleName, macroName,
-                                               getAttachedMacroRoles());
-        if (!macros.empty())
-          return nullptr;
-      }
-    }
-  } else if (auto *identTypeRepr =
-                 dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr())) {
-    auto macros =
-        namelookup::lookupMacros(dc, DeclNameRef(), identTypeRepr->getNameRef(),
-                                 getAttachedMacroRoles());
-    if (!macros.empty())
-      return nullptr;
-  }
+  auto [base, member] = attr->destructureMacroRef();
+  auto moduleName = (base) ? base->getNameRef() : DeclNameRef();
+  auto macroName = (member) ? member->getNameRef() : DeclNameRef();
+  auto macros = namelookup::lookupMacros(dc, moduleName, macroName,
+                                         getAttachedMacroRoles());
+  if (!macros.empty())
+    return nullptr;
 
   // Find the types referenced by the custom attribute.
   auto &ctx = dc->getASTContext();

lib/AST/TypeCheckRequests.cpp

@@ -1792,32 +1792,11 @@ void swift::simple_display(
 // ResolveMacroRequest computation.
 //----------------------------------------------------------------------------//
 
-/// Destructure a type repr for a macro reference.
-///
-/// For a 1-level member type repr whose base and member are both identifier
-/// types, e.g. `Foo.Bar`, return a pair of the base and the member.
-///
-/// For an identifier type repr, return a pair of `nullptr` and the identifier.
-static std::pair<IdentTypeRepr *, IdentTypeRepr *>
-destructureMacroRefTypeRepr(TypeRepr *typeRepr) {
-  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->getBaseComponent()))
-      if (memType->getMemberComponents().size() == 1)
-        if (auto first =
-                dyn_cast<IdentTypeRepr>(memType->getMemberComponents().front()))
-          return {base, first};
-  return {nullptr, nullptr};
-}
-
 DeclNameRef UnresolvedMacroReference::getMacroName() const {
   if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
     return expansion->getMacroName();
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto [_, member] = destructureMacroRefTypeRepr(attr->getTypeRepr());
+    auto [_, member] = attr->destructureMacroRef();
     if (!member)
       return DeclNameRef();
     return member->getNameRef();
@@ -1837,7 +1816,7 @@ DeclNameRef UnresolvedMacroReference::getModuleName() const {
   if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
     return expansion->getModuleName();
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto [base, _] = destructureMacroRefTypeRepr(attr->getTypeRepr());
+    auto [base, _] = attr->destructureMacroRef();
     if (!base)
       return DeclNameRef();
     return base->getNameRef();
@@ -1849,7 +1828,7 @@ DeclNameLoc UnresolvedMacroReference::getModuleNameLoc() const {
   if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
     return expansion->getModuleNameLoc();
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto [base, _] = destructureMacroRefTypeRepr(attr->getTypeRepr());
+    auto [base, _] = attr->destructureMacroRef();
     if (!base)
       return DeclNameLoc();
     return base->getNameLoc();
@@ -1861,7 +1840,7 @@ DeclNameLoc UnresolvedMacroReference::getMacroNameLoc() const {
   if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
     return expansion->getMacroNameLoc();
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto [_, member] = destructureMacroRefTypeRepr(attr->getTypeRepr());
+    auto [_, member] = attr->destructureMacroRef();
     if (!member)
       return DeclNameLoc();
     return member->getNameLoc();
@@ -1874,7 +1853,7 @@ SourceRange UnresolvedMacroReference::getGenericArgsRange() const {
     return expansion->getGenericArgsRange();
 
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto [_, member] = destructureMacroRefTypeRepr(attr->getTypeRepr());
+    auto [_, member] = attr->destructureMacroRef();
     if (!member)
       return SourceRange();
     auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(member);
@@ -1892,7 +1871,7 @@ ArrayRef<TypeRepr *> UnresolvedMacroReference::getGenericArgs() const {
     return expansion->getGenericArgs();
 
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto [_, member] = destructureMacroRefTypeRepr(attr->getTypeRepr());
+    auto [_, member] = attr->destructureMacroRef();
     if (!member)
       return {};
     auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(member);