abi-vtable-macro 0.3.1

//! # abi-vtable-macro — shadow-mode vtable generation
//!
//! **Shadow mode**: annotate a trait and every companion entity of the
//! abi-mode protocol is generated automatically — no manual steps.
//! Usage is as frictionless as the native (dyn) mode.
//!
//! ## When to use which mode
//! - **Rust-only** → native (dyn) mode: fastest and lightest, direct fat
//!   pointers, Arc ref-counting, no vtable struct at all.
//! - **Cross-language (C/C++/Zig/…)** → this (abi) mode: the generated
//!   `#[repr(C)]` vtable is the language-neutral contract.

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, ItemTrait, TraitItem};

/// Map a Rust type to its C-representable counterpart.
/// Only types with a stable C representation are supported.
fn map_type(ty: &syn::Type) -> Option<&'static str> {
    let ty_str = quote!(#ty).to_string().replace(' ', "");
    Some(match ty_str.as_str() {
        "i8" => "i8",
        "i16" => "i16",
        "i32" => "i32",
        "i64" => "i64",
        "u8" => "u8",
        "u16" => "u16",
        "u32" => "u32",
        "u64" => "u64",
        "f32" => "f32",
        "f64" => "f64",
        "bool" => "bool",
        "usize" => "usize",
        "isize" => "isize",
        _ => return None,
    })
}

struct AttrOpts {
    /// Export prefix, e.g. "calc" → `calc_get_vtable`, `CALC_VTABLE`.
    name: String,
}

fn parse_attr_tokens(ts: proc_macro2::TokenStream) -> AttrOpts {
    let mut name = String::new();
    let mut iter = ts.into_iter().peekable();
    while let Some(t) = iter.next() {
        if let proc_macro2::TokenTree::Ident(id) = t {
            if id == "name" {
                for t2 in iter.by_ref() {
                    if let proc_macro2::TokenTree::Literal(lit) = t2 {
                        name = lit.to_string().trim_matches('"').to_string();
                        break;
                    }
                }
            }
        }
    }
    if name.is_empty() {
        panic!("#[abi_vtable] requires `name = \"...\"` (getter symbol prefix)");
    }
    AttrOpts { name }
}

