Add SourceKit support for expanding peer macros

Author: DougGregor

include/swift/AST/TypeCheckRequests.h

@@ -3911,15 +3911,15 @@ class CompilerPluginLoadRequest
 /// Expand peer macros attached to the given declaration.
 class ExpandPeerMacroRequest
     : public SimpleRequest<ExpandPeerMacroRequest,
-                           bool(Decl *),
+                           ArrayRef<unsigned>(Decl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  bool evaluate(Evaluator &evaluator, Decl *decl) const;
+  ArrayRef<unsigned> evaluate(Evaluator &evaluator, Decl *decl) const;
 
 public:
   bool isCached() const { return true; }

lib/Refactoring/Refactoring.cpp

@@ -8514,6 +8514,12 @@ getMacroExpansionBuffers(MacroDecl *macro, const CustomAttr *attr, Decl *decl) {
     allBufferIDs.append(bufferIDs.begin(), bufferIDs.end());
   }
 
+  if (roles.contains(MacroRole::Peer)) {
+    auto bufferIDs = evaluateOrDefault(
+        ctx.evaluator, ExpandPeerMacroRequest{decl}, { });
+    allBufferIDs.append(bufferIDs.begin(), bufferIDs.end());
+  }
+
   // Drop any buffers that come from other macros. We could eliminate this
   // step by adding more fine-grained requests above, which only expand for a
   // single custom attribute.

lib/Sema/TypeCheckMacros.cpp

@@ -391,7 +391,7 @@ static std::string adjustMacroExpansionBufferName(StringRef name) {
 ArrayRef<unsigned> ExpandMemberAttributeMacros::evaluate(Evaluator &evaluator,
                                                          Decl *decl) const {
   if (decl->isImplicit())
-    return { }};
+    return { };
 
   auto *parentDecl = decl->getDeclContext()->getAsDecl();
   if (!parentDecl)
@@ -423,14 +423,16 @@ ArrayRef<unsigned> ExpandSynthesizedMemberMacroRequest::evaluate(
   return decl->getASTContext().AllocateCopy(bufferIDs);
 }
 
-bool ExpandPeerMacroRequest::evaluate(Evaluator &evaluator, Decl *decl) const {
-  bool addedPeers = false;
+ArrayRef<unsigned>
+ExpandPeerMacroRequest::evaluate(Evaluator &evaluator, Decl *decl) const {
+  SmallVector<unsigned, 2> bufferIDs;
   decl->forEachAttachedMacro(MacroRole::Peer,
       [&](CustomAttr *attr, MacroDecl *macro) {
-        addedPeers |= expandPeers(attr, macro, decl);
+        if (auto bufferID = expandPeers(attr, macro, decl))
+          bufferIDs.push_back(*bufferID);
       });
 
-  return addedPeers;
+  return decl->getASTContext().AllocateCopy(bufferIDs);
 }
 
 /// Determine whether the given source file is from an expansion of the given
@@ -1034,7 +1036,7 @@ evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo, CustomAttr *attr,
   case MacroRole::Peer: {
     generatedSourceKind = GeneratedSourceInfo::PeerMacroExpansion;
     SourceLoc afterDeclLoc =
-        Lexer::getLocForEndOfToken(sourceMgr, decl->getEndLoc());
+        Lexer::getLocForEndOfToken(sourceMgr, attachedTo->getEndLoc());
     generatedOriginalSourceRange = CharSourceRange(afterDeclLoc, 0);
     break;
   }
@@ -1175,16 +1177,16 @@ swift::expandMembers(CustomAttr *attr, MacroDecl *macro, Decl *decl) {
   return macroSourceFile->getBufferID();
 }
 
-bool swift::expandPeers(CustomAttr *attr, MacroDecl *macro, Decl *decl) {
+Optional<unsigned>
+swift::expandPeers(CustomAttr *attr, MacroDecl *macro, Decl *decl) {
   auto macroSourceFile = evaluateAttachedMacro(macro, decl, attr,
                                                /*passParentContext*/false,
                                                MacroRole::Peer);
   if (!macroSourceFile)
-    return false;
+    return None;
 
   PrettyStackTraceDecl debugStack("applying expanded peer macro", decl);
 
-  bool addedPeers = false;
   auto *parent = decl->getDeclContext();
   auto topLevelDecls = macroSourceFile->getTopLevelDecls();
   for (auto peer : topLevelDecls) {
@@ -1196,11 +1198,9 @@ bool swift::expandPeers(CustomAttr *attr, MacroDecl *macro, Decl *decl) {
       // TODO: Add peers to global or local contexts.
       continue;
     }
-
-    addedPeers = true;
   }
 
-  return addedPeers;
+  return macroSourceFile->getBufferID();
 }
 
 MacroDecl *

lib/Sema/TypeCheckMacros.h

@@ -70,8 +70,8 @@ expandMembers(CustomAttr *attr, MacroDecl *macro, Decl *decl);
 /// Expand the peer declarations for the given declaration based on
 /// the custom attribute that references the given macro.
 ///
-/// Returns \c true if the macro added new peers, \c false otherwise.
-bool expandPeers(CustomAttr *attr, MacroDecl *macro, Decl *decl);
+/// If expansion occurred, returns the macro expansion buffer ID.
+Optional<unsigned> expandPeers(CustomAttr *attr, MacroDecl *macro, Decl *decl);
 
 } // end namespace swift
 

test/SourceKit/Macros/macro_basic.swift

@@ -34,6 +34,17 @@ struct S2 {
   var y: Int = 17
 }
 
+@attached(peer)
+macro addCompletionHandler() = #externalMacro(module: "MacroDefinition", type: "AddCompletionHandler")
+
+@available(macOS 10.15, *)
+struct S3 {
+  @addCompletionHandler
+  func f(a: Int, for b: String, _ value: Double) async -> String {
+    return b
+  }
+}
+
 // FIXME: Swift parser is not enabled on Linux CI yet.
 // REQUIRES: OS=macosx
 
@@ -113,7 +124,7 @@ struct S2 {
 // ACCESSOR1_EXPAND: source.edit.kind.active:
 // ACCESSOR1_EXPAND:   30:3-30:20 ""
 
-//##-- Refactoring expanding the first accessor macro
+//##-- Refactoring expanding the second accessor macro
 // RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=33:13 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=ACCESSOR2_EXPAND %s
 // ACCESSOR2_EXPAND: source.edit.kind.active:
 // ACCESSOR2_EXPAND:   34:14-34:18 "{
@@ -122,3 +133,14 @@ struct S2 {
 // ACCESSOR2_EXPAND: }"
 // ACCESSOR2_EXPAND: source.edit.kind.active:
 // ACCESSOR2_EXPAND:   33:3-33:20 ""
+
+//##-- Refactoring expanding the second accessor macro
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=42:5 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=PEER_EXPAND %s
+// PEER_EXPAND: source.edit.kind.active:
+// PEER_EXPAND:   45:4-45:4 "func f(a: Int, for b: String, _ value: Double, completionHandler: (String) -> Void) {
+// PEER_EXPAND:  Task {
+// PEER_EXPAND:    completionHandler(await f(a: a, for: b, value))
+// PEER_EXPAND:  }
+// PEER_EXPAND: }"
+// PEER_EXPAND: source.edit.kind.active:
+// PEER_EXPAND:   42:3-42:24 ""