ark-api-macros 0.11.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 = format!("{}", name);

        // 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)?;

    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>);

    static DISALLOWED_TYPES: &[&str] = &["usize", "isize", "u128", "i128"];

    impl<'a> visit::Visit<'a> for FindTypes {
        fn visit_ident(&mut self, node: &'a Ident) {
            if self.0.is_none() {
                DISALLOWED_TYPES
                    .iter()
                    .filter(|ty| node == *ty)
                    .for_each(|ty| {
                        self.0 = Some(Error::new_spanned(
                            node,
                            format!("`{}` is not FFI-safe", ty),
                        ));
                    });
            }
        }

        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 {
                    if path.path.is_ident("bool") {
                        self.0 = Some(Error::new_spanned(path,
                            "Bool ffi parameters are not supported at the moment. Pass them as u32s instead.\n\
                            This is because wasmtime's IntoFunc trait does not support bool parameters.\n\
                            If we did not have this error you would otherwise get a very cryptic error message \
                            about the IntoFunc trait not being implemented, without any line numbers or anything.\n\
                            We could support bool parameters behind the scenes by mapping them to u32s, but this is not done right now.")
                        );
                    }
                }
            }
        }
    }

    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(())
    }
}

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 `Result<T, WasmTrap>`, FFI function returns an `ErrorCode`.
    DeprecatedInfallible,
    /// Shim/export return a `Result<T, WasmTrap>`, 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 ffi_ident = if let Some(ffi_ident) = enum_ffi_ident {
            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 groupe 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!("{}_ptr", ident), ident.span());
                let ident_len = syn::Ident::new(&format!("{}_len", ident), 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!("{}_ptr", ident), ident.span());
                    let ident_len = syn::Ident::new(&format!("{}_len", ident), 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!("{}_ptr", ident), 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 Err;
            type WasmTrap;
        }
    );

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

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

            let err_type: syn::Type = if func.fallible_mode.is_fallible() {
                parse_quote!(Self::Err)
            } else {
                parse_quote!(Self::WasmTrap)
            };

            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 = parse_quote!(Result<#ok_type, #err_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 err_type: Type = if func.fallible_mode.is_fallible() {
                parse_quote!(Self::Err)
            } else {
                parse_quote!(Self::WasmTrap)
            };

            let mut method: TraitItemMethod = parse_quote!(
                fn #method_ident<'a>(
                    memory: &mut Self::Memory,
                    host_context: &mut Self::Context
                ) -> Result<(), #err_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;
    )));
}