#[proc_macro_attribute]
pub fn abi_vtable(attr: TokenStream, item: TokenStream) -> TokenStream {
    let opts = parse_attr_tokens(attr.into());
    let trait_def = parse_macro_input!(item as ItemTrait);
    let trait_name = &trait_def.ident;
    let trait_name_str = trait_name.to_string();
    let vis = &trait_def.vis;

    // ---- collect &self methods with C-representable signatures ----
    struct Method {
        name: syn::Ident,
        params: Vec<proc_macro2::TokenStream>, // "id: CTy"
        ret: Option<syn::Type>,
        call_args: Vec<syn::Ident>,
    }

    let mut methods: Vec<Method> = Vec::new();
    for item in &trait_def.items {
        if let TraitItem::Fn(m) = item {
            let mut params = Vec::new();
            let mut call_args = Vec::new();
            // Receiver: only `&self` is supported. The thunk dispatches on
            // the module-side static instance behind a shared reference —
            // `&mut self` cannot be honored across the C boundary (the host
            // holds no exclusive borrow). Use interior mutability (atomics,
            // locks) in the impl instead. `self` by value makes no sense
            // across the boundary either.
            match m.sig.inputs.first() {
                Some(syn::FnArg::Receiver(r)) => match &r.reference {
                    Some(_) => {
                        if r.mutability.is_some() {
                            panic!(
                                "abi_vtable: `{}::{}` takes `&mut self` — not expressible across the C ABI; use `&self` + interior mutability",
                                trait_name_str, m.sig.ident
                            );
                        }
                        false
                    }
                    None => panic!(
                        "abi_vtable: `{}::{}` takes `self` by value — use `&self`",
                        trait_name_str, m.sig.ident
                    ),
                },
                _ => panic!(
                    "abi_vtable: `{}::{}` has no `self` receiver — associated functions are not vtable methods",
                    trait_name_str, m.sig.ident
                ),
            };
            for input in &m.sig.inputs {
                if let syn::FnArg::Typed(pt) = input {
                    let ident = match pt.pat.as_ref() {
                        syn::Pat::Ident(pi) => pi.ident.clone(),
                        _ => continue,
                    };
                    let ty_str = quote!(#pt.ty).to_string().replace(' ', "");
                    let c_ty = map_type(&pt.ty).unwrap_or_else(|| {
                        panic!(
                            "abi_vtable: unsupported param type `{}` in `{}::{}` — only C-representable scalars cross the boundary",
                            ty_str, trait_name_str, m.sig.ident
                        )
                    });
                    let c_ty: syn::Type = syn::parse_str(c_ty).unwrap();
                    params.push(quote!(#ident: #c_ty));
                    call_args.push(ident);
                }
            }
            let ret = match &m.sig.output {
                syn::ReturnType::Type(_, ty) => {
                    let ty_str = quote!(#ty).to_string().replace(' ', "");
                    let c_ty = map_type(ty).unwrap_or_else(|| {
                        panic!(
                            "abi_vtable: unsupported return type `{}` in `{}::{}`",
                            ty_str, trait_name_str, m.sig.ident
                        )
                    });
                    Some(syn::parse_str::<syn::Type>(c_ty).unwrap())
                }
                syn::ReturnType::Default => None,
            };
            methods.push(Method {
                name: m.sig.ident.clone(),
                params,
                ret,
                call_args,
            });
        }
    }

    // ---- shadow entities ----
    let vtable_struct_name = format_ident!("{}Vtable", trait_name);
    let instance_name = format_ident!("{}_INSTANCE", opts.name.to_uppercase());

    let mut struct_fields = Vec::new();
    // Fully-materialized thunk fns (one token blob per method), consumed by
    // the generated macro_rules! via a flat single-level repetition.
    let mut thunks = Vec::new();
    let mut thunk_assignments = Vec::new();

    for m in &methods {
        let fname = &m.name;
        // Prefix thunks with the export name so two traits in the same crate
        // with same-named methods never collide (e.g. `alpha_value_thunk`).
        let thunk_name = format_ident!("{}_{}_thunk", opts.name, m.name);
        let arg_tys = &m.params;
        let ret = match &m.ret {
            Some(t) => quote!(-> #t),
            None => quote!(),
        };

        struct_fields.push(quote! {
            pub #fname: unsafe extern "C" fn(ctx: *mut ::std::ffi::c_void, #(#arg_tys),*) #ret
        });

        let call_args = &m.call_args;
        // NOTE: no generics here — `$impl_ty` is substituted by the generated
        // macro_rules! at expansion time, so thunks are monomorphic by design.
        // Fully-qualified call: if $impl_ty implements several traits with
        // same-named methods, a bare `.method()` call would be ambiguous.
        //
        // The receiver is the macro-generated static INSTANCE (const-
        // constructible on the module side), NOT the ctx pointer: the host
        // has no way to obtain the module's instance address, and stateless
        // vtables legitimately pass null ctx. `&mut self` methods go through
        // interior mutability on the module side (the static is `&$impl_ty`;
        // mutation is the impl's responsibility, e.g. atomics).
        thunks.push(quote! {
            unsafe extern "C" fn #thunk_name(
                ctx: *mut ::std::ffi::c_void,
                #(#arg_tys),*
            ) #ret {
                let _ = ctx; // instance lives in the module; ctx is reserved
                unsafe { <$impl_ty as #trait_name>::#fname(&*#instance_name, #(#call_args),*) }
            }
        });
        thunk_assignments.push(quote! { #fname: #thunk_name });
    }

    let static_name = format_ident!("{}_VTABLE", opts.name.to_uppercase());
    let getter_name = format_ident!("{}_get_vtable", opts.name);
    let impl_macro = format_ident!("abi_vtable_impl_{}", trait_name_str.to_lowercase());

    // ---- host-side safe wrapper ----
    // A `{Name}Host` struct (`calc` → `CalcHost`) holding the vtable pointer
    // + ctx, exposing one safe method per trait method. The host applies the
    // same `#[abi_vtable]` to the shared trait declaration and gets this
    // wrapper for free — no raw `unsafe` fn-pointer calls required.
    let host_wrapper_name_str = capitalize(&opts.name);
    let host_wrapper = format_ident!("{}Host", host_wrapper_name_str);

    let mut host_methods = Vec::new();
    for m in &methods {
        let fname = &m.name;
        let arg_tys = &m.params;
        let ret = match &m.ret {
            Some(t) => quote!(-> #t),
            None => quote!(),
        };
        let call_args = &m.call_args;
        host_methods.push(quote! {
            pub fn #fname(&self, #(#arg_tys),*) #ret {
                unsafe { ((*self.vtable).#fname)(self.ctx, #(#call_args),*) }
            }
        });
    }

    let expanded = quote! {
        #trait_def

        /// Shadow vtable generated from [`#trait_name`] — the C contract.
        /// Field order IS the protocol; never reorder after release.
        #[repr(C)]
        #[derive(Clone, Copy)]
        #vis struct #vtable_struct_name {
            #(#struct_fields,)*
        }

        /// Host-side safe wrapper generated from [`#trait_name`].
        ///
        /// Holds the vtable pointer plus the instance context; exposes one
        /// safe method per trait method, so host code never writes `unsafe`
        /// fn-pointer calls. Build it from a loaded vtable:
        ///
        /// ```ignore
        /// let module = unsafe { AbiTable::<#vtable_struct_name>::load(path, b"...\0")? };
        /// let host = #host_wrapper::new(module.vtable(), std::ptr::null_mut());
        /// ```
        ///
        /// # Safety contract
        ///
        /// The vtable pointer and ctx must remain valid for as long as this
        /// wrapper exists — keep the library handle alive. Methods are safe
        /// because the generated thunks are plain C functions over C scalars.
        #vis struct #host_wrapper<'a> {
            vtable: &'a #vtable_struct_name,
            ctx: *mut ::std::ffi::c_void,
        }

        // SAFETY: the wrapper only reads the vtable (immutable reference) and
        // passes ctx to extern "C" thunks over C scalars. Thread-safety of the
        // underlying instance is the module's declared guarantee (same policy
        // as the fn-pointer struct itself).
        unsafe impl Send for #host_wrapper<'_> {}
        unsafe impl Sync for #host_wrapper<'_> {}

        impl<'a> #host_wrapper<'a> {
            /// Wrap a loaded vtable. `ctx` is the instance handle received
            /// from the module (`null` for stateless tables).
            pub fn new(
                vtable: &'a #vtable_struct_name,
                ctx: *mut ::std::ffi::c_void,
            ) -> Self {
                Self { vtable, ctx }
            }

            /// The underlying vtable (for hand-written fallback calls).
            pub fn vtable(&self) -> &#vtable_struct_name {
                self.vtable
            }

            /// The instance context pointer (first argument of every call).
            pub fn ctx(&self) -> *mut ::std::ffi::c_void {
                self.ctx
            }

            #(#host_methods)*
        }

        /// One-line binding: impl type → static instance + vtable + getter.
        ///
        /// `$impl_ty` must implement [`#trait_name`]; `$instance` is a
        /// const-constructible expression (e.g. `MyCalc`).
        #[macro_export]
        macro_rules! #impl_macro {
            // (ImplType, instance_expr)
            //
            // `$instance` is a const-constructible expression. The generated
            // static holds a *reference* (`static INSTANCE: &$impl_ty =
            // &$instance;`) — with a unit-struct / const-expression argument
            // Rust's static promotion makes this a single shared instance,
            // so the ctx pointer the host passes and this static alias the
            // same object. Non-Copy types are fine (nothing is moved).
            ($impl_ty:ty, $instance:expr) => {
                static #instance_name: &$impl_ty = &$instance;

                // Concrete thunks — one per method, bound to $impl_ty.
                #(
                    #thunks
                )*

                static #static_name: #vtable_struct_name = #vtable_struct_name {
                    #(#thunk_assignments,)*
                };

                #[no_mangle]
                pub extern "C" fn #getter_name() -> *const #vtable_struct_name {
                    &#static_name
                }
            };
        }
    };

    TokenStream::from(expanded)
}

/// Uppercase the first character of an ASCII identifier (`calc` → `Calc`).
fn capitalize(s: &str) -> String {
    let mut chars = s.chars();
    match chars.next() {
        Some(c) => c.to_ascii_uppercase().to_string() + chars.as_str(),
        None => String::new(),
    }
}