Skip to main content

yog_interop/
lib.rs

1//! Yog inter-mod communication — proc-macros for cross-mod function calls.
2//!
3//! # Export (manual registration, auto-generated binding symbol)
4//!
5//! ```ignore
6//! use yog_api::{yog_export, Registry};
7//!
8//! #[yog_export]
9//! pub fn register_pipe(api: *const yog_api::YogApi, json: *const std::ffi::c_char) {
10//!     // ...
11//! }
12//!
13//! fn register(registry: &mut Registry) {
14//!     // Manual export registration (loader scanning coming later)
15//!     registry.interop().export("register_pipe", register_pipe as *const std::ffi::c_void);
16//! }
17//! ```
18//!
19//! # Import (auto-wired by loader)
20//!
21//! ```ignore
22//! yog_api::import! {
23//!     from "yog-pipes" {
24//!         fn register_pipe(api: *const yog_api::YogApi, json: *const std::ffi::c_char);
25//!     }
26//! }
27//!
28//! // Use like a normal function — loader wires it at init time
29//! register_pipe(api, json_ptr);
30//! ```
31
32use proc_macro::TokenStream;
33use quote::{format_ident, quote};
34
35// ── #[yog_export] ─────────────────────────────────────────────────────────
36
37/// Universal export macro — works on both structs and functions.
38///
39/// **On structs**: auto-derives `rkyv::Archive`, `rkyv::Serialize`,
40/// `rkyv::Deserialize`. Equivalent to:
41/// ```ignore
42/// #[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
43/// #[archive(check_bytes)]
44/// struct MyStruct { ... }
45/// ```
46///
47/// **On functions**: generates a C-ABI wrapper using rkyv zero-copy
48/// serialization + auto-registers the export via `export_mod!`.
49///
50/// Supported fn signatures:
51/// - `fn(InputType) -> OutputType`  (single arg)
52/// - `fn() -> OutputType`          (no arg)
53/// - `fn(InputType)`               (no return)
54/// For multi-arg: `fn((A, B)) -> C`.
55///
56/// No manual `registry.interop().export()` needed.
57#[proc_macro_attribute]
58pub fn yog_export(_attr: TokenStream, item: TokenStream) -> TokenStream {
59    // Try parsing as a struct or enum first (more common for types)
60    if let Ok(s) = syn::parse2::<syn::ItemStruct>(proc_macro2::TokenStream::from(item.clone())) {
61        let vis = &s.vis;
62        let ident = &s.ident;
63        let generics = &s.generics;
64        let fields = &s.fields;
65        let attrs = &s.attrs;
66        return TokenStream::from(quote! {
67            #[derive(::yog_api::rkyv::Archive, ::yog_api::rkyv::Serialize, ::yog_api::rkyv::Deserialize)]
68            #(#attrs)*
69            #vis struct #ident #generics #fields
70        });
71    }
72    if let Ok(e) = syn::parse2::<syn::ItemEnum>(proc_macro2::TokenStream::from(item.clone())) {
73        let vis = &e.vis;
74        let ident = &e.ident;
75        let generics = &e.generics;
76        let variants = &e.variants;
77        let attrs = &e.attrs;
78        return TokenStream::from(quote! {
79            #[derive(::yog_api::rkyv::Archive, ::yog_api::rkyv::Serialize, ::yog_api::rkyv::Deserialize)]
80            #(#attrs)*
81            #vis enum #ident #generics { #variants }
82        });
83    }
84
85    // Fall through to function handling
86    let func = syn::parse_macro_input!(item as syn::ItemFn);
87    let name = &func.sig.ident;
88    let vis = &func.vis;
89    let sig = &func.sig;
90    let block = &func.block;
91    let name_str = name.to_string();
92    let wrap_name = format_ident!("__yog_wrap_{}", name);
93    let _init_name = format_ident!("__yog_export_init_{}", name);
94
95    // Extract input and output types
96    let inputs = &func.sig.inputs;
97    let output = match &func.sig.output {
98        syn::ReturnType::Default => quote! { () },
99        syn::ReturnType::Type(_, ty) => quote! { #ty },
100    };
101
102    // Build the wrapper body — handles both arg and no-arg cases
103    let call_fn = if inputs.is_empty() {
104        quote! { #name() }
105    } else {
106        // Multi-arg: serialize as tuple, then deserialize and destructure
107        let arg_names: Vec<_> = inputs.iter().map(|arg| {
108            match arg {
109                syn::FnArg::Typed(pat_type) => &pat_type.pat,
110                syn::FnArg::Receiver(_) => unreachable!(),
111            }
112        }).collect();
113        quote! { #name(#(#arg_names),*) }
114    };
115
116    // Extract the single input type (or unit for no-arg fns)
117    let input_type = if inputs.is_empty() {
118        quote! { () }
119    } else {
120        // Take the first (and should be only) argument type
121        match inputs.first().unwrap() {
122            syn::FnArg::Typed(pat_type) => {
123                let ty = &pat_type.ty;
124                quote! { #ty }
125            }
126            _ => quote! { () },
127        }
128    };
129
130    let expanded = quote! {
131        #vis #sig #block
132
133        // ctor — runs before main(), pushes entry to export registry
134        #[doc(hidden)]
135        ctor::ctor! {
136            #[ctor(unsafe)]
137            static __yog_export_ctor_{name}: u8 = {
138                ::yog_api::__yog_export_registry()
139                    .lock()
140                    .unwrap()
141                    .push(::yog_api::YogExportEntry {
142                        name: #name_str,
143                        ptr: #wrap_name as usize,
144                    });
145                0
146            };
147        }
148
149        // C-ABI wrapper — deserializes input via rkyv, calls fn, serializes output.
150        #[doc(hidden)]
151        #[no_mangle]
152        pub unsafe extern "C" fn #wrap_name(
153            input_ptr: *const u8,
154            input_len: u32,
155            out_data: *mut *mut u8,
156            out_len: *mut u32,
157            out_cap: *mut u32,
158        ) {
159            let input_slice = std::slice::from_raw_parts(input_ptr, input_len as usize);
160            let args: #input_type = ::yog_api::rkyv::from_bytes::<_, ::yog_api::rkyv::rancor::Error>(input_slice)
161                .expect("yog: deserialization failed in export wrapper");
162            let result: #output = #call_fn;
163            let aligned = ::yog_api::rkyv::to_bytes::<_, 256>(&result)
164                .unwrap_or_default();
165            let bytes: Vec<u8> = aligned.to_vec();
166            *out_data = bytes.as_ptr() as *mut u8;
167            *out_len = bytes.len() as u32;
168            *out_cap = bytes.capacity() as u32;
169            std::mem::forget(bytes); // caller frees via __yog_free_buffer
170        }
171    };
172
173    TokenStream::from(expanded)
174}
175
176// ── import! macro ────────────────────────────────────────────────────────
177
178/// Declare imported functions from another mod. Generates public wrapper
179/// functions and `__yog_bind_NAME` entry points for the loader.
180///
181/// ```ignore
182/// yog_api::import! {
183///     from "yog-pipes" {
184///         fn register_pipe(api: *const YogApi, json: *const c_char) -> bool;
185///     }
186/// }
187/// register_pipe(api, json_ptr); // auto-wired
188/// ```
189#[proc_macro]
190pub fn import(input: TokenStream) -> TokenStream {
191    let import_block = syn::parse_macro_input!(input as ImportBlock);
192    let mod_name = &import_block.mod_name;
193    let mut out = proc_macro2::TokenStream::new();
194
195    for func in &import_block.funcs {
196        let name = &func.sig.ident;
197        let name_str = name.to_string();
198        let slot_name = format_ident!("__yog_slot_{}", name);
199        let bind_name = format_ident!("__yog_bind_{}", name);
200        let fn_sig = &func.sig;
201        let vis = &func.vis;
202        let inputs = &fn_sig.inputs;
203        let output = &fn_sig.output;
204
205        // Extract input type (single arg or unit)
206        let _input_type = if inputs.is_empty() {
207            quote! { () }
208        } else {
209            match inputs.first().unwrap() {
210                syn::FnArg::Typed(pat_type) => {
211                    let ty = &pat_type.ty;
212                    quote! { #ty }
213                }
214                _ => quote! { () },
215            }
216        };
217
218        let output_type = match output {
219            syn::ReturnType::Default => quote! { () },
220            syn::ReturnType::Type(_, ty) => quote! { #ty },
221        };
222
223        // Build arg forwarding for the public function
224        let arg_names: Vec<_> = inputs.iter().map(|arg| {
225            match arg {
226                syn::FnArg::Typed(pat_type) => &pat_type.pat,
227                syn::FnArg::Receiver(_) => unreachable!("self not allowed in imports"),
228            }
229        }).collect();
230
231        // Serialize the arg (or nothing for no-arg fns)
232        let serialize_block = if inputs.is_empty() {
233            quote! { let input_bytes = Vec::new(); }
234        } else {
235            quote! {
236                let aligned = ::yog_api::rkyv::to_bytes::<_, 256>(&(#(#arg_names),*))
237                    .unwrap_or_default();
238                let input_bytes: Vec<u8> = aligned.to_vec();
239            }
240        };
241
242        // C-ABI wrapper type — shared by all exports
243        let wrapper_ty = quote! {
244            unsafe extern "C" fn(
245                input_ptr: *const u8, input_len: u32,
246                out_data: *mut *mut u8, out_len: *mut u32, out_cap: *mut u32,
247            )
248        };
249
250        let _init_name = format_ident!("__yog_import_init_{}", name);
251
252        let wrapper = quote! {
253            // ctor — runs before main(), pushes entry to import registry
254            #[doc(hidden)]
255            ctor::ctor! {
256                #[ctor(unsafe)]
257                static __yog_import_ctor_{name}: u8 = {
258                    ::yog_api::__yog_import_registry()
259                        .lock()
260                        .unwrap()
261                        .push(::yog_api::YogImportEntry {
262                            mod_id: #mod_name,
263                            symbol: #name_str,
264                            bind_fn: #bind_name as usize,
265                        });
266                    0
267                };
268            }
269
270            #[allow(non_snake_case)]
271            #vis fn #name(#inputs) #output {
272                #serialize_block
273                let f = #slot_name.get()
274                    .expect(concat!("yog: import '", #mod_name, ":", #name_str, "' not bound — is mod '", #mod_name, "' loaded?"));
275                let mut out_data: *mut u8 = std::ptr::null_mut();
276                let mut out_len: u32 = 0;
277                let mut out_cap: u32 = 0;
278                unsafe {
279                    f(
280                        input_bytes.as_ptr(), input_bytes.len() as u32,
281                        &mut out_data, &mut out_len, &mut out_cap,
282                    );
283                    let output_slice = std::slice::from_raw_parts(out_data, out_len as usize);
284                    let result: #output_type = ::yog_api::rkyv::from_bytes::<_, ::yog_api::rkyv::rancor::Error>(output_slice)
285                        .expect(concat!("yog: deserialization failed in import '", #mod_name, ":", #name_str, "'"));
286                    // Free the buffer allocated by the export wrapper
287                    let _ = Vec::from_raw_parts(out_data, out_len as usize, out_cap as usize);
288                    result
289                }
290            }
291
292            static #slot_name: ::std::sync::OnceLock<#wrapper_ty>
293                = ::std::sync::OnceLock::new();
294
295            #[doc(hidden)]
296            #[no_mangle]
297            pub unsafe extern "C" fn #bind_name(ptr: *const ::std::os::raw::c_void) {
298                let f: #wrapper_ty = ::std::mem::transmute(ptr);
299                #slot_name.set(f).ok();
300            }
301        };
302
303        out.extend(wrapper);
304    }
305
306    TokenStream::from(out)
307}
308
309// ── Parse helpers ────────────────────────────────────────────────────────
310
311struct ImportBlock {
312    mod_name: String,
313    funcs: Vec<syn::ItemFn>,
314}
315
316impl syn::parse::Parse for ImportBlock {
317    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
318        let from_kw: syn::Ident = input.parse()?;
319        if from_kw != "from" {
320            return Err(syn::Error::new(from_kw.span(), "expected `from`"));
321        }
322        let lit: syn::LitStr = input.parse()?;
323        let mod_name = lit.value();
324
325        let content;
326        syn::braced!(content in input);
327        let mut funcs = Vec::new();
328        while !content.is_empty() {
329            funcs.push(content.parse::<syn::ItemFn>()?);
330            let _ = content.parse::<syn::Token![;]>();
331        }
332
333        Ok(ImportBlock { mod_name, funcs })
334    }
335}