wasm-rquickjs 0.1.1

use crate::javascript::escape_js_ident;
use crate::rust_bindgen::RustWitFunction;
use crate::types::{
    ProcessedParameter, ReturnTypeInformation, WrappedType, get_function_name, get_return_type,
    ident_in_exported_interface, ident_in_exported_interface_or_global, param_refs_as_tuple,
    process_parameter, to_original_func_arg_list, to_wrapped_param_refs, type_borrows_resource,
};
use crate::{EmbeddingMode, GeneratorContext, JsModuleSpec};
use anyhow::{Context, anyhow};
use heck::{ToLowerCamelCase, ToUpperCamelCase};
use proc_macro2::{Ident, Span, TokenStream};
use quote::quote;
use std::collections::BTreeMap;
use syn::{Lit, LitStr};
use wit_parser::{Function, FunctionKind, Interface, TypeId, WorldItem, WorldKey};

/// Generates the `<output>/src/lib.rs` file for the wrapper crate, implementing the component exports
/// and providing the general Rust module declarations.
pub fn generate_export_impls(
    context: &GeneratorContext<'_>,
    js_modules: &[JsModuleSpec],
) -> anyhow::Result<()> {
    let guest_impls = generate_guest_impls(context)?;
    let module_defs = generate_module_defs(js_modules)?;

    let lib_tokens = quote! {
        #[allow(static_mut_refs)]
        #[allow(unsafe_op_in_unsafe_fn)]
        mod bindings;
        mod builtin;
        mod conversions;
        #[allow(unused)]
        mod internal;
        #[allow(unused)]
        mod modules;
        mod wrappers;

        #module_defs

        struct Component;

        #(#guest_impls)*

        bindings::export!(Component with_types_in bindings);
    };

    let lib_ast: syn::File =
        syn::parse2(lib_tokens).context("failed to parse generated lib.rs tokens")?;

    let lib_path = context.output.join("src").join("lib.rs");
    let lib_src = prettier_please::unparse(&lib_ast);

    crate::write_if_changed(&lib_path, lib_src)?;

    Ok(())
}

/// Generates a list of code snippets, each implementing one of the `Guest` traits generated by
/// wit-bindgen-rust for the component's exports.
fn generate_guest_impls(context: &GeneratorContext<'_>) -> anyhow::Result<Vec<TokenStream>> {
    let mut result = Vec::new();

    let world = &context.resolve.worlds[context.world];

    let mut global_exports = Vec::new();
    let mut interface_exports = Vec::new();

    // Enumerating all exports and separating them into global exports and interface exports.
    for (name, export) in &world.exports {
        let name = match name {
            WorldKey::Name(name) => name.clone(),
            WorldKey::Interface(id) => {
                let interface = &context.resolve.interfaces[*id];
                interface
                    .name
                    .clone()
                    .ok_or_else(|| anyhow!("Interface export does not have a name"))?
            }
        };
        match export {
            WorldItem::Interface { id, .. } => {
                let interface = &context.resolve.interfaces[*id];
                interface_exports.push((name, interface));
            }
            WorldItem::Function(function) => {
                global_exports.push((name, function));
            }
            WorldItem::Type(_) => {}
        }
    }

    // Implementing a single Guest trait containing all the global exported functions
    if !global_exports.is_empty() {
        result.extend(generate_guest_impl(
            context,
            quote! { crate::bindings::Guest },
            None,
            &global_exports,
        )?);
    }

    // Implementing a Guest trait per exported interface
    for (name, interface) in interface_exports {
        let interface_exports: Vec<_> = interface
            .functions
            .iter()
            .map(|(name, function)| (name.clone(), function))
            .collect();

        result.extend(generate_guest_impl(
            context,
            ident_in_exported_interface(
                context,
                Ident::new("Guest", Span::call_site()),
                &name,
                interface,
            ),
            Some((&name, interface)),
            &interface_exports,
        )?);
    }

    Ok(result)
}

