Merge pull request #66935 from rintaro/5.9-macros-plugins-collappsed-attached-result

rintaro · web-flow · commit b13d8f6bc6bd · 2023-06-27T09:05:05.000-07:00
[5.9][Macros] Attached macro expansions return single string
diff --git a/include/swift/AST/CASTBridging.h b/include/swift/AST/CASTBridging.h
@@ -308,7 +308,7 @@ typedef const void *PluginCapabilityPtr;
 
 /// Set a capability data to the plugin object. Since the data is just a opaque
 /// pointer, it's not used in AST at all.
-void Plugin_setCapability(PluginHandle handle, PluginCapabilityPtr data);
+void Plugin_setCapability(PluginHandle handle, PluginCapabilityPtr _Nullable data);
 
 /// Get a capability data set by \c Plugin_setCapability .
 PluginCapabilityPtr _Nullable Plugin_getCapability(PluginHandle handle);
diff --git a/lib/ASTGen/Sources/ASTGen/Macros.swift b/lib/ASTGen/Sources/ASTGen/Macros.swift
@@ -407,6 +407,29 @@ func checkMacroDefinition(
   }
 }
 
+// Make an expansion result for '@_cdecl' function caller.
+func makeExpansionOutputResult(
+  expandedSource: String?,
+  outputPointer: UnsafeMutablePointer<UnsafePointer<UInt8>?>,
+  outputLength: UnsafeMutablePointer<Int>
+) -> Int {
+  guard var expandedSource = expandedSource else {
+    return -1
+  }
+
+  // Form the result buffer for our caller.
+  expandedSource.withUTF8 { utf8 in
+    let evaluatedResultPtr = UnsafeMutablePointer<UInt8>.allocate(capacity: utf8.count + 1)
+    if let baseAddress = utf8.baseAddress {
+      evaluatedResultPtr.initialize(from: baseAddress, count: utf8.count)
+    }
+    evaluatedResultPtr[utf8.count] = 0
+
+    outputPointer.pointee = UnsafePointer(evaluatedResultPtr)
+    outputLength.pointee = utf8.count
+  }
+  return 0
+}
 
 @_cdecl("swift_ASTGen_expandFreestandingMacro")
 @usableFromInline
@@ -470,23 +493,11 @@ func expandFreestandingMacro(
       discriminator: discriminator)
   }
 
