alef_backend_java/
trait_bridge.rs

1//! Java (Panama FFM) trait bridge code generation for plugin systems.
2//!
3//! This module generates Java code that wraps C FFI vtables for plugin registration.
4//! Since Java cannot expose method references as raw C function pointers, we use
5//! Java 21+ Foreign Function & Memory API (Panama) upcall stubs to bridge Java
6//! implementations into the C vtable structure.
7//!
8//! For each `[[trait_bridges]]` entry, this module generates:
9//!
10//! 1. A `public interface {TraitName} { ... }` with methods matching the trait's methods
11//!    plus Plugin lifecycle methods (name, version, initialize, shutdown).
12//! 2. A `{TraitName}Bridge` class that:
13//!    - Allocates Panama FFM upcall stubs for each trait method
14//!    - Builds the C vtable as a MemorySegment
15//!    - Manages memory lifecycle with AutoCloseable
16//! 3. Registration helper: `public static void register{TraitName}({TraitName} impl)`
17//!    that builds the vtable and calls the C registration function.
18//! 4. Unregistration helper: `public static void unregister{TraitName}(String name)`.
19
20use heck::{ToPascalCase, ToSnakeCase};
21use std::fmt::Write;
22
23/// Generate all trait bridge code for a single `[[trait_bridges]]` entry.
24/// Returns Java source code as a String.
25///
26/// `has_super_trait`: when `true`, the C vtable includes Plugin lifecycle slots
27/// (name_fn, version_fn, initialize_fn, shutdown_fn) and the bridge emits matching
28/// upcall stubs. When `false`, only the trait-method slots are emitted, matching
29/// the Rust-side vtable layout exactly.
30pub fn gen_trait_bridge(
31    trait_name: &str,
32    prefix: &str,
33    trait_methods: &[(&str, &str)], // [(method_name, return_type), ...]
34    has_super_trait: bool,
35) -> String {
36    let trait_pascal = trait_name.to_pascal_case();
37    let trait_snake = trait_name.to_snake_case();
38    let prefix_upper = prefix.to_uppercase();
39
40    let mut out = String::with_capacity(4096);
41
42    // --- Trait interface ---
43    writeln!(out, "/**").ok();
44    writeln!(out, " * Bridge trait for {} plugin system.", trait_pascal).ok();
45    writeln!(out, " *").ok();
46    writeln!(
47        out,
48        " * Implementations provide methods that are called via upcall stubs"
49    )
50    .ok();
51    writeln!(out, " * into the C vtable during registration.").ok();
52    writeln!(out, " */").ok();
53    writeln!(out, "public interface {} {{", trait_pascal).ok();
54    writeln!(out).ok();
55
56    // Plugin lifecycle methods — only when a super_trait (Plugin) is configured
57    if has_super_trait {
58        writeln!(out, "    /** Return the plugin name. */").ok();
59        writeln!(out, "    String name();").ok();
60        writeln!(out).ok();
61
62        writeln!(out, "    /** Return the plugin version. */").ok();
63        writeln!(out, "    String version();").ok();
64        writeln!(out).ok();
65
66        writeln!(out, "    /** Initialize the plugin. */").ok();
67        writeln!(out, "    void initialize() throws Exception;").ok();
68        writeln!(out).ok();
69
70        writeln!(out, "    /** Shut down the plugin. */").ok();
71        writeln!(out, "    void shutdown() throws Exception;").ok();
72        writeln!(out).ok();
73    }
74
75    // Trait methods
76    for (method_name, return_type) in trait_methods {
77        writeln!(out, "    /** Trait method: {}. */", method_name).ok();
78        writeln!(out, "    {} {}();", return_type, method_name).ok();
79        writeln!(out).ok();
80    }
81
82    writeln!(out, "}}").ok();
83    writeln!(out).ok();
84
85    // --- Bridge class for FFI upcall stubs ---
86    writeln!(out, "/**").ok();
87    writeln!(
88        out,
89        " * Allocates Panama FFM upcall stubs for a {} trait implementation",
90        trait_pascal
91    )
92    .ok();
93    writeln!(out, " * and assembles the C vtable in native memory.").ok();
94    writeln!(out, " */").ok();
95    writeln!(out, "final class {}Bridge implements AutoCloseable {{", trait_pascal).ok();
96    writeln!(out).ok();
97
98    writeln!(out, "    private static final Linker LINKER = Linker.nativeLinker();").ok();
99    writeln!(
100        out,
101        "    private static final MethodHandles.Lookup LOOKUP = MethodHandles.lookup();"
102    )
103    .ok();
104    writeln!(out).ok();
105
106    // Number of vtable fields: optionally name_fn, version_fn, initialize_fn, shutdown_fn,
107    // then trait methods, then free_user_data.
108    let num_methods = trait_methods.len();
109    let num_super_slots = if has_super_trait { 4usize } else { 0usize };
110    let num_vtable_fields = num_super_slots + num_methods + 1; // (super-trait methods +) trait methods + free_user_data
111    writeln!(
112        out,
113        "    // C vtable: {} fields ({} plugin methods + {} trait methods + free_user_data)",
114        num_vtable_fields, num_super_slots, num_methods
115    )
116    .ok();
117    writeln!(
118        out,
119        "    private static final long VTABLE_SIZE = (long) ValueLayout.ADDRESS.byteSize() * {}L;",
120        num_vtable_fields
121    )
122    .ok();
123    writeln!(out).ok();
124
125    writeln!(out, "    private final Arena arena;").ok();
126    writeln!(out, "    private final MemorySegment vtable;").ok();
127    writeln!(out, "    private final {} impl;", trait_pascal).ok();
128    writeln!(out).ok();
129
130    // Constructor
131    writeln!(out, "    {}Bridge({} impl) {{", trait_pascal, trait_pascal).ok();
132    writeln!(out, "        this.impl = impl;").ok();
133    writeln!(out, "        this.arena = Arena.ofConfined();").ok();
134    writeln!(out, "        this.vtable = arena.allocate(VTABLE_SIZE);").ok();
135    writeln!(out).ok();
136    writeln!(out, "        try {{").ok();
137    writeln!(out, "            long offset = 0L;").ok();
138    writeln!(out).ok();
139
140    if has_super_trait {
141        // Register name_fn
142        writeln!(
143            out,
144            "            var stubName = LINKER.upcallStub(LOOKUP.bind(this, \"handleName\","
145        )
146        .ok();
147        writeln!(out, "                MethodType.methodType(MemorySegment.class)),").ok();
148        writeln!(
149            out,
150            "                FunctionDescriptor.of(ValueLayout.ADDRESS, ValueLayout.ADDRESS),"
151        )
152        .ok();
153        writeln!(out, "                arena);").ok();
154        writeln!(out, "            vtable.set(ValueLayout.ADDRESS, offset, stubName);").ok();
155        writeln!(out, "            offset += ValueLayout.ADDRESS.byteSize();").ok();
156        writeln!(out).ok();
157
158        // Register version_fn
159        writeln!(
160            out,
161            "            var stubVersion = LINKER.upcallStub(LOOKUP.bind(this, \"handleVersion\","
162        )
163        .ok();
164        writeln!(out, "                MethodType.methodType(MemorySegment.class)),").ok();
165        writeln!(
166            out,
167            "                FunctionDescriptor.of(ValueLayout.ADDRESS, ValueLayout.ADDRESS),"
168        )
169        .ok();
170        writeln!(out, "                arena);").ok();
171        writeln!(out, "            vtable.set(ValueLayout.ADDRESS, offset, stubVersion);").ok();
172        writeln!(out, "            offset += ValueLayout.ADDRESS.byteSize();").ok();
173        writeln!(out).ok();
174
175        // Register initialize_fn
176        writeln!(
177            out,
178            "            var stubInitialize = LINKER.upcallStub(LOOKUP.bind(this, \"handleInitialize\","
179        )
180        .ok();
181        writeln!(out, "                MethodType.methodType(int.class)),").ok();
182        writeln!(
183            out,
184            "                FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.ADDRESS),"
185        )
186        .ok();
187        writeln!(out, "                arena);").ok();
188        writeln!(
189            out,
190            "            vtable.set(ValueLayout.ADDRESS, offset, stubInitialize);"
191        )
192        .ok();
193        writeln!(out, "            offset += ValueLayout.ADDRESS.byteSize();").ok();
194        writeln!(out).ok();
195
196        // Register shutdown_fn
197        writeln!(
198            out,
199            "            var stubShutdown = LINKER.upcallStub(LOOKUP.bind(this, \"handleShutdown\","
200        )
201        .ok();
202        writeln!(out, "                MethodType.methodType(int.class)),").ok();
203        writeln!(
204            out,
205            "                FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.ADDRESS),"
206        )
207        .ok();
208        writeln!(out, "                arena);").ok();
209        writeln!(
210            out,
211            "            vtable.set(ValueLayout.ADDRESS, offset, stubShutdown);"
212        )
213        .ok();
214        writeln!(out, "            offset += ValueLayout.ADDRESS.byteSize();").ok();
215        writeln!(out).ok();
216    }
217
218    // Register trait methods
219    for (method_name, _) in trait_methods {
220        let handle_name = format!("handle{}", method_name.to_pascal_case());
221        writeln!(
222            out,
223            "            var stub{} = LINKER.upcallStub(LOOKUP.bind(this, \"{}\",",
224            method_name.to_pascal_case(),
225            handle_name
226        )
227        .ok();
228        writeln!(out, "                MethodType.methodType(MemorySegment.class)),").ok();
229        writeln!(
230            out,
231            "                FunctionDescriptor.of(ValueLayout.ADDRESS, ValueLayout.ADDRESS),"
232        )
233        .ok();
234        writeln!(out, "                arena);").ok();
235        writeln!(
236            out,
237            "            vtable.set(ValueLayout.ADDRESS, offset, stub{});",
238            method_name.to_pascal_case()
239        )
240        .ok();
241        writeln!(out, "            offset += ValueLayout.ADDRESS.byteSize();").ok();
242        writeln!(out).ok();
243    }
244
245    // Register free_user_data (NULL for now)
246    writeln!(
247        out,
248        "            vtable.set(ValueLayout.ADDRESS, offset, MemorySegment.NULL);"
249    )
250    .ok();
251    writeln!(out).ok();
252
253    writeln!(out, "        }} catch (ReflectiveOperationException e) {{").ok();
254    writeln!(out, "            arena.close();").ok();
255    writeln!(
256        out,
257        "            throw new RuntimeException(\"Failed to create trait bridge stubs\", e);"
258    )
259    .ok();
260    writeln!(out, "        }}").ok();
261    writeln!(out, "    }}").ok();
262    writeln!(out).ok();
263
264    // Accessor method
265    writeln!(out, "    MemorySegment vtableSegment() {{").ok();
266    writeln!(out, "        return vtable;").ok();
267    writeln!(out, "    }}").ok();
268    writeln!(out).ok();
269
270    // Handle methods
271    writeln!(
272        out,
273        "    // --- Upcall handlers (return MemorySegment pointing to allocated strings) ---"
274    )
275    .ok();
276    writeln!(out).ok();
277
278    if has_super_trait {
279        writeln!(out, "    private MemorySegment handleName() {{").ok();
280        writeln!(out, "        try {{").ok();
281        writeln!(out, "            String name = impl.name();").ok();
282        writeln!(out, "            return arena.allocateFrom(name);").ok();
283        writeln!(out, "        }} catch (Throwable e) {{").ok();
284        writeln!(out, "            return MemorySegment.NULL;").ok();
285        writeln!(out, "        }}").ok();
286        writeln!(out, "    }}").ok();
287        writeln!(out).ok();
288
289        writeln!(out, "    private MemorySegment handleVersion() {{").ok();
290        writeln!(out, "        try {{").ok();
291        writeln!(out, "            String version = impl.version();").ok();
292        writeln!(out, "            return arena.allocateFrom(version);").ok();
293        writeln!(out, "        }} catch (Throwable e) {{").ok();
294        writeln!(out, "            return MemorySegment.NULL;").ok();
295        writeln!(out, "        }}").ok();
296        writeln!(out, "    }}").ok();
297        writeln!(out).ok();
298
299        writeln!(out, "    private int handleInitialize() {{").ok();
300        writeln!(out, "        try {{").ok();
301        writeln!(out, "            impl.initialize();").ok();
302        writeln!(out, "            return 0;").ok();
303        writeln!(out, "        }} catch (Throwable e) {{").ok();
304        writeln!(out, "            return 1;").ok();
305        writeln!(out, "        }}").ok();
306        writeln!(out, "    }}").ok();
307        writeln!(out).ok();
308
309        writeln!(out, "    private int handleShutdown() {{").ok();
310        writeln!(out, "        try {{").ok();
311        writeln!(out, "            impl.shutdown();").ok();
312        writeln!(out, "            return 0;").ok();
313        writeln!(out, "        }} catch (Throwable e) {{").ok();
314        writeln!(out, "            return 1;").ok();
315        writeln!(out, "        }}").ok();
316        writeln!(out, "    }}").ok();
317        writeln!(out).ok();
318    }
319
320    // Trait method handlers
321    for (method_name, return_type) in trait_methods {
322        writeln!(
323            out,
324            "    private MemorySegment handle{}() {{",
325            method_name.to_pascal_case()
326        )
327        .ok();
328        writeln!(out, "        try {{").ok();
329        if *return_type == "void" || *return_type == "Void" {
330            writeln!(out, "            impl.{}();", method_name).ok();
331            writeln!(out, "            return MemorySegment.NULL;").ok();
332        } else {
333            writeln!(out, "            {} result = impl.{}();", return_type, method_name).ok();
334            writeln!(out, "            String json = new com.fasterxml.jackson.databind.ObjectMapper().writeValueAsString(result);").ok();
335            writeln!(out, "            return arena.allocateFrom(json);").ok();
336        }
337        writeln!(out, "        }} catch (Throwable e) {{").ok();
338        writeln!(out, "            return MemorySegment.NULL;").ok();
339        writeln!(out, "        }}").ok();
340        writeln!(out, "    }}").ok();
341        writeln!(out).ok();
342    }
343
344    writeln!(out, "    @Override").ok();
345    writeln!(out, "    public void close() {{").ok();
346    writeln!(out, "        arena.close();").ok();
347    writeln!(out, "    }}").ok();
348    writeln!(out, "}}").ok();
349    writeln!(out).ok();
350
351    // --- Bridge registry (keeps arenas + upcall stubs alive) ---
352    writeln!(
353        out,
354        "/** Registry of live {} bridges — keeps upcall stubs and arenas alive. */",
355        trait_pascal
356    )
357    .ok();
358    writeln!(
359        out,
360        "private static final java.util.concurrent.ConcurrentHashMap<String, {}Bridge> {}_BRIDGES",
361        trait_pascal,
362        trait_snake.to_uppercase()
363    )
364    .ok();
365    writeln!(out, "    = new java.util.concurrent.ConcurrentHashMap<>();").ok();
366    writeln!(out).ok();
367
368    // --- Registration helpers ---
369    writeln!(
370        out,
371        "/** Register a {} implementation via Panama FFM upcall stubs. */",
372        trait_pascal
373    )
374    .ok();
375    writeln!(
376        out,
377        "public static void register{}({} impl) throws Exception {{",
378        trait_pascal, trait_pascal
379    )
380    .ok();
381    // Do NOT use try-with-resources: the bridge arena must stay open for the lifetime
382    // of the plugin. We store it in the static registry so it is not GC'd or closed early.
383    writeln!(out, "    var bridge = new {}Bridge(impl);", trait_pascal).ok();
384    writeln!(out, "    try {{").ok();
385    writeln!(out, "        try (var nameArena = Arena.ofConfined()) {{").ok();
386    writeln!(out, "            var nameCs = nameArena.allocateFrom(impl.name());").ok();
387    writeln!(out, "            var outErrArena = Arena.ofConfined();").ok();
388    writeln!(
389        out,
390        "            MemorySegment outErr = outErrArena.allocate(ValueLayout.ADDRESS);"
391    )
392    .ok();
393    writeln!(
394        out,
395        "            int rc = (int) NativeLib.{}_REGISTER_{}.invoke(nameCs, bridge.vtableSegment(), MemorySegment.NULL, outErr);",
396        prefix_upper,
397        trait_snake.to_uppercase()
398    )
399    .ok();
400    writeln!(out, "            if (rc != 0) {{").ok();
401    writeln!(
402        out,
403        "                MemorySegment errPtr = outErr.get(ValueLayout.ADDRESS, 0);"
404    )
405    .ok();
406    writeln!(
407        out,
408        "                String msg = errPtr.equals(MemorySegment.NULL) ? \"registration failed (rc=\" + rc + \")\" : errPtr.reinterpret(Long.MAX_VALUE).getString(0);"
409    )
410    .ok();
411    writeln!(
412        out,
413        "                throw new RuntimeException(\"register{}: \" + msg);",
414        trait_pascal
415    )
416    .ok();
417    writeln!(out, "            }}").ok();
418    writeln!(out, "        }}").ok();
419    writeln!(out, "    }} catch (Throwable t) {{").ok();
420    writeln!(out, "        bridge.close();").ok();
421    writeln!(out, "        throw t;").ok();
422    writeln!(out, "    }}").ok();
423    writeln!(
424        out,
425        "    {}_BRIDGES.put(impl.name(), bridge);",
426        trait_snake.to_uppercase()
427    )
428    .ok();
429    writeln!(out, "}}").ok();
430    writeln!(out).ok();
431
432    writeln!(out, "/** Unregister a {} implementation. */", trait_pascal).ok();
433    writeln!(
434        out,
435        "public static void unregister{}(String name) throws Exception {{",
436        trait_pascal
437    )
438    .ok();
439    writeln!(out, "    try (var nameArena = Arena.ofConfined()) {{").ok();
440    writeln!(out, "        var nameCs = nameArena.allocateFrom(name);").ok();
441    writeln!(out, "        var outErrArena = Arena.ofConfined();").ok();
442    writeln!(
443        out,
444        "        MemorySegment outErr = outErrArena.allocate(ValueLayout.ADDRESS);"
445    )
446    .ok();
447    writeln!(
448        out,
449        "        int rc = (int) NativeLib.{}_UNREGISTER_{}.invoke(nameCs, outErr);",
450        prefix_upper,
451        trait_snake.to_uppercase()
452    )
453    .ok();
454    writeln!(out, "        if (rc != 0) {{").ok();
455    writeln!(
456        out,
457        "            MemorySegment errPtr = outErr.get(ValueLayout.ADDRESS, 0);"
458    )
459    .ok();
460    writeln!(
461        out,
462        "            String msg = errPtr.equals(MemorySegment.NULL) ? \"unregistration failed (rc=\" + rc + \")\" : errPtr.reinterpret(Long.MAX_VALUE).getString(0);"
463    )
464    .ok();
465    writeln!(
466        out,
467        "            throw new RuntimeException(\"unregister{}: \" + msg);",
468        trait_pascal
469    )
470    .ok();
471    writeln!(out, "        }}").ok();
472    writeln!(out, "    }}").ok();
473    // Close and remove the bridge after the C unregister call succeeds
474    writeln!(
475        out,
476        "    {}Bridge old = {}_BRIDGES.remove(name);",
477        trait_pascal,
478        trait_snake.to_uppercase()
479    )
480    .ok();
481    writeln!(out, "    if (old != null) {{ old.close(); }}").ok();
482    writeln!(out, "}}").ok();
483
484    out
485}
486
487#[cfg(test)]
488mod tests {
489    use super::*;
490
491    #[test]
492    fn test_gen_trait_bridge_basic() {
493        let code = gen_trait_bridge("MyPlugin", "mylib", &[("doWork", "String"), ("getStatus", "int")], true);
494
495        // Basic sanity checks
496        assert!(code.contains("public interface MyPlugin"));
497        assert!(code.contains("String name()"));
498        assert!(code.contains("String version()"));
499        assert!(code.contains("void initialize()"));
500        assert!(code.contains("void shutdown()"));
501        assert!(code.contains("doWork"));
502        assert!(code.contains("getStatus"));
503        assert!(code.contains("MyPluginBridge"));
504        assert!(code.contains("registerMyPlugin"));
505        assert!(code.contains("unregisterMyPlugin"));
506    }
507
508    #[test]
509    fn test_gen_trait_bridge_vtable_stubs() {
510        let code = gen_trait_bridge("Handler", "lib", &[], true);
511
512        // Verify Panama FFM upcall stubs are generated
513        assert!(code.contains("LINKER.upcallStub"));
514        assert!(code.contains("handleName"));
515        assert!(code.contains("handleVersion"));
516        assert!(code.contains("handleInitialize"));
517        assert!(code.contains("handleShutdown"));
518    }
519
520    #[test]
521    fn test_gen_trait_bridge_lifecycle_methods() {
522        let code = gen_trait_bridge("Processor", "pfx", &[("process", "Object")], true);
523
524        // Verify Plugin lifecycle methods are present in Java interface
525        assert!(code.contains("String name()"));
526        assert!(code.contains("String version()"));
527        assert!(code.contains("void initialize()"));
528        assert!(code.contains("void shutdown()"));
529    }
530
531    #[test]
532    fn test_gen_trait_bridge_no_super_trait_omits_lifecycle() {
533        let code = gen_trait_bridge("Transformer", "lib", &[("transform", "String")], false);
534
535        // Without super_trait, no lifecycle slots should be emitted
536        assert!(
537            !code.contains("String name()"),
538            "no lifecycle methods without super_trait"
539        );
540        assert!(
541            !code.contains("handleName"),
542            "no lifecycle handlers without super_trait"
543        );
544        assert!(code.contains("transform"), "trait method must still be emitted");
545    }
546}
alef_backend_java/trait_bridge.rs

alef_backend_java/
trait_bridge.rs
