JNI: non-escaping closures taking or returning jextracted types

Samples/SwiftJavaExtractJNIKotlinSampleApp/Sources/KotlinSampleLibrary/Closures.swift

@@ -34,6 +34,11 @@ public final class Notifier {
     cb("a", 1, nil) && cb("b", 2, Item(id: "b", count: 2))
   }
 
+  /// A non-escaping closure taking and returning a struct.
+  public func transform(id: String, _ f: (Item) -> Item) -> Item {
+    f(Item(id: id, count: 1))
+  }
+
   /// A @Sendable closure called on another thread: the upcall uses that thread's JNIEnv.
   public func deliverInBackground(id: String, _ cb: @escaping @Sendable (Item) -> Void) {
     let group = DispatchGroup()

Samples/SwiftJavaExtractJNIKotlinSampleApp/src/test/kotlin/com/example/kotlin/ClosuresTest.kt

@@ -42,6 +42,13 @@ class ClosuresTest {
         assertEquals(listOf("a1-", "b2b"), calls)
     }
 
+    @Test
+    fun nonEscapingStructClosure() {
+        val out = Notifier.init().transform("a") { item -> Item.init(item.id + "!", item.count + 1) }
+        assertEquals("a!", out.id)
+        assertEquals(2L, out.count)
+    }
+
     @Test
     fun sendableClosureCalledFromAnotherThread() {
         val testThread = Thread.currentThread()

Samples/SwiftJavaExtractJNISampleApp/Sources/MySwiftLibrary/Closures.swift

@@ -123,6 +123,11 @@ public func applyIntClosure(_ f: (Int) -> Int, to x: Int) -> Int {
   f(x)
 }
 
+/// A non-escaping closure taking and returning a struct.
+public func transformStruct(_ value: MySwiftStruct, _ f: (MySwiftStruct) -> MySwiftStruct) -> MySwiftStruct {
+  f(value)
+}
+
 public enum ClosureConstants {
   /// No type annotation: inferred as `Int` from the literal.
   public static let answer = 42

Samples/SwiftJavaExtractJNISampleApp/src/test/java/com/example/swift/ClosuresTest.java

@@ -15,6 +15,7 @@
 package com.example.swift;
 
 import org.junit.jupiter.api.Test;
+import org.swift.swiftkit.core.SwiftArena;
 
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -182,6 +183,19 @@ void intClosure() {
         assertEquals(10, MySwiftLibrary.applyIntClosure(x -> x * 2, 5));
     }
 
+    @Test
+    void nonEscapingStructClosure() {
+        try (var arena = SwiftArena.ofConfined()) {
+            MySwiftStruct result = MySwiftLibrary.transformStruct(
+                MySwiftStruct.init(1, 2, arena),
+                value -> MySwiftStruct.init(value.getCapacity() * 10, value.getLen(), arena),
+                arena
+            );
+            assertEquals(10, result.getCapacity());
+            assertEquals(2, result.getLen());
+        }
+    }
+
     @Test
     void unannotatedStaticLet() {
         assertEquals(42, ClosureConstants.getAnswer());

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+NativeTranslation.swift

@@ -370,56 +370,72 @@ extension JNISwift2JavaGenerator {
       methodName: String,
       parentName: SwiftQualifiedTypeName
     ) throws -> NativeParameter {
-      if functionType.isEscaping {
-        // For escaping closures we e need to create a Swift wrapper around
-        // the passed down Java functional interface because we must keep it
-        // alive with a global ref, that will remain around for as long as the
-        // escaping closure is.
-        //
-        // Prepare the name and shapes of the Java side functional interface
-        // and Swift side @JavaInterface wrapper we'll use to implement that
-        // interface (and keep the Java object reachable).
-
-        // Name it: interface MyClass_myMethod_theClosure { ... }
-        let javaClosureInterfaceName = String.flatName(
-          prefix: "Java",
-          parent: parentName,
-          method: methodName,
-          parameter: parameterName,
-        )
-        let javaBinaryName = String.javaBinaryName(
-          package: self.javaPackage,
-          parent: parentName,
-          method: methodName,
-          qualifier: parameterName,
-        )
-        let generator = JavaInterfaceProtocolWrapperGenerator()
-        let syntheticFunction = try generator.generateSyntheticClosureFunction(
+      if !functionType.isEscaping {
+        // The legacy lowering handles primitive arguments and results only. Other closures go through
+        // the escaping one: its Swift closure can be passed where a non-escaping one is expected.
+        if let legacy = try? translateNonEscapingFunctionParameter(
           functionType: functionType,
-          javaInterfaceName: javaClosureInterfaceName,
-          javaBinaryName: javaBinaryName,
-        )
+          parameterName: parameterName,
+          methodName: methodName,
+          parentName: parentName
+        ) {
+          return legacy
+        }
+      }
+      // For escaping closures we e need to create a Swift wrapper around
+      // the passed down Java functional interface because we must keep it
+      // alive with a global ref, that will remain around for as long as the
+      // escaping closure is.
+      //
+      // Prepare the name and shapes of the Java side functional interface
+      // and Swift side @JavaInterface wrapper we'll use to implement that
+      // interface (and keep the Java object reachable).
+
+      // Name it: interface MyClass_myMethod_theClosure { ... }
+      let javaClosureInterfaceName = String.flatName(
+        prefix: "Java",
+        parent: parentName,
+        method: methodName,
+        parameter: parameterName,
+      )
+      let javaBinaryName = String.javaBinaryName(
+        package: self.javaPackage,
+        parent: parentName,
+        method: methodName,
+        qualifier: parameterName,
+      )
+      let generator = JavaInterfaceProtocolWrapperGenerator()
+      let syntheticFunction = try generator.generateSyntheticClosureFunction(
+        functionType: functionType,
+        javaInterfaceName: javaClosureInterfaceName,
+        javaBinaryName: javaBinaryName,
+      )
 
-        return NativeParameter(
-          parameters: [
-            JavaParameter(
-              name: parameterName,
-              type: .class(
-                package: javaPackage,
-                name: String.javaQualifiedName(parentName.fullName, methodName, parameterName)
-              )
+      return NativeParameter(
+        parameters: [
+          JavaParameter(
+            name: parameterName,
+            type: .class(
+              package: javaPackage,
+              name: String.javaQualifiedName(parentName.fullName, methodName, parameterName)
             )
-          ],
-          conversion: .escapingClosureLowering(
-            syntheticFunction: syntheticFunction
-          ),
-          indirectConversion: nil,
-          conversionCheck: nil,
-          syntheticClosure: syntheticFunction
-        )
-      }
+          )
+        ],
+        conversion: .escapingClosureLowering(
+          syntheticFunction: syntheticFunction
+        ),
+        indirectConversion: nil,
+        conversionCheck: nil,
+        syntheticClosure: syntheticFunction
+      )
+    }
 
-      // Non-escaping closures use the legacy translation
+    private func translateNonEscapingFunctionParameter(
+      functionType: SwiftFunctionType,
+      parameterName: String,
+      methodName: String,
+      parentName: SwiftQualifiedTypeName
+    ) throws -> NativeParameter {
       var parameters = [NativeParameter]()
       for (i, parameter) in functionType.parameters.enumerated() {
         let closureParamName = parameter.parameterName ?? "_\(i)"
@@ -1831,9 +1847,12 @@ extension JNISwift2JavaGenerator {
         // It will cause a new global ref on the javaThis, so no need for explicit global refs.
         // This object will be closed over by the closure we pass to Swift, and therefore keep alive the
         // Java side that we'll make an up-call into when the closure runs.
+        // For a non-escaping parameter, the literal would be inferred as non-escaping and couldn't be
+        // returned: spell out the closure type, which is escaping as a result.
+        let outerHeader = fn.isEscaping ? "{" : "{ () -> \(fn) in"
         printer.print(
           """
-          {
+          \(outerHeader)
             guard let \(placeholder) else {
               fatalError(\"\(placeholder) is null\")
             }

Tests/JExtractSwiftTests/JNI/JNIEscapingClosureJExtractedTypesTests.swift

@@ -87,6 +87,52 @@ struct JNIEscapingClosureJExtractedTypesTests {
     )
   }
 
+  @Test
+  func nonEscapingStructClosure_usesTheEscapingLowering() throws {
+    try assertOutput(
+      input: itemSource + "\npublic func transform(_ item: Item, _ f: (Item) -> Item) -> Item { f(item) }",
+      .jni,
+      .swift,
+      detectChunkByInitialLines: 1,
+      expectedChunks: [
+        """
+        @JavaInterface("com.example.swift.SwiftModule$transform$f")
+        public struct JavaSwiftModule_transform_f {
+          @JavaMethod
+          public func apply(_ _0: JavaItem?) -> JavaItem!
+        }
+        """,
+        // Spelled out: inferred from the non-escaping parameter, the returned closure couldn't escape.
+        """
+        result$.initialize(to: SwiftModule.transform(item$.pointee, { () -> (Item) -> Item in
+        guard let f else {
+          fatalError("f is null")
+        }
+        let javaInterface_f$ = JavaSwiftModule_transform_f(javaThis: f, environment: environment)
+        return { _0 in
+        """,
+      ]
+    )
+  }
+
+  @Test
+  func nonEscapingStructClosure_javaInterface() throws {
+    try assertOutput(
+      input: itemSource + "\npublic func transform(_ item: Item, _ f: (Item) -> Item) -> Item { f(item) }",
+      .jni,
+      .java,
+      detectChunkByInitialLines: 2,
+      expectedChunks: [
+        """
+        @FunctionalInterface
+        public interface f {
+          Item apply(Item _0);
+        }
+        """
+      ]
+    )
+  }
+
   @Test
   func optionalStructParameter_isJavaOptional() throws {
     try assertOutput(