mschema/

lib.rs

use proc_macro::TokenStream;
use quote::{quote, format_ident};
use syn::{ItemFn, Meta, parse::Parse, parse::ParseStream, Token, Lit, Expr, Type, FnArg, PatType, parse_macro_input};
use std::sync::Once;
use aes_types::{TaggedData, Error, ES_ERR_OK};

static INIT: Once = Once::new();
static mut SIGNATURES: Vec<String> = Vec::new();


// Custom parser for attribute arguments
struct AttributeArgs(Vec<Meta>);

impl Parse for AttributeArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut metas = Vec::new();
        while !input.is_empty() {
            let meta = input.parse()?;
            metas.push(meta);
            if !input.is_empty() {
                input.parse::<Token![,]>()?;
            }
        }
        Ok(AttributeArgs(metas))
    }
}

#[proc_macro_attribute]
pub fn ext_ver(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(input as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_block = &input_fn.block;

    let expanded = quote! {
        #[no_mangle]
        pub extern "C" fn ESGetVersion() -> i32 {
            #fn_name()
        }

        fn #fn_name() -> i32 {
            #fn_block
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_attribute]
pub fn ext_init(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(input as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_block = &input_fn.block;

    let expanded = quote! {
        #[no_mangle]
        pub extern "C" fn ESInitialize(
            data: *mut aes_externalobj::TaggedData,
            count: i32
        ) -> *mut i8 {
            // Call user initialization code
            #fn_name();

            // Return signatures
            let signatures = aes_externalobj::internal::get_signatures();
            let c_str = std::ffi::CString::new(signatures).unwrap();
            c_str.into_raw()
        }

        fn #fn_name() {
            #fn_block
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_attribute]
pub fn ext_term(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(input as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_block = &input_fn.block;

    let expanded = quote! {
        #[no_mangle]
        pub extern "C" fn ESTerminate() {
            #fn_name()
        }

        fn #fn_name() {
            #fn_block
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_attribute]
pub fn ext_export(attr: TokenStream, input: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as AttributeArgs);
    let input_fn = parse_macro_input!(input as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_block = &input_fn.block;

    // Get function arguments
    let mut arg_conversions = Vec::new();
    let mut arg_names = Vec::new();
    let mut arg_types = Vec::new();
    let arg_count = input_fn.sig.inputs.len();
    
    for (idx, arg) in input_fn.sig.inputs.iter().enumerate() {
        if let FnArg::Typed(PatType { pat, ty, .. }) = arg {
            let arg_name = quote!(#pat);
            arg_names.push(arg_name);
            
            // Get argument type for signature according to ExtendScript documentation
            let type_str = quote!(#ty).to_string();
            let type_char = match type_str.as_str() {
                "String" => "s",
                "i32" => "d",
                "u32" => "u",
                "f64" => "f",
                "bool" => "b",
                "JsScript" => "a", // JsScript is treated as any type in signatures
                "LiveObject" => "o", // LiveObject is treated as object in signatures
                _ => "a",
            };
            arg_types.push(type_char);
            
            // Generate conversion code for each argument
            let conversion = quote! {
                let #pat: #ty = match args[#idx].try_into() {
                    Ok(val) => val,
                    Err(_) => return Err(aes_externalobj::Error::InvalidArguments),
                };
            };
            
            arg_conversions.push(conversion);
        }
    }

    // Get original return type for the function implementation
    let return_type = if let syn::ReturnType::Type(_, ty) = &input_fn.sig.output {
        quote!(#ty)
    } else {
        quote!(())
    };

    let export_name = if args.0.is_empty() {
        fn_name.to_string()
    } else {
        match &args.0[0] {
            Meta::Path(path) => path.get_ident().map_or_else(
                || fn_name.to_string(),
                |ident| ident.to_string()
            ),
            Meta::NameValue(name_value) => {
                match &name_value.value {
                    Expr::Lit(expr_lit) => {
                        if let Lit::Str(lit_str) = &expr_lit.lit {
                            lit_str.value()
                        } else {
                            fn_name.to_string()
                        }
                    },
                    _ => fn_name.to_string()
                }
            },
            _ => fn_name.to_string()
        }
    };

    // Generate function signature (only with argument types)
    // If there are no arguments, don't add underscore
    let signature = if arg_types.is_empty() {
        export_name.clone()
    } else {
        format!("{}_{}", export_name, arg_types.join(""))
    };
    let signature_str = signature.clone();

    let export_ident = syn::Ident::new(&export_name, fn_name.span());
    let internal_fn_name = syn::Ident::new(&format!("__internal_{}", fn_name), fn_name.span());
    let arg_count_lit = syn::LitInt::new(&arg_count.to_string(), proc_macro2::Span::call_site());

    let expanded = quote! {
        const _: () = {
            #[doc(hidden)]
            #[ctor::ctor]
            fn __register_signature() {
                aes_externalobj::internal::add_signature(#signature_str.to_string());
            }
        };

        #[no_mangle]
        pub extern "C" fn #export_ident(
            args: *mut aes_externalobj::TaggedData,
            argc: i32,
            result: *mut aes_externalobj::TaggedData
        ) -> i32 {
            unsafe {
                let args_slice = std::slice::from_raw_parts(args, argc as usize);
                if args_slice.len() != #arg_count_lit {
                    return aes_externalobj::Error::InvalidArguments.into();
                }
                match #internal_fn_name(args_slice) {
                    Ok(return_value) => {
                        *result = return_value;
                        aes_externalobj::ES_ERR_OK
                    },
                    Err(e) => e.into()
                }
            }
        }

        fn #internal_fn_name(args: &[aes_externalobj::TaggedData]) -> aes_externalobj::Result<aes_externalobj::TaggedData> {
            #(#arg_conversions)*
            
            let user_result = {
                #fn_block
            };

            Ok(aes_externalobj::TaggedData::from(user_result))
        }
    };

    proc_macro::TokenStream::from(expanded)
}