DispatchAny

Author: milyin

prebindgen-ext/src/jni/jni_kotlin_ext.rs

@@ -455,8 +455,9 @@ fn render_wrapper_fn(
         mode: ParamMode,
     }
     enum ParamMode {
-        Borrow,    // &T opaque-handle → withPtr
-        Consume,   // T  opaque-handle → consume
+        Borrow,       // &T opaque-handle → withPtr
+        Consume,      // T  opaque-handle → consume
+        DispatchAny,  // `impl Into<T>` (Kotlin `Any`) → runtime: NativeHandle.withPtr or pass-through
         PassThrough,
     }
 
@@ -472,17 +473,9 @@ fn render_wrapper_fn(
         let entry = registry.input_entry(arg_ty)?;
         let is_opaque = converter_returns_owned_object(&entry.function.sig.output);
 
-        let mode = if is_opaque {
-            if matches!(arg_ty, syn::Type::Reference(_)) {
-                ParamMode::Borrow
-            } else {
-                ParamMode::Consume
-            }
+        let (kt_type_raw, optional) = if is_opaque {
+            ("NativeHandle".to_string(), false)
         } else {
-            ParamMode::PassThrough
-        };
-
-        let (kt_type, optional) = if matches!(mode, ParamMode::PassThrough) {
             // Look up the Kotlin type via the merged type map; fall
             // back to deriving from the wire type when the param
             // isn't pre-registered (e.g. `impl Into<T>` shapes wired
@@ -497,11 +490,26 @@ fn render_wrapper_fn(
                 .or_else(|| kotlin_for_wire(&entry.destination))?;
             let opt = is_option_type(arg_ty);
             (kt, opt)
+        };
+
+        // Mode: opaque → Borrow/Consume by Rust syntactic shape.
+        // Non-opaque → `Any` triggers DispatchAny (`impl Into<T>` wire
+        // = JObject; runtime might be a NativeHandle that we should
+        // unwrap under withPtr to close the race window). Everything
+        // else (primitives, callbacks, data classes) passes through.
+        let mode = if is_opaque {
+            if matches!(arg_ty, syn::Type::Reference(_)) {
+                ParamMode::Borrow
+            } else {
+                ParamMode::Consume
+            }
+        } else if kt_type_raw == "Any" {
+            ParamMode::DispatchAny
         } else {
-            ("NativeHandle".to_string(), false)
+            ParamMode::PassThrough
         };
 
-        let short = register_fqn(&kt_type, imports);
+        let short = register_fqn(&kt_type_raw, imports);
         let suffix = if optional { "?" } else { "" };
         params.push(Param {
             kt_name: name,
@@ -513,63 +521,81 @@ fn render_wrapper_fn(
     // Return type: peel ZResult<...>; detect opaque-handle return.
     let (kt_return, return_is_opaque) = classify_return(&f.sig.output, registry, kotlin_types, imports)?;
 
-    // Body: nest withPtr/consume in declaration order.
-    let mut call = format!("JNINative.{jni_call}(");
-    for (i, p) in params.iter().enumerate() {
-        if i > 0 { call.push_str(", "); }
-        call.push_str(&p.kt_name);
-    }
-    call.push(')');
+    // Helper: build the JNINative call for a given DispatchAny "unwrap mask".
+    // mask bit k = 1 means dispatch_indices[k] is unwrapped (use `<name>_ptr`).
+    let dispatch_indices: Vec<usize> = params
+        .iter()
+        .enumerate()
+        .filter_map(|(i, p)| matches!(p.mode, ParamMode::DispatchAny).then_some(i))
+        .collect();
 
-    // Build innermost expression, wrapping in NativeHandle() if return is opaque.
-    let mut body_expr = if return_is_opaque {
-        format!("NativeHandle({call})")
-    } else {
+    let build_call = |mask: u32| -> String {
+        let mut args: Vec<String> = Vec::with_capacity(params.len());
+        for (i, p) in params.iter().enumerate() {
+            let arg = match p.mode {
+                ParamMode::Borrow | ParamMode::Consume => format!("{}_ptr", p.kt_name),
+                ParamMode::DispatchAny => {
+                    let pos = dispatch_indices.iter().position(|&di| di == i).unwrap();
+                    if (mask >> pos) & 1 == 1 {
+                        format!("{}_ptr", p.kt_name)
+                    } else {
+                        p.kt_name.clone()
+                    }
+                }
+                ParamMode::PassThrough => p.kt_name.clone(),
+            };
+            args.push(arg);
+        }
+        let mut call = format!("JNINative.{jni_call}({})", args.join(", "));
+        if return_is_opaque {
+            call = format!("NativeHandle({call})");
+        }
         call
     };
 
+    // Build the DispatchAny decision tree. At each level we pick one
+    // dispatch_indices[level]: `if (p is NativeHandle) p.withPtr { ... } else { ... }`.
+    // The base case (level == n) emits the JNINative call with the
+    // accumulated unwrap mask. Branches recursively split.
+    fn build_tree(
+        level: usize,
+        mask: u32,
+        dispatch_indices: &[usize],
+        params: &[(String, /*placeholder for type*/ ())],
+        build_call: &dyn Fn(u32) -> String,
+    ) -> String {
+        if level == dispatch_indices.len() {
+            return build_call(mask);
+        }
+        let name = &params[dispatch_indices[level]].0;
+        let with_branch = build_tree(level + 1, mask | (1 << level), dispatch_indices, params, build_call);
+        let else_branch = build_tree(level + 1, mask, dispatch_indices, params, build_call);
+        format!(
+            "if ({name} is NativeHandle) {name}.withPtr {{ {name}_ptr ->\n    {with_branch}\n}} else {{\n    {else_branch}\n}}"
+        )
+    }
+
+    let param_names_for_tree: Vec<(String, ())> = params.iter().map(|p| (p.kt_name.clone(), ())).collect();
+    let mut body_expr = build_tree(0, 0, &dispatch_indices, &param_names_for_tree, &build_call);
+
     // Wrap with nested withPtr/consume from innermost to outermost.
-    let mut indent = String::new();
     for p in params.iter().rev() {
         match p.mode {
             ParamMode::Borrow => {
                 body_expr = format!(
-                    "{name}.withPtr {{ {name}_ptr ->\n{indent}    {expr}\n{indent}}}",
+                    "{name}.withPtr {{ {name}_ptr ->\n    {expr}\n}}",
                     name = p.kt_name,
-                    indent = indent,
                     expr = body_expr,
                 );
             }
             ParamMode::Consume => {
                 body_expr = format!(
-                    "{name}.consume {{ {name}_ptr ->\n{indent}    {expr}\n{indent}}}",
+                    "{name}.consume {{ {name}_ptr ->\n    {expr}\n}}",
                     name = p.kt_name,
-                    indent = indent,
                     expr = body_expr,
                 );
             }
-            ParamMode::PassThrough => {}
-        }
-        // Indentation isn't strictly needed for correctness; keep flat.
-    }
-
-    // Replace `<name>` with `<name>_ptr` inside the call args for opaque params.
-    // (We built `call` with `name`, but for opaque params the underlying
-    // JNI fn takes the raw Long; rewrite arg references accordingly.)
-    let mut fixed = body_expr;
-    for p in &params {
-        if matches!(p.mode, ParamMode::Borrow | ParamMode::Consume) {
-            // Replace the bare parameter name reference inside the JNINative call
-            // with `<name>_ptr`. Word-boundary-safe substitution.
-            let needle = format!(", {}", p.kt_name);
-            let repl = format!(", {}_ptr", p.kt_name);
-            fixed = fixed.replace(&needle, &repl);
-            let head_needle = format!("({}", p.kt_name);
-            let head_repl = format!("({}_ptr", p.kt_name);
-            fixed = fixed.replace(&head_needle, &head_repl);
-            let solo_needle = format!("({})", p.kt_name);
-            let solo_repl = format!("({}_ptr)", p.kt_name);
-            fixed = fixed.replace(&solo_needle, &solo_repl);
+            ParamMode::DispatchAny | ParamMode::PassThrough => {}
         }
     }
 
@@ -581,7 +607,7 @@ fn render_wrapper_fn(
         let _ = write!(out, ": {kt_return}");
     }
     let _ = writeln!(out, " =");
-    let _ = writeln!(out, "    {fixed}");
+    let _ = writeln!(out, "    {body_expr}");
     Some(out)
 }
 

zenoh-jni-runtime/src/commonMain/kotlin/io/zenoh/jni/JNIKeyExpr.kt

@@ -30,17 +30,13 @@ import io.zenoh.exceptions.ZError
 
 /**
  * Pick the declared handle if present, else the raw string. Returns
- * a JVM `Object` (boxed `java.lang.Long` or `java.lang.String`) which
- * the native dispatching converter resolves at runtime.
- *
- * Takes a [NativeHandle] (not a raw `Long`) so callers don't see the
- * inner pointer. The peek is unlocked — the borrow's Rust side runs
- * `Arc::increment_strong_count` and a concurrent close of this
- * keyExpr is the same race window the broader JNI layer already
- * tolerates for borrow-style calls.
+ * the `NativeHandle` (subtype of `Any`) when a handle is declared,
+ * else the `String` (also `Any`). The generator-emitted wrappers in
+ * [JNIWrappers] runtime-dispatch on `is NativeHandle` and acquire
+ * the read lock via `withPtr` before crossing the JNI boundary, so
+ * no `Long` ever escapes outside a held lock.
  */
-fun keyExprArg(handle: NativeHandle?, str: String): Any =
-    handle?.peek()?.takeIf { it != 0L } ?: str
+fun keyExprArg(handle: NativeHandle?, str: String): Any = handle ?: str
 
 @Throws(ZError::class)
 fun keyExprTryFrom(keyExpr: String): String = JNIWrappers.tryFrom(keyExpr)

zenoh-jni-runtime/src/commonMain/kotlin/io/zenoh/jni/JNIQuerier.kt

@@ -37,7 +37,13 @@ public class JNIQuerier(initialPtr: Long) : NativeHandle(initialPtr) {
         payload: ByteArray?,
         encoding: JNIEncoding?,
     ) = withPtr { ptr ->
-        getViaJNI(ptr, keyExprArg(keyExprHandle, keyExprString), parameters, callback, onClose, attachmentBytes, payload, encoding)
+        if (keyExprHandle != null) {
+            keyExprHandle.withPtr { kePtr ->
+                getViaJNI(ptr, kePtr, parameters, callback, onClose, attachmentBytes, payload, encoding)
+            }
+        } else {
+            getViaJNI(ptr, keyExprString, parameters, callback, onClose, attachmentBytes, payload, encoding)
+        }
     }
 
     @Throws(ZError::class)

zenoh-jni-runtime/src/commonMain/kotlin/io/zenoh/jni/JNIQuery.kt

@@ -31,7 +31,13 @@ public class JNIQuery(initialPtr: Long) : NativeHandle(initialPtr) {
         attachment: ByteArray?,
         qosExpress: Boolean,
     ) = withPtr { ptr ->
-        replySuccessViaJNI(ptr, keyExprArg(keyExprHandle, keyExprString), payload, encoding, timestampEnabled, timestampNtp64, attachment, qosExpress)
+        if (keyExprHandle != null) {
+            keyExprHandle.withPtr { kePtr ->
+                replySuccessViaJNI(ptr, kePtr, payload, encoding, timestampEnabled, timestampNtp64, attachment, qosExpress)
+            }
+        } else {
+            replySuccessViaJNI(ptr, keyExprString, payload, encoding, timestampEnabled, timestampNtp64, attachment, qosExpress)
+        }
     }
 
     @Throws(ZError::class)
@@ -48,7 +54,13 @@ public class JNIQuery(initialPtr: Long) : NativeHandle(initialPtr) {
         attachment: ByteArray?,
         qosExpress: Boolean,
     ) = withPtr { ptr ->
-        replyDeleteViaJNI(ptr, keyExprArg(keyExprHandle, keyExprString), timestampEnabled, timestampNtp64, attachment, qosExpress)
+        if (keyExprHandle != null) {
+            keyExprHandle.withPtr { kePtr ->
+                replyDeleteViaJNI(ptr, kePtr, timestampEnabled, timestampNtp64, attachment, qosExpress)
+            }
+        } else {
+            replyDeleteViaJNI(ptr, keyExprString, timestampEnabled, timestampNtp64, attachment, qosExpress)
+        }
     }
 
     fun close() = close { freePtrViaJNI(it) }

zenoh-jni-runtime/src/commonMain/kotlin/io/zenoh/jni/JNISession.kt

@@ -263,7 +263,16 @@ public class JNISession(initialPtr: Long) : NativeHandle(initialPtr) {
     @Throws(ZError::class)
     fun declareLivelinessToken(keyExprHandle: NativeHandle?, keyExprString: String): JNILivelinessToken =
         withPtr { ptr ->
-            JNILivelinessToken(declareLivelinessTokenViaJNI(ptr, keyExprArg(keyExprHandle, keyExprString)))
+            // Mirror JNIWrappers' impl-Into dispatch: hold the keyExpr's
+            // read lock for the JNI call when a declared handle is in
+            // play; otherwise pass the validated string.
+            if (keyExprHandle != null) {
+                keyExprHandle.withPtr { kePtr ->
+                    JNILivelinessToken(declareLivelinessTokenViaJNI(ptr, kePtr))
+                }
+            } else {
+                JNILivelinessToken(declareLivelinessTokenViaJNI(ptr, keyExprString))
+            }
         }
 
     @Throws(ZError::class)
@@ -277,7 +286,13 @@ public class JNISession(initialPtr: Long) : NativeHandle(initialPtr) {
         history: Boolean,
         onClose: JNIOnCloseCallback,
     ): JNISubscriber = withPtr { ptr ->
-        JNISubscriber(declareLivelinessSubscriberViaJNI(ptr, keyExprArg(keyExprHandle, keyExprString), callback, history, onClose))
+        if (keyExprHandle != null) {
+            keyExprHandle.withPtr { kePtr ->
+                JNISubscriber(declareLivelinessSubscriberViaJNI(ptr, kePtr, callback, history, onClose))
+            }
+        } else {
+            JNISubscriber(declareLivelinessSubscriberViaJNI(ptr, keyExprString, callback, history, onClose))
+        }
     }
 
     @Throws(ZError::class)
@@ -297,7 +312,13 @@ public class JNISession(initialPtr: Long) : NativeHandle(initialPtr) {
         timeoutMs: Long,
         onClose: JNIOnCloseCallback,
     ) = withPtr { ptr ->
-        livelinessGetViaJNI(ptr, keyExprArg(keyExprHandle, keyExprString), callback, timeoutMs, onClose)
+        if (keyExprHandle != null) {
+            keyExprHandle.withPtr { kePtr ->
+                livelinessGetViaJNI(ptr, kePtr, callback, timeoutMs, onClose)
+            }
+        } else {
+            livelinessGetViaJNI(ptr, keyExprString, callback, timeoutMs, onClose)
+        }
     }
 
     @Throws(ZError::class)