roc/crates/compiler/test_gen/src/wasm_linking.rs

#![cfg(feature = "gen-wasm")]

use bumpalo::Bump;
use roc_gen_wasm::Env;
use roc_target::TargetInfo;
use std::fs;
use std::process::Command;

use roc_collections::{MutMap, MutSet};
use roc_module::ident::{ForeignSymbol, ModuleName};
use roc_module::low_level::LowLevel;
use roc_module::symbol::{
    IdentIds, IdentIdsByModule, Interns, ModuleId, ModuleIds, PackageModuleIds, PackageQualified,
    Symbol,
};
use roc_mono::ir::{
    Call, CallType, Expr, HostExposedLayouts, Literal, Proc, ProcLayout, SelfRecursive, Stmt,
    UpdateModeId,
};
use roc_mono::layout::{LambdaName, Layout, Niche, STLayoutInterner};
use roc_wasm_interp::{wasi, ImportDispatcher, Instance, WasiDispatcher};
use roc_wasm_module::{Value, WasmModule};

const LINKING_TEST_HOST_WASM: &str = "build/wasm_linking_test_host.wasm";
const LINKING_TEST_HOST_NATIVE: &str = "build/wasm_linking_test_host";

fn create_symbol(home: ModuleId, ident_ids: &mut IdentIds, debug_name: &str) -> Symbol {
    let ident_id = ident_ids.add_str(debug_name);
    Symbol::new(home, ident_id)
}

// Build a fake Roc app in mono IR
// Calls two host functions, one Wasm and one JS
fn build_app_mono<'a>(
    arena: &'a Bump,
    home: ModuleId,
    ident_ids: &mut IdentIds,
) -> (Symbol, MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>) {
    let int_layout = Layout::I32;

    let app_proc = create_symbol(home, ident_ids, "app_proc");
    let js_call_result = create_symbol(home, ident_ids, "js_call_result");
    let host_call_result = create_symbol(home, ident_ids, "host_call_result");
    let bitflag = create_symbol(home, ident_ids, "bitflag");
    let or1 = create_symbol(home, ident_ids, "or1");
    let or2 = create_symbol(home, ident_ids, "or2");

    let js_call = Expr::Call(Call {
        call_type: CallType::Foreign {
            foreign_symbol: ForeignSymbol::from("js_called_directly_from_roc"),
            ret_layout: int_layout,
        },
        arguments: &[],
    });

    let host_call = Expr::Call(Call {
        call_type: CallType::Foreign {
            foreign_symbol: ForeignSymbol::from("host_called_directly_from_roc"),
            ret_layout: int_layout,
        },
        arguments: &[],
    });

    let mut bitflag_bytes = [0; 16];
    bitflag_bytes[0] = 0x20;
    let bitflag_literal = Expr::Literal(Literal::Int(bitflag_bytes));

    let or1_expr = Expr::Call(Call {
        call_type: CallType::LowLevel {
            op: LowLevel::Or,
            update_mode: UpdateModeId::BACKEND_DUMMY,
        },
        arguments: arena.alloc([js_call_result, host_call_result]),
    });

    let or2_expr = Expr::Call(Call {
        call_type: CallType::LowLevel {
            op: LowLevel::Or,
            update_mode: UpdateModeId::BACKEND_DUMMY,
        },
        arguments: arena.alloc([or1, bitflag]),
    });

    let body = Stmt::Let(
        js_call_result,
        js_call,
        int_layout,
        arena.alloc(Stmt::Let(
            host_call_result,
            host_call,
            int_layout,
            arena.alloc(Stmt::Let(
                or1,
                or1_expr,
                int_layout,
                arena.alloc(Stmt::Let(
                    bitflag,
                    bitflag_literal,
                    int_layout,
                    arena.alloc(Stmt::Let(
                        or2,
                        or2_expr,
                        int_layout,
                        //
                        arena.alloc(Stmt::Ret(or2)),
                    )),
                )),
            )),
        )),
    );

    let proc = Proc {
        name: LambdaName::no_niche(app_proc),
        args: &[],
        body,
        closure_data_layout: None,
        ret_layout: int_layout,
        is_self_recursive: SelfRecursive::NotSelfRecursive,
        host_exposed_layouts: HostExposedLayouts::NotHostExposed,
    };

    let proc_layout = ProcLayout {
        arguments: &[],
        result: int_layout,
        niche: Niche::NONE,
    };

    let mut app = MutMap::default();
    app.insert((app_proc, proc_layout), proc);

    (app_proc, app)
}

