[jextract] fix support for objects in async (#428)

Co-authored-by: Konrad `ktoso` Malawski <[email protected]>

Author: madsodgaard

Samples/SwiftJavaExtractJNISampleApp/Sources/MySwiftLibrary/Async.swift

@@ -36,3 +36,7 @@ public func asyncOptional(i: Int64) async throws -> Int64? {
 public func asyncThrows() async throws {
   throw MySwiftError.swiftError
 }
+
+public func asyncString(input: String) async -> String {
+  return input
+}

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

@@ -74,4 +74,10 @@ void asyncOptional() {
         CompletableFuture<OptionalLong> future = MySwiftLibrary.asyncOptional(42);
         assertEquals(OptionalLong.of(42), future.join());
     }
+
+    @Test
+    void asyncString() {
+        CompletableFuture<String> future = MySwiftLibrary.asyncString("hey");
+        assertEquals("hey", future.join());
+    }
 }

Sources/JExtractSwiftLib/Convenience/JavaType+Extensions.swift

@@ -137,4 +137,13 @@ extension JavaType {
       return self
     }
   }
+
+  var requiresBoxing: Bool {
+    switch self {
+    case .boolean, .byte, .char, .short, .int, .long, .float, .double:
+      true
+    default:
+      false
+    }
+  }
 }

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaTranslation.swift

@@ -514,14 +514,12 @@ extension JNISwift2JavaGenerator {
         )
 
         // Update native function
-        nativeFunctionSignature.result.javaType = .void
         nativeFunctionSignature.result.conversion = .asyncCompleteFuture(
-          nativeFunctionSignature.result.conversion,
           swiftFunctionResultType: originalFunctionSignature.result.type,
-          nativeReturnType: nativeFunctionSignature.result.javaType,
-          outParameters: nativeFunctionSignature.result.outParameters,
+          nativeFunctionSignature: nativeFunctionSignature,
           isThrowing: originalFunctionSignature.isThrowing
         )
+        nativeFunctionSignature.result.javaType = .void
         nativeFunctionSignature.result.outParameters.append(.init(name: "result_future", type: nativeFutureType))
       }
     }

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+NativeTranslation.swift

