[Swiftify] enable mutable span (#80387)

* [Swiftify] Emit Mutable[Raw]Span when possible

Previously wrappers would use UnsafeMutable[Raw]Pointer for mutable
pointers, and Span for non-const std::span, to prevent the compiler from
complaining that MutableSpan didn't exist.

Now that MutableSpan has landed we can finally emit MutableSpan without
causing compilation errors. While we had (disabled) support for MutableSpan
syntax already, some unexpected semantic errors required additional
changes:
 - Mutable[Raw]Span parameters need to be inout (for mutation)
 - inout ~Escapable paramters need explicit lifetime annotations
 - MutableSpan cannot be directly bitcast to std::span, because it is
   ~Copyable, so they need unwrapping to UnsafeMutableBufferPointer

rdar://147883022

* [Swiftify] Wrap if-expressions in Immediately Called Closures

When parameters in swiftified wrapper functions are nullable, we use
separate branches for the nil and nonnil cases, because
`withUnsafeBufferPointer` (and similar) cannot be called on nil.
If-expressions have some limitations on where they are allowed in the
grammar, and cannot be passed as arguments to a function. As such, when
the return value is also swiftified, we get an error when trying to
pass the if-expression to the UnsafeBufferPointer/Span constructor.
While it isn't pretty, the best way forward seems to be by wrapping the
if-expressions in Immediately Called Closures.

The closures have the side-effect of acting as a barrier for 'unsafe':
unsafe keywords outside the closure do not "reach" unsafe expressions
inside the closure. We therefore have to emit "unsafe" where unsafe
expressions are used, rather than just when returning.

rdar://148153063

Author: hnrklssn

lib/Macros/Sources/SwiftMacros/SwiftifyImportMacro.swift

@@ -4,9 +4,6 @@ import SwiftSyntax
 import SwiftSyntaxBuilder
 import SwiftSyntaxMacros
 
-// Disable emitting 'MutableSpan' until it has landed
-let enableMutableSpan = false
-
 // avoids depending on SwiftifyImport.swift
 // all instances are reparsed and reinstantiated by the macro anyways,
 // so linking is irrelevant
@@ -279,36 +276,49 @@ func getUnqualifiedStdName(_ type: String) -> String? {
 func getSafePointerName(mut: Mutability, generateSpan: Bool, isRaw: Bool) -> TokenSyntax {
   switch (mut, generateSpan, isRaw) {
   case (.Immutable, true, true): return "RawSpan"
-  case (.Mutable, true, true): return if enableMutableSpan {
-    "MutableRawSpan"
-  } else {
-    "RawSpan"
-  }
+  case (.Mutable, true, true): return "MutableRawSpan"
   case (.Immutable, false, true): return "UnsafeRawBufferPointer"
   case (.Mutable, false, true): return "UnsafeMutableRawBufferPointer"
 
   case (.Immutable, true, false): return "Span"
-  case (.Mutable, true, false): return if enableMutableSpan {
-    "MutableSpan"
-  } else {
-    "Span"
-  }
+  case (.Mutable, true, false): return "MutableSpan"
   case (.Immutable, false, false): return "UnsafeBufferPointer"
   case (.Mutable, false, false): return "UnsafeMutableBufferPointer"
   }
 }
 
-func transformType(_ prev: TypeSyntax, _ generateSpan: Bool, _ isSizedBy: Bool) throws -> TypeSyntax {
+func hasOwnershipSpecifier(_ attrType: AttributedTypeSyntax) -> Bool {
+  return attrType.specifiers.contains(where: { e in
+    guard let simpleSpec = e.as(SimpleTypeSpecifierSyntax.self) else {
+      return false
+    }
+    let specifierText = simpleSpec.specifier.text
+    switch specifierText {
+    case "borrowing":
+      return true
+    case "inout":
+      return true
+    case "consuming":
+      return true
+    default:
+      return false
+    }
+  })
+}
+
+func transformType(_ prev: TypeSyntax, _ generateSpan: Bool, _ isSizedBy: Bool, _ setMutableSpanInout: Bool) throws -> TypeSyntax {
   if let optType = prev.as(OptionalTypeSyntax.self) {
     return TypeSyntax(
-      optType.with(\.wrappedType, try transformType(optType.wrappedType, generateSpan, isSizedBy)))
+      optType.with(\.wrappedType, try transformType(optType.wrappedType, generateSpan, isSizedBy, setMutableSpanInout)))
   }
   if let impOptType = prev.as(ImplicitlyUnwrappedOptionalTypeSyntax.self) {
-    return try transformType(impOptType.wrappedType, generateSpan, isSizedBy)
+    return try transformType(impOptType.wrappedType, generateSpan, isSizedBy, setMutableSpanInout)
   }
   if let attrType = prev.as(AttributedTypeSyntax.self) {
+    // We insert 'inout' by default for MutableSpan, but it shouldn't override existing ownership
+    let setMutableSpanInoutNext = setMutableSpanInout && !hasOwnershipSpecifier(attrType)
     return TypeSyntax(
-      attrType.with(\.baseType, try transformType(attrType.baseType, generateSpan, isSizedBy)))
+      attrType.with(\.baseType, try transformType(attrType.baseType, generateSpan, isSizedBy, setMutableSpanInoutNext)))
   }
   let name = try getTypeName(prev)
   let text = name.text
@@ -326,10 +336,15 @@ func transformType(_ prev: TypeSyntax, _ generateSpan: Bool, _ isSizedBy: Bool)
         + " - first type token is '\(text)'", node: name)
   }
   let token = getSafePointerName(mut: kind, generateSpan: generateSpan, isRaw: isSizedBy)
