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

Author: hborla

include/swift/AST/Decl.h

@@ -8319,10 +8319,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;
 
@@ -8341,6 +8351,25 @@ 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);
+
+    missing->unexpandedPeer.macroAttr = macro;
+    missing->unexpandedPeer.attachedTo = attachedTo;
+
+    return missing;
+  }
+
+  using ExpandedPeerCallback = llvm::function_ref<void(ValueDecl *)>;
+  void forEachExpandedPeer(ExpandedPeerCallback callback);
+
+  SourceRange getSourceRange() const {
+    return SourceRange();
+  }
+
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::Missing;
   }
@@ -8446,6 +8475,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,68 @@ 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:
+      // Only freestanding and member macros can have arbitrary names.
+      assert(role == MacroRole::Member || role == MacroRole::Declaration);
+      break;
+    }
+  }
 }
 
 MacroDefinition MacroDecl::getDefinition() const {

lib/AST/Module.cpp

@@ -170,6 +170,12 @@ class swift::SourceLookupCache {
   void addToUnqualifiedLookupCache(Range decls, bool onlyOperators);
   template<typename Range>
   void addToMemberCache(Range decls);
+
+  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 +240,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);
+        }
       }
     }
 
@@ -309,6 +319,56 @@ void SourceLookupCache::addToMemberCache(Range decls) {
   }
 }
 
+void SourceLookupCache::populateAuxiliaryDeclCache() {
+  for (auto *decl : MayHaveAuxiliaryDecls) {
+    auto *dc = decl->getDeclContext();
+
+    // 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 moduleScopeDC = dc->getModuleScopeContext();
+      ASTContext &ctx = moduleScopeDC->getASTContext();
+      UnqualifiedLookupDescriptor descriptor(macroRef.getMacroName(), moduleScopeDC);
+      auto lookup = evaluateOrDefault(
+          ctx.evaluator, UnqualifiedLookupRequest{descriptor}, {});
+      for (const auto &found : lookup.allResults()) {
+        if (auto macro = dyn_cast<MacroDecl>(found.getValueDecl())) {
+          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,
@@ -336,9 +396,35 @@ void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,
   auto I = TopLevelValues.find(Name);
   if (I == TopLevelValues.end()) return;
 
+  bool foundMacros = false;
   Result.reserve(I->second.size());
-  for (ValueDecl *Elt : I->second)
+  for (ValueDecl *Elt : I->second) {
     Result.push_back(Elt);
+    foundMacros = isa<MacroDecl>(Elt);
+  }
+
+  // If none of the results are macro decls, look into peers.
+  //
+  // FIXME: This approach is an approximation for whether lookup is attempting
+  // to find macro declarations, and it may cause cycles for invalid macro
+  // references. We need a more principled way to determine whether we're looking
+  // for a macro and should not look into auxiliary decls during lookup.
+  //
+  // We also need to not consider auxiliary decls if we're doing lookup from
+  // inside a macro argument at module scope.
+  if (!foundMacros) {
+    populateAuxiliaryDeclCache();
+    auto I = TopLevelAuxiliaryDecls.find(Name);
+    if (I == TopLevelAuxiliaryDecls.end()) return;
+
+    for (auto *unexpandedDecl : I->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")