jextract: pass arena from Swift instead of default interface for callbacks (#489)

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

Author: madsodgaard

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

@@ -77,7 +77,7 @@ void protocolMethod() {
     void protocolClassMethod() {
         try (var arena = SwiftArena.ofConfined()) {
             ProtocolA proto1 = ConcreteProtocolAB.init(10, 5, arena);
-            assertEquals(10, proto1.makeClass().getX());
+            assertEquals(10, proto1.makeClass(arena).getX());
         }
     }
 

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+InterfaceWrapperGeneration.swift

@@ -42,6 +42,7 @@ extension JNISwift2JavaGenerator {
     let protocolType: SwiftNominalType
     let functions: [Function]
     let variables: [Variable]
+    let importedType: ImportedNominalType
 
     var wrapperName: String {
       protocolType.nominalTypeDecl.javaInterfaceSwiftProtocolWrapperName
@@ -99,7 +100,8 @@ extension JNISwift2JavaGenerator {
       return JavaInterfaceSwiftWrapper(
         protocolType: SwiftNominalType(nominalTypeDecl: type.swiftNominal),
         functions: functions,
-        variables: variables
+        variables: variables,
+        importedType: type
       )
     }
 

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaBindingsPrinting.swift

@@ -214,7 +214,7 @@ extension JNISwift2JavaGenerator {
         }
         
         public static \(decl.swiftNominal.name) wrapMemoryAddressUnsafe(long selfPointer) {
-          return new \(decl.swiftNominal.name)(selfPointer, SwiftMemoryManagement.GLOBAL_SWIFT_JAVA_ARENA);
+          return new \(decl.swiftNominal.name)(selfPointer, SwiftMemoryManagement.DEFAULT_SWIFT_JAVA_AUTO_ARENA);
         }
         """
       )
@@ -531,17 +531,9 @@ extension JNISwift2JavaGenerator {
     // 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 shouldGenerateGlobalArenaVariation = config.effectiveMemoryManagementMode.requiresGlobalArena && translatedSignature.requiresSwiftArena
     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)
@@ -554,7 +546,7 @@ extension JNISwift2JavaGenerator {
       }
 
       printer.printBraceBlock("\(annotationsStr)\(modifiers.joined(separator: " ")) \(resultType) \(translatedDecl.name)(\(parametersStr))\(throwsClause)") { printer in
-        let globalArenaName = "SwiftMemoryManagement.GLOBAL_SWIFT_JAVA_ARENA"
+        let globalArenaName = "SwiftMemoryManagement.DEFAULT_SWIFT_JAVA_AUTO_ARENA"
         let arguments = translatedDecl.translatedFunctionSignature.parameters.map(\.parameter.name) + [globalArenaName]
         let call = "\(translatedDecl.name)(\(arguments.joined(separator: ", ")))"
         if translatedDecl.translatedFunctionSignature.resultType.javaType.isVoid {

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaTranslation.swift

@@ -1317,5 +1317,9 @@ extension JNISwift2JavaGenerator {
 
     // FIXME: Remove once we support protocol variables
     case protocolVariablesNotSupported
+
+    /// We cannot generate interface wrappers for
+    /// protocols that we unable to be jextracted.
+    case protocolWasNotExtracted
   }
 }

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+SwiftThunkPrinting.swift

@@ -124,9 +124,9 @@ extension JNISwift2JavaGenerator {
       printer.print("var \(translatedWrapper.javaInterfaceVariableName): \(translatedWrapper.javaInterfaceName) { get }")
     }
     printer.println()
-    printer.printBraceBlock("extension \(translatedWrapper.wrapperName)") { printer in
+    try printer.printBraceBlock("extension \(translatedWrapper.wrapperName)") { printer in
       for function in translatedWrapper.functions {
-        printInterfaceWrapperFunctionImpl(&printer, function, inside: translatedWrapper)
+        try printInterfaceWrapperFunctionImpl(&printer, function, inside: translatedWrapper)
         printer.println()
       }
 
@@ -142,16 +142,29 @@ extension JNISwift2JavaGenerator {
     _ printer: inout CodePrinter,
     _ function: JavaInterfaceSwiftWrapper.Function,
     inside wrapper: JavaInterfaceSwiftWrapper
-  ) {
+  ) throws {
+    guard let protocolMethod = wrapper.importedType.methods.first(where: { $0.functionSignature == function.originalFunctionSignature }) else {
+      fatalError("Failed to find protocol method")
+    }
+    guard let translatedDecl = self.translatedDecl(for: protocolMethod) else {
+      throw JavaTranslationError.protocolWasNotExtracted
+    }
+
     printer.printBraceBlock(function.swiftDecl.signatureString) { printer in
-      let upcallArguments = zip(
+      var upcallArguments = zip(
         function.originalFunctionSignature.parameters,
         function.parameterConversions
       ).map { param, conversion in
         // Wrap-java does not extract parameter names, so no labels
         conversion.render(&printer, param.parameterName!)
       }
 
+      // If the underlying translated method requires
+      // a SwiftArena, we pass in the global arena
+      if translatedDecl.translatedFunctionSignature.requiresSwiftArena {
+        upcallArguments.append("JNI.shared.defaultAutoArena")
+      }
+
       let tryClause = function.originalFunctionSignature.isThrowing ? "try " : ""
       let javaUpcall = "\(tryClause)\(wrapper.javaInterfaceVariableName).\(function.swiftFunctionName)(\(upcallArguments.joined(separator: ", ")))"
 

Sources/SwiftJava/AnyJavaObject.swift

@@ -77,6 +77,11 @@ extension AnyJavaObject {
     return String(seq)
   }
 
+  /// The mangled name for this java class
+  public static var mangledName: String {
+    "L\(fullJavaClassNameWithSlashes);"
+  }
+
   /// Initialize a Java object from its instance.
   public init(javaThis: jobject, environment: JNIEnvironment) {
     self.init(javaHolder: JavaObjectHolder(object: javaThis, environment: environment))

Sources/SwiftJavaRuntimeSupport/JNI.swift

@@ -15,13 +15,38 @@
 import SwiftJava
 import CSwiftJavaJNI
 
-final class JNI {
-  static var shared: JNI!
+/// A type that represents the shared JNI environment
+/// used to share any global JNI variables.
+///
+/// This is initialized when the `JNI_OnLoad` is triggered,
+/// which happens when you call `System.loadLibrary(...)`
+/// from Java.
+public final class JNI {
+  /// The shared JNI object, initialized by `JNI_OnLoad`
+  public fileprivate(set) static var shared: JNI!
 
-  let applicationClassLoader: JavaClassLoader
+  /// The default application class loader
+  public let applicationClassLoader: JavaClassLoader
+
+  /// The default auto arena of SwiftKitCore
+  public let defaultAutoArena: JavaSwiftArena
 
   init(fromVM javaVM: JavaVirtualMachine) {
-    self.applicationClassLoader = try! JavaClass<JavaThread>(environment: javaVM.environment()).currentThread().getContextClassLoader()
+    let environment = try! javaVM.environment()
+
+    self.applicationClassLoader = try! JavaClass<JavaThread>(environment: environment).currentThread().getContextClassLoader()
+
+    // Find global arena
+    let swiftMemoryClass = environment.interface.FindClass(environment, "org/swift/swiftkit/core/SwiftMemoryManagement")!
+    let arenaFieldID = environment.interface.GetStaticFieldID(
+        environment,
+        swiftMemoryClass,
+        "DEFAULT_SWIFT_JAVA_AUTO_ARENA",
+        JavaSwiftArena.mangledName
+    )
+    let localObject = environment.interface.GetStaticObjectField(environment, swiftMemoryClass, arenaFieldID)!
+    self.defaultAutoArena = JavaSwiftArena(javaThis: localObject, environment: environment)
+    environment.interface.DeleteLocalRef(environment, localObject)
   }
 }
 

Sources/SwiftJavaRuntimeSupport/generated/JavaSwiftArena.swift

@@ -0,0 +1,18 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftJava
+
+@JavaInterface("org.swift.swiftkit.core.SwiftArena")
+public struct JavaSwiftArena {}

Sources/SwiftJavaRuntimeSupport/swift-java.config

@@ -0,0 +1,6 @@
+{
+  "classes" : {
+    "org.swift.swiftkit.core.JNISwiftInstance" : "JavaJNISwiftInstance",
+    "org.swift.swiftkit.core.SwiftArena" : "JavaSwiftArena"
+  }
+}

SwiftKitCore/src/main/java/org/swift/swiftkit/core/SwiftMemoryManagement.java

@@ -15,5 +15,5 @@
 package org.swift.swiftkit.core;
 
 public class SwiftMemoryManagement {
-    public static final SwiftArena GLOBAL_SWIFT_JAVA_ARENA = SwiftArena.ofAuto();
+    public static final SwiftArena DEFAULT_SWIFT_JAVA_AUTO_ARENA = SwiftArena.ofAuto();
 }