jextract: fix protocols that return java classes. (#479)

* fix protocols

* remove comments

* comments

* cleanup

Author: madsodgaard

Package.swift

@@ -460,6 +460,7 @@ let package = Package(
         .product(name: "SwiftSyntax", package: "swift-syntax"),
         .product(name: "SwiftSyntaxBuilder", package: "swift-syntax"),
         .product(name: "ArgumentParser", package: "swift-argument-parser"),
+        .product(name: "OrderedCollections", package: "swift-collections"),
         "JavaTypes",
         "SwiftJavaShared",
         "SwiftJavaConfigurationShared",

Samples/SwiftJavaExtractJNISampleApp/Sources/MySwiftLibrary/ConcreteProtocolAB.swift

@@ -21,6 +21,10 @@ public class ConcreteProtocolAB: ProtocolA, ProtocolB {
     return "ConcreteProtocolAB"
   }
 
+  public func makeClass() -> MySwiftClass {
+    return MySwiftClass(x: 10, y: 50)
+  }
+
   public init(constantA: Int64, constantB: Int64) {
     self.constantA = constantA
     self.constantB = constantB

Samples/SwiftJavaExtractJNISampleApp/Sources/MySwiftLibrary/ProtocolA.swift

@@ -17,6 +17,7 @@ public protocol ProtocolA {
   var mutable: Int64 { get set }
 
   func name() -> String
+  func makeClass() -> MySwiftClass
 }
 
 public func takeProtocol(_ proto1: any ProtocolA, _ proto2: some ProtocolA) -> Int64 {

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

@@ -74,7 +74,7 @@ public String withString(String input) {
         public void withVoid() {}
 
         @Override
-        public MySwiftClass withObject(MySwiftClass input) {
+        public MySwiftClass withObject(MySwiftClass input, SwiftArena swiftArena$) {
             return input;
         }
 
@@ -84,7 +84,7 @@ public OptionalLong withOptionalInt64(OptionalLong input) {
         }
 
         @Override
-        public Optional<MySwiftClass> withOptionalObject(Optional<MySwiftClass> input) {
+        public Optional<MySwiftClass> withOptionalObject(Optional<MySwiftClass> input, SwiftArena swiftArena$) {
             return input;
         }
     }

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

@@ -73,6 +73,14 @@ void protocolMethod() {
         }
     }
 
+    @Test
+    void protocolClassMethod() {
+        try (var arena = SwiftArena.ofConfined()) {
+            ProtocolA proto1 = ConcreteProtocolAB.init(10, 5, arena);
+            assertEquals(10, proto1.makeClass().getX());
+        }
+    }
+
     static class JavaStorage implements Storage {
         StorageItem item;
 
@@ -81,7 +89,7 @@ static class JavaStorage implements Storage {
         }
 
         @Override
-        public StorageItem load() {
+        public StorageItem load(SwiftArena swiftArena$) {
             return item;
         }
 

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaBindingsPrinting.swift

@@ -144,17 +144,17 @@ extension JNISwift2JavaGenerator {
 
     printer.printBraceBlock("public interface \(decl.swiftNominal.name)\(extendsString)") { printer in
       for initializer in decl.initializers {
-        printFunctionDowncallMethods(&printer, initializer, skipMethodBody: true, skipArenas: true)
+        printFunctionDowncallMethods(&printer, initializer, skipMethodBody: true)
         printer.println()
       }
 
       for method in decl.methods {
-        printFunctionDowncallMethods(&printer, method, skipMethodBody: true, skipArenas: true)
+        printFunctionDowncallMethods(&printer, method, skipMethodBody: true)
         printer.println()
       }
 
       for variable in decl.variables {
-        printFunctionDowncallMethods(&printer, variable, skipMethodBody: true, skipArenas: true)
+        printFunctionDowncallMethods(&printer, variable, skipMethodBody: true)
         printer.println()
       }
     }
@@ -420,16 +420,15 @@ extension JNISwift2JavaGenerator {
         printer.print("record _NativeParameters(\(nativeParameters.joined(separator: ", "))) {}")
       }
 
-      self.printJavaBindingWrapperMethod(&printer, translatedCase.getAsCaseFunction, skipMethodBody: false, skipArenas: false)
+      self.printJavaBindingWrapperMethod(&printer, translatedCase.getAsCaseFunction, skipMethodBody: false)
       printer.println()
     }
   }
 
   private func printFunctionDowncallMethods(
     _ printer: inout CodePrinter,
     _ decl: ImportedFunc,
-    skipMethodBody: Bool = false,
-    skipArenas: Bool = false
+    skipMethodBody: Bool = false
   ) {
     guard translatedDecl(for: decl) != nil else {
       // Failed to translate. Skip.
@@ -440,7 +439,7 @@ extension JNISwift2JavaGenerator {
 
     printJavaBindingWrapperHelperClass(&printer, decl)
 
-    printJavaBindingWrapperMethod(&printer, decl, skipMethodBody: skipMethodBody, skipArenas: skipArenas)
+    printJavaBindingWrapperMethod(&printer, decl, skipMethodBody: skipMethodBody)
   }
 
   /// Print the helper type container for a user-facing Java API.
@@ -486,21 +485,19 @@ extension JNISwift2JavaGenerator {
   private func printJavaBindingWrapperMethod(
     _ printer: inout CodePrinter,
     _ decl: ImportedFunc,
-    skipMethodBody: Bool,
-    skipArenas: Bool
+    skipMethodBody: Bool
   ) {
     guard let translatedDecl = translatedDecl(for: decl) else {
       fatalError("Decl was not translated, \(decl)")
     }
-    printJavaBindingWrapperMethod(&printer, translatedDecl, importedFunc: decl, skipMethodBody: skipMethodBody, skipArenas: skipArenas)
+    printJavaBindingWrapperMethod(&printer, translatedDecl, importedFunc: decl, skipMethodBody: skipMethodBody)
   }
 
   private func printJavaBindingWrapperMethod(
     _ printer: inout CodePrinter,
     _ translatedDecl: TranslatedFunctionDecl,
     importedFunc: ImportedFunc? = nil,
-    skipMethodBody: Bool,
-    skipArenas: Bool
+    skipMethodBody: Bool
   ) {
     var modifiers = ["public"]
     if translatedDecl.isStatic {
@@ -531,14 +528,28 @@ extension JNISwift2JavaGenerator {
     let parametersStr = parameters.joined(separator: ", ")
 
     // Print default global arena variation
+    // If we have enabled javaCallbacks we must emit default
+    // arena methods for protocols, as this is what
+    // Swift will call into, when you call a interface from Swift.
+    let shouldGenerateGlobalArenaVariation: Bool
+    let isParentProtocol = importedFunc?.parentType?.asNominalType?.isProtocol ?? false
+
     if config.effectiveMemoryManagementMode.requiresGlobalArena && translatedSignature.requiresSwiftArena {
+      shouldGenerateGlobalArenaVariation = true
+    } else if isParentProtocol, translatedSignature.requiresSwiftArena, config.effectiveEnableJavaCallbacks {
+      shouldGenerateGlobalArenaVariation = true
+    } else {
+      shouldGenerateGlobalArenaVariation = false
+    }
+
+    if shouldGenerateGlobalArenaVariation {
       if let importedFunc {
         printDeclDocumentation(&printer, importedFunc)
       }
       var modifiers = modifiers
 
       // If we are a protocol, we emit this as default method
-      if importedFunc?.parentType?.asNominalTypeDeclaration?.kind == .protocol {
+      if isParentProtocol {
         modifiers.insert("default", at: 1)
       }
 
@@ -555,7 +566,7 @@ extension JNISwift2JavaGenerator {
       printer.println()
     }
 
-    if translatedSignature.requiresSwiftArena, !skipArenas {
+    if translatedSignature.requiresSwiftArena {
       parameters.append("SwiftArena swiftArena$")
     }
     if let importedFunc {

Sources/SwiftJavaConfigurationShared/Configuration.swift

@@ -65,7 +65,10 @@ public struct Configuration: Codable {
     asyncFuncMode ?? .default
   }
 
-  public var enableJavaCallbacks: Bool? // FIXME: default it to false, but that plays not nice with Codable
+  public var enableJavaCallbacks: Bool?
+  public var effectiveEnableJavaCallbacks: Bool {
+    enableJavaCallbacks ?? false
+  }
 
   public var generatedJavaSourcesListFileOutput: String?
 

Tests/JExtractSwiftTests/JNI/JNIProtocolTests.swift

@@ -32,7 +32,9 @@ struct JNIProtocolTests {
   
     public protocol B {}
   
-    public class SomeClass: SomeProtocol {}
+    public class SomeClass: SomeProtocol {
+      public func makeClass() -> SomeClass {}
+    }
   
     public func takeProtocol(x: some SomeProtocol, y: any SomeProtocol)
     public func takeGeneric<S: SomeProtocol>(s: S)
@@ -61,7 +63,29 @@ struct JNIProtocolTests {
           ...
           public void method();
           ...
-          public SomeClass withObject(SomeClass c);
+          public SomeClass withObject(SomeClass c, SwiftArena swiftArena$);
+          ...
+        }
+        """
+      ])
+  }
+
+  @Test
+  func emitsDefault() throws {
+    try assertOutput(
+      input: source,
+      config: config,
+      .jni, .java,
+      detectChunkByInitialLines: 1,
+      expectedChunks: [
+        """
+        public interface SomeProtocol {
+          ...
+          public default SomeClass withObject(SomeClass c) {
+            return withObject(c, SwiftMemoryManagement.GLOBAL_SWIFT_JAVA_ARENA);
+          }
+          ...
+          public SomeClass withObject(SomeClass c, SwiftArena swiftArena$);
           ...
         }
         """
@@ -78,7 +102,11 @@ struct JNIProtocolTests {
       expectedChunks: [
         """
         public final class SomeClass implements JNISwiftInstance, SomeProtocol {
-        """
+          ...
+          public SomeClass makeClass(SwiftArena swiftArena$) {
+          ...
+        }
+        """,
       ])
   }
 

Tests/JExtractSwiftTests/MemoryManagementModeTests.swift

@@ -99,7 +99,7 @@ struct MemoryManagementModeTests {
         }
         """,
         """
-        public MyClass f();
+        public MyClass f(SwiftArena swiftArena$);
         """
       ]
     )