Skip to content

Commit 5767e8f

Browse files
ApolloZhuApolloZhu
authored
[Macros] Support module-qualified attached macro lookup (#69457)
Allow attached macro expansion syntax to have a module qualifier, `@Foo.Bar`. rdar: //108621205
1 parent 310fb8c commit 5767e8f

16 files changed

+239
-80
lines changed

‎include/swift/AST/Attr.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ class ModuleDecl;
6464
class PatternBindingInitializer;
6565
class TrailingWhereClause;
6666
class TypeExpr;
67+
class IdentTypeRepr;
6768

6869
class alignas(1 << AttrAlignInBits) AttributeBase
6970
: public ASTAllocated<AttributeBase> {
@@ -1769,6 +1770,15 @@ class CustomAttr final : public DeclAttribute {
17691770
TypeRepr *getTypeRepr() const;
17701771
Type getType() const;
17711772

1773+
/// Destructure an attribute's type repr for a macro reference.
1774+
///
1775+
/// For a 1-level member type repr whose base and member are both identifier
1776+
/// types, e.g. `Foo.Bar`, return a pair of the base and the member.
1777+
///
1778+
/// For an identifier type repr, return a pair of `nullptr` and the
1779+
/// identifier.
1780+
std::pair<IdentTypeRepr *, IdentTypeRepr *> destructureMacroRef();
1781+
17721782
/// Whether the attribute has any arguments.
17731783
bool hasArgs() const { return argList != nullptr; }
17741784

‎include/swift/AST/FreestandingMacroExpansion.h‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ class ArgumentList;
3434
/// declaration/expression nodes.
3535
struct MacroExpansionInfo : ASTAllocated<MacroExpansionInfo> {
3636
SourceLoc SigilLoc;
37+
DeclNameRef ModuleName;
38+
DeclNameLoc ModuleNameLoc;
3739
DeclNameRef MacroName;
3840
DeclNameLoc MacroNameLoc;
3941
SourceLoc LeftAngleLoc, RightAngleLoc;
@@ -43,13 +45,16 @@ struct MacroExpansionInfo : ASTAllocated<MacroExpansionInfo> {
4345
/// The referenced macro.
4446
ConcreteDeclRef macroRef;
4547

46-
MacroExpansionInfo(SourceLoc sigilLoc, DeclNameRef macroName,
48+
MacroExpansionInfo(SourceLoc sigilLoc, DeclNameRef moduleName,
49+
DeclNameLoc moduleNameLoc, DeclNameRef macroName,
4750
DeclNameLoc macroNameLoc, SourceLoc leftAngleLoc,
4851
SourceLoc rightAngleLoc, ArrayRef<TypeRepr *> genericArgs,
4952
ArgumentList *argList)
50-
: SigilLoc(sigilLoc), MacroName(macroName), MacroNameLoc(macroNameLoc),
51-
LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
52-
GenericArgs(genericArgs), ArgList(argList) {}
53+
: SigilLoc(sigilLoc), ModuleName(moduleName),
54+
ModuleNameLoc(moduleNameLoc), MacroName(macroName),
55+
MacroNameLoc(macroNameLoc), LeftAngleLoc(leftAngleLoc),
56+
RightAngleLoc(rightAngleLoc), GenericArgs(genericArgs),
57+
ArgList(argList) {}
5358

5459
SourceLoc getLoc() const { return SigilLoc; }
5560
SourceRange getGenericArgsRange() const {
@@ -85,6 +90,10 @@ class FreestandingMacroExpansion {
8590

8691
SourceLoc getPoundLoc() const { return getExpansionInfo()->SigilLoc; }
8792

93+
DeclNameLoc getModuleNameLoc() const {
94+
return getExpansionInfo()->ModuleNameLoc;
95+
}
96+
DeclNameRef getModuleName() const { return getExpansionInfo()->ModuleName; }
8897
DeclNameLoc getMacroNameLoc() const {
8998
return getExpansionInfo()->MacroNameLoc;
9099
}

‎include/swift/AST/NameLookup.h‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,8 @@ enum class UnqualifiedLookupFlags {
245245
ExcludeMacroExpansions = 1 << 6,
246246
/// This lookup should only return macros.
247247
MacroLookup = 1 << 7,
248+
/// This lookup should only return modules
249+
ModuleLookup = 1 << 8,
248250
};
249251

250252
using UnqualifiedLookupOptions = OptionSet<UnqualifiedLookupFlags>;
@@ -550,8 +552,10 @@ void filterForDiscriminator(SmallVectorImpl<Result> &results,
550552

551553
/// \returns The set of macro declarations with the given name that
552554
/// fulfill any of the given macro roles.
553-
SmallVector<MacroDecl *, 1>
554-
lookupMacros(DeclContext *dc, DeclNameRef macroName, MacroRoles roles);
555+
SmallVector<MacroDecl *, 1> lookupMacros(DeclContext *dc,
556+
DeclNameRef moduleName,
557+
DeclNameRef macroName,
558+
MacroRoles roles);
555559

556560
/// \returns Whether the given source location is inside an attached
557561
/// or freestanding macro argument.

‎include/swift/AST/TypeCheckRequests.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3458,6 +3458,8 @@ class UnresolvedMacroReference {
34583458
}
34593459

34603460
SourceLoc getSigilLoc() const;
3461+
DeclNameRef getModuleName() const;
3462+
DeclNameLoc getModuleNameLoc() const;
34613463
DeclNameRef getMacroName() const;
34623464
DeclNameLoc getMacroNameLoc() const;
34633465
SourceRange getGenericArgsRange() const;

‎lib/AST/Attr.cpp‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2587,6 +2587,21 @@ CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
25872587
CustomAttr(atLoc, range, type, initContext, argList, implicit);
25882588
}
25892589

2590+
std::pair<IdentTypeRepr *, IdentTypeRepr *> CustomAttr::destructureMacroRef() {
2591+
TypeRepr *typeRepr = getTypeRepr();
2592+
if (!typeRepr)
2593+
return {nullptr, nullptr};
2594+
if (auto *identType = dyn_cast<IdentTypeRepr>(typeRepr))
2595+
return {nullptr, identType};
2596+
if (auto *memType = dyn_cast<MemberTypeRepr>(typeRepr))
2597+
if (auto *base = dyn_cast<IdentTypeRepr>(memType->getBaseComponent()))
2598+
if (memType->getMemberComponents().size() == 1)
2599+
if (auto first =
2600+
dyn_cast<IdentTypeRepr>(memType->getMemberComponents().front()))
2601+
return {base, first};
2602+
return {nullptr, nullptr};
2603+
}
2604+
25902605
TypeRepr *CustomAttr::getTypeRepr() const { return typeExpr->getTypeRepr(); }
25912606
Type CustomAttr::getType() const { return typeExpr->getInstanceType(); }
25922607

‎lib/AST/Decl.cpp‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11534,11 +11534,16 @@ MacroExpansionDecl::create(
1153411534
ArgumentList *args
1153511535
) {
1153611536
ASTContext &ctx = dc->getASTContext();
11537-
MacroExpansionInfo *info = new (ctx) MacroExpansionInfo{
11538-
poundLoc, macro, macroLoc,
11539-
leftAngleLoc, rightAngleLoc, genericArgs,
11540-
args ? args : ArgumentList::createImplicit(ctx, {})
11541-
};
11537+
MacroExpansionInfo *info = new (ctx)
11538+
MacroExpansionInfo{poundLoc,
11539+
/*moduleName*/ DeclNameRef(),
11540+
/*moduleNameLoc*/ DeclNameLoc(),
11541+
macro,
11542+
macroLoc,
11543+
leftAngleLoc,
11544+
rightAngleLoc,
11545+
genericArgs,
11546+
args ? args : ArgumentList::createImplicit(ctx, {})};
1154211547
return new (ctx) MacroExpansionDecl(dc, info);
1154311548
}
1154411549

‎lib/AST/Expr.cpp‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2769,10 +2769,15 @@ MacroExpansionExpr *MacroExpansionExpr::create(
27692769
) {
27702770
ASTContext &ctx = dc->getASTContext();
27712771
MacroExpansionInfo *info = new (ctx) MacroExpansionInfo{
2772-
sigilLoc, macroName, macroNameLoc,
2773-
leftAngleLoc, rightAngleLoc, genericArgs,
2774-
argList ? argList : ArgumentList::createImplicit(ctx, {})
2775-
};
2772+
sigilLoc,
2773+
/*moduleName*/ DeclNameRef(),
2774+
/*moduleNameLoc*/ DeclNameLoc(),
2775+
macroName,
2776+
macroNameLoc,
2777+
leftAngleLoc,
2778+
rightAngleLoc,
2779+
genericArgs,
2780+
argList ? argList : ArgumentList::createImplicit(ctx, {})};
27762781
return new (ctx) MacroExpansionExpr(dc, info, roles, isImplicit, ty);
27772782
}
27782783

‎lib/AST/NameLookup.cpp‎

Lines changed: 56 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ void swift::simple_display(llvm::raw_ostream &out,
113113
{UnqualifiedLookupFlags::IncludeOuterResults, "IncludeOuterResults"},
114114
{UnqualifiedLookupFlags::TypeLookup, "TypeLookup"},
115115
{UnqualifiedLookupFlags::MacroLookup, "MacroLookup"},
116+
{UnqualifiedLookupFlags::ModuleLookup, "ModuleLookup"},
116117
};
117118

118119
auto flagsToPrint = llvm::make_filter_range(
@@ -1621,41 +1622,59 @@ static DeclName adjustLazyMacroExpansionNameKey(
16211622
return name;
16221623
}
16231624

1624-
SmallVector<MacroDecl *, 1>
1625-
namelookup::lookupMacros(DeclContext *dc, DeclNameRef macroName,
1626-
MacroRoles roles) {
1625+
SmallVector<MacroDecl *, 1> namelookup::lookupMacros(DeclContext *dc,
1626+
DeclNameRef moduleName,
1627+
DeclNameRef macroName,
1628+
MacroRoles roles) {
16271629
SmallVector<MacroDecl *, 1> choices;
16281630
auto moduleScopeDC = dc->getModuleScopeContext();
16291631
ASTContext &ctx = moduleScopeDC->getASTContext();
16301632

1631-
// When performing lookup for freestanding macro roles, only consider
1632-
// macro names, ignoring types.
1633-
bool onlyMacros = static_cast<bool>(roles & getFreestandingMacroRoles()) &&
1634-
!(roles - getFreestandingMacroRoles());
1635-
1636-
// Macro lookup should always exclude macro expansions; macro
1637-
// expansions cannot introduce new macro declarations. Note that
1638-
// the source location here doesn't matter.
1639-
UnqualifiedLookupDescriptor descriptor{
1640-
macroName, moduleScopeDC, SourceLoc(),
1641-
UnqualifiedLookupFlags::ExcludeMacroExpansions
1642-
};
1643-
1644-
if (onlyMacros)
1645-
descriptor.Options |= UnqualifiedLookupFlags::MacroLookup;
1646-
1647-
auto lookup = evaluateOrDefault(
1648-
ctx.evaluator, UnqualifiedLookupRequest{descriptor}, {});
1649-
1650-
for (const auto &found : lookup.allResults()) {
1651-
if (auto macro = dyn_cast<MacroDecl>(found.getValueDecl())) {
1633+
auto addChoiceIfApplicable = [&](ValueDecl *decl) {
1634+
if (auto macro = dyn_cast<MacroDecl>(decl)) {
16521635
auto candidateRoles = macro->getMacroRoles();
16531636
if ((candidateRoles && roles.contains(candidateRoles)) ||
16541637
// FIXME: `externalMacro` should have all roles.
16551638
macro->getBaseIdentifier().str() == "externalMacro") {
16561639
choices.push_back(macro);
16571640
}
16581641
}
1642+
};
1643+
1644+
// When a module is specified, it's a module-qualified lookup.
1645+
if (moduleName) {
1646+
UnqualifiedLookupDescriptor moduleLookupDesc(
1647+
moduleName, moduleScopeDC, SourceLoc(),
1648+
UnqualifiedLookupFlags::ModuleLookup);
1649+
auto moduleLookup = evaluateOrDefault(
1650+
ctx.evaluator, UnqualifiedLookupRequest{moduleLookupDesc}, {});
1651+
auto foundTypeDecl = moduleLookup.getSingleTypeResult();
1652+
auto *moduleDecl = dyn_cast_or_null<ModuleDecl>(foundTypeDecl);
1653+
if (!moduleDecl)
1654+
return {};
1655+
1656+
ModuleQualifiedLookupRequest req{moduleScopeDC, moduleDecl, macroName,
1657+
SourceLoc(),
1658+
NL_ExcludeMacroExpansions | NL_OnlyMacros};
1659+
auto lookup = evaluateOrDefault(ctx.evaluator, req, {});
1660+
for (auto *found : lookup)
1661+
addChoiceIfApplicable(found);
1662+
}
1663+
// Otherwise it's an unqualified lookup.
1664+
else {
1665+
// Macro lookup should always exclude macro expansions; macro
1666+
// expansions cannot introduce new macro declarations. Note that
1667+
// the source location here doesn't matter.
1668+
UnqualifiedLookupDescriptor descriptor{
1669+
macroName, moduleScopeDC, SourceLoc(),
1670+
UnqualifiedLookupFlags::ExcludeMacroExpansions |
1671+
UnqualifiedLookupFlags::MacroLookup};
1672+
1673+
auto lookup = evaluateOrDefault(ctx.evaluator,
1674+
UnqualifiedLookupRequest{descriptor}, {});
1675+
1676+
for (const auto &found : lookup.allResults())
1677+
addChoiceIfApplicable(found.getValueDecl());
16591678
}
16601679

16611680
return choices;
@@ -1683,8 +1702,9 @@ namelookup::isInMacroArgument(SourceFile *sourceFile, SourceLoc loc) {
16831702
inMacroArgument = true;
16841703
} else if (auto *attr = macro.getAttr()) {
16851704
auto *moduleScope = sourceFile->getModuleScopeContext();
1686-
auto results = lookupMacros(moduleScope, macro.getMacroName(),
1687-
getAttachedMacroRoles());
1705+
auto results =
1706+
lookupMacros(moduleScope, macro.getModuleName(),
1707+
macro.getMacroName(), getAttachedMacroRoles());
16881708
inMacroArgument = !results.empty();
16891709
}
16901710

@@ -1705,9 +1725,9 @@ void namelookup::forEachPotentialResolvedMacro(
17051725
) {
17061726
ASTContext &ctx = moduleScopeCtx->getASTContext();
17071727
UnqualifiedLookupDescriptor lookupDesc{
1708-
macroName, moduleScopeCtx, SourceLoc(),
1709-
UnqualifiedLookupFlags::ExcludeMacroExpansions
1710-
};
1728+
macroName, moduleScopeCtx, SourceLoc(),
1729+
UnqualifiedLookupFlags::ExcludeMacroExpansions |
1730+
UnqualifiedLookupFlags::MacroLookup};
17111731

17121732
auto lookup = evaluateOrDefault(
17131733
ctx.evaluator, UnqualifiedLookupRequest{lookupDesc}, {});
@@ -3568,13 +3588,13 @@ CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
35683588
// Look for names at module scope, so we don't trigger name lookup for
35693589
// nested scopes. At this point, we're looking to see whether there are
35703590
// any suitable macros.
3571-
if (auto *identTypeRepr =
3572-
dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr())) {
3573-
auto macros = namelookup::lookupMacros(
3574-
dc, identTypeRepr->getNameRef(), getAttachedMacroRoles());
3575-
if (!macros.empty())
3576-
return nullptr;
3577-
}
3591+
auto [module, macro] = attr->destructureMacroRef();
3592+
auto moduleName = (module) ? module->getNameRef() : DeclNameRef();
3593+
auto macroName = (macro) ? macro->getNameRef() : DeclNameRef();
3594+
auto macros = namelookup::lookupMacros(dc, moduleName, macroName,
3595+
getAttachedMacroRoles());
3596+
if (!macros.empty())
3597+
return nullptr;
35783598

35793599
// Find the types referenced by the custom attribute.
35803600
auto &ctx = dc->getASTContext();

‎lib/AST/TypeCheckRequests.cpp‎

Lines changed: 38 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1817,10 +1817,10 @@ DeclNameRef UnresolvedMacroReference::getMacroName() const {
18171817
if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
18181818
return expansion->getMacroName();
18191819
if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
1820-
auto *identTypeRepr = dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr());
1821-
if (!identTypeRepr)
1820+
auto [_, macro] = attr->destructureMacroRef();
1821+
if (!macro)
18221822
return DeclNameRef();
1823-
return identTypeRepr->getNameRef();
1823+
return macro->getNameRef();
18241824
}
18251825
llvm_unreachable("Unhandled case");
18261826
}
@@ -1833,14 +1833,38 @@ SourceLoc UnresolvedMacroReference::getSigilLoc() const {
18331833
llvm_unreachable("Unhandled case");
18341834
}
18351835

1836+
DeclNameRef UnresolvedMacroReference::getModuleName() const {
1837+
if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
1838+
return expansion->getModuleName();
1839+
if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
1840+
auto [module, _] = attr->destructureMacroRef();
1841+
if (!module)
1842+
return DeclNameRef();
1843+
return module->getNameRef();
1844+
}
1845+
llvm_unreachable("Unhandled case");
1846+
}
1847+
1848+
DeclNameLoc UnresolvedMacroReference::getModuleNameLoc() const {
1849+
if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
1850+
return expansion->getModuleNameLoc();
1851+
if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
1852+
auto [module, _] = attr->destructureMacroRef();
1853+
if (!module)
1854+
return DeclNameLoc();
1855+
return module->getNameLoc();
1856+
}
1857+
llvm_unreachable("Unhandled case");
1858+
}
1859+
18361860
DeclNameLoc UnresolvedMacroReference::getMacroNameLoc() const {
18371861
if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
18381862
return expansion->getMacroNameLoc();
18391863
if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
1840-
auto *identTypeRepr = dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr());
1841-
if (!identTypeRepr)
1864+
auto [_, macro] = attr->destructureMacroRef();
1865+
if (!macro)
18421866
return DeclNameLoc();
1843-
return identTypeRepr->getNameLoc();
1867+
return macro->getNameLoc();
18441868
}
18451869
llvm_unreachable("Unhandled case");
18461870
}
@@ -1850,8 +1874,10 @@ SourceRange UnresolvedMacroReference::getGenericArgsRange() const {
18501874
return expansion->getGenericArgsRange();
18511875

18521876
if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
1853-
auto *typeRepr = attr->getTypeRepr();
1854-
auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(typeRepr);
1877+
auto [_, macro] = attr->destructureMacroRef();
1878+
if (!macro)
1879+
return SourceRange();
1880+
auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(macro);
18551881
if (!genericTypeRepr)
18561882
return SourceRange();
18571883

@@ -1866,8 +1892,10 @@ ArrayRef<TypeRepr *> UnresolvedMacroReference::getGenericArgs() const {
18661892
return expansion->getGenericArgs();
18671893

18681894
if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
1869-
auto *typeRepr = attr->getTypeRepr();
1870-
auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(typeRepr);
1895+
auto [_, macro] = attr->destructureMacroRef();
1896+
if (!macro)
1897+
return {};
1898+
auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(macro);
18711899
if (!genericTypeRepr)
18721900
return {};
18731901

‎lib/AST/UnqualifiedLookup.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,6 +281,12 @@ void UnqualifiedLookupFactory::performUnqualifiedLookup() {
281281
"performUnqualifiedLookup",
282282
DC->getParentSourceFile());
283283

284+
if (options.contains(UnqualifiedLookupFlags::ModuleLookup)) {
285+
lookForAModuleWithTheGivenName(DC->getModuleScopeContext());
286+
recordCompletionOfAScope();
287+
return;
288+
}
289+
284290
if (Loc.isValid() && DC->getParentSourceFile()) {
285291
// Operator lookup is always global, for the time being.
286292
if (!Name.isOperator())

0 commit comments

Comments
 (0)