Merge pull request #64065 from hborla/global-peer-macro-lookup

[Macros] Implement unqualified lookup for global macro-expanded peer declarations.

Author: hborla

include/swift/AST/Decl.h

@@ -8306,10 +8306,20 @@ class PostfixOperatorDecl : public OperatorDecl {
   }
 };
 
-/// Represents a missing declaration in the source code. This
-/// is used for parser recovery, e.g. when parsing a floating
-/// attribute list.
-class MissingDecl: public Decl {
+/// Represents a missing declaration in the source code.
+///
+/// This is used for parser recovery, e.g. when parsing a floating
+/// attribute list, and to represent placeholders for unexpanded
+/// declarations generated by macros.
+class MissingDecl : public Decl {
+  /// If this missing decl represents an unexpanded peer generated
+  /// by a macro, \c unexpandedPeer contains the macro custom attribute
+  /// and the declaration the macro is attached to.
+  struct {
+    CustomAttr *macroAttr = nullptr;
+    ValueDecl *attachedTo = nullptr;
+  } unexpandedPeer;
+
   /// The location that the decl would be if it wasn't missing.
   SourceLoc Loc;
 
@@ -8328,6 +8338,21 @@ class MissingDecl: public Decl {
 
   SourceRange getSourceRange() const { return SourceRange(Loc); }
 
+  static MissingDecl *
+  forUnexpandedPeer(CustomAttr *macro, ValueDecl *attachedTo) {
+    auto &ctx = attachedTo->getASTContext();
+    auto *dc = attachedTo->getDeclContext();
+    auto *missing = new (ctx) MissingDecl(dc, SourceLoc());
+
+    missing->unexpandedPeer.macroAttr = macro;
+    missing->unexpandedPeer.attachedTo = attachedTo;
+
+    return missing;
+  }
+
+  using ExpandedPeerCallback = llvm::function_ref<void(ValueDecl *)>;
+  void forEachExpandedPeer(ExpandedPeerCallback callback);
+
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::Missing;
   }
@@ -8433,6 +8458,11 @@ class MacroDecl : public GenericContext, public ValueDecl {
   /// Retrieve the attribute that declared the given macro role.
   const MacroRoleAttr *getMacroRoleAttr(MacroRole role) const;
 
+  /// Populate the \c names vector with the decl names introduced
+  /// by a given role of this macro.
+  void getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
+                          SmallVectorImpl<DeclName> &names) const;
+
   /// Retrieve the definition of this macro.
   MacroDefinition getDefinition() const;
 

lib/AST/Decl.cpp

@@ -9974,6 +9974,28 @@ MacroRoles swift::getAttachedMacroRoles() {
   return attachedMacroRoles;
 }
 
+void
+MissingDecl::forEachExpandedPeer(ExpandedPeerCallback callback) {
+  auto *macro = unexpandedPeer.macroAttr;
+  auto *attachedTo = unexpandedPeer.attachedTo;
+  if (!macro || !attachedTo)
+    return;
+
+  attachedTo->visitAuxiliaryDecls(
+      [&](Decl *auxiliaryDecl) {
+        auto *sf = auxiliaryDecl->getInnermostDeclContext()->getParentSourceFile();
+        auto *macroAttr = sf->getAttachedMacroAttribute();
+        if (macroAttr != unexpandedPeer.macroAttr)
+          return;
+
+        auto *value = dyn_cast<ValueDecl>(auxiliaryDecl);
+        if (!value)
+          return;
+
+        callback(value);
+      });
+}
+
 MacroDecl::MacroDecl(
     SourceLoc macroLoc, DeclName name, SourceLoc nameLoc,
     GenericParamList *genericParams,
@@ -10027,7 +10049,71 @@ const MacroRoleAttr *MacroDecl::getMacroRoleAttr(MacroRole role) const {
   for (auto attr : getAttrs().getAttributes<MacroRoleAttr>())
     if (attr->getMacroRole() == role)
       return attr;
-  llvm_unreachable("Macro role not declared for this MacroDecl");
+
+  return nullptr;
+}
+
+void MacroDecl::getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
+                                   SmallVectorImpl<DeclName> &names) const {
+  ASTContext &ctx = getASTContext();
+  auto *attr = getMacroRoleAttr(role);
+  if (!attr)
+    return;
+
+  for (auto expandedName : attr->getNames()) {
+    switch (expandedName.getKind()) {
+    case MacroIntroducedDeclNameKind::Named: {
+      names.push_back(DeclName(expandedName.getIdentifier()));
+      break;
+    }
+
+    case MacroIntroducedDeclNameKind::Overloaded: {
+      if (!attachedTo)
+        break;
+
+      names.push_back(attachedTo->getBaseName());
+      break;
+    }
+
+    case MacroIntroducedDeclNameKind::Prefixed: {
+      if (!attachedTo)
+        break;
+
+      auto baseName = attachedTo->getBaseName();
+      std::string prefixedName;
+      {
+        llvm::raw_string_ostream out(prefixedName);
+        out << expandedName.getIdentifier();
+        out << baseName.getIdentifier();
+      }
+
+      Identifier nameId = ctx.getIdentifier(prefixedName);
+      names.push_back(DeclName(nameId));
+      break;
+    }
+
+    case MacroIntroducedDeclNameKind::Suffixed: {
+      if (!attachedTo)
+        break;
+
+      auto baseName = attachedTo->getBaseName();
+      std::string suffixedName;
+      {
+        llvm::raw_string_ostream out(suffixedName);
+        out << baseName.getIdentifier();
+        out << expandedName.getIdentifier();
+      }
+
+      Identifier nameId = ctx.getIdentifier(suffixedName);
+      names.push_back(DeclName(nameId));
+      break;
+    }
+
+    case MacroIntroducedDeclNameKind::Arbitrary:
+      // FIXME: Indicate that the macro covers arbitrary names.
+      break;
+    }
+  }
 }
 
 MacroDefinition MacroDecl::getDefinition() const {

lib/AST/Module.cpp

@@ -170,6 +170,15 @@ class swift::SourceLookupCache {
   void addToUnqualifiedLookupCache(Range decls, bool onlyOperators);
   template<typename Range>
   void addToMemberCache(Range decls);
+
+  using MacroDeclMap = llvm::DenseMap<Identifier, TinyPtrVector<MacroDecl *>>;
+  MacroDeclMap MacroDecls;
+
+  using AuxiliaryDeclMap = llvm::DenseMap<DeclName, TinyPtrVector<MissingDecl *>>;
+  AuxiliaryDeclMap TopLevelAuxiliaryDecls;
+  SmallVector<ValueDecl *, 4> MayHaveAuxiliaryDecls;
+  void populateAuxiliaryDeclCache();
+
 public:
   SourceLookupCache(const SourceFile &SF);
   SourceLookupCache(const ModuleDecl &Mod);
@@ -234,6 +243,10 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
       if (onlyOperators ? VD->isOperator() : VD->hasName()) {
         // Cache the value under both its compound name and its full name.
         TopLevelValues.add(VD);
+
+        if (VD->getAttrs().hasAttribute<CustomAttr>()) {
+          MayHaveAuxiliaryDecls.push_back(VD);
+        }
       }
     }
 
@@ -256,6 +269,9 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
 
     if (auto *PG = dyn_cast<PrecedenceGroupDecl>(D))
       PrecedenceGroups[PG->getName()].push_back(PG);
+
+    if (auto *macro = dyn_cast<MacroDecl>(D))
+      MacroDecls[macro->getBaseIdentifier()].push_back(macro);
   }
 }
 
@@ -309,6 +325,53 @@ void SourceLookupCache::addToMemberCache(Range decls) {
   }
 }
 
