add caching

Author: madsodgaard

Plugins/JExtractSwiftPlugin/JExtractSwiftPlugin.swift

@@ -114,6 +114,13 @@ struct JExtractSwiftBuildToolPlugin: SwiftJavaPluginProtocol, BuildToolPlugin {
       outputSwiftDirectory.appending(path: "\(sourceModule.name)Module+SwiftJava.swift")
     ]
 
+    // Append any JNI cache files
+    if configuration?.mode == .jni {
+      outputSwiftFiles += [
+        outputSwiftDirectory.appending(path: "\(sourceModule.name)+JNICaches.swift")
+      ]
+    }
+
     // If the module uses 'Data' type, the thunk file is emitted as if 'Data' is declared
     // in that module. Declare the thunk file as the output.
     // FIXME: Make this conditional.

Samples/JExtractJNISampleApp/src/jmh/java/com/example/swift/EnumBenchmark.java

@@ -22,6 +22,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.concurrent.TimeUnit;
 
 @BenchmarkMode(Mode.AverageTime)
@@ -39,7 +40,7 @@ public static class BenchmarkState {
         @Setup(Level.Trial)
         public void beforeAll() {
             arena = new ConfinedSwiftMemorySession();
-            vehicle = Vehicle.motorbike("Yamaha", 900, arena);
+            vehicle = Vehicle.motorbike("Yamaha", 900, OptionalInt.empty(), arena);
         }
 
         @TearDown(Level.Trial)
@@ -49,7 +50,7 @@ public void afterAll() {
     }
 
     @Benchmark
-    public Vehicle.Motorbike java_copy(BenchmarkState state, Blackhole bh) {
+    public Vehicle.Motorbike getAssociatedValues(BenchmarkState state, Blackhole bh) {
         Vehicle.Motorbike motorbike = state.vehicle.getAsMotorbike().orElseThrow();
         bh.consume(motorbike.arg0());
         bh.consume(motorbike.horsePower());

Sources/JExtractSwiftLib/Convenience/String+Extensions.swift

@@ -24,6 +24,14 @@ extension String {
     return "\(f.uppercased())\(String(dropFirst()))"
   }
 
+  var firstCharacterLowercased: String {
+    guard let f = first else {
+      return self
+    }
+
+    return "\(f.lowercased())\(String(dropFirst()))"
+  }
+
   /// Returns whether the string is of the format `isX`
   var hasJavaBooleanNamingConvention: Bool {
     guard self.hasPrefix("is"), self.count > 2 else {

Sources/JExtractSwiftLib/JNI/JNICaching.swift

@@ -0,0 +1,19 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2025 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
+//
+//===----------------------------------------------------------------------===//
+
+enum JNICaching {
+  static func cacheName(for enumCase: ImportedEnumCase) -> String {
+    "\(enumCase.enumType.nominalTypeDecl.name.firstCharacterLowercased)\(enumCase.name.firstCharacterUppercased)Cache"
+  }
+}

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaBindingsPrinting.swift

@@ -259,14 +259,31 @@ extension JNISwift2JavaGenerator {
 
       // Print record
       printer.printBraceBlock("public record \(caseName)(\(members.joined(separator: ", "))) implements Case") { printer in
+        printer.printBraceBlock("static class $JNI") { printer in
+          printer.print("private static native void $nativeInit();")
+        }
+
+        // Used to ensure static initializer has been calling to trigger caching.
+        printer.print("static void $ensureInitialized() {}")
+
+        printer.printBraceBlock("static") { printer in
+          printer.print("$JNI.$nativeInit();")
+        }
+
         let nativeParameters = zip(translatedCase.translatedValues, translatedCase.parameterConversions).flatMap { value, conversion in
           ["\(conversion.native.javaType) \(value.parameter.name)"]
         }
 
         printer.print("record $NativeParameters(\(nativeParameters.joined(separator: ", "))) {}")
       }
 
-      self.printJavaBindingWrapperMethod(&printer, translatedCase.getAsCaseFunction)
+      self.printJavaBindingWrapperMethod(
+        &printer,
+        translatedCase.getAsCaseFunction,
+        prefix: { printer in
+          printer.print("\(caseName).$ensureInitialized();")
+        }
+      )
       printer.println()
     }
   }
@@ -335,7 +352,11 @@ extension JNISwift2JavaGenerator {
     printJavaBindingWrapperMethod(&printer, translatedDecl)
   }
 
-  private func printJavaBindingWrapperMethod(_ printer: inout CodePrinter, _ translatedDecl: TranslatedFunctionDecl) {
+  private func printJavaBindingWrapperMethod(
+    _ printer: inout CodePrinter,
+    _ translatedDecl: TranslatedFunctionDecl,
+    prefix: (inout CodePrinter) -> Void = { _ in }
+  ) {
     var modifiers = ["public"]
     if translatedDecl.isStatic {
       modifiers.append("static")
@@ -358,6 +379,7 @@ extension JNISwift2JavaGenerator {
     printer.printBraceBlock(
       "\(annotationsStr)\(modifiersStr) \(resultType) \(translatedDecl.name)(\(parametersStr))\(throwsClause)"
     ) { printer in
+      prefix(&printer)
       printDowncall(&printer, translatedDecl)
     }
 

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+SwiftThunkPrinting.swift

@@ -56,6 +56,8 @@ extension JNISwift2JavaGenerator {
         self.expectedOutputSwiftFiles.remove(moduleFilename)
       }
 
+      try self.writeJNICacheSource(&printer)
+
       for (_, ty) in self.analysis.importedTypes.sorted(by: { (lhs, rhs) in lhs.key < rhs.key }) {
         let fileNameBase = "\(ty.swiftNominal.qualifiedName)+SwiftJava"
         let filename = "\(fileNameBase).swift"
@@ -80,6 +82,34 @@ extension JNISwift2JavaGenerator {
     }
   }
 
+  private func writeJNICacheSource(_ printer: inout CodePrinter) throws {
+    printer.print("import JavaKit")
+    printer.println()
+    printer.printBraceBlock("enum JNICaches") { printer in
+      let enumCases = self.analysis.importedTypes.values.filter { $0.swiftNominal.kind == .enum }.flatMap(\.cases)
+      for enumCase in enumCases {
+        printer.print("static var \(JNICaching.cacheName(for: enumCase)): _JNICache!")
+      }
+      printer.println()
+      printer.printBraceBlock("func cleanup()") { printer in
+        for enumCase in enumCases {
+          printer.print("JNICaches.\(JNICaching.cacheName(for: enumCase)) = nil")
+        }
+      }
+    }
+
+    let fileName = "\(self.swiftModuleName)+JNICaches.swift"
+
+    if let outputFile = try printer.writeContents(
+      outputDirectory: self.swiftOutputDirectory,
+      javaPackagePath: nil,
+      filename: fileName
+    ) {
+      print("[swift-java] Generated: \(fileName.bold) (at \(outputFile))")
+      self.expectedOutputSwiftFiles.remove(fileName)
+    }
+  }
+
   private func printGlobalSwiftThunkSources(_ printer: inout CodePrinter) throws {
     printHeader(&printer)
 
@@ -156,12 +186,32 @@ extension JNISwift2JavaGenerator {
     printSwiftFunctionThunk(&printer, enumCase.caseFunction)
     printer.println()
 
+    // Print enum case native init
+    printEnumNativeInit(&printer, translatedCase)
+
     // Print getAsCase method
     if !translatedCase.translatedValues.isEmpty {
       printEnumGetAsCaseThunk(&printer, translatedCase)
     }
   }
 
+  private func printEnumNativeInit(_ printer: inout CodePrinter, _ enumCase: TranslatedEnumCase) {
+    printCDecl(
+      &printer,
+      javaMethodName: "$nativeInit",
+      parentName: "\(enumCase.original.enumType.nominalTypeDecl.name)$\(enumCase.name)$$JNI",
+      parameters: [],
+      resultType: .void
+    ) { printer in
+      // Setup caching
+      let nativeParametersClassName = "\(javaPackagePath)/\(enumCase.enumName)$\(enumCase.name)$$NativeParameters"
+      let methodSignature = MethodSignature(resultType: .void, parameterTypes: enumCase.parameterConversions.map(\.native.javaType))
+      let methods = #"[.init(name: "<init>", signature: "\#(methodSignature.mangledName)")]"#
+
+      printer.print(#"JNICaches.\#(JNICaching.cacheName(for: enumCase.original)) = _JNICache(environment: environment, className: "\#(nativeParametersClassName)", methods: \#(methods))"#)
+    }
+  }
+
   private func printEnumGetAsCaseThunk(
     _ printer: inout CodePrinter,
     _ enumCase: TranslatedEnumCase
@@ -176,14 +226,15 @@ extension JNISwift2JavaGenerator {
       }
       let caseNamesWithLet = caseNames.map { "let \($0)" }
       let methodSignature = MethodSignature(resultType: .void, parameterTypes: enumCase.parameterConversions.map(\.native.javaType))
-      // TODO: Caching of class and static method ID.
       printer.print(
         """
         guard case .\(enumCase.original.name)(\(caseNamesWithLet.joined(separator: ", "))) = \(selfPointer).pointee else {
           fatalError("Expected enum case '\(enumCase.original.name)', but was '\\(\(selfPointer).pointee)'!")
         }
-        let class$ = environment.interface.FindClass(environment, "\(javaPackagePath)/\(enumCase.enumName)$\(enumCase.name)$$NativeParameters")!
-        let constructorID$ = environment.interface.GetMethodID(environment, class$, "<init>", "\(methodSignature.mangledName)")!
+        let cache$ = JNICaches.\(JNICaching.cacheName(for: enumCase.original))!
+        let class$ = cache$.javaClass
+        let method$ = _JNICache.Method(name: "<init>", signature: "\(methodSignature.mangledName)")
+        let constructorID$ = cache$[method$]
         """
       )
       let upcallArguments = zip(enumCase.parameterConversions, caseNames).map { conversion, caseName in

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator.swift

@@ -72,6 +72,7 @@ package class JNISwift2JavaGenerator: Swift2JavaGenerator {
         return String(filePathPart.replacing(".swift", with: "+SwiftJava.swift"))
       })
       self.expectedOutputSwiftFiles.insert("\(translator.swiftModuleName)Module+SwiftJava.swift")
+      self.expectedOutputSwiftFiles.insert("\(translator.swiftModuleName)+JNICaches.swift")
 
       // FIXME: Can we avoid this?
       self.expectedOutputSwiftFiles.insert("Data+SwiftJava.swift")

Sources/JavaKit/Helpers/_JNICache.swift

@@ -0,0 +1,71 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+/// A cache used to hold references for JNI method and classes.
+///
+/// This type is used internally in by the outputted JExtract wrappers
+/// to improve performance of any JNI lookups.
+///
+/// - Warning: This type is inherently thread-unsafe.
+///            We assume that there exists **at most** one of these
+///            per Java `class`, and it must only be initialized and modified in the
+///            static initializer of that Java class.
+public final class _JNICache: @unchecked Sendable {
+  public struct Method: Hashable {
+    public let name: String
+    public let signature: String
+
+    public init(name: String, signature: String) {
+      self.name = name
+      self.signature = signature
+    }
+  }
+
+  var _class: jclass?
+  let environment: JNIEnvironment
+  let methods: [Method: jmethodID]
+
+  public var javaClass: jclass {
+    self._class!
+  }
+
+  public init(environment: UnsafeMutablePointer<JNIEnv?>!, className: String, methods: [Method]) {
+    guard let clazz = environment.interface.FindClass(environment, className) else {
+      fatalError("Class \(className) could not be found!")
+    }
+    self.environment = environment
+    self._class = environment.interface.NewGlobalRef(environment, clazz)!
+    self.methods = methods.reduce(into: [:]) { (result, method) in
+      if let methodID = environment.interface.GetMethodID(environment, clazz, method.name, method.signature) {
+        result[method] = methodID
+      } else {
+        fatalError("Method \(method.signature) with signature \(method.signature) not found in class \(className)")
+      }
+    }
+  }
+
+
+  public subscript(_ method: Method) -> jmethodID? {
+    methods[method]
+  }
+
+  func cleanup() {
+    environment.interface.DeleteGlobalRef(environment, self._class)
+  }
+
+  deinit {
+    cleanup()
+    self._class = nil
+  }
+}