[Swiftify] Emit Span for counted_by return values with lifetime info

__counted_by return values with .lifetimeDependence are now mapped to
Span instead of UnsafeBufferPointer. Also fixes bug where std::span
return values would map to Span even if lifetime dependence info was
missing.

Author: hnrklssn

lib/Macros/Sources/SwiftMacros/SwiftifyImportMacro.swift

@@ -88,6 +88,9 @@ struct CxxSpan: ParamInfo {
       return CxxSpanThunkBuilder(base: base, index: i - 1, signature: funcDecl.signature,
         typeMappings: typeMappings, node: original, nonescaping: nonescaping)
     case .return:
+      if dependencies.isEmpty {
+        return base
+      }
       return CxxSpanReturnThunkBuilder(base: base, signature: funcDecl.signature,
         typeMappings: typeMappings, node: original)
     case .self:
@@ -125,7 +128,7 @@ struct CountedBy: ParamInfo {
       return CountedOrSizedReturnPointerThunkBuilder(
         base: base, countExpr: count,
         signature: funcDecl.signature,
-        nonescaping: nonescaping, isSizedBy: sizedBy)
+        nonescaping: nonescaping, isSizedBy: sizedBy, dependencies: dependencies)
     case .self:
       return base
     }
@@ -184,11 +187,13 @@ func getTypeName(_ type: TypeSyntax) throws -> TokenSyntax {
   }
 }
 