-  if isSizedBy {
-    return TypeSyntax(IdentifierTypeSyntax(name: token))
+  let mainType = if isSizedBy {
+    TypeSyntax(IdentifierTypeSyntax(name: token))
+  } else {
+    try replaceTypeName(prev, token)
   }
-  return try replaceTypeName(prev, token)
+  if setMutableSpanInout && generateSpan && kind == .Mutable {
+    return TypeSyntax("inout \(mainType)")
+  }
+  return mainType
 }
 
 func isMutablePointerType(_ type: TypeSyntax) -> Bool {
@@ -431,10 +446,11 @@ struct FunctionCallBuilder: BoundsCheckedThunkBuilder {
       let colon: TokenSyntax? = label != nil ? .colonToken() : nil
       return LabeledExprSyntax(label: label, colon: colon, expression: arg, trailingComma: comma)
     }
-    return ExprSyntax(
+    let call = ExprSyntax(
       FunctionCallExprSyntax(
         calledExpression: functionRef, leftParen: .leftParenToken(),
         arguments: LabeledExprListSyntax(labeledArgs), rightParen: .rightParenToken()))
+    return "unsafe \(call)"
   }
 }
 
@@ -446,6 +462,7 @@ struct CxxSpanThunkBuilder: SpanBoundsThunkBuilder, ParamBoundsThunkBuilder {
   public let node: SyntaxProtocol
   public let nonescaping: Bool
   let isSizedBy: Bool = false
+  let isParameter: Bool = true
 
   func buildBoundsChecks() throws -> [CodeBlockItemSyntax.Item] {
     return try base.buildBoundsChecks()
@@ -462,8 +479,26 @@ struct CxxSpanThunkBuilder: SpanBoundsThunkBuilder, ParamBoundsThunkBuilder {
     var args = pointerArgs
     let typeName = getUnattributedType(oldType).description
     assert(args[index] == nil)
-    args[index] = ExprSyntax("\(raw: typeName)(\(raw: name))")
-    return try base.buildFunctionCall(args)
+
+    let (_, isConst) = dropCxxQualifiers(try genericArg)
+    if isConst {
+      args[index] = ExprSyntax("\(raw: typeName)(\(raw: name))")
+      return try base.buildFunctionCall(args)
+    } else {
+      let unwrappedName = TokenSyntax("_\(name)Ptr")
+      args[index] = ExprSyntax("\(raw: typeName)(\(unwrappedName))")
+      let call = try base.buildFunctionCall(args)
+
+      // MutableSpan - unlike Span - cannot be bitcast to std::span due to being ~Copyable,
+      // so unwrap it to an UnsafeMutableBufferPointer that we can cast
+      let unwrappedCall = ExprSyntax(
+        """
+          unsafe \(name).withUnsafeMutableBufferPointer { \(unwrappedName) in
+            return \(call)
+          }
+        """)
+      return unwrappedCall
+    }
   }
 }
 
@@ -472,6 +507,7 @@ struct CxxSpanReturnThunkBuilder: SpanBoundsThunkBuilder {
   public let signature: FunctionSignatureSyntax
   public let typeMappings: [String: String]
   public let node: SyntaxProtocol
+  let isParameter: Bool = false
 
   var oldType: TypeSyntax {
     return signature.returnClause!.type
@@ -490,12 +526,12 @@ struct CxxSpanReturnThunkBuilder: SpanBoundsThunkBuilder {
   func buildFunctionCall(_ pointerArgs: [Int: ExprSyntax]) throws -> ExprSyntax {
     let call = try base.buildFunctionCall(pointerArgs)
     let (_, isConst) = dropCxxQualifiers(try genericArg)
-    let cast = if isConst || !enableMutableSpan {
+    let cast = if isConst {
       "Span"
     } else {
       "MutableSpan"
     }
-    return "_cxxOverrideLifetime(\(raw: cast)(_unsafeCxxSpan: \(call)), copying: ())"
+    return "unsafe _cxxOverrideLifetime(\(raw: cast)(_unsafeCxxSpan: \(call)), copying: ())"
   }
 }
 
@@ -508,11 +544,12 @@ protocol BoundsThunkBuilder: BoundsCheckedThunkBuilder {
 protocol SpanBoundsThunkBuilder: BoundsThunkBuilder {
   var typeMappings: [String: String] { get }
   var node: SyntaxProtocol { get }
+  var isParameter: Bool { get }
 }
 extension SpanBoundsThunkBuilder {
   var desugaredType: TypeSyntax {
     get throws {
-      let typeName = try getUnattributedType(oldType).description
+      let typeName = getUnattributedType(oldType).description
       guard let desugaredTypeName = typeMappings[typeName] else {
         throw DiagnosticError(
           "unable to desugar type with name '\(typeName)'", node: node)
@@ -547,14 +584,18 @@ extension SpanBoundsThunkBuilder {
   var newType: TypeSyntax {
     get throws {
       let (strippedArg, isConst) = dropCxxQualifiers(try genericArg)
-      let mutablePrefix = if isConst || !enableMutableSpan {
+      let mutablePrefix = if isConst {
         ""
       } else {
         "Mutable"
       }
-      return replaceBaseType(
+      let mainType = replaceBaseType(
         oldType,
         TypeSyntax("\(raw: mutablePrefix)Span<\(raw: strippedArg)>"))
+      if !isConst && isParameter {
+        return TypeSyntax("inout \(mainType)")
+      }
+      return mainType
     }
   }
 }
@@ -563,13 +604,14 @@ protocol PointerBoundsThunkBuilder: BoundsThunkBuilder {
   var nullable: Bool { get }
   var isSizedBy: Bool { get }
   var generateSpan: Bool { get }
+  var isParameter: Bool { get }
 }
 
 extension PointerBoundsThunkBuilder {
   var nullable: Bool { return oldType.is(OptionalTypeSyntax.self) }
 
   var newType: TypeSyntax { get throws {
-    return try transformType(oldType, generateSpan, isSizedBy) }
+    return try transformType(oldType, generateSpan, isSizedBy, isParameter) }
   }
 }
 
@@ -599,8 +641,9 @@ struct CountedOrSizedReturnPointerThunkBuilder: PointerBoundsThunkBuilder {
   public let nonescaping: Bool
   public let isSizedBy: Bool
   public let dependencies: [LifetimeDependence]
+  let isParameter: Bool = false
 
-  var generateSpan: Bool { !dependencies.isEmpty && (!isMutablePointerType(oldType) || enableMutableSpan)}
+  var generateSpan: Bool { !dependencies.isEmpty }
 
   var oldType: TypeSyntax {
     return signature.returnClause!.type
@@ -623,9 +666,25 @@ struct CountedOrSizedReturnPointerThunkBuilder: PointerBoundsThunkBuilder {
     } else {
       "start"
     }
+    var cast = try newType
+    if nullable {
+      if let optType = cast.as(OptionalTypeSyntax.self) {
+        cast = optType.wrappedType
+      }
+      return """
+      { () in
+        let _resultValue = \(call)
+        if unsafe _resultValue == nil {
+          return nil
+        } else {
+          return unsafe \(raw: try cast)(\(raw: startLabel): _resultValue!, count: Int(\(countExpr)))
+        }
+      }()
+      """
+    }
     return
       """
-      \(raw: try newType)(\(raw: startLabel): \(call), count: Int(\(countExpr)))
+      unsafe \(raw: try cast)(\(raw: startLabel): \(call), count: Int(\(countExpr)))
       """
   }
 }
@@ -639,8 +698,9 @@ struct CountedOrSizedPointerThunkBuilder: ParamBoundsThunkBuilder, PointerBounds
   public let nonescaping: Bool
   public let isSizedBy: Bool
   public let skipTrivialCount: Bool
+  let isParameter: Bool = true
 
-  var generateSpan: Bool { nonescaping && (!isMutablePointerType(oldType) || enableMutableSpan) }
+  var generateSpan: Bool { nonescaping }
 
   func buildFunctionSignature(_ argTypes: [Int: TypeSyntax?], _ returnType: TypeSyntax?) throws
     -> (FunctionSignatureSyntax, Bool) {
@@ -702,11 +762,16 @@ struct CountedOrSizedPointerThunkBuilder: ParamBoundsThunkBuilder, PointerBounds
     let call = try base.buildFunctionCall(args)
     let ptrRef = unwrapIfNullable(ExprSyntax(DeclReferenceExprSyntax(baseName: name)))
 
-    let funcName = isSizedBy ? "withUnsafeBytes" : "withUnsafeBufferPointer"
+    let funcName = switch (isSizedBy, isMutablePointerType(oldType)) {
+    case (true, true): "withUnsafeMutableBytes"
+    case (true, false): "withUnsafeBytes"
+    case (false, true): "withUnsafeMutableBufferPointer"
+    case (false, false): "withUnsafeBufferPointer"
+    }
     let unwrappedCall = ExprSyntax(
       """
-        \(ptrRef).\(raw: funcName) { \(unwrappedName) in
-          return unsafe \(call)
+        unsafe \(ptrRef).\(raw: funcName) { \(unwrappedName) in
+          return \(call)
         }
       """)
     return unwrappedCall
@@ -766,11 +831,11 @@ struct CountedOrSizedPointerThunkBuilder: ParamBoundsThunkBuilder, PointerBounds
         nullArgs[index] = ExprSyntax(NilLiteralExprSyntax(nilKeyword: .keyword(.nil)))
         return ExprSyntax(
           """
-            if \(name) == nil {
-              unsafe \(try base.buildFunctionCall(nullArgs))
-            } else {
-              \(unwrappedCall)
-            }
+            { () in return if \(name) == nil {
+                \(try base.buildFunctionCall(nullArgs))
+              } else {
+                \(unwrappedCall)
+              } }()
           """)
       }
       return unwrappedCall
@@ -1161,7 +1226,7 @@ public struct SwiftifyImportMacro: PeerMacro {
     }
   }
 
-  static func lifetimeAttributes(_ funcDecl: FunctionDeclSyntax,
+  static func getReturnLifetimeAttribute(_ funcDecl: FunctionDeclSyntax,
     _ dependencies: [SwiftifyExpr: [LifetimeDependence]]) -> [AttributeListSyntax.Element] {
     let returnDependencies = dependencies[.`return`, default: []]
     if returnDependencies.isEmpty {
@@ -1190,6 +1255,66 @@ public struct SwiftifyImportMacro: PeerMacro {
               rightParen: .rightParenToken()))]
   }
 
+  static func isMutableSpan(_ type: TypeSyntax) -> Bool {
+    if let optType = type.as(OptionalTypeSyntax.self) {
+      return isMutableSpan(optType.wrappedType)
+    }
+    if let impOptType = type.as(ImplicitlyUnwrappedOptionalTypeSyntax.self) {
+      return isMutableSpan(impOptType.wrappedType)
+    }
+    if let attrType = type.as(AttributedTypeSyntax.self) {
+      return isMutableSpan(attrType.baseType)
+    }
+    guard let identifierType = type.as(IdentifierTypeSyntax.self) else {
+      return false
+    }
+    let name = identifierType.name.text
+    return name == "MutableSpan" || name == "MutableRawSpan"
+  }
+
+  static func containsLifetimeAttr(_ attrs: AttributeListSyntax, for paramName: TokenSyntax) -> Bool {
+    for elem in attrs {
+      guard let attr = elem.as(AttributeSyntax.self) else {
+        continue
+      }
+      if attr.attributeName != "lifetime" {
+        continue
+      }
+      guard let args = attr.arguments?.as(LabeledExprListSyntax.self) else {
+        continue
+      }
+      for arg in args {
+        if arg.label == paramName {
+          return true
+        }
+      }
+    }
+    return false
+  }
+
+  // Mutable[Raw]Span parameters need explicit @lifetime annotations since they are inout
+  static func paramLifetimeAttributes(_ newSignature: FunctionSignatureSyntax, _ oldAttrs: AttributeListSyntax) -> [AttributeListSyntax.Element] {
+    var defaultLifetimes: [AttributeListSyntax.Element] = []
+    for param in newSignature.parameterClause.parameters {
+      if !isMutableSpan(param.type) {
+        continue
+      }
+      let paramName = param.secondName ?? param.firstName
+      if containsLifetimeAttr(oldAttrs, for: paramName) {
+        continue
+      }
+      let expr = ExprSyntax("\(paramName): copy \(paramName)")
+
+      defaultLifetimes.append(.attribute(AttributeSyntax(
+              atSign: .atSignToken(),
+              attributeName: IdentifierTypeSyntax(name: "lifetime"),
+              leftParen: .leftParenToken(),
+              arguments: .argumentList(LabeledExprListSyntax([LabeledExprSyntax(expression: expr)])),
+              rightParen: .rightParenToken())))
+    }
+    return defaultLifetimes
+  }
+
   public static func expansion(
     of node: AttributeSyntax,
     providingPeersOf declaration: some DeclSyntaxProtocol,
@@ -1255,9 +1380,10 @@ public struct SwiftifyImportMacro: PeerMacro {
         item: CodeBlockItemSyntax.Item(
           ReturnStmtSyntax(
             returnKeyword: .keyword(.return, trailingTrivia: " "),
-            expression: ExprSyntax("unsafe \(try builder.buildFunctionCall([:]))"))))
+            expression: try builder.buildFunctionCall([:]))))
       let body = CodeBlockSyntax(statements: CodeBlockItemListSyntax(checks + [call]))
-      let lifetimeAttrs = lifetimeAttributes(funcDecl, lifetimeDependencies)
+      let returnLifetimeAttribute = getReturnLifetimeAttribute(funcDecl, lifetimeDependencies)
+      let lifetimeAttrs = returnLifetimeAttribute + paramLifetimeAttributes(newSignature, funcDecl.attributes)
       let disfavoredOverload : [AttributeListSyntax.Element] = (onlyReturnTypeChanged ? [
         .attribute(
           AttributeSyntax(