Merge pull request #63598 from hborla/global-peer-macros

hborla · web-flow · commit a18eebf47c9f · 2023-02-14T20:15:45.000-08:00
[Macros] Enable global peer macros.
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -867,6 +867,15 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   /// attribute macro expansion.
   DeclAttributes getSemanticAttrs() const;
 
+  using AuxiliaryDeclCallback = llvm::function_ref<void(Decl *)>;
+
+  /// Iterate over the auxiliary declarations for this declaration,
+  /// invoking the given callback with each auxiliary decl.
+  ///
+  /// Auxiliary declarations can be property wrapper backing variables,
+  /// backing variables for 'lazy' vars, or peer macro expansions.
+  void visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const;
+
   using MacroCallback = llvm::function_ref<void(CustomAttr *, MacroDecl *)>;
 
   /// Iterate over each attached macro with the given role, invoking the
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -441,7 +441,7 @@ SWIFT_REQUEST(TypeChecker, ExpandSynthesizedMemberMacroRequest,
               ArrayRef<unsigned>(Decl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandPeerMacroRequest,
-              bool(Decl *),
+              ArrayRef<unsigned>(Decl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SynthesizeRuntimeMetadataAttrGenerator,
               Expr *(CustomAttr *, ValueDecl *),
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -376,6 +376,27 @@ DeclAttributes Decl::getSemanticAttrs() const {
   return getAttrs();
 }
 
+void Decl::visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const {
+  auto &ctx = getASTContext();
+  auto *mutableThis = const_cast<Decl *>(this);
+  auto peerBuffers =
+      evaluateOrDefault(ctx.evaluator,
+                        ExpandPeerMacroRequest{mutableThis},
+                        {});
+
+  SourceManager &sourceMgr = ctx.SourceMgr;
+  auto *moduleDecl = getModuleContext();
+  for (auto bufferID : peerBuffers) {
+    auto startLoc = sourceMgr.getLocForBufferStart(bufferID);
+    auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
+    for (auto *peer : sourceFile->getTopLevelDecls()) {
+      callback(peer);
+    }
+  }
+
+  // FIXME: fold VarDecl::visitAuxiliaryDecls into this.
+}
+
 void Decl::forEachAttachedMacro(MacroRole role,
                                 MacroCallback macroCallback) const {
   auto *dc = getDeclContext();
diff --git a/lib/ASTGen/Sources/ASTGen/Macros.swift b/lib/ASTGen/Sources/ASTGen/Macros.swift
@@ -293,6 +293,8 @@ private func findSyntaxNodeInSourceFile<Node: SyntaxProtocol>(
 func expandAttachedMacro(
   diagEnginePtr: UnsafeMutablePointer<UInt8>,
   macroPtr: UnsafeRawPointer,
+  discriminatorText: UnsafePointer<UInt8>,
+  discriminatorTextLength: Int,
   rawMacroRole: UInt8,
   customAttrSourceFilePtr: UnsafeRawPointer,
   customAttrSourceLocPointer: UnsafePointer<UInt8>?,
@@ -344,7 +346,13 @@ func expandAttachedMacro(
   sourceManager.insert(declarationSourceFilePtr)
 
   // Create an expansion context
-  let context = sourceManager.createMacroExpansionContext()
+  let discriminatorBuffer = UnsafeBufferPointer(
+    start: discriminatorText, count: discriminatorTextLength
+  )
+  let discriminator = String(decoding: discriminatorBuffer, as: UTF8.self)
+  let context = sourceManager.createMacroExpansionContext(
+    discriminator: discriminator
+  )
 
   let macroName = customAttrNode.attributeName.description
   var evaluatedSyntaxStr: String
diff --git a/lib/Parse/ParseRequests.cpp b/lib/Parse/ParseRequests.cpp
@@ -200,8 +200,12 @@ SourceFileParsingResult ParseSourceFileRequest::evaluate(Evaluator &evaluator,
     }
 
     case GeneratedSourceInfo::PeerMacroExpansion: {
-      // FIXME: this does not work for top-level or local peer expansions.
-      parser.parseExpandedMemberList(items);
+      if (parser.CurDeclContext->isTypeContext()) {
+        parser.parseExpandedMemberList(items);
+      } else {
+        parser.parseTopLevelItems(items);
+      }
+
       break;
     }
     }
diff --git a/lib/SILGen/SILGen.cpp b/lib/SILGen/SILGen.cpp
@@ -2188,6 +2188,13 @@ class SILGenModuleRAII {
 
     SourceFileScope scope(SGM, sf);
     for (auto *D : sf->getTopLevelDecls()) {
+      // Emit auxiliary decls.
+      D->visitAuxiliaryDecls([&](Decl *auxiliaryDecl) {
+        FrontendStatsTracer StatsTracer(SGM.getASTContext().Stats,
+                                        "SILgen-decl", auxiliaryDecl);
+        SGM.visit(auxiliaryDecl);
+      });
+
       FrontendStatsTracer StatsTracer(SGM.getASTContext().Stats,
                                       "SILgen-decl", D);
       SGM.visit(D);
diff --git a/lib/Sema/LookupVisibleDecls.cpp b/lib/Sema/LookupVisibleDecls.cpp
@@ -616,7 +616,7 @@ static void synthesizeMemberDeclsForLookup(NominalTypeDecl *NTD,
     (void)evaluateOrDefault(
         ctx.evaluator,
         ExpandPeerMacroRequest{member},
-        false);
+        {});
   }
 
   synthesizePropertyWrapperVariables(NTD);
diff --git a/lib/Sema/TypeCheckDecl.cpp b/lib/Sema/TypeCheckDecl.cpp
@@ -2806,7 +2806,7 @@ static ArrayRef<Decl *> evaluateMembersRequest(
     (void)evaluateOrDefault(
         ctx.evaluator,
         ExpandPeerMacroRequest{member},
-        false);
+        {});
 
     if (auto *var = dyn_cast<VarDecl>(member)) {
       // The projected storage wrapper ($foo) might have
diff --git a/lib/Sema/TypeCheckDeclPrimary.cpp b/lib/Sema/TypeCheckDeclPrimary.cpp
@@ -1852,6 +1852,11 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
   }
 
   void visit(Decl *decl) {
+    // Visit auxiliary decls first.
+    decl->visitAuxiliaryDecls([&](Decl *auxiliaryDecl) {
+      this->visit(auxiliaryDecl);
+    });
+
     if (auto *Stats = getASTContext().Stats)
       ++Stats->getFrontendCounters().NumDeclsTypechecked;
 
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
@@ -58,6 +58,8 @@ extern "C" ptrdiff_t swift_ASTGen_expandFreestandingMacro(
 
 extern "C" ptrdiff_t swift_ASTGen_expandAttachedMacro(
     void *diagEngine, void *macro,
+    const char *discriminator,
+    ptrdiff_t discriminatorLength,
     uint32_t rawMacroRole,
     void *customAttrSourceFile,
     const void *customAttrSourceLocation,
@@ -867,6 +869,7 @@ evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo, CustomAttr *attr,
   // Evaluate the macro.
   NullTerminatedStringRef evaluatedSource;
 
+  std::string discriminator;
   auto macroDef = macro->getDefinition();
   switch (macroDef.kind) {
   case MacroDefinition::Kind::Undefined:
@@ -932,11 +935,18 @@ evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo, CustomAttr *attr,
     if (auto var = dyn_cast<VarDecl>(attachedTo))
       searchDecl = var->getParentPatternBinding();
 
+    {
+      Mangle::ASTMangler mangler;
+      discriminator =
+        mangler.mangleAttachedMacroExpansion(attachedTo, attr, role);
+    }
+
     const char *evaluatedSourceAddress;
     ptrdiff_t evaluatedSourceLength;
     swift_ASTGen_expandAttachedMacro(
         &ctx.Diags,
         externalDef.opaqueHandle,
+        discriminator.data(), discriminator.size(),
         static_cast<uint32_t>(role),
         astGenAttrSourceFile, attr->AtLoc.getOpaquePointerValue(),
         astGenDeclSourceFile, searchDecl->getStartLoc().getOpaquePointerValue(),
@@ -955,14 +965,7 @@ evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo, CustomAttr *attr,
   }
 
   // Figure out a reasonable name for the macro expansion buffer.
-  std::string discriminator;
-  std::string bufferName;
-  {
-    Mangle::ASTMangler mangler;
-    discriminator =
-      mangler.mangleAttachedMacroExpansion(attachedTo, attr, role);
-    bufferName = adjustMacroExpansionBufferName(discriminator);
-  }
+  std::string bufferName = adjustMacroExpansionBufferName(discriminator);
 
   // Dump macro expansions to standard output, if requested.
   if (ctx.LangOpts.DumpMacroExpansions) {
@@ -1194,9 +1197,6 @@ swift::expandPeers(CustomAttr *attr, MacroDecl *macro, Decl *decl) {
       nominal->addMember(peer);
     } else if (auto *extension = dyn_cast<ExtensionDecl>(parent)) {
       extension->addMember(peer);
-    } else {
-      // TODO: Add peers to global or local contexts.
-      continue;
     }
   }
 
diff --git a/lib/Sema/TypeCheckStorage.cpp b/lib/Sema/TypeCheckStorage.cpp
@@ -121,7 +121,7 @@ static void computeLoweredStoredProperties(NominalTypeDecl *decl,
     (void)evaluateOrDefault(
         ctx.evaluator,
         ExpandPeerMacroRequest{member},
-        false);
+        {});
 
     auto *var = dyn_cast<VarDecl>(member);
     if (!var || var->isStatic())
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/test/Macros/Inputs/syntax_macro_definitions.swift
@@ -637,7 +637,7 @@ public struct AddCompletionHandler: PeerMacro {
       FunctionParameterSyntax(
         firstName: .identifier("completionHandler"),
         colon: .colonToken(trailingTrivia: .space),
-        type: "(\(resultType ?? "")) -> Void" as TypeSyntax
+        type: "@escaping (\(resultType ?? "")) -> Void" as TypeSyntax
       )
 
     // Add the completion handler parameter to the parameter list.
@@ -732,6 +732,84 @@ public struct AddCompletionHandler: PeerMacro {
   }
 }
 
+public struct WrapInType: PeerMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard let funcDecl = declaration.as(FunctionDeclSyntax.self) else {
+      throw CustomError.message("@wrapInType only applies to functions")
+    }
+
+    // Build a new function with the same signature that forwards arguments
+    // to the the original function.
+    let parameterList = funcDecl.signature.input.parameterList
+    let callArguments: [String] = try parameterList.map { param in
+      guard let argName = param.secondName ?? param.firstName else {
+        throw CustomError.message("@wrapInType argument must have a name")
+      }
+
+      if let paramName = param.firstName, paramName.text != "_" {
+        return "\(paramName.text): \(argName.text)"
+      }
+
+      return "\(argName.text)"
+    }
+
+    let call: ExprSyntax =
+      """
+      \(funcDecl.identifier)(\(raw: callArguments.joined(separator: ", ")))
+      """
+
+    // Drop the peer macro attribute from the new declaration.
+    let newAttributeList = AttributeListSyntax(
+      funcDecl.attributes?.filter {
+        guard case let .attribute(attribute) = $0,
+              let attributeType = attribute.attributeName.as(SimpleTypeIdentifierSyntax.self),
+              let nodeType = node.attributeName.as(SimpleTypeIdentifierSyntax.self)
+        else {
+          return true
+        }
+
+        return attributeType.name.text != nodeType.name.text
+      } ?? []
+    )
+
+    let method =
+      funcDecl
+      .with(
+        \.identifier,
+         "\(context.createUniqueName(funcDecl.identifier.text))"
+      )
+      .with(
+        \.signature,
+        funcDecl.signature
+      )
+      .with(
+        \.body,
+        CodeBlockSyntax(
+          leftBrace: .leftBraceToken(leadingTrivia: .space),
+          statements: CodeBlockItemListSyntax(
+            [CodeBlockItemSyntax(item: .expr(call))]
+          )
+          .with(\.leadingTrivia, [.newlines(1), .spaces(2)]),
+          rightBrace: .rightBraceToken(leadingTrivia: .newline)
+        )
+      )
+      .with(\.attributes, newAttributeList)
+
+    let structType: DeclSyntax =
+      """
+      struct \(context.createUniqueName(funcDecl.identifier.text)) {
+        \(method)
+      }
+      """
+
+    return [structType]
+  }
+}
+
 private extension DeclSyntaxProtocol {
   var isObservableStoredProperty: Bool {
     if let property = self.as(VariableDeclSyntax.self),
diff --git a/test/Macros/macro_expand_peers.swift b/test/Macros/macro_expand_peers.swift
@@ -18,9 +18,24 @@ struct S {
   }
 
   // CHECK-DUMP: @__swiftmacro_18macro_expand_peers1SV1f1a3for_SSSi_SSSdtYaF20addCompletionHandlerfMp_.swift
-  // CHECK-DUMP: func f(a: Int, for b: String, _ value: Double, completionHandler: (String) -> Void) {
+  // 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: }
 }
+
+@attached(peer)
+macro wrapInType() = #externalMacro(module: "MacroDefinition", type: "WrapInType")
+
+@wrapInType
+func global(a: Int, b: String) {
+  print(a, b)
+}
+
+// CHECK-DUMP: @__swiftmacro_18macro_expand_peers6global1a1bySi_SStF10wrapInTypefMp_.swift
+// CHECK-DUMP: struct $s18macro_expand_peers6global1a1bySi_SStF10wrapInTypefMp_6globalfMu0_ {
+// CHECK-DUMP:   func $s18macro_expand_peers6global1a1bySi_SStF10wrapInTypefMp_6globalfMu_(a: Int, b: String)  {
+// CHECK-DUMP:     global(a: a, b: b)
+// CHECK-DUMP:   }
+// CHECK-DUMP: }
diff --git a/test/SourceKit/Macros/macro_basic.swift b/test/SourceKit/Macros/macro_basic.swift
@@ -139,7 +139,7 @@ struct S3 {
 // 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 "
-// PEER_EXPAND: func f(a: Int, for b: String, _ value: Double, completionHandler: (String) -> Void) {
+// PEER_EXPAND: func f(a: Int, for b: String, _ value: Double, completionHandler: @escaping (String) -> Void) {
 // PEER_EXPAND:  Task {
 // PEER_EXPAND:    completionHandler(await f(a: a, for: b, value))
 // PEER_EXPAND:  }

Original file line number	Diff line number	Diff line change
`@@ -200,8 +200,12 @@ SourceFileParsingResult ParseSourceFileRequest::evaluate(Evaluator &evaluator,`
`200`	`200`	`}`
`201`	`201`
`202`	`202`	`case GeneratedSourceInfo::PeerMacroExpansion: {`
`203`		`- // FIXME: this does not work for top-level or local peer expansions.`
`204`		`- parser.parseExpandedMemberList(items);`
	`203`	`+ if (parser.CurDeclContext->isTypeContext()) {`
	`204`	`+ parser.parseExpandedMemberList(items);`
	`205`	`+ } else {`
	`206`	`+ parser.parseTopLevelItems(items);`
	`207`	`+ }`
	`208`	`+`
`205`	`209`	`break;`
`206`	`210`	`}`
`207`	`211`	`}`
Original file line number	Diff line number	Diff line change
`@@ -616,7 +616,7 @@ static void synthesizeMemberDeclsForLookup(NominalTypeDecl *NTD,`
`616`	`616`	`(void)evaluateOrDefault(`
`617`	`617`	`ctx.evaluator,`
`618`	`618`	`ExpandPeerMacroRequest{member},`
`619`		`- false);`
	`619`	`+ {});`
`620`	`620`	`}`
`621`	`621`
`622`	`622`	`synthesizePropertyWrapperVariables(NTD);`