whamm 0.1.0

#![allow(clippy::too_many_arguments)]

use crate::common::error::ErrorGen;
use crate::emitter::memory_allocator::MemoryAllocator;
use crate::emitter::tag_handler::get_tag_for;
use crate::generator::ast::Script;
use crate::lang_features::libraries::core::utils::utils_adapter::UtilsAdapter;
use crate::lang_features::libraries::core::utils::UtilsPackage;
use crate::lang_features::libraries::core::{
    LibPackage, ASSUMED_LIB_MEM_NAME, WHAMM_CORE_LIB_NAME,
};
use crate::parser::types::Location;
use crate::verifier::types::{Record, SymbolTable};
use std::collections::HashSet;
use wirm::ir::id::FunctionID;
use wirm::wasmparser::{ExternalKind, MemoryType};
use wirm::{DataType, Module};

// Some documentation on why it's difficult to only import the *used* functions.
//
// TLDR; Rust ownership.
// If I pass in a reference to both the application module (to conditionally import
// if not already done) AND a function modifier to the library adapter, I'll have
// two mutable references to the app module.
// This means that the caller will have to check for needed imports BEFORE actually
// delegating to the adapter...which means that I'd have to break the practice of
// information hiding. Or, I could create a 'check_OPERATION()' per 'OPERATION()',
// which would just be more clunky.
//
// So for now, we'll do this, but it can be changed later if I get a better idea.

pub fn link_core_lib(
    ast: &[Script],
    app_wasm: &mut Module,
    core_lib: &[u8],
    mem_allocator: &mut MemoryAllocator,
    utils: &mut UtilsPackage,
    packages: &mut [&mut dyn LibPackage],
    err: &mut ErrorGen,
) -> Vec<FunctionID> {
    let mut injected_funcs = vec![];

    let mut used_a_pkg = false;
    let mut should_use = vec![];
    for package in packages.iter_mut() {
        package.visit_ast(ast);
        package.set_adapter_usage(package.is_used());
        package.set_global_adapter_usage(package.is_used_in_global_scope());
        used_a_pkg |= package.is_used();

        should_use.push(package.is_used());
    }

    if used_a_pkg {
        let core_lib = Module::parse(core_lib, false, false).unwrap();
        import_lib_package(
            app_wasm,
            &None,
            WHAMM_CORE_LIB_NAME.to_string(),
            &None,
            &core_lib,
            utils,
        )
    }

    for (package, is_used) in packages.iter_mut().zip(should_use.iter()) {
        if *is_used {
            let core_lib = Module::parse(core_lib, false, false).unwrap();
            if package.import_memory() {
                let lib_mem_id = import_lib_memory(
                    app_wasm,
                    &None,
                    WHAMM_CORE_LIB_NAME,
                    &None,
                    get_lib_mem_name(&core_lib),
                    None,
                );
                package.set_lib_mem_id(lib_mem_id);
            }
            package.set_instr_mem_id(mem_allocator.mem_id as i32);
            import_lib_package(
                app_wasm,
                &None,
                WHAMM_CORE_LIB_NAME.to_string(),
                &None,
                &core_lib,
                *package,
            );
            injected_funcs.extend(gen_package_helpers(
                app_wasm,
                mem_allocator,
                &utils.adapter,
                *package,
                err,
            ));
        }
    }
    injected_funcs
}

pub fn link_user_lib(
    app_wasm: &mut Module,
    loc: &Option<Location>,
    lib_wasm: &Module,
    lib_name: String,
    lib_name_import_override: &Option<String>,
    used_mem: bool,
    used_lib_fns: &HashSet<String>,
    table: &mut SymbolTable,
) -> Vec<FunctionID> {
    let added = import_lib_fn_names(
        app_wasm,
        loc,
        &lib_name,
        lib_name_import_override,
        lib_wasm,
        used_lib_fns,
        Some(table),
    );

    if used_mem {
        import_lib_memory(
            app_wasm,
            loc,
            &lib_name,
            lib_name_import_override,
            get_lib_mem_name(lib_wasm),
            Some(table),
        );
    }

    let mut injected_funcs = vec![];
    for (_, fid) in added.iter() {
        injected_funcs.push(FunctionID(*fid));
    }

    injected_funcs
}

pub(crate) fn get_lib_mem_name<'a>(lib_wasm: &'a Module) -> &'a str {
    let mut lib_mem_name = ASSUMED_LIB_MEM_NAME;
    for export in lib_wasm.exports.iter() {
        if export.kind == ExternalKind::Memory {
            // assume it's the FIRST lib memory
            lib_mem_name = export.name.as_str();
            break;
        }
    }
    lib_mem_name
}