/// Generates the implementation of a `Guest` trait for the component, implementing the exported functions.
///
/// The `guest_trait` parameter is a Rust snippet containing the fully-qualified path to the `Guest` trait to
/// be implemented.
///
/// If there are resources in the interface, the return contains all the trait implementations, for the interface
/// and the resources as well.
fn generate_guest_impl(
    context: &GeneratorContext<'_>,
    guest_trait: TokenStream,
    interface: Option<(&str, &Interface)>,
    exports: &[(String, &Function)],
) -> anyhow::Result<Vec<TokenStream>> {
    let mut func_impls = Vec::new();
    let mut resource_impls = Vec::new();
    let mut resource_functions = BTreeMap::new();

    for (name, function) in exports {
        match &function.kind {
            FunctionKind::Freestanding => {
                if name == "wizer-initialize" {
                    // wizer-initialize calls directly into the skeleton's
                    // pre-init function instead of dispatching to JS
                    func_impls.push(quote! {
                        fn wizer_initialize() {
                            crate::internal::wizer_initialize();
                        }
                    });
                } else {
                    let func_impl =
                        generate_exported_function_impl(context, interface, name, function)?;
                    func_impls.push(func_impl);
                }
            }
            FunctionKind::AsyncFreestanding
            | FunctionKind::AsyncMethod(_)
            | FunctionKind::AsyncStatic(_) => {
                Err(anyhow!("Async exported functions are not supported yet"))?
            }
            FunctionKind::Method(type_id)
            | FunctionKind::Static(type_id)
            | FunctionKind::Constructor(type_id) => {
                resource_functions
                    .entry(type_id)
                    .or_insert_with(Vec::new)
                    .push((name, function));
            }
        }
    }

    let mut resource_types = Vec::new();
    for (resource_type_id, resource_funcs) in resource_functions {
        let typ = context
            .resolve
            .types
            .get(*resource_type_id)
            .ok_or_else(|| anyhow!("Unknown resource type id"))?;

        let resource_name = typ
            .name
            .as_ref()
            .ok_or_else(|| anyhow!("Resource type has no name"))?;
        let resource_name_ident =
            Ident::new(&resource_name.to_upper_camel_case(), Span::call_site());
        let resource_name_borrow_ident = Ident::new(
            &format!("{}Borrow", resource_name.to_upper_camel_case()),
            Span::call_site(),
        );
        let guest_name_ident = Ident::new(
            &format!("Guest{}", resource_name.to_upper_camel_case()),
            Span::call_site(),
        );
        let guest_trait =
            ident_in_exported_interface_or_global(context, guest_name_ident, interface);

        let borrow_wrapper =
            ident_in_exported_interface_or_global(context, resource_name_borrow_ident, interface);
        let owned_wrapper =
            ident_in_exported_interface_or_global(context, resource_name_ident.clone(), interface);

        let mut resource_func_impls = Vec::new();
        for (name, resource_function) in resource_funcs {
            let func_impl = generate_exported_resource_function_impl(
                context,
                interface,
                resource_type_id,
                name,
                resource_function,
            )?;
            resource_func_impls.push(func_impl);
        }

        resource_impls.push(quote! {
            struct #resource_name_ident {
                resource_id: usize
            }

            impl #guest_trait for #resource_name_ident {
                #(#resource_func_impls)*
            }

            impl Drop for #resource_name_ident {
                fn drop(&mut self) {
                    crate::internal::enqueue_drop_js_resource(self.resource_id);
                }
            }

            impl<'js> rquickjs::IntoJs<'js> for #borrow_wrapper<'_> {
                fn into_js(self, ctx: &rquickjs::Ctx<'js>) -> rquickjs::Result<rquickjs::Value<'js>> {
                    let inner: &#resource_name_ident = self.get();
                    let resource_table: rquickjs::Object = ctx.globals().get(crate::internal::RESOURCE_TABLE_NAME)
                        .expect("Failed to get the resource table");
                    let resource_instance: rquickjs::Object = resource_table.get(inner.resource_id.to_string())
                        .expect(&format!("Failed to get resource instance with id {}", inner.resource_id));
                    Ok(resource_instance.into_value())
                }
            }

            impl<'js> rquickjs::IntoJs<'js> for #owned_wrapper {
                fn into_js(self, ctx: &rquickjs::Ctx<'js>) -> rquickjs::Result<rquickjs::Value<'js>> {
                    let inner: &#resource_name_ident = self.get();
                    let resource_table: rquickjs::Object = ctx.globals().get(crate::internal::RESOURCE_TABLE_NAME)
                        .expect("Failed to get the resource table");
                    let resource_instance: rquickjs::Object = resource_table.get(inner.resource_id.to_string())
                        .expect(&format!("Failed to get resource instance with id {}", inner.resource_id));
                    Ok(resource_instance.into_value())
                }
            }

            impl<'js> rquickjs::FromJs<'js> for #owned_wrapper {
                fn from_js(ctx: &rquickjs::Ctx<'js>, value: rquickjs::Value<'js>) -> rquickjs::Result<Self> {
                    let resource = value.into_object().ok_or_else(|| {
                        rquickjs::Error::new_from_js_message(
                            "JS Resource instance",
                            "WASM resource instance",
                            "The value is not an object",
                        )
                    })?;

                    let already_registered = resource.contains_key(crate::internal::RESOURCE_ID_KEY)?;
                    let resource_id: usize = if already_registered {
                        // This resource instance is already registered in the resource table
                        resource.get(crate::internal::RESOURCE_ID_KEY)?
                    } else {
                        // This is a new resource instance, we need to store it in the resource table
                        let resource_table: rquickjs::Object = ctx.globals().get(crate::internal::RESOURCE_TABLE_NAME)?;
                        let resource_id = crate::internal::get_free_resource_id();
                        resource_table.set(resource_id.to_string(), resource)?;
                        resource_id
                    };

                    Ok(#owned_wrapper::new(#resource_name_ident { resource_id }))
                }
            }
        });

        resource_types.push(quote! {
            type #resource_name_ident = #resource_name_ident;
        })
    }

    let mut guest_impls = Vec::new();
    guest_impls.extend(resource_impls);
    guest_impls.push(quote! {
        impl #guest_trait for Component {
            #(#resource_types)*
            #(#func_impls)*
        }
    });

    Ok(guest_impls)
}

