ark-api-macros 0.13.0

use crate::utils;
use proc_macro2::Ident;
use proc_macro2::Span;
use proc_macro2::TokenStream;
use quote::ToTokens;
use syn::parse::Error;
use syn::parse::Parse;
use syn::parse::ParseStream;
use syn::parse::Result;
use syn::parse_quote;
use syn::spanned::Spanned;
use syn::Attribute;
use syn::FnArg;
use syn::ItemFn;
use syn::ItemMod;
use syn::ItemTrait;
use syn::Pat;
use syn::Signature;
use syn::Token;
use syn::TraitItemMethod;
use syn::Type;

#[derive(Clone)]
pub struct Args {
    /// imports = "<name>"
    imports: String,
}

impl Parse for Args {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        match try_parse(input) {
            Ok(args) if input.is_empty() => Ok(args),
            Ok(_) | Err(_) => Err(Error::new(
                Span::call_site(),
                "expected #[ark_bindgen(imports = \"<name>\"]",
            )),
        }
    }
}

mod kw {
    syn::custom_keyword!(imports);
}

fn try_parse(input: ParseStream<'_>) -> Result<Args> {
    if input.peek(kw::imports) {
        input.parse::<kw::imports>()?;
        input.parse::<Token![=]>()?;
        let name = input.parse::<proc_macro2::Literal>()?;
        let name = name.to_string();

        // Remove the surrounding quotes
        let name = &name[1..name.len() - 1];
        Ok(Args {
            imports: name.to_owned(),
        })
    } else {
        Err(Error::new(Span::call_site(), "no imports"))
    }
}

pub struct Item(ItemMod);

impl Parse for Item {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let mut lookahead = input.lookahead1();
        if lookahead.peek(Token![pub]) {
            let ahead = input.fork();
            ahead.parse::<Token![pub]>()?;
            lookahead = ahead.lookahead1();
        }
        if lookahead.peek(Token![mod]) {
            let mut item: ItemMod = input.parse()?;
            item.attrs = attrs;
            Ok(Self(item))
        } else {
            Err(lookahead.error())
        }
    }
}

impl ToTokens for Item {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.0.to_tokens(tokens);
    }
}

struct Context {
    mod_name: Ident,
}

pub fn expand(input: &mut Item, args: Args) -> Result<()> {
    validate_input_and_output_types(&input.0)?;
    validate_types(&input.0)?;

    let ctx = Context {
        mod_name: input.0.ident.clone(),
    };
    let items = extract_items_from_input(input)?;

    let extern_enums = translate_enums(items.1)?;
    let extern_funcs = expand_extern_functions(items.0, &extern_enums, &ctx, &args)?;
    expand_safe_mod(&extern_funcs, &extern_enums, input);
    expand_host_shim(&extern_funcs, &extern_enums, input, &ctx, &args)?;
    expand_use(&extern_enums, input);

    Ok(())
}

fn validate_input_and_output_types(item_mod: &ItemMod) -> Result<()> {
    use syn::visit;

    struct FindSigs(FindTypes);

    impl<'a> visit::Visit<'a> for FindSigs {
        fn visit_signature(&mut self, node: &'a Signature) {
            visit::visit_signature(&mut self.0, node);
        }
    }

    struct FindTypes(Option<Error>);

    impl<'a> visit::Visit<'a> for FindTypes {
        fn visit_fn_arg(&mut self, node: &'a FnArg) {
            // Base implementation
            visit::visit_fn_arg(self, node);

            if let syn::FnArg::Typed(pt) = node {
                if let syn::Type::Path(path) = &*pt.ty {
                    for (invalid_type, suggested_type) in [
                        ("bool", "u32"),
                        ("u16", "u32"),
                        ("u8", "u32"),
                        ("i16", "i32"),
                        ("i8", "i32"),
                    ] {
                        if path.path.is_ident(invalid_type) {
                            self.0 = Some(Error::new_spanned(
                                path,
                                format!("FFI parameters of type `{invalid_type}` are not supported at the moment. Pass it as a `{suggested_type}` instead.")
                            ));
                            break; // NOTE[TSolberg]: We cannot have multiple spanned errors so we only check until we hit an invalid type.
                        }
                    }
                }
            }
        }
    }

    let mut sigs = FindSigs(FindTypes(None));
    visit::visit_item_mod(&mut sigs, item_mod);

    if let Some(error) = (sigs.0).0 {
        Err(error)
    } else {
        Ok(())
    }
}

fn validate_types(item_mod: &ItemMod) -> Result<()> {
    use syn::visit;
    use syn::visit::Visit;
    const DISALLOWED_TYPES: &[&str] = &["usize", "isize", "u128", "i128"];

    struct TypeVisitor(Option<Error>);
    impl<'a> visit::Visit<'a> for TypeVisitor {
        fn visit_type(&mut self, t: &'a Type) {
            visit::visit_type(self, t);
            if let Type::Path(p) = t {
                for ident in DISALLOWED_TYPES {
                    if p.path.is_ident(ident) {
                        let e = Error::new_spanned(p, format!("`{ident}` is not FFI-safe"));
                        match &mut self.0 {
                            Some(e2) => e2.combine(e),
                            None => self.0 = Some(e),
                        }
                    }
                }
            }
        }
    }

    let mut visitor = TypeVisitor(None);
    visitor.visit_item_mod(item_mod);

    match visitor.0 {
        Some(e) => Err(e),
        None => Ok(()),
    }
}