struct BackendInputs<'a> {
    env: Env<'a>,
    interns: Interns,
    host_module: WasmModule<'a>,
    procedures: MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
}

impl<'a> BackendInputs<'a> {
    fn new(arena: &'a Bump) -> Self {
        // Compile the host from an external source file
        let host_bytes = fs::read(LINKING_TEST_HOST_WASM).unwrap();
        let host_module: WasmModule = roc_gen_wasm::parse_host(arena, &host_bytes).unwrap();

        // Identifier stuff to build the mono IR
        let module_name = ModuleName::from("UserApp");
        let pkg_qualified_module_name = PackageQualified::Unqualified(module_name);
        let mut package_module_ids = PackageModuleIds::default();
        let module_id: ModuleId = package_module_ids.get_or_insert(&pkg_qualified_module_name);
        let mut ident_ids = IdentIds::default();

        // IR for the app
        let (roc_main_sym, procedures) = build_app_mono(arena, module_id, &mut ident_ids);
        let mut exposed_to_host = MutSet::default();
        exposed_to_host.insert(roc_main_sym);
        let env = Env {
            arena,
            module_id,
            exposed_to_host,
            stack_bytes: Env::DEFAULT_STACK_BYTES,
        };

        // Identifier stuff for the backend
        let module_ids = ModuleIds::default();
        let mut all_ident_ids: IdentIdsByModule = IdentIds::exposed_builtins(1);
        all_ident_ids.insert(module_id, ident_ids);
        let interns = Interns {
            module_ids,
            all_ident_ids,
        };

        BackendInputs {
            env,
            interns,
            host_module,
            procedures,
        }
    }
}

struct TestDispatcher<'a> {
    wasi: WasiDispatcher<'a>,
}

impl ImportDispatcher for TestDispatcher<'_> {
    fn dispatch(
        &mut self,
        module_name: &str,
        function_name: &str,
        arguments: &[Value],
        memory: &mut [u8],
    ) -> Option<Value> {
        if module_name == wasi::MODULE_NAME {
            self.wasi.dispatch(function_name, arguments, memory)
        } else if module_name == "env" {
            match function_name {
                "js_called_directly_from_roc" => Some(Value::I32(0x01)),
                "js_called_indirectly_from_roc" => Some(Value::I32(0x02)),
                "js_called_directly_from_main" => Some(Value::I32(0x04)),
                "js_called_indirectly_from_main" => Some(Value::I32(0x08)),
                "js_unused" => Some(Value::I32(0x10)),
                _ => panic!("Unknown import env.{}", function_name),
            }
        } else {
            panic!(
                "TestDispatcher does not implement {}.{}",
                module_name, function_name
            );
        }
    }
}

fn execute_wasm_module<'a>(arena: &'a Bump, orig_module: WasmModule<'a>) -> Result<i32, String> {
    // FIXME: see if we can skip serializing and re-parsing!
    // Some metadata seems to be left over from the host file. e.g. CodeSection::section_start
    let module = {
        let mut buffer = Vec::with_capacity(orig_module.size());
        orig_module.serialize(&mut buffer);
        WasmModule::preload(arena, &buffer, false).map_err(|e| format!("{:?}", e))?
    };

    let dispatcher = TestDispatcher {
        wasi: wasi::WasiDispatcher::default(),
    };
    let is_debug_mode = false;
    let mut inst = Instance::for_module(&arena, &module, dispatcher, is_debug_mode)?;

    // In Zig, main can only return u8 or void, but our result is too wide for that.
    // But I want to use main so that I can test that _start is created for it!
    // So return void from main, and call another function to get the result.
    inst.call_export("_start", [])?;

    // FIXME: read_host_result does not actually appear as an export!
    // The interpreter has to look it up in debug info! (Apparently Wasm3 did this!)
    // If we change gen_wasm to export it, then it does the same for js_unused,
    // so we can't test import elimination and function reordering.
    // We should to come back to this and fix it.
    inst.call_export("read_host_result", [])?
        .ok_or(String::from("expected a return value"))?
        .expect_i32()
        .map_err(|type_err| format!("{:?}", type_err))
}