fn import_lib_memory(
    app_wasm: &mut Module,
    loc: &Option<Location>,
    lib_name: &str,
    lib_name_import_override: &Option<String>,
    lib_mem_name: &str,
    mut table: Option<&mut SymbolTable>,
) -> i32 {
    let import_module_name = if let Some(name_override) = lib_name_import_override {
        name_override.as_str()
    } else {
        lib_name
    };
    let id = if let Some(table) = table.as_mut() {
        let Some(Record::Library { mem_id, .. }) = table.lookup_lib_mut(lib_name) else {
            panic!("unexpected type");
        };

        match mem_id {
            Some(id) => *id,
            None => {
                // memory for this library hasn't been imported yet, fix that!
                let id = import_memory(
                    import_module_name,
                    lib_mem_name,
                    &format!("{lib_name}_lib_mem"),
                    loc,
                    app_wasm,
                );

                // save the MEM_ID to the symbol table
                *mem_id = Some(id);
                id
            }
        }
    } else {
        // memory for this library hasn't been imported yet, fix that!
        let id = import_memory(
            import_module_name,
            ASSUMED_LIB_MEM_NAME,
            &format!("{lib_name}_lib_mem"),
            loc,
            app_wasm,
        );

        id
    };

    id as i32
}

fn gen_package_helpers(
    app_wasm: &mut Module,
    mem_allocator: &mut MemoryAllocator,
    utils: &UtilsAdapter,
    package: &mut dyn LibPackage,
    err: &mut ErrorGen,
) -> Vec<FunctionID> {
    // enable the library to define in-module helper functions
    package.define_helper_funcs(utils, mem_allocator, app_wasm, err)
}

fn import_lib_package(
    app_wasm: &mut Module,
    loc: &Option<Location>,
    lib_name: String,
    lib_name_import_override: &Option<String>,
    lib_wasm: &Module,
    package: &mut dyn LibPackage,
) {
    // should only import the EXPORTED contents of the lib_wasm
    let added = import_lib_fn_names(
        app_wasm,
        loc,
        &lib_name,
        lib_name_import_override,
        lib_wasm,
        &HashSet::from_iter(package.get_fn_names().iter().cloned()),
        None,
    );

    for (name, fid) in added.iter() {
        // save the FID
        package.add_fid_to_adapter(name.as_str(), *fid);
    }
}

fn import_lib_fn_names(
    app_wasm: &mut Module,
    loc: &Option<Location>,
    lib_name: &str,
    lib_name_import_override: &Option<String>,
    lib_wasm: &Module,
    lib_fns: &HashSet<String>,
    mut table: Option<&mut SymbolTable>,
) -> Vec<(String, u32)> {
    let mut injected_fns = vec![];
    for export in lib_wasm.exports.iter() {
        // we don't care about non-function exports
        if let ExternalKind::Func = export.kind {
            if lib_fns.contains(&export.name) {
                let func = lib_wasm.functions.get(FunctionID(export.index));
                if let Some(ty) = lib_wasm.types.get(func.get_type_id()) {
                    let import_name = if let Some(name_override) = lib_name_import_override {
                        name_override.as_str()
                    } else {
                        lib_name
                    };
                    let fn_name = export.name.as_str();

                    let fid = import_func(
                        import_name,
                        fn_name,
                        &ty.params().clone(),
                        &ty.results().clone(),
                        loc,
                        app_wasm,
                    );
                    // save the FID to the symbol table
                    if let Some(table) = table.as_mut() {
                        let Some(Record::LibFn { addr, .. }) =
                            table.lookup_lib_fn_mut(lib_name, fn_name)
                        else {
                            panic!("unexpected type");
                        };

                        *addr = Some(fid);
                    }

                    // save the FID as an injected function
                    injected_fns.push((export.name.clone(), fid));
                } else {
                    unreachable!(
                        "ImportLib: Could not add function \"{}\" as application import",
                        export.name
                    );
                }
            }
        }
    }
    injected_fns
}

fn import_memory(
    module_name: &str,
    mem_name: &str,
    use_name: &str,
    loc: &Option<Location>,
    app_wasm: &mut Module,
) -> u32 {
    let (mem_id, imp_id) = app_wasm.add_import_memory_with_tag(
        module_name.to_string(),
        mem_name.to_string(),
        MemoryType {
            memory64: false,
            shared: false,
            initial: 0,
            maximum: None,
            page_size_log2: None,
        },
        get_tag_for(loc),
    );
    app_wasm.imports.set_name(use_name.to_string(), imp_id);

    *mem_id
}

pub fn import_func(
    module_name: &str,
    fname: &str,
    params: &[DataType],
    results: &[DataType],
    loc: &Option<Location>,
    app_wasm: &mut Module,
) -> u32 {
    let ty_id = app_wasm
        .types
        .add_func_type_with_tag(params, results, get_tag_for(loc));
    let (fid, imp_id) = app_wasm.add_import_func_with_tag(
        module_name.to_string(),
        fname.to_string(),
        ty_id,
        get_tag_for(loc),
    );
    app_wasm.imports.set_name(fname.to_string(), imp_id);

    *fid
}