perf: use nonvirtual JNI calls for known Kotlin implementations

packages/nitrogen/src/syntax/kotlin/FbjniHybridObject.ts

@@ -8,15 +8,36 @@ import { getKotlinPropertyAccessors, KotlinInstanceMethod } from './KotlinInstan
 import { wrapKotlinType } from './wrapKotlinType.js'
 import { getImportsCode } from '../getImportCode.js'
 import { createKotlinHybridClassBindings } from './KotlinHybridClassBindings.js'
+import type { KotlinHybridObjectType } from './types/KotlinHybridObjectType.js'
 
-export function createFbjniHybridObject(spec: HybridObjectSpec<'kotlin'>): SourceFile[] {
+interface FbjniHybridObjectOptions {
+  spec: HybridObjectSpec<'kotlin'>
+  hybridObject: KotlinHybridObjectType
+}
+
+export function createFbjniHybridObject({ spec, hybridObject }: FbjniHybridObjectOptions): SourceFile[] {
   const name = getHybridObjectName(spec.name)
   const properties = spec.properties
   const nativeMethods = spec.methods
+  const nonvirtual = hybridObject.canUseNonvirtualJniCalls
   const propertyAccessors = properties.flatMap(getKotlinPropertyAccessors)
-  const propertyDefinitions = propertyAccessors.flatMap((accessor) => accessor.createBridge().cppMethods)
-  const methodDefinitions = nativeMethods.flatMap((method) => new KotlinInstanceMethod({ method: method }).createBridge().cppMethods)
+  const propertyDefinitions = propertyAccessors.flatMap((accessor) => accessor.createBridge({ nonvirtual: nonvirtual }).cppMethods)
+  const methodDefinitions = nativeMethods.flatMap(
+    (method) => new KotlinInstanceMethod({ method: method }).createBridge({ nonvirtual: nonvirtual }).cppMethods
+  )
   const bindings = createKotlinHybridClassBindings(spec)
+  let getClassBody: string
+  if (nonvirtual) {
+    getClassBody = `
+static const auto __field = CxxPart::javaClassStatic()->getStaticField<jni::JClass>("implementationClass_cxx");
+static const auto __class = jni::make_global(CxxPart::javaClassStatic()->getStaticFieldValue(__field));
+return __class;
+`.trim()
+  } else {
+    getClassBody = `
+return ${name.JHybridTSpec}::JavaPart::javaClassStatic();
+`.trim()
+  }
   const registrations: string[] = []
   for (const property of properties) {
     const accessors = [`&${name.JHybridTSpec}::${property.getGetterName('other')}`]
@@ -96,6 +117,7 @@ namespace ${cxxNamespace} {
       static constexpr auto kJavaDescriptor = "L${cxxPartJniClassDescriptor};";
       static jni::local_ref<jhybriddata> initHybrid(jni::alias_ref<jhybridobject> jThis);
       static void registerNatives();
+      static jni::alias_ref<jni::JClass> getClass();
       using HybridBase::HybridBase;
     protected:
       std::shared_ptr<JHybridObject> createHybridObject(const jni::local_ref<JHybridObject::JavaPart>& javaPart) override;
@@ -203,6 +225,10 @@ namespace ${cxxNamespace} {
     return makeCxxInstance(jThis);
   }
 
+  jni::alias_ref<jni::JClass> ${name.JHybridTSpec}::CxxPart::getClass() {
+    ${indent(getClassBody, '    ')}
+  }
+
   std::shared_ptr<JHybridObject> ${name.JHybridTSpec}::CxxPart::createHybridObject(const jni::local_ref<JHybridObject::JavaPart>& javaPart) {
     auto castJavaPart = jni::dynamic_ref_cast<${name.JHybridTSpec}::JavaPart>(javaPart);
     if (castJavaPart == nullptr) [[unlikely]] {

packages/nitrogen/src/syntax/kotlin/KotlinConstructor.ts

@@ -6,7 +6,7 @@ import { HybridObjectType } from '../types/HybridObjectType.js'
 import { KotlinNativeHybridObjectType } from './types/KotlinNativeHybridObjectType.js'
 import {
   createCppJniCallBody,
-  createKotlinForwarderBody,
+  createKotlinBridgeBody,
   getKotlinImplementation,
   prepareKotlinMethod,
   type KotlinMethodBridge,
@@ -60,7 +60,7 @@ export class KotlinConstructor extends Constructor {
       jniMethodLookup: `${this.jniClass}->getStaticMethod`,
       jniReceiver: this.jniClass,
     })
-    const kotlinBody = createKotlinForwarderBody({
+    const kotlinBody = createKotlinBridgeBody({
       callable: this.callable,
       getKotlinInvocation: (args) => `${this.implementationClassName}(${args.join(', ')})`,
     })

packages/nitrogen/src/syntax/kotlin/KotlinHybridObject.ts

@@ -19,6 +19,7 @@ import type { KotlinStaticProperty } from './KotlinStaticProperty.js'
 import { wrapKotlinType } from './wrapKotlinType.js'
 import { getImportCode } from '../getImportCode.js'
 import { getKotlinNativeType } from './getKotlinNativeType.js'
+import { KotlinHybridObjectType } from './types/KotlinHybridObjectType.js'
 
 export function createKotlinHybridObject(publicSpec: HybridObjectSpec<'kotlin'>): SourceFile[] {
   const nativeProperties = publicSpec.properties.map((property) => {
@@ -29,6 +30,7 @@ export function createKotlinHybridObject(publicSpec: HybridObjectSpec<'kotlin'>)
     ...publicSpec,
     properties: nativeProperties,
   }
+  const hybridObject = new KotlinHybridObjectType(new HybridObjectType(spec))
   const name = getHybridObjectName(spec.name)
   const properties = spec.properties.map(getKotlinPropertyCode).join('\n\n')
   const methods = spec.methods.map(getKotlinMethodCode).join('\n\n')
@@ -40,6 +42,7 @@ export function createKotlinHybridObject(publicSpec: HybridObjectSpec<'kotlin'>)
     })
     .join('\n\n')
   const bridgeCompanion = getKotlinCompanion(spec)
+  const cxxPartMembers = getKotlinCxxPartMembers({ spec: spec, hybridObject: hybridObject })
 
   const extraImports: SourceImport[] = [
     ...spec.properties.flatMap((p) => p.getRequiredImports('kotlin')),
@@ -48,6 +51,7 @@ export function createKotlinHybridObject(publicSpec: HybridObjectSpec<'kotlin'>)
     ...spec.staticMethods.flatMap((m) => m.getRequiredImports('kotlin')),
     ...(spec.constructor?.getRequiredImports('kotlin') ?? []),
     ...bridgeCompanion.imports,
+    ...cxxPartMembers.imports,
     ...spec.baseTypes.flatMap((b) => new HybridObjectType(b).getRequiredImports('kotlin')),
     {
       name: 'com.margelo.nitro.core.HybridObject',
@@ -124,9 +128,7 @@ abstract class ${name.HybridTSpec}: ${supertypes} {
   @DoNotStrip
   @Keep
   protected open class CxxPart(javaPart: ${name.HybridTSpec}): ${cxxPartBase}(javaPart) {
-    // C++ ${name.JHybridTSpec}::CxxPart::initHybrid(...)
-    @FastNative
-    external override fun initHybrid(): HybridData
+    ${indent(cxxPartMembers.code, '    ')}
   }
   override fun createCxxPart(): CxxPart {
     return CxxPart(this)
@@ -147,7 +149,7 @@ abstract class ${name.HybridTSpec}: ${supertypes} {
   `.trim()
 
   // 2. Create C++ (fbjni) bindings
-  const cppFiles = createFbjniHybridObject(spec)
+  const cppFiles = createFbjniHybridObject({ spec: spec, hybridObject: hybridObject })
 
   // 3. Create enums or structs in Kotlin
   const nativeTypes = getAllTypes(spec)
@@ -181,8 +183,39 @@ interface KotlinCompanionGroup {
   callables: (KotlinConstructor | KotlinStaticMethod | KotlinStaticProperty)[]
 }
 
+interface KotlinCxxPartMembersOptions {
+  spec: HybridObjectSpec<'kotlin'>
+  hybridObject: KotlinHybridObjectType
+}
+
+function getKotlinCxxPartMembers({ spec, hybridObject }: KotlinCxxPartMembersOptions): KotlinCompanion {
+  const name = getHybridObjectName(spec.name)
+  const initHybridCode = `
+// C++ ${name.JHybridTSpec}::CxxPart::initHybrid(...)
+@FastNative
+external override fun initHybrid(): HybridData
+`.trim()
+  const sections = [initHybridCode]
+  const imports: SourceImport[] = []
+  if (hybridObject.canUseNonvirtualJniCalls) {
+    const implementation = getKotlinImplementation(spec)
+    imports.push(implementation.kotlinImport)
+    // A class literal keeps the implementation reachable after R8 renaming.
+    const implementationClassCode = `
+companion object {
+  @DoNotStrip
+  @Keep
+  @JvmField
+  val implementationClass_cxx: Class<out ${name.HybridTSpec}> = ${implementation.className}::class.java
+}
+`.trim()
+    sections.push(implementationClassCode)
+  }
+  return { code: sections.join('\n\n'), imports: imports }
+}
+
 /**
- * The companion object holds the forwarders JNI calls to construct the class or call its static methods.
+ * The companion object holds the bridges JNI calls to construct the class or call its static methods.
  * They let Kotlin check the implementation at compile time, keep it from being stripped by R8,
  * and expose actual JVM static methods through `@JvmStatic`.
  * TODO: Remove the Kotlin-side conversions (`convert(..., { in: 'kotlin' })`, for `UInt64`
@@ -211,56 +244,56 @@ protected const val TAG = "${name.HybridTSpec}"
   // Only static members require a companion; instance-only specs may have no configured implementation.
   if (spec.staticMethods.length > 0 || spec.staticProperties.length > 0) {
     const implementation = getKotlinImplementation(spec)
-    // Check the configured companion's contract; forwarders still call it directly.
+    // Check the configured companion's contract; bridges still call it directly.
     const companionRequirement = `
 private val implementationCompanion: Companion
   get() = ${implementation.className}
 `.trim()
     sections.push(companionRequirement)
   }
   for (const group of groups) {
-    const forwarders = group.callables
+    const bridges = group.callables
       .flatMap((callable) => callable.createBridge().kotlinMethods)
       .map(({ method, body }) => method.getDefinition('kotlin', { body: body }))
     const groupCode = `
 // ${group.comment}
-${forwarders.join('\n\n')}
+${bridges.join('\n\n')}
 `.trim()
     sections.push(groupCode)
   }
   const imports: SourceImport[] = []
   if (groups.length > 0) {
-    // The forwarders call the implementation by its class name.
+    // The bridges call the implementation by its class name.
     imports.push(getKotlinImplementation(spec).kotlinImport)
   }
   return { code: sections.join('\n\n'), imports: imports }
 }
 
 function getKotlinMethodCode(method: InstanceMethod): string {
   const kotlinMethod = new KotlinInstanceMethod({ method: method })
-  // Without a forwarder, JNI calls the abstract member directly.
+  // Without a bridge, JNI calls the abstract member directly.
   const declaration = new InstanceMethod(method.name, method.returnType, method.parameters, {
-    doNotStrip: !kotlinMethod.needsKotlinForwarder,
+    doNotStrip: !kotlinMethod.needsKotlinBridge,
     pureVirtual: true,
   })
-  const forwarders = kotlinMethod
+  const bridges = kotlinMethod
     .createBridge()
-    .kotlinMethods.map(({ method: forwarder, body }) => forwarder.getDefinition('kotlin', { body: body }))
-  return [declaration.getDeclaration('kotlin'), ...forwarders].join('\n\n')
+    .kotlinMethods.map(({ method: bridge, body }) => bridge.getDefinition('kotlin', { body: body }))
+  return [declaration.getDeclaration('kotlin'), ...bridges].join('\n\n')
 }
 
 function getKotlinPropertyCode(property: Property): string {
   const accessors = getKotlinPropertyAccessors(property)
-  // Without forwarders, JNI calls the abstract accessors directly.
-  const isCalledFromJni = accessors.every((accessor) => !accessor.needsKotlinForwarder)
+  // Without bridges, JNI calls the abstract accessors directly.
+  const isCalledFromJni = accessors.every((accessor) => !accessor.needsKotlinBridge)
   const declaration = property.getCode('kotlin', {
     doNotStrip: isCalledFromJni,
     virtual: true,
   })
-  const forwarders = accessors
+  const bridges = accessors
     .flatMap((accessor) => accessor.createBridge().kotlinMethods)
     .map(({ method, body }) => method.getDefinition('kotlin', { body: body }))
-  return [declaration, ...forwarders].join('\n\n')
+  return [declaration, ...bridges].join('\n\n')
 }
 
 function getImplementedInterfaceImport(implementedInterface: KotlinImplementedInterface): SourceImport<'kotlin'> {

packages/nitrogen/src/syntax/kotlin/KotlinInstanceMethod.ts

@@ -3,7 +3,7 @@ import { Parameter } from '../Parameter.js'
 import type { Property } from '../Property.js'
 import {
   createCppJniCallBody,
-  createKotlinForwarderBody,
+  createKotlinBridgeBody,
   prepareKotlinMethod,
   type KotlinMethodBridge,
   type KotlinMethodDefinition,
@@ -21,9 +21,13 @@ interface KotlinInstanceMethodOptions {
   accessor?: KotlinAccessor
 }
 
+interface KotlinInstanceBridgeOptions {
+  nonvirtual?: boolean
+}
+
 /**
  * The Kotlin member that C++ calls through JNI.
- * If values need conversion, this is a private `_cxx` forwarder;
+ * If values need conversion, this is a private `_cxx` bridge;
  * otherwise JNI calls the public abstract member directly.
  */
 export class KotlinInstanceMethod extends InstanceMethod {
@@ -70,35 +74,40 @@ export class KotlinInstanceMethod extends InstanceMethod {
     this.callable = callable
   }
 
-  /** Whether JNI calls a private forwarder instead of the public member. */
-  get needsKotlinForwarder(): boolean {
+  /** Whether JNI calls a private bridge instead of the public member. */
+  get needsKotlinBridge(): boolean {
     return this.callable.needsKotlinConversion
   }
 
-  createBridge(): KotlinMethodBridge {
+  createBridge({ nonvirtual = false }: KotlinInstanceBridgeOptions = {}): KotlinMethodBridge {
     const nativeMethod = this.nativeMethod
     const cppMethod = new InstanceMethod(nativeMethod.name, nativeMethod.returnType, nativeMethod.parameters, {
       generated: true,
     })
+    const isNonvirtual = nonvirtual || this.needsKotlinBridge
+    const jniMethod = isNonvirtual ? 'getNonvirtualMethod' : 'getMethod'
+    const jniMethodLookup = `CxxPart::getClass()->${jniMethod}`
+    const jniDeclaringClass = isNonvirtual ? 'CxxPart::getClass()' : undefined
     const cppBody = createCppJniCallBody({
       callable: this.callable,
       jvmName: this.name,
-      jniMethodLookup: '_javaPart->javaClassStatic()->getMethod',
+      jniMethodLookup: jniMethodLookup,
       jniReceiver: '_javaPart',
+      jniDeclaringClass: jniDeclaringClass,
     })
     const kotlinMethods: KotlinMethodDefinition[] = []
-    if (this.needsKotlinForwarder) {
-      kotlinMethods.push({ method: this, body: this.getKotlinForwarderBody() })
+    if (this.needsKotlinBridge) {
+      kotlinMethods.push({ method: this, body: this.getKotlinBridgeBody() })
     }
     return {
       kotlinMethods: kotlinMethods,
       cppMethods: [{ method: cppMethod, body: cppBody }],
     }
   }
 
-  private getKotlinForwarderBody(): string {
+  private getKotlinBridgeBody(): string {
     if (this.accessor == null) {
-      return createKotlinForwarderBody({
+      return createKotlinBridgeBody({
         callable: this.callable,
         getKotlinInvocation: (args) => `${this.nativeMethod.name}(${args.join(', ')})`,
       })

packages/nitrogen/src/syntax/kotlin/KotlinMethodHelpers.ts

@@ -52,13 +52,13 @@ export function prepareKotlinMethod(nativeMethod: Method): KotlinMethodInfo {
   }
 }
 
-interface KotlinForwarderBodyOptions {
+interface KotlinBridgeBodyOptions {
   callable: KotlinMethodInfo
   getKotlinInvocation: (args: string[]) => string
 }
 
-/** Kotlin forwarders convert JNI values to the public API and back. */
-export function createKotlinForwarderBody({ callable, getKotlinInvocation }: KotlinForwarderBodyOptions): string {
+/** Kotlin bridges convert JNI values to the public API and back. */
+export function createKotlinBridgeBody({ callable, getKotlinInvocation }: KotlinBridgeBodyOptions): string {
   const kotlinArguments = callable.parameters.map((parameter) =>
     parameter.type.convert(parameter.native.type.escapedName, {
       from: 'c++',
@@ -87,18 +87,26 @@ interface CppJniCallBodyOptions {
   callable: KotlinMethodInfo
   /** The JVM method C++ calls, e.g. `addNumbers_cxx`. */
   jvmName: string
-  /** Looks up the fbjni method, e.g. `_javaPart->javaClassStatic()->getMethod`. */
+  /** Looks up the fbjni method, e.g. `CxxPart::getClass()->getMethod`. */
   jniMethodLookup: string
   /** The Java instance, or the class for static methods. */
   jniReceiver: string
+  /** The lookup class passed to fbjni's nonvirtual invocation. */
+  jniDeclaringClass?: string
 }
 
 /**
  * Calls Kotlin through a cached fbjni method, converting values between C++ and JNI.
  * The caller must already be attached to the JVM.
  * Generated locals use the reserved `__` prefix, so they cannot shadow parameters.
  */
-export function createCppJniCallBody({ callable, jvmName, jniMethodLookup, jniReceiver }: CppJniCallBodyOptions): string {
+export function createCppJniCallBody({
+  callable,
+  jvmName,
+  jniMethodLookup,
+  jniReceiver,
+  jniDeclaringClass,
+}: CppJniCallBodyOptions): string {
   const jniReturnType = callable.returnType.getCode('c++', {
     context: 'return',
   })
@@ -107,16 +115,19 @@ export function createCppJniCallBody({ callable, jvmName, jniMethodLookup, jniRe
     return `${type} /* ${parameter.native.name} */`
   })
   const jniSignature = `${jniReturnType}(${jniParameterTypes.join(', ')})`
-  const jniArguments = [
-    jniReceiver,
+  const jniArguments = [jniReceiver]
+  if (jniDeclaringClass != null) {
+    jniArguments.push(jniDeclaringClass)
+  }
+  jniArguments.push(
     ...callable.parameters.map((parameter) =>
       parameter.type.convert(parameter.native.type.escapedName, {
         from: 'c++',
         to: 'kotlin',
         in: 'c++',
       })
-    ),
-  ]
+    )
+  )
   if (callable.nativeMethod.returnType.kind === 'void') {
     return `
 static const auto __method = ${jniMethodLookup}<${jniSignature}>("${jvmName}");