trait_ffi/

lib.rs

#![doc = include_str!("../README.md")]

use convert_case::Casing;
use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{format_ident, quote};
use syn::{Ident, ItemImpl, ItemTrait, parse_macro_input, spanned::Spanned};

macro_rules! bail {
    ($i:expr, $msg:expr) => {
        return syn::parse::Error::new($i, $msg).to_compile_error().into();
    };
}

fn get_crate_name() -> String {
    std::env::var("CARGO_PKG_NAME").unwrap_or_else(|_| "unknown".to_string())
}

fn get_crate_version() -> String {
    std::env::var("CARGO_PKG_VERSION").unwrap_or_else(|_| "0.1.0".to_string())
}

fn prefix_version() -> String {
    let version = lenient_semver::parse(&get_crate_version()).unwrap();
    let major = version.major;
    let minor = version.minor;
    if major == 0 {
        format!("0_{minor}")
    } else {
        major.to_string()
    }
}

fn extern_fn_name(crate_name: &str, fn_name: &Ident) -> Ident {
    let crate_name = crate_name.to_lowercase().replace("-", "_");
    let version = prefix_version();

    format_ident!("__{crate_name}_{version}_{fn_name}")
}

fn parse_def_extern_trait_args(args: TokenStream) -> Result<String, String> {
    if args.is_empty() {
        return Ok("rust".to_string()); // 默认使用 Rust ABI
    }

    let args_str = args.to_string();
    let mut abi = None;

    // 简单解析 abi="value" 形式
    let parts: Vec<&str> = args_str.split(',').collect();

    for part in parts {
        let part = part.trim();
        if part.starts_with("abi") {
            if let Some(start) = part.find('"') {
                if let Some(end) = part.rfind('"') {
                    if start < end {
                        abi = Some(part[start + 1..end].to_string());
                    }
                }
            }
        }
    }

    let abi = abi.unwrap_or_else(|| "rust".to_string());

    if abi != "c" && abi != "rust" {
        return Err("Invalid abi parameter. Supported values: \"c\", \"rust\"".to_string());
    }

    Ok(abi)
}

/// Defines an extern trait that can be called across FFI boundaries.
///
/// This macro converts a regular Rust trait into a trait that can be called through FFI.
/// It generates:
/// 1. The original trait definition
/// 2. A module containing wrapper functions that call external implementations
/// 3. A helper macro `impl_trait!` for implementing the trait
/// 4. A checker function to ensure the trait is properly implemented
///
/// # Arguments
/// - `abi`: Optional parameter specifying ABI type ("c" or "rust"), defaults to "rust"
///
/// # Example
/// ```rust
/// #[def_extern_trait(abi = "c")]
/// trait Calculator {
///     fn add(&self, a: i32, b: i32) -> i32;
///     fn multiply(&self, a: i32, b: i32) -> i32;
/// }
/// ```
///
/// This will generate a `calculator` module containing functions that can call external implementations.
#[proc_macro_attribute]
pub fn def_extern_trait(args: TokenStream, input: TokenStream) -> TokenStream {
    let abi = match parse_def_extern_trait_args(args) {
        Ok(abi) => abi,
        Err(error_msg) => {
            bail!(Span::call_site(), error_msg);
        }
    };

    let input = parse_macro_input!(input as ItemTrait);
    let vis = input.vis.clone();
    let mod_name = format_ident!(
        "{}",
        input.ident.to_string().to_case(convert_case::Case::Snake)
    );
    let crate_name_str = get_crate_name();

    let mut fn_list = vec![];

    for item in &input.items {
        if let syn::TraitItem::Fn(func) = item {
            let fn_name = func.sig.ident.clone();
            let extern_fn_name = extern_fn_name(&crate_name_str, &fn_name);

            let attrs = &func.attrs;
            let inputs = &func.sig.inputs;
            let output = &func.sig.output;
            let generics = &func.sig.generics;

            let mut param_names = vec![];
            let mut param_types = vec![];

            for input in inputs {
                if let syn::FnArg::Typed(pat_type) = input {
                    param_names.push(&pat_type.pat);
                    param_types.push(&pat_type.ty);
                }
            }

            let extern_abi = if abi == "rust" { "Rust" } else { "C" };

            fn_list.push(quote! {
                #(#attrs)*
                pub fn #fn_name #generics (#inputs) #output {
                    unsafe extern #extern_abi {
                        fn #extern_fn_name #generics (#inputs) #output;
                    }
                    unsafe{ #extern_fn_name(#(#param_names),*) }
                }
            });
        } else {
            bail!(
                item.span(),
                "Only function items are allowed in extern traits"
            );
        }
    }

    let crate_name = format_ident!("{}", crate_name_str.replace("-", "_"));

    let warn_fn_name = format_ident!(
        "Trait_{}_in_crate_{}_{}_need_impl",
        input.ident,
        crate_name_str.replace("-", "_"),
        prefix_version()
    );

    let generated_macro = quote! {
        #[macro_export]
        macro_rules! impl_trait {
            (impl $trait:ident for $type:ty { $($body:tt)* }) => {
                #[#crate_name::impl_extern_trait(name = #crate_name_str, abi = #abi)]
                impl $trait for $type {
                    $($body)*
                }

                #[allow(snake_case)]
                #[unsafe(no_mangle)]
                extern "C" fn #warn_fn_name() { }
            };
        }
    };

    quote! {
        pub use trait_ffi::impl_extern_trait;

        #input

        #vis mod #mod_name {
            use super::*;
            /// `trait-ffi` generated.
            pub fn ____checker_do_not_use(){
                unsafe extern "C" {
                    fn #warn_fn_name();
                }
                unsafe { #warn_fn_name() };
            }
            #(#fn_list)*
        }

        #generated_macro
    }
    .into()
}