/// Generates one trait method implementation for an exported freestanding function
fn generate_exported_function_impl(
    context: &GeneratorContext<'_>,
    interface: Option<(&str, &Interface)>,
    name: &str,
    function: &Function,
) -> anyhow::Result<TokenStream> {
    let rust_fn = RustWitFunction::new(context, name, function);
    let func_name = rust_fn.function_name_ident();

    let param_ident_type: Vec<_> = function
        .params
        .iter()
        .zip(rust_fn.export_parameters.clone())
        .zip(rust_fn.import_parameters.clone())
        .map(
            |(((param_name, param_type), export_parameter), import_parameter)| {
                process_parameter(
                    context,
                    param_name,
                    param_type,
                    &export_parameter,
                    &import_parameter,
                )
            },
        )
        .collect::<anyhow::Result<Vec<_>>>()?;

    let func_arg_list = to_original_func_arg_list(&param_ident_type);
    let return_types = get_return_type(context, function, name, &rust_fn)?;

    let param_refs = to_wrapped_param_refs(&param_ident_type);
    let param_refs_tuple = param_refs_as_tuple(&param_refs);

    let js_func_name_str = Lit::Str(LitStr::new(
        &escape_js_ident(name.to_lower_camel_case()),
        func_name.span(),
    ));
    let (js_func_path, wit_package_lit) = match interface {
        Some((iface_name, iface)) => {
            let if_name_str = LitStr::new(
                &escape_js_ident(iface_name.to_lower_camel_case()),
                func_name.span(),
            );

            let owner_package_name = match iface.package {
                Some(package_id) => {
                    let package = context.resolve.packages.get(package_id).ok_or_else(|| {
                        anyhow!("Unknown owner package of interface: {iface_name}")
                    })?;
                    package.name.to_string()
                }
                None => context.root_package_name().to_string(),
            };

            (
                quote! { &[#if_name_str, #js_func_name_str] },
                Lit::Str(LitStr::new(&owner_package_name, Span::call_site())),
            )
        }
        None => (
            quote! { &[#js_func_name_str] },
            Lit::Str(LitStr::new(&context.root_package_name(), Span::call_site())),
        ),
    };

    let original_result = &return_types.wit_level_ret.original_type_ref;
    let wrapped_result = &return_types.wit_level_ret.wrapped_type_ref;
    let unwrap = &return_types.wit_level_ret.unwrap;
    let unwrap_result = unwrap.run(quote! { result });
    let call = if return_types.expected_exception.is_some() {
        quote! { call_js_export_returning_result }
    } else {
        quote! { call_js_export }
    };
    let func_impl = quote! {
       fn #func_name(#(#func_arg_list),*) -> #original_result {
           crate::internal::async_exported_function(async move {
               let result: #wrapped_result = crate::internal::#call(
                   #wit_package_lit,
                   #js_func_path,
                   #param_refs_tuple
               ).await;
               #unwrap_result
           })
       }
    };
    Ok(func_impl)
}

/// Generates one trait method implementation for an exported freestanding function
fn generate_exported_resource_function_impl(
    context: &GeneratorContext<'_>,
    interface: Option<(&str, &Interface)>,
    resource_type_id: &TypeId,
    name: &str,
    function: &Function,
) -> anyhow::Result<TokenStream> {
    let func_name = get_function_name(name, function)?;

    let rust_fn = RustWitFunction::new(context, &func_name, function);
    let func_name_ident = rust_fn.function_name_ident();

    let param_ident_type: Vec<_> = function
        .params
        .iter()
        .zip(rust_fn.export_parameters.clone())
        .zip(rust_fn.import_parameters.clone())
        .map(|(((param_name, param_type), export_param), import_param)| {
            if matches!(
                function.kind,
                FunctionKind::Method(_) | FunctionKind::AsyncMethod(_)
            ) && type_borrows_resource(context, param_type, resource_type_id)?
            {
                Ok(ProcessedParameter {
                    ident: Ident::new(param_name, Span::call_site()),
                    wrapped_type: None,
                    export_parameter: export_param,
                    import_parameter: import_param,
                })
            } else {
                process_parameter(
                    context,
                    param_name,
                    param_type,
                    &export_param,
                    &import_param,
                )
            }
        })
        .collect::<anyhow::Result<Vec<_>>>()?;

    let func_arg_list = to_original_func_arg_list(&param_ident_type);
    let return_types = if matches!(function.kind, FunctionKind::Constructor(_)) {
        ReturnTypeInformation {
            wit_level_ret: WrappedType::no_wrapping(quote! { Self }),
            func_ret: WrappedType::no_wrapping(quote! { Self }),
            expected_exception: None,
        }
    } else {
        get_return_type(context, function, name, &rust_fn)?
    };

    let param_refs = to_wrapped_param_refs(&param_ident_type);

    let resource_name = context
        .resolve
        .types
        .get(*resource_type_id)
        .ok_or_else(|| anyhow::anyhow!("Unknown resource type id"))?
        .name
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("Resource type has no name"))?;

    let js_resource_name_str = Lit::Str(LitStr::new(
        &resource_name.to_upper_camel_case(),
        Span::call_site(),
    ));
    let (js_resource_path, wit_package_lit) = match interface {
        Some((iface_name, iface)) => {
            let if_name_str = LitStr::new(
                &escape_js_ident(iface_name.to_lower_camel_case()),
                Span::call_site(),
            );

            let owner_package_name = match iface.package {
                Some(package_id) => {
                    let package = context.resolve.packages.get(package_id).ok_or_else(|| {
                        anyhow!("Unknown owner package of interface: {iface_name}")
                    })?;
                    package.name.to_string()
                }
                None => context.root_package_name().to_string(),
            };

            (
                quote! { &[#if_name_str, #js_resource_name_str] },
                Lit::Str(LitStr::new(&owner_package_name, Span::call_site())),
            )
        }
        None => (
            quote! { &[#js_resource_name_str] },
            Lit::Str(LitStr::new(&context.root_package_name(), Span::call_site())),
        ),
    };

    let js_func_name_str = Lit::Str(LitStr::new(
        &escape_js_ident(func_name.to_lower_camel_case()),
        Span::call_site(),
    ));
    let js_static_func_path = match interface {
        Some((iface_name, _)) => {
            let if_name_str = LitStr::new(
                &escape_js_ident(iface_name.to_lower_camel_case()),
                Span::call_site(),
            );
            quote! { &[#if_name_str, #js_resource_name_str, #js_func_name_str] }
        }
        None => quote! { &[#js_func_name_str] },
    };

    let func_impl = match &function.kind {
        FunctionKind::Constructor(_) => {
            let param_refs_tuple = param_refs_as_tuple(&param_refs);

            quote! {
              fn #func_name_ident(#(#func_arg_list),*) -> Self {
                  crate::internal::async_exported_function(async move {
                    let resource_id = crate::internal::call_js_resource_constructor(
                         #wit_package_lit,
                         #js_resource_path,
                         #param_refs_tuple,
                    ).await;
                    Self {
                        resource_id
                    }
                  })
              }
            }
        }
        FunctionKind::Method(_) => {
            let param_refs = param_refs[1..].to_vec();
            let param_refs_tuple = param_refs_as_tuple(&param_refs);
            let original_result = &return_types.func_ret.original_type_ref;
            let wrapped_result = &return_types.func_ret.wrapped_type_ref;
            let unwrap = &return_types.func_ret.unwrap;
            let unwrap_result = unwrap.run(quote! { result });
            let call = if return_types.expected_exception.is_some() {
                quote! { call_js_resource_method_returning_result }
            } else {
                quote! { call_js_resource_method }
            };
            quote! {
               fn #func_name_ident(#(#func_arg_list),*) -> #original_result {
                   crate::internal::async_exported_function(async move {
                       let result: #wrapped_result = crate::internal::#call(
                            #wit_package_lit,
                            #js_resource_path,
                            self.resource_id,
                            #js_func_name_str,
                            #param_refs_tuple,
                       ).await;
                       #unwrap_result
                   })
               }
            }
        }
        FunctionKind::Static(_) => {
            let param_refs_tuple = param_refs_as_tuple(&param_refs);
            let original_result = &return_types.wit_level_ret.original_type_ref;
            let wrapped_result = &return_types.wit_level_ret.wrapped_type_ref;
            let unwrap = &return_types.wit_level_ret.unwrap;
            let unwrap_result = unwrap.run(quote! { result });
            let call = if return_types.expected_exception.is_some() {
                quote! { call_js_export_returning_result }
            } else {
                quote! { call_js_export }
            };
            quote! {
               fn #func_name_ident(#(#func_arg_list),*) -> #original_result {
                   crate::internal::async_exported_function(async move {
                       let result: #wrapped_result = crate::internal::#call(
                           #wit_package_lit,
                           #js_static_func_path,
                           #param_refs_tuple,
                       ).await;
                       #unwrap_result
                   })
               }
            }
        }
        FunctionKind::AsyncMethod(_) | FunctionKind::AsyncStatic(_) => Err(anyhow::anyhow!(
            "Async exported functions are not supported yet",
        ))?,
        FunctionKind::Freestanding | FunctionKind::AsyncFreestanding => Err(anyhow::anyhow!(
            "Freestanding functions are not expected in resource methods",
        ))?,
    };

    Ok(func_impl)
}

fn generate_module_defs(js_modules: &[JsModuleSpec]) -> anyhow::Result<TokenStream> {
    if let Some((export_module, additional_modules)) = js_modules.split_first() {
        let export_module_name = LitStr::new(&export_module.name, Span::call_site());
        let export_module_file_name = LitStr::new(&export_module.file_name(), Span::call_site());

        let mut additional_module_pairs = Vec::new();
        for module in additional_modules {
            match module.mode {
                EmbeddingMode::EmbedFile(_) => {
                    let name = LitStr::new(&module.name, Span::call_site());
                    let file_name = LitStr::new(&module.file_name(), Span::call_site());
                    additional_module_pairs
                        .push(quote! { (#name, Box::new(|| { include_str!(#file_name) })) });
                }
                EmbeddingMode::Composition => {
                    let name = LitStr::new(&module.name, Span::call_site());
                    additional_module_pairs
                        .push(quote! { (#name, Box::new(|| { crate::bindings::get_script() })) });
                }
            }
        }

        Ok(quote! {
            static JS_EXPORT_MODULE_NAME: &str = #export_module_name;
            static JS_EXPORT_MODULE: &str = include_str!(#export_module_file_name);

            static JS_ADDITIONAL_MODULES: std::sync::LazyLock<Vec<(&str, Box<dyn (Fn() -> String) + Send + Sync>)>> =
              std::sync::LazyLock::new(|| { vec![
                 #(#additional_module_pairs),*
              ]});
        })
    } else {
        Err(anyhow!("No JS modules provided."))?
    }
}