make SwiftJava a dynamic lib and add class lookup fallbacks

Package.swift

@@ -100,6 +100,7 @@ let package = Package(
     // ==== SwiftJava (i.e. calling Java directly Swift utilities)
     .library(
       name: "SwiftJava",
+      type: .dynamic,
       targets: ["SwiftJava"]
     ),
 

Sources/JExtractSwiftLib/FFM/FFMSwift2JavaGenerator.swift

@@ -195,6 +195,7 @@ extension FFMSwift2JavaGenerator {
         private static final boolean INITIALIZED_LIBS = initializeLibs();
         static boolean initializeLibs() {
             System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_CORE);
+            System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_JAVA);
             System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_RUNTIME_FUNCTIONS);
             System.loadLibrary(LIB_NAME);
             return true;
@@ -345,6 +346,7 @@ extension FFMSwift2JavaGenerator {
         private static SymbolLookup getSymbolLookup() {
             if (SwiftLibraries.AUTO_LOAD_LIBS) {
                 System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_CORE);
+                System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_JAVA);
                 System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_RUNTIME_FUNCTIONS);
                 System.loadLibrary(LIB_NAME);
             }

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaBindingsPrinting.swift

@@ -99,6 +99,7 @@ extension JNISwift2JavaGenerator {
         static final String LIB_NAME = "\(swiftModuleName)";
         
         static {
+          System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_JAVA);
           System.loadLibrary(LIB_NAME);
         }
         """
@@ -169,6 +170,7 @@ extension JNISwift2JavaGenerator {
         @SuppressWarnings("unused")
         private static final boolean INITIALIZED_LIBS = initializeLibs();
         static boolean initializeLibs() {
+            System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_JAVA);
             System.loadLibrary(LIB_NAME);
             return true;
         }

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+SwiftThunkPrinting.swift

@@ -162,7 +162,7 @@ extension JNISwift2JavaGenerator {
       // If the underlying translated method requires
       // a SwiftArena, we pass in the global arena
       if translatedDecl.translatedFunctionSignature.requiresSwiftArena {
-        upcallArguments.append("JNI.shared.defaultAutoArena")
+        upcallArguments.append("JavaSwiftArena.defaultAutoArena")
       }
 
       let tryClause = function.originalFunctionSignature.isThrowing ? "try " : ""
@@ -201,8 +201,6 @@ extension JNISwift2JavaGenerator {
   private func printGlobalSwiftThunkSources(_ printer: inout CodePrinter) throws {
     printHeader(&printer)
 
-    printJNIOnLoad(&printer)
-
     for decl in analysis.importedGlobalFuncs {
       printSwiftFunctionThunk(&printer, decl)
       printer.println()
@@ -214,18 +212,6 @@ extension JNISwift2JavaGenerator {
     }
   }
 
-  private func printJNIOnLoad(_ printer: inout CodePrinter) {
-    printer.print(
-      """
-      @_cdecl("JNI_OnLoad")
-      func JNI_OnLoad(javaVM: JavaVMPointer, reserved: UnsafeMutableRawPointer) -> jint {
-        SwiftJavaRuntimeSupport._JNI_OnLoad(javaVM, reserved)
-        return JNI_VERSION_1_6
-      }
-      """
-    )
-  }
-
   private func printNominalTypeThunks(_ printer: inout CodePrinter, _ type: ImportedNominalType) throws {
     printHeader(&printer)
 

Sources/SwiftJava/AnyJavaObject.swift

@@ -102,23 +102,39 @@ extension AnyJavaObject {
     in environment: JNIEnvironment,
     _ body: (jclass) throws -> Result
   ) throws -> Result {
-    let resolvedClass = try environment.translatingJNIExceptions {
-      environment.interface.FindClass(
-        environment,
-        fullJavaClassNameWithSlashes
-      )
-    }!
-    return try body(resolvedClass)
+    do {
+      let resolvedClass = try environment.translatingJNIExceptions {
+        environment.interface.FindClass(
+          environment,
+          fullJavaClassNameWithSlashes
+        )
+      }!
+      return try body(resolvedClass)
+    } catch {
+      // If we are in a Java environment where we have loaded
+      // SwiftJava dynamically, we have access to the application class loader
+      // so lets try that as as a fallback
+      if let applicationClassLoader = JNI.shared?.applicationClassLoader {
+        return try _withJNIClassFromCustomClassLoader(
+          applicationClassLoader,
+          in: environment
+        ) { applicationLoadedClass in
+          return try body(applicationLoadedClass)
+        }
+      } else {
+        throw error
+      }
+    }
   }
 
   /// Retrieve the Java class for this type using a specific class loader.
   private static func _withJNIClassFromCustomClassLoader<Result>(
     _ classLoader: JavaClassLoader,
     in environment: JNIEnvironment,
-    _ body: (jclass?) throws -> Result
+    _ body: (jclass) throws -> Result
   ) throws -> Result {
-    let resolvedClass = try? classLoader.findClass(fullJavaClassName)
-    return try body(resolvedClass?.javaThis)
+    let resolvedClass = try classLoader.findClass(fullJavaClassName)
+    return try body(resolvedClass!.javaThis)
   }
 
   /// Retrieve the Java class for this type and execute body().
@@ -129,16 +145,15 @@ extension AnyJavaObject {
   ) throws -> Result {
     if let AnyJavaObjectWithCustomClassLoader = self as? AnyJavaObjectWithCustomClassLoader.Type,
        let customClassLoader = try AnyJavaObjectWithCustomClassLoader.getJavaClassLoader(in: environment) {
-      try _withJNIClassFromCustomClassLoader(customClassLoader, in: environment) { clazz in
-        guard let clazz else {
-          // If the custom class loader did not find the class
-          // let's look in the default class loader.
-          return try _withJNIClassFromDefaultClassLoader(in: environment, body)
+      do {
+        return try _withJNIClassFromCustomClassLoader(customClassLoader, in: environment) { clazz in
+          return try body(clazz)
         }
-        return try body(clazz)
+      } catch {
+        return try _withJNIClassFromDefaultClassLoader(in: environment, body)
       }
     } else {
-      try _withJNIClassFromDefaultClassLoader(in: environment, body)
+      return try _withJNIClassFromDefaultClassLoader(in: environment, body)
     }
   }
 }

Sources/SwiftJavaRuntimeSupport/JNI.swift → Sources/SwiftJava/JNI.swift

@@ -12,7 +12,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-import SwiftJava
 import CSwiftJavaJNI
 
 /// A type that represents the shared JNI environment
@@ -21,36 +20,28 @@ import CSwiftJavaJNI
 /// This is initialized when the `JNI_OnLoad` is triggered,
 /// which happens when you call `System.loadLibrary(...)`
 /// from Java.
-public final class JNI {
+package final class JNI {
   /// The shared JNI object, initialized by `JNI_OnLoad`
-  public fileprivate(set) static var shared: JNI!
+  ///
+  /// This may be `nil` in the case where `SwiftJava` is not loaded as a dynamic lib
+  /// by the Java sources.
+  package fileprivate(set) static var shared: JNI?
 
   /// The default application class loader
-  public let applicationClassLoader: JavaClassLoader
-
-  /// The default auto arena of SwiftKitCore
-  public let defaultAutoArena: JavaSwiftArena
+  package let applicationClassLoader: JavaClassLoader
 
   init(fromVM javaVM: JavaVirtualMachine) {
+    // Update the global JavaVM
+    JavaVirtualMachine.sharedJVM.withLock {
+      $0 = javaVM
+    }
     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)
   }
 }
 
-// Called by generated code, and not automatically by Java.
-public func _JNI_OnLoad(_ javaVM: JavaVMPointer, _ reserved: UnsafeMutableRawPointer) {
+@_cdecl("JNI_OnLoad")
+func SwiftJava_JNI_OnLoad(javaVM: JavaVMPointer, reserved: UnsafeMutableRawPointer) -> jint {
   JNI.shared = JNI(fromVM: JavaVirtualMachine(adoptingJVM: javaVM))
+  return JNI_VERSION_1_6
 }

Sources/SwiftJava/JVM/JavaVirtualMachine.swift

@@ -220,7 +220,7 @@ extension JavaVirtualMachine {
   /// TODO: If the use of the lock itself ends up being slow, we could
   /// use an atomic here instead because our access pattern is fairly
   /// simple.
-  private static let sharedJVM: LockedState<JavaVirtualMachine?> = .init(initialState: nil)
+  static let sharedJVM: LockedState<JavaVirtualMachine?> = .init(initialState: nil)
 
   /// Access the shared Java Virtual Machine instance.
   ///

Sources/SwiftJavaRuntimeSupport/_JNIMethodIDCache.swift

@@ -55,7 +55,8 @@ public final class _JNIMethodIDCache: Sendable {
       // Clear any ClassNotFound exceptions from FindClass
       environment.interface.ExceptionClear(environment)
 
-      if let javaClass = try? JNI.shared.applicationClassLoader.loadClass(
+      // OK to force unwrap, we are in a jextract environment.
+      if let javaClass = try? JNI.shared!.applicationClassLoader.loadClass(
         className.replacingOccurrences(of: "/", with: ".")
       ) {
         clazz = javaClass.javaThis

Sources/SwiftJavaRuntimeSupport/generated/JavaJNISwiftInstance.swift

@@ -20,6 +20,7 @@ public struct JavaJNISwiftInstance: AnyJavaObjectWithCustomClassLoader {
   public func memoryAddress() -> Int64
 
   public static func getJavaClassLoader(in environment: JNIEnvironment) throws -> JavaClassLoader! {
-    JNI.shared.applicationClassLoader
+    // OK to force unwrap, we are in a jextract environment.
+    JNI.shared!.applicationClassLoader
   }
 }

Sources/SwiftJavaRuntimeSupport/generated/JavaSwiftArena.swift

@@ -16,3 +16,11 @@ import SwiftJava
 
 @JavaInterface("org.swift.swiftkit.core.SwiftArena")
 public struct JavaSwiftArena {}
+
+extension JavaSwiftArena {
+  /// A cache for the default auto arena found in SwiftKitCore
+  public static internal(set) var defaultAutoArena: JavaSwiftArena = {
+    let swiftMemoryClass = try! JavaClass<SwiftJavaRuntimeSupport.JavaSwiftMemoryManagement>()
+    return swiftMemoryClass.defaultSwiftJavaAutoArena
+  }()
+}

Sources/SwiftJavaRuntimeSupport/generated/JavaSwiftMemoryManagement.swift

@@ -0,0 +1,23 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+
+@JavaClass("org.swift.swiftkit.core.SwiftMemoryManagement")
+public class JavaSwiftMemoryManagement: JavaObject {}
+
+extension JavaClass<JavaSwiftMemoryManagement> {
+  @JavaStaticField("DEFAULT_SWIFT_JAVA_AUTO_ARENA", isFinal: true)
+  var defaultSwiftJavaAutoArena: JavaSwiftArena!
+}

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

@@ -28,6 +28,7 @@ public final class SwiftLibraries {
     public static final String LIB_NAME_SWIFT_CORE = "swiftCore";
     public static final String LIB_NAME_SWIFT_CONCURRENCY = "swift_Concurrency";
     public static final String LIB_NAME_SWIFT_RUNTIME_FUNCTIONS = "SwiftRuntimeFunctions";
+    public static final String LIB_NAME_SWIFT_JAVA = "SwiftJava";
 
     /** 
      * Allows for configuration if jextracted types should automatically attempt to load swiftCore and the library type is from.
@@ -44,6 +45,7 @@ public final class SwiftLibraries {
 
     public static boolean loadLibraries(boolean loadSwiftRuntimeFunctions) {
         System.loadLibrary(LIB_NAME_SWIFT_CORE);
+        System.loadLibrary(LIB_NAME_SWIFT_JAVA);
         if (loadSwiftRuntimeFunctions) {
             System.loadLibrary(LIB_NAME_SWIFT_RUNTIME_FUNCTIONS);
         }

Tests/JExtractSwiftTests/JNI/JNIClassTests.swift

@@ -66,6 +66,7 @@ struct JNIClassTests {
           @SuppressWarnings("unused")
           private static final boolean INITIALIZED_LIBS = initializeLibs();
           static boolean initializeLibs() {
+            System.loadLibrary(SwiftLibraries.LIB_NAME_SWIFT_JAVA);
             System.loadLibrary(LIB_NAME);
             return true;
           }

Tests/JExtractSwiftTests/JNI/JNIProtocolTests.swift

@@ -307,7 +307,7 @@ struct JNIProtocolTests {
             let cClass = try! JavaClass<JavaSomeClass>(environment: JavaVirtualMachine.shared().environment())
             let cPointer = UnsafeMutablePointer<SomeClass>.allocate(capacity: 1)
             cPointer.initialize(to: c)
-            guard let unwrapped$ = _javaSomeProtocolInterface.withObject(cClass.wrapMemoryAddressUnsafe(Int64(Int(bitPattern: cPointer))), JNI.shared.defaultAutoArena) else {
+            guard let unwrapped$ = _javaSomeProtocolInterface.withObject(cClass.wrapMemoryAddressUnsafe(Int64(Int(bitPattern: cPointer))), JavaSwiftArena.defaultAutoArena) else {
               fatalError("Upcall to withObject unexpectedly returned nil")
             }
             let result$MemoryAddress$ = unwrapped$.as(JavaJNISwiftInstance.self)!.memoryAddress()