fn parse_extern_trait_args(args: TokenStream) -> Result<(String, String), String> {
    if args.is_empty() {
        return Err(
            "Missing parameters. Usage: #[impl_extern_trait(name=\"crate_name\", abi=\"c\")]"
                .to_string(),
        );
    }

    let args_str = args.to_string();
    let mut name = None;
    let mut abi = None;

    let parts: Vec<&str> = args_str.split(',').collect();

    for part in parts {
        let part = part.trim();
        if part.starts_with("name") {
            if let Some(start) = part.find('"') {
                if let Some(end) = part.rfind('"') {
                    if start < end {
                        name = Some(part[start + 1..end].to_string());
                    }
                }
            }
        } else if part.starts_with("abi") {
            if let Some(start) = part.find('"') {
                if let Some(end) = part.rfind('"') {
                    if start < end {
                        abi = Some(part[start + 1..end].to_string());
                    }
                }
            }
        }
    }

    let name = name.ok_or_else(|| {
        "Missing name parameter. Usage: #[impl_extern_trait(name=\"crate_name\", abi=\"c\")]"
            .to_string()
    })?;
    let abi = abi.unwrap_or_else(|| "c".to_string());

    if abi != "c" && abi != "rust" {
        return Err("Invalid abi parameter. Supported values: \"c\", \"rust\"".to_string());
    }

    Ok((name, abi))
}

/// Implements an extern trait for a type and generates corresponding C function exports.
///
/// This macro takes a trait implementation and generates extern "C" functions that can be
/// called from other languages. Each method in the trait implementation gets a corresponding
/// extern function with a mangled name based on the crate name and version.
///
/// # Arguments
/// - `name`: The name of the crate that defines the extern trait
/// - `abi`: The ABI to use for the extern functions ("c" or "rust"), defaults to "c"
///
/// # Example
/// ```rust
/// struct Calculator;
///
/// #[impl_extern_trait(name = "calculator_crate", abi = "c")]
/// impl MyTrait for Calculator {
///     fn add(&self, a: i32, b: i32) -> i32 {
///         a + b
///     }
/// }
/// ```
///
/// This will generate extern "C" functions that can be called from other languages.
#[proc_macro_attribute]
pub fn impl_extern_trait(args: TokenStream, input: TokenStream) -> TokenStream {
    let (crate_name_str, abi) = match parse_extern_trait_args(args) {
        Ok((name, abi)) => (name, abi),
        Err(error_msg) => {
            bail!(Span::call_site(), error_msg);
        }
    };
    let input = parse_macro_input!(input as ItemImpl);
    let mut extern_fn_list = vec![];

    let struct_name = input.self_ty.clone();
    let trait_name = input.clone().trait_.unwrap().1;

    for item in &input.items {
        if let syn::ImplItem::Fn(func) = item {
            let fn_name_raw = &func.sig.ident;
            let fn_name = extern_fn_name(&crate_name_str, fn_name_raw);

            let inputs = &func.sig.inputs;
            let output = &func.sig.output;
            let generics = &func.sig.generics;

            let extern_abi = if abi == "rust" { "Rust" } else { "C" };

            let mut param_names = vec![];
            let mut param_types = vec![];

            for input in inputs {
                if let syn::FnArg::Typed(pat_type) = input {
                    param_names.push(&pat_type.pat);
                    param_types.push(&pat_type.ty);
                }
            }

            extern_fn_list.push(quote! {
                /// `trait-ffi` generated extern function.
                #[unsafe(no_mangle)]
                pub extern #extern_abi fn #fn_name #generics (#inputs) #output {
                    <#struct_name as #trait_name>::#fn_name_raw(#(#param_names),*)
                }
            });
        }
    }

    quote! {
        #input
        #(#extern_fn_list)*
    }
    .into()
}