refactor: reorganize code generation functions and add utility methods for IR handling

Author: Arnaud Riess

crates/herkos/src/codegen/mod.rs

@@ -22,31 +22,6 @@ fn module_name_to_trait_name(module_name: &str) -> String {
     format!("{}Imports", module_name.to_upper_camel_case())
 }
 
-/// Group items by module name using a key function.
-fn group_by_module<'a, T, F>(
-    items: &'a [T],
-    key: F,
-) -> std::collections::BTreeMap<String, Vec<&'a T>>
-where
-    F: Fn(&T) -> &str,
-{
-    let mut grouped: std::collections::BTreeMap<String, Vec<&'a T>> =
-        std::collections::BTreeMap::new();
-    for item in items {
-        grouped.entry(key(item).to_string()).or_default().push(item);
-    }
-    grouped
-}
-
-/// Collect all unique module names from function and global imports.
-fn all_import_module_names(info: &ModuleInfo) -> std::collections::BTreeSet<String> {
-    info.func_imports
-        .iter()
-        .map(|i| i.module_name.clone())
-        .chain(info.imported_globals.iter().map(|g| g.module_name.clone()))
-        .collect()
-}
-
 /// Build a call args vector by conditionally adding globals/memory/table/host.
 fn build_inner_call_args(
     base_args: &[String],
@@ -197,22 +172,16 @@ impl<'a, B: Backend> CodeGenerator<'a, B> {
         let mut code = String::new();
         code.push_str("// Generated by herkos\n");
         code.push_str("// DO NOT EDIT\n\n");
+        code.push_str("use herkos_runtime::*;\n\n");
         code
     }
 
     /// Generate standalone functions (no module wrapper).
     fn generate_standalone_module(&self, info: &ModuleInfo) -> Result<String> {
         let mut rust_code = Self::rust_code_preamble();
 
-        // Imports
-        rust_code.push_str("#[allow(unused_imports)]\n");
         if info.has_memory {
-            rust_code.push_str("use herkos_runtime::{WasmResult, WasmTrap, IsolatedMemory};\n\n");
             rust_code.push_str(&format!("const MAX_PAGES: usize = {};\n\n", info.max_pages));
-        } else if info.has_memory_import {
-            rust_code.push_str("use herkos_runtime::{WasmResult, WasmTrap, IsolatedMemory};\n\n");
-        } else {
-            rust_code.push_str("use herkos_runtime::{WasmResult, WasmTrap};\n\n");
         }
 
         // Host trait definitions
@@ -239,28 +208,10 @@ impl<'a, B: Backend> CodeGenerator<'a, B> {
         let mut rust_code = Self::rust_code_preamble();
         let has_mut_globals = info.has_mutable_globals();
 
-        // Imports
-        rust_code.push_str("#[allow(unused_imports)]\n");
-        let funcref_import = if !info.element_segments.is_empty() {
-            ", FuncRef"
-        } else {
-            ""
-        };
         if info.has_memory {
-            rust_code.push_str(&format!(
-                "use herkos_runtime::{{WasmResult, WasmTrap, IsolatedMemory, Module, Table{funcref_import}}};\n\n",
-            ));
-            rust_code.push_str(&format!("const MAX_PAGES: usize = {};\n", info.max_pages));
-        } else if info.has_memory_import {
-            // Memory imported — use LibraryModule, include IsolatedMemory, no MAX_PAGES constant
-            rust_code.push_str(&format!(
-                "use herkos_runtime::{{WasmResult, WasmTrap, IsolatedMemory, LibraryModule, Table{funcref_import}}};\n\n",
-            ));
-        } else {
-            rust_code.push_str(&format!(
-                "use herkos_runtime::{{WasmResult, WasmTrap, LibraryModule, Table{funcref_import}}};\n\n",
-            ));
+            rust_code.push_str(&format!("const MAX_PAGES: usize = {};\n\n", info.max_pages));
         }
+
         if info.has_table() {
             rust_code.push_str(&format!("const TABLE_MAX: usize = {};\n", info.table_max));
         }
@@ -542,32 +493,6 @@ impl<'a, B: Backend> CodeGenerator<'a, B> {
         code
     }
 
-    /// Check if an IR function contains any CallImport instructions.
-    fn has_import_calls(ir_func: &IrFunction) -> bool {
-        ir_func.blocks.iter().any(|block| {
-            block
-                .instructions
-                .iter()
-                .any(|instr| matches!(instr, IrInstr::CallImport { .. }))
-        })
-    }
-
-    /// Check if an IR function accesses any imported globals.
-    fn has_global_import_access(ir_func: &IrFunction, num_imported_globals: usize) -> bool {
-        if num_imported_globals == 0 {
-            return false;
-        }
-        ir_func.blocks.iter().any(|block| {
-            block.instructions.iter().any(|instr| {
-                matches!(
-                    instr,
-                    IrInstr::GlobalGet { index, .. } | IrInstr::GlobalSet { index, .. }
-                    if *index < num_imported_globals
-                )
-            })
-        })
-    }
-
     /// Generate a complete Rust function from IR with module info.
     ///
     /// `is_public` controls whether the function is `pub fn` or `fn`.
@@ -731,12 +656,7 @@ impl<'a, B: Backend> CodeGenerator<'a, B> {
 
         for (var, ty) in sorted_vars {
             let rust_ty = self.wasm_type_to_rust(ty);
-            let default = match ty {
-                WasmType::I32 => "0i32",
-                WasmType::I64 => "0i64",
-                WasmType::F32 => "0.0f32",
-                WasmType::F64 => "0.0f64",
-            };
+            let default = ty.default_value_literal();
             output.push_str(&format!("    let mut {var}: {rust_ty} = {default};\n"));
         }
 
@@ -810,8 +730,8 @@ impl<'a, B: Backend> CodeGenerator<'a, B> {
         let visibility = if is_public { "pub " } else { "" };
 
         // Check if function needs host parameter (imports or global imports)
-        let needs_host = Self::has_import_calls(ir_func)
-            || Self::has_global_import_access(ir_func, info.imported_globals.len());
+        let needs_host = has_import_calls(ir_func)
+            || has_global_import_access(ir_func, info.imported_globals.len());
         let trait_bounds_opt = if needs_host {
             self.build_trait_bounds(info)
         } else {

crates/herkos/src/ir/types.rs

@@ -61,6 +61,20 @@ impl WasmType {
             _ => panic!("Unsupported value type: {:?}", vt),
         }
     }
+
+    /// Get the default/zero value literal string for this type (used in code generation).
+    ///
+    /// Examples:
+    /// - `I32` → `"0i32"`
+    /// - `F64` → `"0.0f64"`
+    pub fn default_value_literal(&self) -> &'static str {
+        match self {
+            WasmType::I32 => "0i32",
+            WasmType::I64 => "0i64",
+            WasmType::F32 => "0.0f32",
+            WasmType::F64 => "0.0f64",
+        }
+    }
 }
 
 /// IR representation of a complete function.
@@ -810,3 +824,54 @@ impl ModuleInfo {
         }
     }
 }
+
+/// Group items by module name using a key function.
+pub fn group_by_module<'a, T, F>(
+    items: &'a [T],
+    key: F,
+) -> std::collections::BTreeMap<String, Vec<&'a T>>
+where
+    F: Fn(&T) -> &str,
+{
+    let mut grouped: std::collections::BTreeMap<String, Vec<&'a T>> =
+        std::collections::BTreeMap::new();
+    for item in items {
+        grouped.entry(key(item).to_string()).or_default().push(item);
+    }
+    grouped
+}
+
+/// Collect all unique module names from function and global imports.
+pub fn all_import_module_names(info: &ModuleInfo) -> std::collections::BTreeSet<String> {
+    info.func_imports
+        .iter()
+        .map(|i| i.module_name.clone())
+        .chain(info.imported_globals.iter().map(|g| g.module_name.clone()))
+        .collect()
+}
+
+/// Check if an IR function contains any CallImport instructions.
+pub fn has_import_calls(ir_func: &IrFunction) -> bool {
+    ir_func.blocks.iter().any(|block| {
+        block
+            .instructions
+            .iter()
+            .any(|instr| matches!(instr, IrInstr::CallImport { .. }))
+    })
+}
+
+/// Check if an IR function accesses any imported globals.
+pub fn has_global_import_access(ir_func: &IrFunction, num_imported_globals: usize) -> bool {
+    if num_imported_globals == 0 {
+        return false;
+    }
+    ir_func.blocks.iter().any(|block| {
+        block.instructions.iter().any(|instr| {
+            matches!(
+                instr,
+                IrInstr::GlobalGet { index, .. } | IrInstr::GlobalSet { index, .. }
+                if *index < num_imported_globals
+            )
+        })
+    })
+}