fn get_native_result() -> i32 {
    let output = Command::new(LINKING_TEST_HOST_NATIVE)
        .output()
        .expect(&format!("failed to run {}", LINKING_TEST_HOST_NATIVE));

    let result_str = std::str::from_utf8(&output.stdout).unwrap().trim();
    result_str.parse().unwrap()
}

fn test_help(
    eliminate_dead_code: bool,
    expected_host_import_names: &[&str],
    expected_final_import_names: &[&str],
    expected_name_section_start: &[(u32, &str)],
    dump_filename: &str,
) {
    let arena = Bump::new();
    let mut layout_interner = STLayoutInterner::with_capacity(4, TargetInfo::default_wasm32());

    let BackendInputs {
        env,
        mut interns,
        host_module,
        procedures,
    } = BackendInputs::new(&arena);

    let host_import_names = Vec::from_iter(host_module.import.imports.iter().map(|imp| imp.name));
    assert_eq!(&host_import_names, expected_host_import_names);

    assert!(&host_module.names.function_names.is_empty());

    let (mut final_module, called_fns, _roc_main_index) = roc_gen_wasm::build_app_module(
        &env,
        &mut layout_interner,
        &mut interns,
        host_module,
        procedures,
    );

    if eliminate_dead_code {
        final_module.eliminate_dead_code(env.arena, called_fns);
    }

    if std::env::var("DEBUG_WASM").is_ok() {
        let mut buffer = Vec::with_capacity(final_module.size());
        final_module.serialize(&mut buffer);
        fs::write(dump_filename, &buffer).unwrap();
    }

    let final_import_names = Vec::from_iter(final_module.import.imports.iter().map(|i| i.name));

    assert_eq!(&final_import_names, expected_final_import_names);

    let name_count = expected_name_section_start.len();
    assert_eq!(
        &final_module.names.function_names[0..name_count],
        expected_name_section_start
    );

    let wasm_result = execute_wasm_module(&arena, final_module).unwrap();
    assert_eq!(wasm_result, get_native_result());
}

const EXPECTED_HOST_IMPORT_NAMES: [&'static str; 9] = [
    "__linear_memory",
    "__stack_pointer",
    "js_called_indirectly_from_roc",
    "js_unused",
    "js_called_directly_from_roc",
    "js_called_directly_from_main",
    "roc__app_proc_1_exposed",
    "js_called_indirectly_from_main",
    "__indirect_function_table",
];

#[test]
fn test_linking_without_dce() {
    let expected_final_import_names = &[
        "js_called_indirectly_from_roc",
        "js_unused", // not eliminated
        "js_called_directly_from_roc",
        "js_called_directly_from_main",
        "js_called_indirectly_from_main",
    ];

    let expected_name_section_start = &[
        (0, "js_called_indirectly_from_roc"),
        (1, "js_unused"), // not eliminated
        (2, "js_called_directly_from_roc"),
        (3, "js_called_directly_from_main"),
        (4, "js_called_indirectly_from_main"),
    ];

    let eliminate_dead_code = false;
    let dump_filename = "build/without_dce.wasm";

    test_help(
        eliminate_dead_code,
        &EXPECTED_HOST_IMPORT_NAMES,
        expected_final_import_names,
        expected_name_section_start,
        dump_filename,
    );
}

#[test]
fn test_linking_with_dce() {
    let expected_final_import_names = &[
        "js_called_indirectly_from_roc",
        // "js_unused", // eliminated
        "js_called_directly_from_roc",
        "js_called_directly_from_main",
        "js_called_indirectly_from_main",
    ];

    let expected_name_section_start = &[
        (0, "js_called_indirectly_from_roc"),
        (1, "js_called_directly_from_roc"),    // index changed
        (2, "js_called_directly_from_main"),   // index changed
        (3, "js_called_indirectly_from_main"), // index changed
        (4, "js_unused"), // still exists, but now an internal dummy, with index changed
    ];

    let eliminate_dead_code = true;
    let dump_filename = "build/with_dce.wasm";

    test_help(
        eliminate_dead_code,
        &EXPECTED_HOST_IMPORT_NAMES,
        expected_final_import_names,
        expected_name_section_start,
        dump_filename,
    );
}