-  guard var expandedSource = expandedSource else {
-    return -1
-  }
-
-  // Form the result buffer for our caller.
-  expandedSource.withUTF8 { utf8 in
-    let evaluatedResultPtr = UnsafeMutablePointer<UInt8>.allocate(capacity: utf8.count + 1)
-    if let baseAddress = utf8.baseAddress {
-      evaluatedResultPtr.initialize(from: baseAddress, count: utf8.count)
-    }
-    evaluatedResultPtr[utf8.count] = 0
-
-    expandedSourcePointer.pointee = UnsafePointer(evaluatedResultPtr)
-    expandedSourceLength.pointee = utf8.count
-  }
-
-  return 0
+  return makeExpansionOutputResult(
+    expandedSource: expandedSource,
+    outputPointer: expandedSourcePointer,
+    outputLength: expandedSourceLength
+  )
 }
 
 func expandFreestandingMacroIPC(
@@ -516,7 +527,7 @@ func expandFreestandingMacroIPC(
     preconditionFailure("unhandled macro role for freestanding macro")
 
   case .expression: pluginMacroRole = .expression
-  case .declaration: pluginMacroRole = .freeStandingDeclaration
+  case .declaration: pluginMacroRole = .declaration
   case .codeItem: pluginMacroRole = .codeItem
   }
 
@@ -528,9 +539,15 @@ func expandFreestandingMacroIPC(
     syntax: PluginMessage.Syntax(syntax: Syntax(expansionSyntax), in: sourceFilePtr)!)
   do {
     let result = try macro.plugin.sendMessageAndWait(message)
-    guard
-      case .expandFreestandingMacroResult(let expandedSource, let diagnostics) = result
-    else {
+    let expandedSource: String?
+    let diagnostics: [PluginMessage.Diagnostic]
+    switch result {
+    case
+        .expandMacroResult(let _expandedSource, let _diagnostics),
+        .expandFreestandingMacroResult(let _expandedSource, let _diagnostics):
+      expandedSource = _expandedSource
+      diagnostics = _diagnostics
+    default:
       throw PluginError.invalidReponseKind
     }
 
@@ -707,10 +724,10 @@ func expandAttachedMacro(
   )
   let discriminator = String(decoding: discriminatorBuffer, as: UTF8.self)
 
-  let expandedSources: [String]?
+  let expandedSource: String?
   switch MacroPluginKind(rawValue: macroKind)! {
   case .Executable:
-    expandedSources = expandAttachedMacroIPC(
+    expandedSource = expandAttachedMacroIPC(
       diagEnginePtr: diagEnginePtr,
       macroPtr: macroPtr,
       rawMacroRole: rawMacroRole,
@@ -722,7 +739,7 @@ func expandAttachedMacro(
       parentDeclSourceFilePtr: parentDeclSourceFilePtr,
       parentDeclNode: parentDeclNode)
   case .InProcess:
-    expandedSources = expandAttachedMacroInProcess(
+    expandedSource = expandAttachedMacroInProcess(
       diagEnginePtr: diagEnginePtr,
       macroPtr: macroPtr,
       rawMacroRole: rawMacroRole,
@@ -735,26 +752,11 @@ func expandAttachedMacro(
       parentDeclNode: parentDeclNode)
   }
 
-  guard let expandedSources = expandedSources else {
-    return -1
-  }
-
-  // Fixup the source.
-  var expandedSource: String = collapse(expansions: expandedSources, for: MacroRole(rawMacroRole: rawMacroRole), attachedTo: declarationNode)
-
-  // Form the result buffer for our caller.
-  expandedSource.withUTF8 { utf8 in
-    let evaluatedResultPtr = UnsafeMutablePointer<UInt8>.allocate(capacity: utf8.count + 1)
-    if let baseAddress = utf8.baseAddress {
-      evaluatedResultPtr.initialize(from: baseAddress, count: utf8.count)
-    }
-    evaluatedResultPtr[utf8.count] = 0
-
-    expandedSourcePointer.pointee = UnsafePointer(evaluatedResultPtr)
-    expandedSourceLength.pointee = utf8.count
-  }
-
-  return 0
+  return makeExpansionOutputResult(
+    expandedSource: expandedSource,
+    outputPointer: expandedSourcePointer,
+    outputLength: expandedSourceLength
+  )
 }
 
 func expandAttachedMacroIPC(
@@ -768,7 +770,7 @@ func expandAttachedMacroIPC(
   attachedTo declarationNode: DeclSyntax,
   parentDeclSourceFilePtr: UnsafePointer<ExportedSourceFile>?,
   parentDeclNode: DeclSyntax?
-) -> [String]? {
+) -> String? {
   let macroName: String = customAttrNode.attributeName.description
   let macro = macroPtr.assumingMemoryBound(to: ExportedExecutableMacro.self).pointee
 
@@ -810,10 +812,27 @@ func expandAttachedMacroIPC(
     declSyntax: declSyntax,
     parentDeclSyntax: parentDeclSyntax)
   do {
-    let result = try macro.plugin.sendMessageAndWait(message)
-    guard
-      case .expandAttachedMacroResult(let expandedSources, let diagnostics) = result
-    else {
+    let expandedSource: String?
+    let diagnostics: [PluginMessage.Diagnostic]
+    switch try macro.plugin.sendMessageAndWait(message) {
+    case .expandMacroResult(let _expandedSource, let _diagnostics):
+      expandedSource = _expandedSource
+      diagnostics = _diagnostics
+
+    // Handle legacy result message.
+    case .expandAttachedMacroResult(let _expandedSources, let _diagnostics):
+      if let _expandedSources = _expandedSources {
+        expandedSource = SwiftSyntaxMacroExpansion.collapse(
+          expansions: _expandedSources,
+          for: MacroRole(rawMacroRole: rawMacroRole),
+          attachedTo: declarationNode
+        )
+      } else {
+        expandedSource = nil
+      }
+      diagnostics = _diagnostics
+      break
+    default:
       throw PluginError.invalidReponseKind
     }
 
@@ -827,7 +846,7 @@ func expandAttachedMacroIPC(
       }
       diagEngine.emit(diagnostics, messageSuffix: " (from macro '\(macroName)')")
     }
-    return expandedSources
+    return expandedSource
 
   } catch let error {
     let srcMgr = SourceManager(cxxDiagnosticEngine: diagEnginePtr)
@@ -861,7 +880,7 @@ func expandAttachedMacroInProcess(
   attachedTo declarationNode: DeclSyntax,
   parentDeclSourceFilePtr: UnsafePointer<ExportedSourceFile>?,
   parentDeclNode: DeclSyntax?
-) -> [String]? {
+) -> String? {
   // Get the macro.
   let macroPtr = macroPtr.bindMemory(to: ExportedMacro.self, capacity: 1)
   let macro = macroPtr.pointee.macro
@@ -925,54 +944,3 @@ fileprivate extension SyntaxProtocol {
     return formatted.trimmedDescription(matching: { $0.isWhitespace })
   }
 }
-
-fileprivate func collapse<Node: SyntaxProtocol>(
-  expansions: [String],
-  for role: MacroRole,
-  attachedTo declarationNode: Node
-) -> String {
-  if expansions.isEmpty {
-    return ""
-  }
-
-  var expansions = expansions
-  var separator: String = "\n\n"
-
-  if role == .accessor,
-     let varDecl = declarationNode.as(VariableDeclSyntax.self),
-     let binding = varDecl.bindings.first,
-     binding.accessor == nil {
-    let indentation = String(repeating: " ", count: 4)
-
-    expansions = expansions.map({ indent($0, with: indentation) })
-    expansions[0] = "{\n" + expansions[0]
-    expansions[expansions.count - 1] += "\n}"
-  } else if role == .memberAttribute {
-    separator = " "
-  }
-
-  return expansions.joined(separator: separator)
-}
-
-fileprivate func indent(_ source: String, with indentation: String) -> String {
-  if source.isEmpty || indentation.isEmpty {
-    return source
-  }
-
-  var indented = ""
-  var remaining = source[...]
-  while let nextNewline = remaining.firstIndex(where: { $0.isNewline }) {
-    if nextNewline != remaining.startIndex {
-      indented += indentation
-    }
-    indented += remaining[...nextNewline]
-    remaining = remaining[remaining.index(after: nextNewline)...]
-  }
-
-  if !remaining.isEmpty {
-    indented += indentation
-    indented += remaining
-  }
-
-  return indented
-}
diff --git a/lib/ASTGen/Sources/ASTGen/PluginHost.swift b/lib/ASTGen/Sources/ASTGen/PluginHost.swift
@@ -154,24 +154,38 @@ struct CompilerPlugin {
 
   /// Initialize the plugin. This should be called inside lock.
   func initialize() throws {
-    // Get capability.
-    let response = try self.sendMessageAndWaitWithoutLock(.getCapability)
+    // Send host capability and get plugin capability.
+    let hostCapability = PluginMessage.HostCapability(
+      protocolVersion: PluginMessage.PROTOCOL_VERSION_NUMBER
+    )
+    let request = HostToPluginMessage.getCapability(capability: hostCapability)
+    let response = try self.sendMessageAndWaitWithoutLock(request)
     guard case .getCapabilityResult(let capability) = response else {
       throw PluginError.invalidReponseKind
     }
+
+    deinitializePluginCapabilityIfExist()
+
     let ptr = UnsafeMutablePointer<Capability>.allocate(capacity: 1)
     ptr.initialize(to: .init(capability))
     Plugin_setCapability(opaqueHandle, UnsafeRawPointer(ptr))
   }
 
+  /// Deinitialize and unset the plugin capability stored in C++
+  /// 'LoadedExecutablePlugin'. This should be called inside lock.
+  func deinitializePluginCapabilityIfExist() {
+    if let ptr = Plugin_getCapability(opaqueHandle) {
+      let capabilityPtr = UnsafeMutableRawPointer(mutating: ptr)
+        .assumingMemoryBound(to: PluginMessage.PluginCapability.self)
+      capabilityPtr.deinitialize(count: 1)
+      capabilityPtr.deallocate()
+      Plugin_setCapability(opaqueHandle, nil)
+    }
+  }
+
   func deinitialize() {
     self.withLock {
-      if let ptr = Plugin_getCapability(opaqueHandle) {
-        let capabilityPtr = UnsafeMutableRawPointer(mutating: ptr)
-          .assumingMemoryBound(to: PluginMessage.PluginCapability.self)
-        capabilityPtr.deinitialize(count: 1)
-        capabilityPtr.deallocate()
-      }
+      deinitializePluginCapabilityIfExist()
     }
   }
 
diff --git a/lib/ASTGen/Sources/ASTGen/PluginMessages.swift b/lib/ASTGen/Sources/ASTGen/PluginMessages.swift
@@ -14,8 +14,10 @@
 // NOTE: Types in this file should be self-contained and should not depend on any non-stdlib types.
 
 internal enum HostToPluginMessage: Codable {
-  /// Get capability of this plugin.
-  case getCapability
+  /// Send capability of the host, and get capability of the plugin.
+  case getCapability(
+    capability: PluginMessage.HostCapability?
+  )
 
   /// Expand a '@freestanding' macro.
   case expandFreestandingMacro(
@@ -49,11 +51,19 @@ internal enum PluginToHostMessage: Codable {
     capability: PluginMessage.PluginCapability
   )
 
+  /// Unified response for freestanding/attached macro expansion.
+  case expandMacroResult(
+    expandedSource: String?,
+    diagnostics: [PluginMessage.Diagnostic]
+  )
+
+  // @available(*, deprecated: "use expandMacroResult() instead")
   case expandFreestandingMacroResult(
     expandedSource: String?,
     diagnostics: [PluginMessage.Diagnostic]
   )
 
+  // @available(*, deprecated: "use expandMacroResult() instead")
   case expandAttachedMacroResult(
     expandedSources: [String]?,
     diagnostics: [PluginMessage.Diagnostic]
@@ -66,7 +76,11 @@ internal enum PluginToHostMessage: Codable {
 }
 
 /*namespace*/ internal enum PluginMessage {
-  static var PROTOCOL_VERSION_NUMBER: Int { 4 }  // Added 'loadPluginLibrary'.
+  static var PROTOCOL_VERSION_NUMBER: Int { 5 }  // Added 'expandMacroResult'.
+
+  struct HostCapability: Codable {
+    var protocolVersion: Int
+  }
 
   struct PluginCapability: Codable {
     var protocolVersion: Int
@@ -86,7 +100,7 @@ internal enum PluginToHostMessage: Codable {
 
   enum MacroRole: String, Codable {
     case expression
-    case freeStandingDeclaration
+    case declaration
     case accessor
     case memberAttribute
     case member
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
@@ -292,6 +292,15 @@ initializeExecutablePlugin(ASTContext &ctx,
     if (!swift_ASTGen_initializePlugin(executablePlugin, &ctx.Diags)) {
       return nullptr;
     }
+
+    // Resend the compiler capability on reconnect.
+    auto *callback = new std::function<void(void)>(
+        [executablePlugin]() {
+          (void)swift_ASTGen_initializePlugin(
+              executablePlugin, /*diags=*/nullptr);
+        });
+    executablePlugin->addOnReconnect(callback);
+
     executablePlugin->setCleanup([executablePlugin] {
       swift_ASTGen_deinitializePlugin(executablePlugin);
     });
diff --git a/test/Macros/macro_plugin_basic.swift b/test/Macros/macro_plugin_basic.swift
@@ -21,7 +21,7 @@
 
 // RUN: %FileCheck -strict-whitespace %s < %t/macro-expansions.txt
 
-// CHECK: ->(plugin:[[#PID:]]) {"getCapability":{}}
+// CHECK: ->(plugin:[[#PID:]]) {"getCapability":{"capability":{"protocolVersion":[[#PROTOCOL_VERSION:]]}}}
 // CHECK: <-(plugin:[[#PID]]) {"getCapabilityResult":{"capability":{"protocolVersion":1}}}
 // CHECK: ->(plugin:[[#PID]]) {"expandFreestandingMacro":{"discriminator":"$s{{.+}}","macro":{"moduleName":"TestPlugin","name":"testString","typeName":"TestStringMacro"},"macroRole":"expression","syntax":{"kind":"expression","location":{"column":19,"fileID":"MyApp/test.swift","fileName":"BUILD_DIR{{.+}}test.swift","line":5,"offset":301},"source":"#testString(123)"}}}
 // CHECK: <-(plugin:[[#PID]]) {"expandFreestandingMacroResult":{"diagnostics":[],"expandedSource":"\"123\"\n  +   \"foo  \""}}
diff --git a/test/Macros/macro_plugin_error.swift b/test/Macros/macro_plugin_error.swift
diff --git a/test/Macros/macro_plugin_server.swift b/test/Macros/macro_plugin_server.swift