@@ -570,10 +570,8 @@ extension JNISwift2JavaGenerator {
     indirect case unwrapOptional(NativeSwiftConversionStep, name: String, fatalErrorMessage: String)
 
     indirect case asyncCompleteFuture(
-      NativeSwiftConversionStep,
       swiftFunctionResultType: SwiftType,
-      nativeReturnType: JavaType,
-      outParameters: [JavaParameter],
+      nativeFunctionSignature: NativeFunctionSignature,
       isThrowing: Bool
     )
 
@@ -806,15 +804,22 @@ extension JNISwift2JavaGenerator {
         return unwrappedName
 
       case .asyncCompleteFuture(
-        let inner,
         let swiftFunctionResultType,
-        let nativeReturnType,
-        let outParameters,
+        let nativeFunctionSignature,
         let isThrowing
       ):
+        var globalRefs: [String] = ["globalFuture"]
+
         // Global ref all indirect returns
-        for outParameter in outParameters {
+        for outParameter in nativeFunctionSignature.result.outParameters {
           printer.print("let \(outParameter.name) = environment.interface.NewGlobalRef(environment, \(outParameter.name))")
+          globalRefs.append(outParameter.name)
+        }
+
+        // We also need to global ref any objects passed in
+        for parameter in nativeFunctionSignature.parameters.flatMap(\.parameters) where !parameter.type.isPrimitive {
+          printer.print("let \(parameter.name) = environment.interface.NewGlobalRef(environment, \(parameter.name))")
+          globalRefs.append(parameter.name)
         }
 
         printer.print(
@@ -826,26 +831,28 @@ extension JNISwift2JavaGenerator {
         func printDo(printer: inout CodePrinter) {
           printer.print("let swiftResult$ = await \(placeholder)")
           printer.print("environment = try! JavaVirtualMachine.shared().environment()")
-          let inner = inner.render(&printer, "swiftResult$")
+          let inner = nativeFunctionSignature.result.conversion.render(&printer, "swiftResult$")
           if swiftFunctionResultType.isVoid {
             printer.print("environment.interface.CallBooleanMethodA(environment, globalFuture, _JNIMethodIDCache.CompletableFuture.complete, [jvalue(l: nil)])")
           } else {
-            printer.print(
-              """
-              let boxedResult$ = SwiftJavaRuntimeSupport._JNIBoxedConversions.box(\(inner), in: environment)
-              environment.interface.CallBooleanMethodA(environment, globalFuture, _JNIMethodIDCache.CompletableFuture.complete, [jvalue(l: boxedResult$)])
-              """
-            )
+            let result: String
+            if nativeFunctionSignature.result.javaType.requiresBoxing {
+              printer.print("let boxedResult$ = SwiftJavaRuntimeSupport._JNIBoxedConversions.box(\(inner), in: environment)")
+              result = "boxedResult$"
+            } else {
+              result = inner
+            }
+
+            printer.print("environment.interface.CallBooleanMethodA(environment, globalFuture, _JNIMethodIDCache.CompletableFuture.complete, [jvalue(l: \(result))])")
           }
         }
 
         func printTaskBody(printer: inout CodePrinter) {
           printer.printBraceBlock("defer") { printer in
             // Defer might on any thread, so we need to attach environment.
             printer.print("let deferEnvironment = try! JavaVirtualMachine.shared().environment()")
-            printer.print("environment.interface.DeleteGlobalRef(deferEnvironment, globalFuture)")
-            for outParameter in outParameters {
-              printer.print("environment.interface.DeleteGlobalRef(deferEnvironment, \(outParameter.name))")
+            for globalRef in globalRefs {
+              printer.print("environment.interface.DeleteGlobalRef(deferEnvironment, \(globalRef))")
             }
           }
           if isThrowing {

Sources/JExtractSwiftLib/JavaParameter.swift

@@ -35,6 +35,15 @@ struct JavaParameter {
       }
     }
 
+    var isPrimitive: Bool {
+      switch self {
+      case .concrete(let javaType):
+        javaType.isPrimitive
+      case .generic(let name, let extends):
+        false
+      }
+    }
+
     var jniTypeName: String {
       switch self {
       case .concrete(let type): type.jniTypeName

Tests/JExtractSwiftTests/JNI/JNIAsyncTests.swift

@@ -353,4 +353,59 @@ struct JNIAsyncTests {
       ]
     )
   }
+
+  @Test("Import: (String) async -> String (Java, CompletableFuture)")
+  func completableFuture_asyncStringToString_java() throws {
+    try assertOutput(
+      input: """
+      public func async(s: String) async -> String
+      """,
+      .jni, .java,
+      detectChunkByInitialLines: 2,
+      expectedChunks: [
+        """
+        public static java.util.concurrent.CompletableFuture<java.lang.String> async(java.lang.String s) {
+          java.util.concurrent.CompletableFuture<java.lang.String> $future = new java.util.concurrent.CompletableFuture<java.lang.String>();
+          SwiftModule.$async(s, $future);
+          return $future.thenApply((futureResult$) -> {
+            return futureResult$;
+          }
+          );
+        }
+        """,
+        """
+        private static native void $async(java.lang.String s, java.util.concurrent.CompletableFuture<java.lang.String> result_future);
+        """,
+      ]
+    )
+  }
+
+  @Test("Import: (String) async -> String (Swift, CompletableFuture)")
+  func completableFuture_asyncStringToString_swift() throws {
+    try assertOutput(
+      input: """
+      public func async(s: String) async -> String
+      """,
+      .jni, .swift,
+      detectChunkByInitialLines: 1,
+      expectedChunks: [
+        """
+        @_cdecl("Java_com_example_swift_SwiftModule__00024async__Ljava_lang_String_2Ljava_util_concurrent_CompletableFuture_2")
+        func Java_com_example_swift_SwiftModule__00024async__Ljava_lang_String_2Ljava_util_concurrent_CompletableFuture_2(environment: UnsafeMutablePointer<JNIEnv?>!, thisClass: jclass, s: jstring?, result_future: jobject?) {
+          let s = environment.interface.NewGlobalRef(environment, s)
+          let globalFuture = environment.interface.NewGlobalRef(environment, result_future)
+          ...
+          defer {
+            let deferEnvironment = try! JavaVirtualMachine.shared().environment()
+            environment.interface.DeleteGlobalRef(deferEnvironment, globalFuture)
+            environment.interface.DeleteGlobalRef(deferEnvironment, s)
+          }
+          ...
+          environment.interface.CallBooleanMethodA(environment, globalFuture, _JNIMethodIDCache.CompletableFuture.complete, [jvalue(l: swiftResult$.getJNIValue(in: environment))])
+          ...
+        }
+        """
+      ]
+    )
+  }
 }