[Swiftify] Add _SwiftifyImportProtocol for safe overloads for protocols

The existing _SwiftifyImport macro is a peer macro, limiting it to only
emitting function wrappers in the same scope as the original function.
Protocols cannot contain function implementations, so these need to be
placed in a separate protocol extension instead. _SwiftifyImportProtocol
is an extension macro rather than a peer macro, to enable this
functionality.

Rather than operating on a single function, like _SwiftifyImport,
_SwiftifyImportProtocol takes information about multiple methods and
creates a single protocol extension with all wrappers in a one-shot
operation.

rdar://144335990

Author: hnrklssn

lib/Macros/Sources/SwiftMacros/SwiftifyImportMacro.swift

@@ -1730,6 +1730,106 @@ public struct SwiftifyImportMacro: PeerMacro {
   }
 }
 
+func parseProtocolMacroParam(
+  _ paramAST: LabeledExprSyntax,
+  methods: [String: FunctionDeclSyntax]
+) throws -> (FunctionDeclSyntax, [ExprSyntax]) {
+  let paramExpr = paramAST.expression
+  guard let enumConstructorExpr = paramExpr.as(FunctionCallExprSyntax.self) else {
+    throw DiagnosticError(
+      "expected _SwiftifyProtocolMethodInfo enum literal as argument, got '\(paramExpr)'", node: paramExpr)
+  }
+  let enumName = try parseEnumName(paramExpr)
+  if enumName != "method" {
+    throw DiagnosticError(
+      "expected 'method', got '\(enumName)'",
+      node: enumConstructorExpr)
+  }
+  let argumentList = enumConstructorExpr.arguments
+  let methodSignatureArg = try getArgumentByName(argumentList, "signature")
+  guard let methodSignatureStringLit = methodSignatureArg.as(StringLiteralExprSyntax.self) else {
+    throw DiagnosticError(
+      "expected string literal for 'signature' parameter, got \(methodSignatureArg)", node: methodSignatureArg)
+  }
+  let methodSignature = methodSignatureStringLit.representedLiteralValue!
+  guard let methodSyntax = methods[methodSignature] else {
+    var notes: [Note] = []
+    var name = methodSignature
+    if let methodSyntax = DeclSyntax("\(raw: methodSignature)").as(FunctionDeclSyntax.self) {
+      name = methodSyntax.name.trimmed.text
+    }
+    for (tmp, method) in methods where method.name.trimmed.text == name {
+      notes.append(Note(node: Syntax(method.name), message: MacroExpansionNoteMessage("did you mean '\(method.trimmed.description)'?")))
+    }
+    throw DiagnosticError(
+      "method with signature '\(methodSignature)' not found in protocol", node: methodSignatureArg, notes: notes)
+  }
+  let paramInfoArg = try getArgumentByName(argumentList, "paramInfo")
+  guard let paramInfoArgList = paramInfoArg.as(ArrayExprSyntax.self) else {
+    throw DiagnosticError("expected array literal for 'paramInfo' parameter, got \(paramInfoArg)", node: paramInfoArg)
+  }
+  return (methodSyntax, paramInfoArgList.elements.map { ExprSyntax($0.expression) })
+}
+
+/// Similar to SwiftifyImportMacro, but for providing overloads to methods in
+/// protocols using an extension, rather than in the same scope as the original.
+public struct SwiftifyImportProtocolMacro: ExtensionMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    attachedTo declaration: some DeclGroupSyntax,
+    providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
+    in context: some MacroExpansionContext
+  ) throws -> [ExtensionDeclSyntax] {
+    do {
+      guard let protocolDecl = declaration.as(ProtocolDeclSyntax.self) else {
+        throw DiagnosticError("@_SwiftifyImportProtocol only works on protocols", node: declaration)
+      }
+      let argumentList = node.arguments!.as(LabeledExprListSyntax.self)!
+      var arguments = [LabeledExprSyntax](argumentList)
+      let typeMappings = try parseTypeMappingParam(arguments.last)
+      if typeMappings != nil {
+        arguments = arguments.dropLast()
+      }
+      let spanAvailability = try parseSpanAvailabilityParam(arguments.last)
+      if spanAvailability != nil {
+        arguments = arguments.dropLast()
+      }
+
+      var methods: [String: FunctionDeclSyntax] = [:]
+      for member in protocolDecl.memberBlock.members {
+        guard let methodDecl = member.decl.as(FunctionDeclSyntax.self) else {
+          continue
+        }
+        let trimmedDecl = methodDecl.with(\.body, nil)
+                                    .with(\.attributes, [])
+                                    .trimmed
+        methods[trimmedDecl.description] = methodDecl
+      }
+      let overloads = try arguments.map {
+        let (method, args) = try parseProtocolMacroParam($0, methods: methods)
+        let function = try constructOverloadFunction(
+          forDecl: method, leadingTrivia: Trivia(), args: args,
+          spanAvailability: spanAvailability,
+          typeMappings: typeMappings)
+        return MemberBlockItemSyntax(decl: function)
+      }
+
+      return [ExtensionDeclSyntax(extensionKeyword: .identifier("extension"), extendedType: type,
+                                  memberBlock: MemberBlockSyntax(leftBrace: .leftBraceToken(),
+                                                                 members: MemberBlockItemListSyntax(overloads),
+                                                                 rightBrace: .rightBraceToken())
+      )]
+    } catch let error as DiagnosticError {
+      context.diagnose(
+        Diagnostic(
+          node: error.node, message: MacroExpansionErrorMessage(error.description),
+          notes: error.notes))
+      return []
+    }
+  }
+}
+
 // MARK: syntax utils
 extension TypeSyntaxProtocol {
   public var isSwiftCoreModule: Bool {

stdlib/public/core/SwiftifyImport.swift

@@ -66,6 +66,30 @@ public macro _SwiftifyImport(_ paramInfo: _SwiftifyInfo...,
     #externalMacro(module: "SwiftMacros", type: "SwiftifyImportMacro")
 #endif
 
+/// Allows annotating pointer parameters in a protocol method using the `@_SwiftifyImportProtocol` macro.
+///
+/// This is not marked @available, because _SwiftifyImportProtocolMethod is available for any target.
+/// Instances of _SwiftifyProtocolMethodInfo should ONLY be passed as arguments directly to
+/// _SwiftifyImportProtocolMethod, so they should not affect linkage since there are never any instances
+/// at runtime.
+public enum _SwiftifyProtocolMethodInfo {
+  case method(signature: String, paramInfo: [_SwiftifyInfo])
+}
+
+/// Like _SwiftifyImport, but since protocols cannot contain function implementations they need to
+/// be placed in a separate extension instead. Unlike _SwiftifyImport, which applies to a single
+/// function, this macro supports feeding info about multiple methods and generating safe overloads
+/// for all of them at once.
+#if hasFeature(Macros)
+  @attached(extension, names: arbitrary)
+  public macro _SwiftifyImportProtocol(
+    _ methodInfo: _SwiftifyProtocolMethodInfo...,
+    spanAvailability: String? = nil,
+    typeMappings: [String: String] = [:]
+  ) =
+    #externalMacro(module: "SwiftMacros", type: "SwiftifyImportProtocolMacro")
+#endif
+
 /// Unsafely discard any lifetime dependency on the `dependent` argument. Return
 /// a value identical to `dependent` with a lifetime dependency on the caller's
 /// borrow scope of the `source` argument.

test/Macros/SwiftifyImport/CountedBy/Protocol.swift

@@ -0,0 +1,106 @@
+// REQUIRES: swift_swift_parser
+// REQUIRES: swift_feature_Lifetimes
+
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %target-swift-frontend %t/test.swift -emit-module -plugin-path %swift-plugin-dir -enable-experimental-feature Lifetimes -verify
+// RUN: %target-swift-frontend %t/test.swift -typecheck -plugin-path %swift-plugin-dir -enable-experimental-feature Lifetimes -dump-macro-expansions 2> %t/expansions.out
+// RUN: %diff %t/expansions.out %t/expansions.expected
+
+//--- test.swift
+@_SwiftifyImportProtocol(.method(signature: "func myFunc(_ ptr: UnsafePointer<CInt>, _ len: CInt)", paramInfo: [.countedBy(pointer: .param(1), count: "len")]))
+protocol SimpleProtocol {
+  func myFunc(_ ptr: UnsafePointer<CInt>, _ len: CInt)
+}
+
+@_SwiftifyImportProtocol(.method(signature: "func foo(_ ptr: UnsafePointer<CInt>, _ len: CInt)", paramInfo: [.countedBy(pointer: .param(1), count: "len"), .nonescaping(pointer: .param(1))]),
+                         .method(signature: "func bar(_ len: CInt) -> UnsafePointer<CInt>", paramInfo: [.countedBy(pointer: .return, count: "len"), .nonescaping(pointer: .return), .lifetimeDependence(dependsOn: .self, pointer: .return, type: .borrow)]))
+protocol SpanProtocol {
+  func foo(_ ptr: UnsafePointer<CInt>, _ len: CInt)
+  func bar(_ len: CInt) -> UnsafePointer<CInt>
+}
+
+@_SwiftifyImportProtocol(.method(signature: "func foo(_ ptr: UnsafePointer<CInt>, _ len: CInt)", paramInfo: [.countedBy(pointer: .param(1), count: "len"), .nonescaping(pointer: .param(1))]),
+                         .method(signature: "func bar(_ ptr: UnsafePointer<CInt>, _ len: CInt)", paramInfo: [.countedBy(pointer: .param(1), count: "len")]))
+protocol MixedProtocol {
+  /// Some doc comment
+  func foo(_ ptr: UnsafePointer<CInt>, _ len: CInt)
+  func bar(_ ptr: UnsafePointer<CInt>, _ len: CInt)
+}
+
+@_SwiftifyImportProtocol(.method(signature: "func foo(_ ptr: UnsafePointer<CInt>, _ len1: CInt)", paramInfo: [.countedBy(pointer: .param(1), count: "len1")]),
+                         .method(signature: "func foo(bar: UnsafePointer<CInt>, _ len2: CInt)", paramInfo: [.countedBy(pointer: .param(1), count: "len2")]))
+protocol OverloadedProtocol {
+  func foo(_ ptr: UnsafePointer<CInt>, _ len1: CInt)
+  func foo(bar: UnsafePointer<CInt>, _ len2: CInt)
+  func foo()
+}
+
+//--- expansions.expected
+@__swiftmacro_4test14SimpleProtocol015_SwiftifyImportC0fMe_.swift
+------------------------------
+extension SimpleProtocol {
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload
+      func myFunc(_ ptr: UnsafeBufferPointer<CInt>) {
+        let len = CInt(exactly: ptr.count)!
+        return unsafe myFunc(ptr.baseAddress!, len)
+    }
+}
+------------------------------
+@__swiftmacro_4test12SpanProtocol015_SwiftifyImportC0fMe_.swift
+------------------------------
+extension SpanProtocol {
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload
+      func foo(_ ptr: Span<CInt>) {
+        let len = CInt(exactly: ptr.count)!
+        return unsafe ptr.withUnsafeBufferPointer { _ptrPtr in
+          return unsafe foo(_ptrPtr.baseAddress!, len)
+        }
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_lifetime(borrow self) @_disfavoredOverload
+      func bar(_ len: CInt) -> Span<CInt> {
+        return unsafe _swiftifyOverrideLifetime(Span<CInt>(_unsafeStart: unsafe bar(len), count: Int(len)), copying: ())
+    }
+}
+------------------------------
+@__swiftmacro_4test13MixedProtocol015_SwiftifyImportC0fMe_.swift
+------------------------------
+extension MixedProtocol {
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload
+      /// Some doc comment
+      func foo(_ ptr: Span<CInt>) {
+        let len = CInt(exactly: ptr.count)!
+        return unsafe ptr.withUnsafeBufferPointer { _ptrPtr in
+          return unsafe foo(_ptrPtr.baseAddress!, len)
+        }
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload
+      func bar(_ ptr: UnsafeBufferPointer<CInt>) {
+        let len = CInt(exactly: ptr.count)!
+        return unsafe bar(ptr.baseAddress!, len)
+    }
+}
+------------------------------
+@__swiftmacro_4test18OverloadedProtocol015_SwiftifyImportC0fMe_.swift
+------------------------------
+extension OverloadedProtocol {
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload
+      func foo(_ ptr: UnsafeBufferPointer<CInt>) {
+        let len1 = CInt(exactly: ptr.count)!
+        return unsafe foo(ptr.baseAddress!, len1)
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload
+      func foo(bar: UnsafeBufferPointer<CInt>) {
+        let len2 = CInt(exactly: bar.count)!
+        return unsafe foo(bar: bar.baseAddress!, len2)
+    }
+}
+------------------------------

test/abi/Inputs/macOS/arm64/stdlib/baseline

@@ -8937,6 +8937,10 @@ _$ss27_BidirectionalCollectionBoxCMu
 _$ss27_BidirectionalCollectionBoxCfD
 _$ss27_BidirectionalCollectionBoxCfd
 _$ss27_BidirectionalCollectionBoxCy7ElementQzs09_AnyIndexC0_pcig
+_$ss27_SwiftifyProtocolMethodInfoO6methodyABSS_Says01_aD0OGtcABmFWC
+_$ss27_SwiftifyProtocolMethodInfoOMa
+_$ss27_SwiftifyProtocolMethodInfoOMn
+_$ss27_SwiftifyProtocolMethodInfoON
 _$ss27_allocateUninitializedArrayySayxG_BptBwlF
 _$ss27_bridgeAnythingToObjectiveCyyXlxlF
 _$ss27_debuggerTestingCheckExpectyySS_SStF

test/abi/Inputs/macOS/arm64/stdlib/baseline-asserts

@@ -8949,6 +8949,10 @@ _$ss27_BidirectionalCollectionBoxCMu
 _$ss27_BidirectionalCollectionBoxCfD
 _$ss27_BidirectionalCollectionBoxCfd
 _$ss27_BidirectionalCollectionBoxCy7ElementQzs09_AnyIndexC0_pcig
+_$ss27_SwiftifyProtocolMethodInfoO6methodyABSS_Says01_aD0OGtcABmFWC
+_$ss27_SwiftifyProtocolMethodInfoOMa
+_$ss27_SwiftifyProtocolMethodInfoOMn
+_$ss27_SwiftifyProtocolMethodInfoON
 _$ss27_allocateUninitializedArrayySayxG_BptBwlF
 _$ss27_bridgeAnythingToObjectiveCyyXlxlF
 _$ss27_debuggerTestingCheckExpectyySS_SStF

test/abi/Inputs/macOS/x86_64/stdlib/baseline

@@ -8962,6 +8962,10 @@ _$ss27_BidirectionalCollectionBoxCMu
 _$ss27_BidirectionalCollectionBoxCfD
 _$ss27_BidirectionalCollectionBoxCfd
 _$ss27_BidirectionalCollectionBoxCy7ElementQzs09_AnyIndexC0_pcig
+_$ss27_SwiftifyProtocolMethodInfoO6methodyABSS_Says01_aD0OGtcABmFWC
+_$ss27_SwiftifyProtocolMethodInfoOMa
+_$ss27_SwiftifyProtocolMethodInfoOMn
+_$ss27_SwiftifyProtocolMethodInfoON
 _$ss27_allocateUninitializedArrayySayxG_BptBwlF
 _$ss27_bridgeAnythingToObjectiveCyyXlxlF
 _$ss27_debuggerTestingCheckExpectyySS_SStF

test/abi/Inputs/macOS/x86_64/stdlib/baseline-asserts

@@ -8974,6 +8974,10 @@ _$ss27_BidirectionalCollectionBoxCMu
 _$ss27_BidirectionalCollectionBoxCfD
 _$ss27_BidirectionalCollectionBoxCfd
 _$ss27_BidirectionalCollectionBoxCy7ElementQzs09_AnyIndexC0_pcig
+_$ss27_SwiftifyProtocolMethodInfoO6methodyABSS_Says01_aD0OGtcABmFWC
+_$ss27_SwiftifyProtocolMethodInfoOMa
+_$ss27_SwiftifyProtocolMethodInfoOMn
+_$ss27_SwiftifyProtocolMethodInfoON
 _$ss27_allocateUninitializedArrayySayxG_BptBwlF
 _$ss27_bridgeAnythingToObjectiveCyyXlxlF
 _$ss27_debuggerTestingCheckExpectyySS_SStF