From 9964203b01ee171e0e7fc83eb20712ef8437508c Mon Sep 17 00:00:00 2001
From: wasm-forge <122647775+wasm-forge@users.noreply.github.com>
Date: Sat, 7 Oct 2023 15:15:32 +0200
Subject: [PATCH] add tests

---
 src/main.rs  |  2 +-
 src/tests.rs | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/main.rs b/src/main.rs
index 5349a91..abc6fc6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -22,6 +22,7 @@ fn get_replacement_module_id(module: &walrus::Module, import_item: &walrus::Impo
                 log::debug!("Function replacement found: {:?} -> {:?}.", 
                     module.funcs.get(fn_id).name, 
                     module.funcs.get(fun.id()).name);
+                    
                 assert_eq!(module.funcs.get(fn_id).ty(), module.funcs.get(fun.id()).ty());
 
                 return Some(fun.id());
@@ -77,7 +78,6 @@ fn gather_replacement_ids(m: &walrus::Module) -> HashMap<FunctionId, FunctionId>
             },
 
             walrus::ImportKind::Table(_) => todo!(),
-            
             walrus::ImportKind::Memory(_) => {},
             walrus::ImportKind::Global(_) => {},
         }
diff --git a/src/tests.rs b/src/tests.rs
index b873c63..b466bd8 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -111,3 +111,62 @@ fn test_remove_start_export() {
     assert!(None == export_found);
 
 }
+
+
+
+#[test]
+fn test_gather_replacement_ids() {
+
+    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 $__ic_custom_random_get
+          i32.const 1
+          i32.const 2
+          call $_wasi_snapshot_preview_random_get
+          i32.const 4
+          i32.const 5
+          call $__ic_custom_fd_write
+          drop
+        )
+
+        (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 module = walrus::Module::from_buffer(&binary).unwrap();
+
+    let id_reps: HashMap<usize, usize> = gather_replacement_ids(&module).iter().map(|(x, y)| (x.index(), y.index())).collect();
+
+    assert!(id_reps[&2] == 8);
+    assert!(id_reps[&3] == 7);
+
+}