-func replaceTypeName(_ type: TypeSyntax, _ name: TokenSyntax) -> TypeSyntax {
+func replaceTypeName(_ type: TypeSyntax, _ name: TokenSyntax) throws -> TypeSyntax {
   if let memberType = type.as(MemberTypeSyntax.self) {
     return TypeSyntax(memberType.with(\.name, name))
   }
-  let idType = type.as(IdentifierTypeSyntax.self)!
+  guard let idType = type.as(IdentifierTypeSyntax.self) else {
+    throw DiagnosticError("unexpected type \(type) with kind \(type.kind)", node: type)
+  }
   return TypeSyntax(idType.with(\.name, name))
 }
 
@@ -264,6 +269,10 @@ func transformType(_ prev: TypeSyntax, _ generateSpan: Bool, _ isSizedBy: Bool)
   if let impOptType = prev.as(ImplicitlyUnwrappedOptionalTypeSyntax.self) {
     return try transformType(impOptType.wrappedType, generateSpan, isSizedBy)
   }
+  if let attrType = prev.as(AttributedTypeSyntax.self) {
+    return TypeSyntax(
+      attrType.with(\.baseType, try transformType(attrType.baseType, generateSpan, isSizedBy)))
+  }
   let name = try getTypeName(prev)
   let text = name.text
   let isRaw = isRawPointerType(text: text)
@@ -283,7 +292,7 @@ func transformType(_ prev: TypeSyntax, _ generateSpan: Bool, _ isSizedBy: Bool)
   if isSizedBy {
     return TypeSyntax(IdentifierTypeSyntax(name: token))
   }
-  return replaceTypeName(prev, token)
+  return try replaceTypeName(prev, token)
 }
 
 protocol BoundsCheckedThunkBuilder {
@@ -485,8 +494,9 @@ struct CountedOrSizedReturnPointerThunkBuilder: PointerBoundsThunkBuilder {
   public let signature: FunctionSignatureSyntax
   public let nonescaping: Bool
   public let isSizedBy: Bool
+  public let dependencies: [LifetimeDependence]
 
-  var generateSpan: Bool = false // needs more lifetime information
+  var generateSpan: Bool { !dependencies.isEmpty }
 
   var oldType: TypeSyntax {
     return signature.returnClause!.type
@@ -504,9 +514,14 @@ struct CountedOrSizedReturnPointerThunkBuilder: PointerBoundsThunkBuilder {
 
   func buildFunctionCall(_ pointerArgs: [Int: ExprSyntax]) throws -> ExprSyntax {
     let call = try base.buildFunctionCall(pointerArgs)
+    let startLabel = if generateSpan {
+      "_unsafeStart"
+    } else {
+      "start"
+    }
     return
       """
-      \(raw: try newType)(start: \(call), count: Int(\(countExpr)))
+      \(raw: try newType)(\(raw: startLabel): \(call), count: Int(\(countExpr)))
       """
   }
 }
@@ -1102,6 +1117,11 @@ public struct SwiftifyImportMacro: PeerMacro {
         }
       }
       try checkArgs(parsedArgs, funcDecl)
+      parsedArgs.sort { a, b in
+        // make sure return value cast to Span happens last so that withUnsafeBufferPointer
+        // doesn't return a ~Escapable type
+        (a.pointerIndex != .return && b.pointerIndex == .return) || paramOrReturnIndex(a.pointerIndex) < paramOrReturnIndex(b.pointerIndex)
+      }
       let baseBuilder = FunctionCallBuilder(funcDecl)
 
       let skipTrivialCount = hasTrivialCountVariants(parsedArgs)

test/Interop/C/swiftify-import/Inputs/counted-by-noescape.h

@@ -2,6 +2,7 @@
 
 #define __counted_by(x) __attribute__((__counted_by__(x)))
 #define __noescape __attribute__((noescape))
+#define __lifetimebound __attribute__((lifetimebound))
 
 void simple(int len, int * __counted_by(len) __noescape p);
 
@@ -20,3 +21,4 @@ void nonnull(int len, int * __counted_by(len) _Nonnull __noescape p);
 
 int * __counted_by(len) __noescape returnPointer(int len);
 
+int * __counted_by(len1) returnLifetimeBound(int len1, int len2, int * __counted_by(len2) p __lifetimebound);

test/Interop/C/swiftify-import/counted-by-noescape.swift

@@ -1,6 +1,7 @@
 // REQUIRES: swift_feature_SafeInteropWrappers
+// REQUIRES: swift_feature_LifetimeDependence
 
-// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByNoEscapeClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers | %FileCheck %s
+// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByNoEscapeClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers -enable-experimental-feature LifetimeDependence | %FileCheck %s
 
 // swift-ide-test doesn't currently typecheck the macro expansions, so run the compiler as well
 // RUN: %empty-directory(%t)
@@ -13,6 +14,8 @@ import CountedByNoEscapeClang
 // CHECK:      @_alwaysEmitIntoClient public func complexExpr(_ len: Int{{.*}}, _ offset: Int{{.*}}, _ p: MutableSpan<Int{{.*}}>)
 // CHECK-NEXT: @_alwaysEmitIntoClient public func nonnull(_  p: MutableSpan<Int{{.*}}>)
 // CHECK-NEXT: @_alwaysEmitIntoClient public func nullUnspecified(_  p: MutableSpan<Int{{.*}}>)
+// CHECK-NEXT: @lifetime(p)
+// CHECK-NEXT: @_alwaysEmitIntoClient public func returnLifetimeBound(_ len1: Int32, _ p: MutableSpan<Int32>) -> MutableSpan<Int32>
 // CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func returnPointer(_  len: Int{{.*}}) -> UnsafeMutableBufferPointer<Int{{.*}}>
 // CHECK-NEXT: @_alwaysEmitIntoClient public func shared(_ len: Int{{.*}}, _ p1: MutableSpan<Int{{.*}}>, _ p2: MutableSpan<Int{{.*}}>)
 // CHECK-NEXT: @_alwaysEmitIntoClient public func simple(_  p: MutableSpan<Int{{.*}}>)

test/Interop/Cxx/stdlib/Inputs/std-span.h

@@ -109,4 +109,8 @@ inline void mixedFuncWithSafeWrapper5(ConstSpanOfInt s,
 inline void mixedFuncWithSafeWrapper6(ConstSpanOfInt s,
                                       int * __counted_by(len) p, int len) {}
 
+inline ConstSpanOfInt mixedFuncWithSafeWrapper7(const int * __counted_by(len) p, int len) {
+  return ConstSpanOfInt(p, len);
+}
+
 #endif // TEST_INTEROP_CXX_STDLIB_INPUTS_STD_SPAN_H

test/Interop/Cxx/stdlib/std-span-interface.swift

@@ -35,8 +35,9 @@ import CxxStdlib
 // CHECK-NEXT: @lifetime(p)
 // CHECK-NEXT: @_alwaysEmitIntoClient public func mixedFuncWithSafeWrapper1(_ p: Span<Int32>) -> Span<CInt>
 // CHECK-NEXT: @lifetime(borrow v)
-// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func mixedFuncWithSafeWrapper2(_ v: borrowing VecOfInt, _ len: Int32) -> UnsafeBufferPointer<Int32>
+// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func mixedFuncWithSafeWrapper2(_ v: borrowing VecOfInt, _ len: Int32) -> Span<Int32>
 // CHECK-NEXT: @_alwaysEmitIntoClient public func mixedFuncWithSafeWrapper3(_ s: Span<CInt>, _ p: UnsafeMutableBufferPointer<Int32>)
 // CHECK-NEXT: @_alwaysEmitIntoClient public func mixedFuncWithSafeWrapper4(_ s: Span<CInt>, _ p: Span<Int32>)
 // CHECK-NEXT: @_alwaysEmitIntoClient public func mixedFuncWithSafeWrapper5(_ s: ConstSpanOfInt, _ p: Span<Int32>)
 // CHECK-NEXT: @_alwaysEmitIntoClient public func mixedFuncWithSafeWrapper6(_ s: ConstSpanOfInt, _ p: UnsafeMutableBufferPointer<Int32>)
+// CHECK-NEXT: @_alwaysEmitIntoClient public func mixedFuncWithSafeWrapper7(_ p: UnsafeBufferPointer<Int32>) -> ConstSpanOfInt

test/Macros/SwiftifyImport/CountedBy/PointerReturn.swift

@@ -1,6 +1,6 @@
 // REQUIRES: swift_swift_parser
 
-// RUN: %target-swift-frontend %s -swift-version 5 -module-name main -disable-availability-checking -typecheck -plugin-path %swift-plugin-dir -dump-macro-expansions 2>&1 | %FileCheck --match-full-lines %s
+// RUN: %target-swift-frontend %s -swift-version 5 -module-name main -disable-availability-checking -typecheck -plugin-path %swift-plugin-dir -enable-experimental-feature LifetimeDependence -dump-macro-expansions 2>&1 | %FileCheck --match-full-lines %s
 
 @_SwiftifyImport(.countedBy(pointer: .return, count: "len"))
 func myFunc(_ len: CInt) -> UnsafeMutablePointer<CInt> {
@@ -10,6 +10,14 @@ func myFunc(_ len: CInt) -> UnsafeMutablePointer<CInt> {
 func nonEscaping(_ len: CInt) -> UnsafePointer<CInt> {
 }
 
+@_SwiftifyImport(.countedBy(pointer: .return, count: "len2"), .countedBy(pointer: .param(1), count: "len1"), .lifetimeDependence(dependsOn: .param(1), pointer: .return, type: .copy))
+func lifetimeDependentCopy(_ p: UnsafePointer<CInt>, _ len1: CInt, _ len2: CInt) -> UnsafePointer<CInt> {
+}
+
+@_SwiftifyImport(.countedBy(pointer: .return, count: "len2"), .countedBy(pointer: .param(1), count: "len1"), .lifetimeDependence(dependsOn: .param(1), pointer: .return, type: .borrow))
+func lifetimeDependentBorrow(_ p: borrowing UnsafePointer<CInt>, _ len1: CInt, _ len2: CInt) -> UnsafePointer<CInt> {
+}
+
 // CHECK:      @_alwaysEmitIntoClient @_disfavoredOverload
 // CHECK-NEXT: func myFunc(_ len: CInt) -> UnsafeMutableBufferPointer<CInt> {
 // CHECK-NEXT:     return UnsafeMutableBufferPointer<CInt> (start: myFunc(len), count: Int(len))
@@ -18,5 +26,16 @@ func nonEscaping(_ len: CInt) -> UnsafePointer<CInt> {
 // CHECK:      @_alwaysEmitIntoClient @_disfavoredOverload
 // CHECK-NEXT: func nonEscaping(_ len: CInt) -> UnsafeBufferPointer<CInt> {
 // CHECK-NEXT:     return UnsafeBufferPointer<CInt> (start: nonEscaping(len), count: Int(len))
+
+// CHECK:      @_alwaysEmitIntoClient @lifetime(p)
+// CHECK-NEXT: func lifetimeDependentCopy(_ p: Span<CInt>, _ len2: CInt) -> Span<CInt> {
+// CHECK-NEXT:     return Span<CInt> (_unsafeStart:   p.withUnsafeBufferPointer { _pPtr in
+// CHECK-NEXT:         return lifetimeDependentCopy(_pPtr.baseAddress!, CInt(exactly: p.count)!, len2)
+// CHECK-NEXT:       }, count: Int(len2))
+// CHECK-NEXT: }
+
+// CHECK:      @_alwaysEmitIntoClient @lifetime(borrow p)
+// CHECK-NEXT: func lifetimeDependentBorrow(_ p: borrowing UnsafeBufferPointer<CInt>, _ len2: CInt) -> Span<CInt> {
+// CHECK-NEXT:     return Span<CInt> (_unsafeStart: lifetimeDependentBorrow(p.baseAddress!, CInt(exactly: p.count)!, len2), count: Int(len2))
 // CHECK-NEXT: }