enum PreciseType {
    /// Returning a `std::vec::Vec<u8>`.
    ByteVec,
    /// Returning a `std::string::String`. Encoding and decoding are handled by the proc macros.
    String,
}

impl PreciseType {
    fn make_type_asserts(
        &self,
        fallible_mode: FallibleMode,
        original_result_type: &syn::Type,
    ) -> syn::Stmt {
        let expected_type = match self {
            Self::ByteVec => quote::quote!(::std::vec::Vec<u8>),
            Self::String => quote::quote!(::std::string::String),
        };
        let expected_type = if fallible_mode.is_fallible() {
            quote::quote!(crate::error_code::FFIResult<#expected_type>)
        } else {
            expected_type
        };
        parse_quote!(
            static_assertions::assert_type_eq_all!(
                #original_result_type,
                #expected_type
            );
        )
    }
}

enum ReturnType {
    /// Function returns a very specific type, that is type-checked at compile-time.
    PreciseType(PreciseType),
    /// Function return `Result<T, ...>` in fallible mode or `T`, where `T` is a plain-old data
    /// type (in particular, it must implement `Copy`).
    GenericPod(syn::TypePath),
    /// Function returns `Result<(), ...>` in fallible mode, or nothing.
    UnitType,
}

impl ReturnType {
    fn make_type_asserts(
        &self,
        fallible_mode: FallibleMode,
        original_return_type: &syn::Type,
    ) -> syn::Stmt {
        match self {
            Self::PreciseType(byte_vec) => {
                byte_vec.make_type_asserts(fallible_mode, original_return_type)
            }

            Self::GenericPod(plain_ty) => {
                // If the function is fallible, ensure the `Result` type is specialized with our
                // custom error type.
                let expected_type_assert = if fallible_mode.is_fallible() {
                    Some(quote::quote!(
                        static_assertions::assert_type_eq_all!(
                            crate::error_code::FFIResult<#plain_ty>,
                            #original_return_type
                        );
                    ))
                } else {
                    // No need to check that the type is equal to itself!
                    None
                };
                // Ensure the type implements `Copy`.
                parse_quote!({
                    trait ValidReturn {}
                    impl<T: Copy + Clone> ValidReturn for T {}
                    static_assertions::assert_impl_all!(#plain_ty: ValidReturn);
                    #expected_type_assert
                })
            }

            Self::UnitType => {
                // Sanity-check: infallible functions returning () don't even have an
                // `original_return_type` set.
                assert!(fallible_mode.is_fallible());
                // Just check the `Result` type is the one we expect.
                parse_quote!({
                    static_assertions::assert_type_eq_all!(
                        crate::error_code::FFIResult<()>,
                        #original_return_type
                    );
                })
            }
        }
    }
}

#[derive(Clone, Copy)]
enum FallibleMode {
    /// Shim/export return a `Result<T, ApiError>`, FFI function returns an `ErrorCode`.
    Fallible,
    /// Shim/export return a `anyhow::Result<T>`, FFI function returns an `ErrorCode`.
    DeprecatedInfallible,
    /// Shim/export return a `anyhow::Result<T>`, FFI function returns nothing.
    Infallible,
}

impl FallibleMode {
    fn is_fallible(self) -> bool {
        matches!(self, Self::Fallible)
    }

    fn check_ffi_return(self, res_code: syn::Ident) -> Option<syn::Stmt> {
        match self {
            Self::Fallible => Some(parse_quote!(if #res_code != crate::ErrorCode::Success {
                return Err(#res_code);
            })),
            Self::DeprecatedInfallible => Some(parse_quote!(assert_eq!(
                #res_code,
                crate::ErrorCode::Success,
                "unexpected error in deprecated infallible function"
            );)),
            Self::Infallible => None,
        }
    }

    fn ensure_ffi_success(self, res_code: syn::Ident) -> syn::Stmt {
        match self {
            Self::Fallible => {
                parse_quote!(if #res_code != crate::ErrorCode::Success {
                    return Err(#res_code);
                })
            }
            Self::DeprecatedInfallible | Self::Infallible => {
                parse_quote!(assert_eq!(
                    #res_code,
                    crate::ErrorCode::Success,
                    "unexpected error in deprecated infallible function"
                );)
            }
        }
    }

