new tests

src/main.rs

@@ -253,33 +253,37 @@ fn remove_start_export(module: &mut walrus::Module) {
 
 }
 
-fn main() -> anyhow::Result<()> {
+fn do_module_replacements(module: &mut walrus::Module) {
+    // find corresponding IDs for replacements
+    let fn_replacement_ids = gather_replacement_ids(&module);
 
-    let exe_name = std::env::current_exe().unwrap().file_name().unwrap().to_str().unwrap().to_owned();
+    // do recursive call replacement
+    replace_calls(module, &fn_replacement_ids);
+
+    // add start entry (this is needed to do initialization)
+    add_start_entry(module);
 
+    // remove the _start export to clean up the module exports
+    remove_start_export(module);
+
+    // clean-up unused imports
+    walrus::passes::gc::run(module);
+}
+
+fn main() -> anyhow::Result<()> {
     env_logger::init();
 
+    let exe_name = std::env::current_exe().unwrap().file_name().unwrap().to_str().unwrap().to_owned();
+
     let input_wasm = std::env::args()
         .nth(1)
         .ok_or_else(|| anyhow::anyhow!("The launch parameters are incorrect, try: {} <input.wasm> [output.wasm]", {exe_name}))?;
 
     // load Wasm module from file
     let mut module = walrus::Module::from_file(&input_wasm)?;
 
-    // find corresponding IDs for replacements
-    let fn_replacement_ids = gather_replacement_ids(&module);
-
-    // do recursive call replacement
-    replace_calls(&mut module, &fn_replacement_ids);
-
-    // add start entry (this is needed to do initialization)
-    add_start_entry(&mut module);
-
-    // remove the _start export to clean up the module exports
-    remove_start_export(&mut module);
-
-    // clean-up unused imports
-    walrus::passes::gc::run(&mut module);
+    // do the substitution workflow
+    do_module_replacements(&mut module);
 
     // try store as binary
     let wasm = module.emit_wasm();

src/tests.rs

@@ -1,7 +1,7 @@
 
 use crate::*;
 
-#[warn(dead_code)]
+#[allow(dead_code)]
 fn print_wasm(module: &mut walrus::Module) {
     let wasm = module.emit_wasm();
     let text = wasmprinter::print_bytes(wasm).unwrap();
@@ -113,12 +113,11 @@ fn test_remove_start_export() {
 }
 
 
-
 #[test]
 fn test_gather_replacement_ids() {
 
     let wat = r#"
-    (module
+      (module
         (type (;0;) (func))
         (type (;1;) (func (param i32)))
         (type (;2;) (func (param i32 i32)))
@@ -170,3 +169,99 @@ fn test_gather_replacement_ids() {
     assert!(id_reps[&3] == 7);
 
 }
+
+#[test]
+fn test_do_module_replacements() {
+
+    let wat = r#"
+    (module
+        (type (;0;) (func))
+        (type (;1;) (func (param i32)))
+        (type (;2;) (func (param i32 i32)))
+        (type (;3;) (func (param i32 i32) (result i32)))
+        (type (;4;) (func (param i32 i32 i32)))
+        (type (;5;) (func (param i32 i32 i32 i32) (result i32)))
+        (import "ic0" "debug_print" (func $_dprint (;0;) (type 2)))
+        (import "ic0" "msg_reply" (func $_msg_reply (;1;) (type 0)))
+        (import "wasi_snapshot_preview1" "fd_write" (func $_wasi_snapshot_preview_fd_write (;2;) (type 5)))
+        (import "wasi_snapshot_preview1" "random_get" (func $_wasi_snapshot_preview_random_get (;3;) (type 3)))
+        (import "wasi_snapshot_preview1" "environ_get" (func $__imported_wasi_snapshot_preview1_environ_get (;4;) (type 3)))
+        (import "wasi_snapshot_preview1" "proc_exit" (func $__imported_wasi_snapshot_preview1_proc_exit (;5;) (type 1)))
+
+        (func $_start (;6;) (type 0)
+            i32.const 1
+            i32.const 2
+            call $_wasi_snapshot_preview_random_get
+
+            (block $test_block
+                i32.const 0
+                (if 
+                    (then
+                        (block $test_block3
+                            i32.const 1
+                            i32.const 2
+                
+                            call $_wasi_snapshot_preview_random_get
+    
+                            drop
+                        )
+                    )
+                    (else 
+                        (loop $test_loop
+                            i32.const 1
+                            i32.const 2
+                
+                            call $_wasi_snapshot_preview_random_get
+    
+                            br_if $test_loop
+                        )
+                    )
+                )
+            )
+
+            i32.const 3
+            i32.const 4
+            i32.const 5
+            call $_wasi_snapshot_preview_fd_write
+
+            call $__imported_wasi_snapshot_preview1_proc_exit
+        )
+
+        (func $__ic_custom_random_get (;7;) (type 3) (param i32 i32) (result i32)
+            call $_msg_reply
+            i32.const 421
+        )
+
+        (func $__ic_custom_fd_write (;8;) (type 5) (param i32 i32 i32 i32) (result i32)
+            i32.const 0
+            i32.const 0
+            call $_dprint
+            i32.const 42
+        )
+
+        (export "_start" (func $_start))
+    )
+    "#;
+
+    let binary = wat::parse_str(wat).unwrap();
+    let mut module = walrus::Module::from_buffer(&binary).unwrap();
+
+    do_module_replacements(&mut module);
+
+    // we expect random_get and fd_write to be replaced, environ_get to be removed and the calls to the proc_exit to remain
+    let imports = module.imports;
+
+    let result = imports.find("ic0", "debug_print");
+    assert!(result.is_some());
+    let result = imports.find("ic0", "msg_reply");
+    assert!(result.is_some());
+    let result = imports.find("wasi_snapshot_preview1", "proc_exit");
+    assert!(result.is_some());
+    let result = imports.find("wasi_snapshot_preview1", "fd_write");
+    assert!(None == result);
+    let result = imports.find("wasi_snapshot_preview1", "random_get");
+    assert!(None == result);
+    let result = imports.find("wasi_snapshot_preview1", "environ_get");
+    assert!(None == result);
+
+}