+void SourceLookupCache::populateAuxiliaryDeclCache() {
+  for (auto *decl : MayHaveAuxiliaryDecls) {
+    // Gather macro-introduced peer names.
+    llvm::SmallDenseMap<CustomAttr *, llvm::SmallVector<DeclName, 2>> introducedNames;
+
+    // This code deliberately avoids `forEachAttachedMacro`, because it
+    // will perform overload resolution and possibly invoke unqualified
+    // lookup for macro arguments, which will recursively populate the
+    // auxiliary decl cache and cause request cycles.
+    //
+    // We do not need a fully resolved macro until expansion. Instead, we
+    // conservatively consider peer names for all macro declarations with a
+    // custom attribute name. Unqualified lookup for that name will later
+    // invoke expansion of the macro, and will yield no results if the resolved
+    // macro does not produce the requested name, so the only impact is possibly
+    // expanding earlier than needed / unnecessarily looking in the top-level
+    // auxiliary decl cache.
+    for (auto attrConst : decl->getSemanticAttrs().getAttributes<CustomAttr>()) {
+      auto *attr = const_cast<CustomAttr *>(attrConst);
+      UnresolvedMacroReference macroRef(attr);
+      auto macroName = macroRef.getMacroName().getBaseIdentifier();
+
+      auto found = MacroDecls.find(macroName);
+      if (found == MacroDecls.end())
+        continue;
+
+      for (const auto *macro : found->second) {
+        macro->getIntroducedNames(MacroRole::Peer,
+                                  decl, introducedNames[attr]);
+      }
+    }
+
+    // Add macro-introduced names to the top-level auxiliary decl cache as
+    // unexpanded peer decls represented by a MissingDecl.
+    for (auto macroNames : introducedNames) {
+      auto *macroAttr = macroNames.getFirst();
+      for (auto name : macroNames.getSecond()) {
+        auto *placeholder =
+            MissingDecl::forUnexpandedPeer(macroAttr, decl);
+        name.addToLookupTable(TopLevelAuxiliaryDecls, placeholder);
+      }
+    }
+  }
+
+  MayHaveAuxiliaryDecls.clear();
+}
+
 /// Populate our cache on the first name lookup.
 SourceLookupCache::SourceLookupCache(const SourceFile &SF) {
   FrontendStatsTracer tracer(SF.getASTContext().Stats,
@@ -339,6 +402,23 @@ void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,
   Result.reserve(I->second.size());
   for (ValueDecl *Elt : I->second)
     Result.push_back(Elt);
+
+  // Add top-level auxiliary decls to the result.
+  //
+  // FIXME: We need to not consider auxiliary decls if we're doing lookup
+  // from inside a macro argument at module scope.
+  populateAuxiliaryDeclCache();
+  auto auxDecls = TopLevelAuxiliaryDecls.find(Name);
+  if (auxDecls == TopLevelAuxiliaryDecls.end())
+    return;
+
+  for (auto *unexpandedDecl : auxDecls->second) {
+    // Add expanded peers to the result.
+    unexpandedDecl->forEachExpandedPeer(
+        [&](ValueDecl *expandedPeer) {
+          Result.push_back(expandedPeer);
+        });
+  }
 }
 
 void SourceLookupCache::getPrecedenceGroups(

test/Macros/macro_expand_peers.swift

@@ -8,7 +8,7 @@
 // FIXME: Swift parser is not enabled on Linux CI yet.
 // REQUIRES: OS=macosx
 
-@attached(peer)
+@attached(peer, names: overloaded)
 macro addCompletionHandler() = #externalMacro(module: "MacroDefinition", type: "AddCompletionHandler")
 
 struct S {
@@ -23,8 +23,33 @@ struct S {
   // CHECK-DUMP:     completionHandler(await f(a: a, for: b, value))
   // CHECK-DUMP:   }
   // CHECK-DUMP: }
+
+  func useOverload() {
+    f(a: 1, for: "", 2.0) {
+      print($0)
+    }
+  }
+}
+
+// CHECK-DUMP: @__swiftmacro_18macro_expand_peers1f1a3for_SSSi_SSSdtYaF20addCompletionHandlerfMp_.swift
+// CHECK-DUMP: func f(a: Int, for b: String, _ value: Double, completionHandler: @escaping (String) -> Void) {
+// CHECK-DUMP:   Task {
+// CHECK-DUMP:     completionHandler(await f(a: a, for: b, value))
+// CHECK-DUMP:   }
+// CHECK-DUMP: }
+
+@addCompletionHandler
+func f(a: Int, for b: String, _ value: Double) async -> String {
+  return b
+}
+
+func useOverload() {
+  f(a: 1, for: "", 2.0) {
+    print($0)
+  }
 }
 
+
 @attached(peer)
 macro wrapInType() = #externalMacro(module: "MacroDefinition", type: "WrapInType")