    fn return_result(self, result: syn::Ident) -> syn::Stmt {
        match self {
            Self::Fallible => parse_quote!(return Ok(#result);),
            Self::DeprecatedInfallible | Self::Infallible => {
                parse_quote!(return #result;)
            }
        }
    }

    fn return_unit(self) -> Option<syn::Stmt> {
        match self {
            Self::Fallible => Some(parse_quote!(return Ok(());)),
            Self::DeprecatedInfallible | Self::Infallible => None,
        }
    }
}

struct ExternFn {
    /// The original, Rusty, signature, for the high-level `safe` variant.
    sig: Signature,

    /// The attributes (e.g. #[cfg]) on the original function.
    attrs: Vec<Attribute>,

    /// Should the shim method take a `Memory` parameter?
    with_memory: bool,

    fallible_mode: FallibleMode,

    return_type: ReturnType,
    params: Vec<FfiParam>,

    /// The identifier for the FFI function generated from the original (e.g. `module__func_name`).
    ffi_ident: Ident,
}

impl ExternFn {
    /// Create type assertions to ensure that the AST-declared types are actually the ones we
    /// expect.
    /// Grep "Note about fraudulent type paths" in this file to understand where these are coming
    /// from.
    fn make_type_asserts(&self) -> Option<syn::Stmt> {
        // What was the original return type, used in the high-level declaration?
        let original_return_type = if let syn::ReturnType::Type(_, typ) = &self.sig.output {
            &**typ
        } else {
            // Returns ().
            assert!(matches!(self.return_type, ReturnType::UnitType));
            return None;
        };
        Some(
            self.return_type
                .make_type_asserts(self.fallible_mode, original_return_type),
        )
    }
}

pub(crate) struct ExternEnum {
    /// The identifier of the original enum
    pub ident: Ident,
    /// The identifier to be used as a FFI function argument
    /// Extracted from #[repr(primitive_type)]
    ffi_ty: Type,
}

fn extract_items_from_input(
    input: &mut Item,
) -> Result<(
    &mut syn::ItemForeignMod,
    Vec<&mut syn::ItemEnum>,
    Vec<&syn::ItemStruct>,
)> {
    let foreign_item;
    let mut enum_items = Vec::new();
    let mut struct_items = Vec::new();

    // Dodge the borrow checker here for error messages.
    let cloned_input = input.0.clone();

    match &mut input.0.content {
        Some(items) => {
            let mut foreign = None;
            for item in &mut items.1 {
                match item {
                    // We expect exactly **1** item, an extern "C" block
                    syn::Item::ForeignMod(fmod) => match foreign {
                        Some(_) => {
                            return Err(Error::new(
                                fmod.span(),
                                "an extern module was already declared in this module",
                            ));
                        }
                        None => foreign = Some(fmod),
                    },
                    syn::Item::Enum(ienum) => enum_items.push(ienum),
                    syn::Item::Struct(istruct) => struct_items.push(istruct as &syn::ItemStruct),
                    _ => {}
                }
            }

            if let Some(fmod) = foreign {
                foreign_item = fmod;
            } else {
                return Err(Error::new(
                    cloned_input.span(),
                    "the module doesn't contain an extern module",
                ));
            }
        }
        None => {
            return Err(Error::new(
                cloned_input.span(),
                "can't bindgen an empty module",
            ))
        }
    }

    Ok((foreign_item, enum_items, struct_items))
}

fn translate_enums(enum_items: Vec<&mut syn::ItemEnum>) -> Result<Vec<ExternEnum>> {
    let mut extern_enums = Vec::with_capacity(enum_items.len());

    for enum_item in enum_items {
        let mut enum_ffi_ident = None;

        // We expect exactly **1** #[repr(primitive_type)] attribute
        for attribute in &enum_item.attrs {
            let ident = utils::search_enum_attribute(
                attribute,
                "repr",
                &[
                    "u8", "i8", "u16", "i16", "u32", "i32", "u64", "i64", "u128", "i128",
                ],
            );

            if ident.is_some() {
                if enum_ffi_ident.is_none() {
                    enum_ffi_ident = ident;
                } else {
                    return Err(Error::new(
                        attribute.span(),
                        "#[repr(primitive_type)] was already declared for this enum",
                    ));
                }
            }

            // #[derive(IntoPrimitive, TryFromPrimitive)] are generated automatically!
            if utils::search_enum_attribute(
                attribute,
                "derive",
                &["TryFromPrimitive", "IntoPrimitive"],
            )
            .is_some()
            {
                return Err(Error::new(attribute.span(),
                "IntoPrimitive and TryFromPrimitive are generated by the ark_bindgen macro, don't declare them yourself!"));
            }

            // #[derive(NoUninit, CheckedBitPattern)] are generated automatically!
            if utils::search_enum_attribute(attribute, "derive", &["NoUninit", "CheckedBitPattern"])
                .is_some()
            {
                return Err(Error::new(attribute.span(),
                "For enums, NoUninit and CheckedBitPattern are generated by the ark_bindgen macro, don't declare them yourself!"));
            }
        }

        let Some(ffi_ident) = enum_ffi_ident else {
            return Err(Error::new(
                enum_item.span(),
                format!(
                    "{} enum is missing the '#[repr(primitive_type)]' attribute",
                    enum_item.ident
                ),
            ));
        };

        enum_item.attrs.push(
            parse_quote!(#[derive(IntoPrimitive, TryFromPrimitive, NoUninit, CheckedBitPattern)]),
        );

        extern_enums.push(ExternEnum {
            ident: enum_item.ident.clone(),
            ffi_ty: parse_quote!(#ffi_ident),
        });
    }

    Ok(extern_enums)
}

/// Converts functions in the extern "C" block into their FFI forms, and returns the list of those
/// functions for further processing
fn expand_extern_functions(
    foreign_mod: &mut syn::ItemForeignMod,
    extern_enums: &[ExternEnum],
    ctx: &Context,
    args: &Args,
) -> Result<Vec<ExternFn>> {
    let imports = &args.imports;
    foreign_mod
        .attrs
        .push(parse_quote!(#[link(wasm_import_module = #imports)]));

    let extern_items = &mut foreign_mod.items;

    let mut externs = Vec::new();

    let with_memory_attr: Attribute = parse_quote!(#[with_memory]);
    let deprecated_infallible_attr: Attribute = parse_quote!(#[deprecated_infallible]);

    for func in extern_items.iter_mut().filter_map(|fi| {
        if let syn::ForeignItem::Fn(func) = fi {
            Some(func)
        } else {
            None
        }
    }) {
        // Parse custom attributes, and remove them before Rust sees them.
        let mut with_memory = false;
        let mut deprecated_infallible = false;
        func.attrs.retain(|attr| {
            if utils::has_custom_attribute(attr, &with_memory_attr) {
                with_memory = true;
                false
            } else if utils::has_custom_attribute(attr, &deprecated_infallible_attr) {
                deprecated_infallible = true;
                false
            } else {
                true
            }
        });

        // Keep a copy of the original signature, the "safe" method we generate will have exactly
        // the same.
        let original_sig = func.sig.clone();

        let (params, return_type, fallible_mode) = to_ffi_sig(
            &ctx.mod_name,
            &mut func.sig,
            extern_enums,
            deprecated_infallible,
        )?;

        externs.push(ExternFn {
            sig: original_sig,
            attrs: func.attrs.clone(),
            with_memory,
            fallible_mode,
            return_type,
            params,
            ffi_ident: func.sig.ident.clone(),
        });
    }

    Ok(externs)
}

/// Modifies in-place a 'normal' rust function signature into a FFI compatible one.
/// Returns the components of the signature of the shim wrapper.
///
/// Input:
///
/// ```ignore
/// pub fn launch(
///     fn_name: &str,
///     input: &[u8],
/// ) -> FFIResult<RequestHandle>;
/// ```
///
/// Output:
///
/// ```ignore
/// pub fn module_run__launch(
///     fn_name_ptr: *const u8,
///     fn_name_len: u32,
///     input_ptr: *const u8,
///     input_len: u32,
///     __ark_ffi_output: *mut RequestHandle,
/// ) -> ErrorCode;
/// ```
fn to_ffi_sig(
    namespace: &Ident,
    sig: &mut Signature,
    extern_enums: &[ExternEnum],
    deprecated_infallible: bool,
) -> Result<(Vec<FfiParam>, ReturnType, FallibleMode)> {
    let mut params = Vec::new();

    sig.ident = syn::Ident::new(&format!("{namespace}__{}", sig.ident), sig.ident.span());

    for param in &sig.inputs {
        match param {
            FnArg::Typed(pt) => {
                if let Pat::Ident(pat) = &*pt.pat {
                    convert_ffi_param(&mut params, &pat.ident, pt.ty.as_ref(), extern_enums)?;
                } else {
                    return Err(Error::new(pt.span(), "parameter is missing an identifier"));
                }
            }
            other @ FnArg::Receiver(_) => {
                return Err(Error::new(
                    other.span(),
                    "you're trying to pass a self via FFI, this will not work",
                ));
            }
        }
    }

    let (return_type, infallible) = return_to_ffi(sig, &mut params)?;

    if deprecated_infallible && !infallible {
        return Err(Error::new(
            sig.output.span(),
            "#[deprecated_infallible] only applies to plain, non-Result return types",
        ));
    }

    let fallible_mode = if deprecated_infallible {
        FallibleMode::DeprecatedInfallible
    } else if infallible {
        FallibleMode::Infallible
    } else {
        FallibleMode::Fallible
    };

    // Tweak the FFI's return type: fallible and deprecated infallible return a status code, while
    // infallible don't return anything.
    sig.output = match fallible_mode {
        FallibleMode::Fallible | FallibleMode::DeprecatedInfallible => {
            parse_quote!(-> crate::ErrorCode)
        }
        FallibleMode::Infallible => syn::ReturnType::Default,
    };

    // Nuke the existing parameters and recreate them from the new, actual, FFI parameters
    sig.inputs = params.iter().map(|fp| fp.param.clone()).collect();

    Ok((params, return_type, fallible_mode))
}

/// Create a FFI-compatible signature, possibly by adding one or more output parameters.  Also
/// returns the `Result` type used in the shim function.
fn return_to_ffi(sig: &Signature, params: &mut Vec<FfiParam>) -> Result<(ReturnType, bool)> {
    let ty = match &sig.output {
        syn::ReturnType::Default => {
            return Ok((ReturnType::UnitType, true));
        }
        syn::ReturnType::Type(_, ty) => match ty.as_ref() {
            Type::Path(tp) => tp.clone(),
            _ => {
                return Err(Error::new(
                    ty.span(),
                    "unhandled return kind: can only return plain types",
                ));
            }
        },
    };

    if utils::type_path_ends_with(&ty, "ErrorCode") {
        return Err(Error::new(
            ty.span(),
            "Returning ErrorCode is deprecated, return FFIResult<()> instead",
        ));
    }

    // Note about fraudulent type paths: the user of this API could try to mistake us by using a
    // type alias that's named `FFIResult` (e.g. `use any_crate::any_type as FFIResult`), and pass
    // something else here. Macros are only operating at the AST level, so we don't have type
    // information yet, so what we'll do instead is add a static assertion in the code later, that
    // the type was precisely the one we expected.
    let (ret_type_path, infallible) = if utils::type_path_ends_with(&ty, "FFIResult") {
        let ok_type = utils::extract_single_generic_type(&ty)?;
        match ok_type {
            Some(type_path) => (type_path, false),
            None => {
                return Ok((ReturnType::UnitType, false));
            }
        }
    } else if utils::type_path_ends_with(&ty, "Result") {
        // Error type must be "ErrorCode".
        match utils::extract_generic_type(1, Some(2), &ty)? {
            Some(err_ty) if utils::type_path_ends_with(&err_ty, "ErrorCode") => {}
            _ => return Err(Error::new(ty.span(), "Result error type must be ErrorCode")),
        };
        let ok_type = utils::extract_first_generic_type(&ty)?;
        match ok_type {
            Some(type_path) => (type_path, false),
            None => {
                return Ok((ReturnType::UnitType, false));
            }
        }
    } else {
        (ty, true)
    };

    // Add an extra parameter for the return value.
    let (out_param, ret_type, reserved_param_name) = convert_ffi_result(ret_type_path)?;

    // We've already converted the original parameters, so check to make sure none of them were
    // called the same as the reserved name used for the new out-parameter. Otherwise, things will
    // get confusing and it's easier to get the user to change the parameter name than it is to
    // generate a different identifier if they've already named one of theirs the same
    let should_rename = params.iter().find(|p| {
        if let FnArg::Typed(pat) = &p.param {
            if let Pat::Ident(pat) = &*pat.pat {
                return pat.ident == reserved_param_name;
            }
        }
        false
    });
    if let Some(should_rename) = should_rename {
        return Err(Error::new(
            should_rename.param.span(),
            "this name is reserved, please rename!",
        ));
    }

    params.push(out_param);

    Ok((ret_type, infallible))
}

/// A parameter or return value in an FFI function.
struct FfiParam {
    /// Formal parameter in the FFI function's signature.
    param: syn::FnArg,
    /// Formal parameter in the _export function: none if it it's the same as the type
    /// used in the `param`, `Some(other_type)` if it's different.
    export_type: Option<Type>,
    /// Actual parameter in FFI function's call sites: any block that converts from a regular Rust
    /// item into the FFI expected type.
    to_ffi: syn::Expr,
}

fn is_str(ty: &syn::Type) -> bool {
    if let Type::Path(tp) = ty {
        match tp.path.get_ident() {
            None => false,
            Some(id) => {
                let idents = id.to_string();
                idents == "str"
            }
        }
    } else {
        false
    }
}

/// Converts a unique return value to a single FFI out-parameter that will be appended to the
/// export function signatures.
fn convert_ffi_result(tp: syn::TypePath) -> Result<(FfiParam, ReturnType, &'static str)> {
    // Grep "Note about fraudulent type paths" in this file.
    if utils::type_path_ends_with(&tp, "String") {
        // Special casing! We'll write the size of the vector here, and the safe
        // wrapper will use it to consume the host return vector.
        Ok((
            FfiParam {
                param: parse_quote!(__ark_byte_size: *mut u32),
                to_ffi: parse_quote!(&mut __ark_byte_size),
                export_type: Some(parse_quote!(u32)),
            },
            ReturnType::PreciseType(PreciseType::String),
            "__ark_byte_size",
        ))
    } else if utils::type_path_ends_with(&tp, "Vec") {
        // Special casing! We'll write the size of the vector here, and the safe
        // wrapper will use it to consume the host return vector.
        match utils::extract_single_generic_type(&tp)? {
            Some(tp) if utils::type_path_ends_with(&tp, "u8") => {}
            _ => {
                return Err(Error::new(
                    tp.span(),
                    "only supported return Vec type is Vec<u8>",
                ))
            }
        };
        Ok((
            FfiParam {
                param: parse_quote!(__ark_byte_size: *mut u32),
                export_type: Some(parse_quote!(u32)),
                to_ffi: parse_quote!(&mut __ark_byte_size),
            },
            ReturnType::PreciseType(PreciseType::ByteVec),
            "__ark_byte_size",
        ))
    } else {
        Ok((
            FfiParam {
                param: parse_quote!(__ark_ffi_output: *mut #tp),
                export_type: Some(parse_quote!(u32)),
                to_ffi: parse_quote!(&mut __ark_ffi_output),
            },
            ReturnType::GenericPod(parse_quote!(#tp)),
            "__ark_ffi_output",
        ))
    }
}

/// Converts a parameter (not returns) to a single or group of FFI parameters that will be
/// appended to the export function signature.
fn convert_ffi_param(
    params: &mut Vec<FfiParam>,
    ident: &syn::Ident,
    ty: &syn::Type,
    extern_enums: &[ExternEnum],
) -> Result<()> {
    match ty {
        syn::Type::Path(tp) => {
            let (param, to_ffi) = match utils::type_path_is_enum(tp, extern_enums) {
                Some(ee) => {
                    let enum_ty = &ee.ffi_ty;
                    (parse_quote!(#ident: #enum_ty), parse_quote!(#ident.into()))
                }
                None => {
                    if utils::type_path_ends_with(tp, "Vec") {
                        return Err(syn::Error::new(
                            ty.span(),
                            "Vec not supported in function parameter position, use slices instead",
                        ));
                    } else {
                        (parse_quote!(#ident: #ty), parse_quote!(#ident))
                    }
                }
            };

            params.push(FfiParam {
                param,
                to_ffi,
                export_type: None,
            });
        }

        syn::Type::Reference(tr) => {
            let is_mut = tr.mutability.is_some();

            // We only support 3 reference types, scalar types, slices of
            // scalar types, str, and cstr
            if is_str(tr.elem.as_ref()) {
                if is_mut {
                    return Err(syn::Error::new(
                        tr.span(),
                        "&mut str is not allowed, consider returning a String instead!",
                    ));
                }

                let ident_ptr = syn::Ident::new(&format!("{ident}_ptr"), ident.span());
                let ident_len = syn::Ident::new(&format!("{ident}_len"), ident.span());

                params.push(FfiParam {
                    param: parse_quote!(#ident_ptr: *const u8),
                    to_ffi: parse_quote!(#ident.as_ptr()),
                    export_type: Some(parse_quote!(u32)),
                });
                params.push(FfiParam {
                    param: parse_quote!(#ident_len: u32),
                    to_ffi: parse_quote!(#ident.len() as u32),
                    export_type: None,
                });
            } else if let syn::Type::Slice(inner) = tr.elem.as_ref() {
                if let syn::Type::Path(tp) = inner.elem.as_ref() {
                    let ident_ptr = syn::Ident::new(&format!("{ident}_ptr"), ident.span());
                    let ident_len = syn::Ident::new(&format!("{ident}_len"), ident.span());

                    params.push(FfiParam {
                        param: if is_mut {
                            parse_quote!(#ident_ptr: *mut #tp)
                        } else {
                            parse_quote!(#ident_ptr: *const #tp)
                        },
                        to_ffi: if is_mut {
                            parse_quote!(#ident.as_mut_ptr())
                        } else {
                            parse_quote!(#ident.as_ptr())
                        },
                        export_type: Some(parse_quote!(u32)),
                    });
                    params.push(FfiParam {
                        param: parse_quote!(#ident_len: u32),
                        to_ffi: parse_quote!(#ident.len() as u32),
                        export_type: None,
                    });
                } else {
                    return Err(Error::new(tr.elem.span(), "not a simple type path"));
                }
            } else if let syn::Type::Path(tp) = tr.elem.as_ref() {
                let ident_ptr = syn::Ident::new(&format!("{ident}_ptr"), ident.span());

                params.push(FfiParam {
                    param: if is_mut {
                        parse_quote!(#ident_ptr: *mut #tp)
                    } else {
                        parse_quote!(#ident_ptr: *const #tp)
                    },
                    to_ffi: parse_quote!(#ident),
                    export_type: Some(parse_quote!(u32)),
                });
            } else {
                return Err(Error::new(tr.span(), "this type is not supported"));
            }
        }

        _ => return Err(Error::new(ty.span(), "this type is not supported")),
    }

    Ok(())
}

/// Introduces a new submodule called `safe` which contains a "safe" wrapper function for
/// each of the extern functions, whose sole responsibility is converting the rust types
/// to the FFI types and invoking the extern function. They are all inlined.
fn expand_safe_mod(functions: &[ExternFn], enums: &[ExternEnum], input: &mut Item) {
    let mut safe_funcs = Vec::with_capacity(functions.len());

    for func in functions {
        let sig = &func.sig;
        let mut safe_func: ItemFn = parse_quote!(
            #[inline]
            pub #sig {
            }
        );
        safe_func.attrs.extend_from_slice(&func.attrs);

        // Generate the args that we'll be passing to the FFI function.
        let args = {
            let mut args: syn::punctuated::Punctuated<syn::Expr, syn::token::Comma> =
                syn::punctuated::Punctuated::new();
            for param in &func.params {
                args.push(param.to_ffi.clone());
            }
            args
        };

        let ffi_ident = &func.ffi_ident;

        let type_asserts = func.make_type_asserts();

        safe_func.block = match &func.return_type {
            ReturnType::PreciseType(inner) => {
                // Note: in theory, the `core__take_host_return_vec` shouldn't be used directly
                // here, but remember that this macro is invoked in the `ark-api-ffi` crate. The
                // high-level helper would live in `ark-api`, which can't be depended upon by the
                // FFI crate, as the former already depends upon the latter (and there'd be a
                // dependency cycle).

                // Tail expression invoked after retrieving the Vec<u8> `buffer`.
                let convert_buffer: Option<syn::Stmt> = match inner {
                    // Raw bytes are meant to just be returned as is
                    PreciseType::ByteVec => None,
                    // Reinterpret the raw bytes as a string.
                    PreciseType::String => Some(if func.fallible_mode.is_fallible() {
                        parse_quote!(
                            let buffer = String::from_utf8(buffer)
                                .map_err(|_decode_err| crate::ErrorCode::InternalError)?;
                        )
                    } else {
                        parse_quote!(
                            let buffer = String::from_utf8(buffer)
                                .expect("invalid utf8 bytes returned by infallible host function");
                        )
                    }),
                };

                let check_ffi_return = func.fallible_mode.check_ffi_return(parse_quote!(res_code));
                let check_core_result = func
                    .fallible_mode
                    .ensure_ffi_success(parse_quote!(res_code));
                let return_expr = func.fallible_mode.return_result(parse_quote!(buffer));

                parse_quote!({
                    #type_asserts;

                    // First, call the API that will produce the vector.
                    let mut __ark_byte_size = 0;
                    let res_code = unsafe { #ffi_ident(#args) };
                    #check_ffi_return

                    // Then, retrieve the vector produced by the above call.
                    let mut buffer = vec![0; __ark_byte_size as usize];
                    let res_code = unsafe {
                        crate::core_v4::core__take_host_return_vec(
                            buffer.as_mut_slice().as_mut_ptr(),
                            __ark_byte_size
                        )
                    };

                    #check_core_result

                    // Then return the buffer.
                    #convert_buffer
                    #return_expr
                })
            }

            ReturnType::GenericPod(type_path) => {
                // Grep "Note about fraudulent type paths" in this file to understand where this
                // assertion is coming from.
                let mut output: syn::Expr = parse_quote!(#type_path::default());
                if let Some(ee) = utils::type_path_is_enum(type_path, enums) {
                    if let Type::Path(tp) = &ee.ffi_ty {
                        let ident = &ee.ident;
                        output = if func.fallible_mode.is_fallible() {
                            parse_quote!(
                                #ident::try_from(#tp::default())
                                    .map_err(|_| crate::ErrorCode::InvalidArguments)?
                            )
                        } else {
                            parse_quote!(
                                #ident::try_from(#tp::default())
                                    .expect("invalid enum value")
                            )
                        };
                    }
                }

                let check_ffi_return = func.fallible_mode.check_ffi_return(parse_quote!(res_code));
                let return_expr = func
                    .fallible_mode
                    .return_result(parse_quote!(__ark_ffi_output));

                parse_quote!({
                    #type_asserts;
                    let mut __ark_ffi_output = #output;
                    let res_code = unsafe { #ffi_ident(#args) };
                    #check_ffi_return
                    #return_expr
                })
            }

            ReturnType::UnitType => {
                let check_ffi_return = func.fallible_mode.check_ffi_return(parse_quote!(res_code));
                let return_unit = func.fallible_mode.return_unit();
                parse_quote!({
                    #type_asserts;
                    let res_code = unsafe { #ffi_ident(#args) };
                    #check_ffi_return
                    #return_unit
                })
            }
        };

        safe_funcs.push(syn::Item::Fn(safe_func));
    }

    input.0.content.as_mut().unwrap().1.append(&mut safe_funcs);
}

/// Introduces a new submodule called `host` which contains a single trait called
/// `Shim` that contains safe methods for each of the exports.
fn expand_host_shim(
    functions: &[ExternFn],
    extern_enums: &[ExternEnum],
    input: &mut Item,
    ctx: &Context,
    args: &Args,
) -> Result<()> {
    let mut shim_trait: ItemTrait = parse_quote!(
        #[cfg(not(target_arch = "wasm32"))]
        pub trait HostShim<'t> {
            type Memory;
            type Context;
            type WasmLinker;
            type ImportError;
            type ApiError;
        }
    );

    for func in functions {
        let sig = &func.sig;

        // Add the shim method
        {
            let method_ident = Ident::new(&format!("{}_shim", sig.ident), sig.span());

            let ok_type: syn::Type = match &func.return_type {
                ReturnType::PreciseType(inner) => match inner {
                    PreciseType::ByteVec => parse_quote!(Vec<u8>),
                    PreciseType::String => parse_quote!(String),
                },
                ReturnType::GenericPod(ty) => parse_quote!(#ty),
                ReturnType::UnitType => parse_quote!(()),
            };

            let shim_return: syn::Type = if func.fallible_mode.is_fallible() {
                parse_quote!(Result<#ok_type, Self::ApiError>)
            } else {
                parse_quote!(anyhow::Result<#ok_type>)
            };

            let mut method: TraitItemMethod = parse_quote!(
                fn #method_ident(&mut self) -> #shim_return;
            );
            method.attrs.extend_from_slice(&func.attrs);

            let params = &mut method.sig.inputs;

            if func.with_memory {
                let arg: FnArg = parse_quote!(memory: &mut Self::Memory);
                params.push(arg);
            }

            for arg in &sig.inputs {
                params.push(arg.clone());
            }

            // Allow `clippy::too_many_arguments` if needed.
            if method.sig.inputs.len() > 7 {
                let clippy_attr: Attribute = parse_quote!(#[allow(clippy::too_many_arguments)]);
                method.attrs.push(clippy_attr);
            }

            shim_trait.items.push(syn::TraitItem::Method(method));
        }

        // Add the export method
        {
            let method_ident = Ident::new(&format!("{}_export", sig.ident), sig.span());

            let res_type: Type = if func.fallible_mode.is_fallible() {
                parse_quote!(Result<(), Self::ApiError>)
            } else {
                parse_quote!(anyhow::Result<()>)
            };

            let mut method: TraitItemMethod = parse_quote!(
                fn #method_ident<'a>(
                    memory: &mut Self::Memory,
                    host_context: &mut Self::Context
                ) -> #res_type;
            );

            method.attrs.extend_from_slice(&func.attrs);

            let formal_params = &mut method.sig.inputs;

            for param in &func.params {
                match &param.param {
                    FnArg::Typed(pat_type) => {
                        if let Pat::Ident(pat) = &*pat_type.pat {
                            let ty = match &param.export_type {
                                Some(ty) => ty,
                                None => pat_type.ty.as_ref(),
                            };
                            let ident = &pat.ident;
                            formal_params.push(parse_quote!(#ident: #ty));
                        } else {
                            return Err(Error::new(
                                pat_type.span(),
                                "parameter is missing an identifier",
                            ));
                        }
                    }

                    other @ FnArg::Receiver(_) => {
                        return Err(Error::new(
                            other.span(),
                            "you're trying to pass a self via FFI, this will not work",
                        ));
                    }
                }
            }

            // Allow `clippy::too_many_arguments` if needed.
            if method.sig.inputs.len() > 7 {
                let clippy_attr: Attribute = parse_quote!(#[allow(clippy::too_many_arguments)]);
                method.attrs.push(clippy_attr);
            }

            shim_trait.items.push(syn::TraitItem::Method(method));
        }
    }

    // Add the import table methods
    {
        {
            let imports_ident = Ident::new("imports", shim_trait.span());

            let method: TraitItemMethod = parse_quote!(
                fn #imports_ident(linker: &mut Self::WasmLinker) -> Result<(), Self::ImportError>;
            );

            shim_trait.items.push(syn::TraitItem::Method(method));
        }

        {
            let ns_ident = Ident::new("namespace", shim_trait.span());

            let mut method: TraitItemMethod = parse_quote!(
                fn #ns_ident() -> (&'static str, &'static str);
            );

            let imports = &args.imports;
            let prefix = &ctx.mod_name;

            method.default = Some(parse_quote!({
                (#imports, stringify!(#prefix))
            }));

            shim_trait.items.push(syn::TraitItem::Method(method));
        }
    }

    // Add enum types
    {
        for extern_enum in extern_enums {
            let enum_ident = Ident::new(&format!("{}_Repr", extern_enum.ident), shim_trait.span());
            shim_trait
                .items
                .push(syn::TraitItem::Type(parse_quote!(type #enum_ident;)));
        }
    }

    input
        .0
        .content
        .as_mut()
        .unwrap()
        .1
        .push(syn::Item::Trait(shim_trait));

    Ok(())
}

/// Adds use statements for traits used when doing enum conversions
fn expand_use(extern_enums: &[ExternEnum], input: &mut Item) {
    if extern_enums.is_empty() {
        return;
    }

    // Enums present! add the necessary use items.
    let items = &mut input.0.content.as_mut().unwrap().1;
    items.push(syn::Item::Use(parse_quote!(
        use num_enum::{IntoPrimitive, TryFromPrimitive};
    )));
    items.push(syn::Item::Use(parse_quote!(
        use bytemuck::{NoUninit, CheckedBitPattern};
    )));
    items.push(syn::Item::Use(parse_quote!(
        use std::convert::TryFrom;
    )));
}