Refactor host import handling and update related tests

crates/herkos-core/src/backend/mod.rs

@@ -47,14 +47,12 @@ pub trait Backend {
     ) -> Result<String>;
 
     /// Emit Rust code for a function call (local function).
-    #[allow(clippy::too_many_arguments)]
+    /// All functions uniformly take env, memory, and table parameters.
     fn emit_call(
         &self,
         dest: Option<VarId>,
         func_idx: usize,
         args: &[VarId],
-        has_host: bool,
-        has_globals: bool,
         has_memory: bool,
         has_table: bool,
     ) -> String;

crates/herkos-core/src/backend/safe.rs

@@ -476,18 +476,12 @@ impl Backend for SafeBackend {
         dest: Option<VarId>,
         func_idx: usize,
         args: &[VarId],
-        has_host: bool,
-        has_globals: bool,
         has_memory: bool,
         has_table: bool,
     ) -> String {
         let mut call_args: Vec<String> = args.iter().map(|a| a.to_string()).collect();
-        if has_host {
-            call_args.push("host".to_string());
-        }
-        if has_globals {
-            call_args.push("globals".to_string());
-        }
+        // All functions uniformly receive env
+        call_args.push("env".to_string());
         if has_memory {
             call_args.push("memory".to_string());
         }
@@ -505,23 +499,23 @@ impl Backend for SafeBackend {
         func_name: &str,
         args: &[VarId],
     ) -> String {
-        // Generate: host.func_name(args)?
+        // Generate: env.host.func_name(args)?
         // Note: module_name is ignored for now (Milestone 3 will use it for trait names)
         let args_str: Vec<String> = args.iter().map(|a| a.to_string()).collect();
-        let call_expr = format!("host.{}({})?", func_name, args_str.join(", "));
+        let call_expr = format!("env.host.{}({})?", func_name, args_str.join(", "));
         emit_call_result(dest, &call_expr)
     }
 
     fn emit_global_get(&self, dest: VarId, index: usize, is_mutable: bool) -> String {
         if is_mutable {
-            format!("                {dest} = globals.g{index};")
+            format!("                {dest} = env.globals.g{index};")
         } else {
             format!("                {dest} = G{index};")
         }
     }
 
     fn emit_global_set(&self, index: usize, value: VarId) -> String {
-        format!("                globals.g{index} = {value};")
+        format!("                env.globals.g{index} = {value};")
     }
 
     fn emit_assign(&self, dest: VarId, src: VarId) -> String {

crates/herkos-core/src/codegen/constructor.rs

@@ -89,14 +89,14 @@ pub fn generate_constructor<B: Backend>(
         && info.element_segments.is_empty()
     {
         code.push_str("pub fn new() -> Result<WasmModule, ConstructionError> {\n");
-        code.push_str("    Ok(WasmModule(LibraryModule::new((), Table::try_new(0)?)))\n");
+        code.push_str("    Ok(WasmModule(LibraryModule::new(Globals {}, Table::try_new(0)?)))\n");
         code.push_str("}\n");
         return Ok(code);
     }
 
     code.push_str("pub fn new() -> WasmResult<WasmModule> {\n");
 
-    // Build globals initializer
+    // Build globals initializer (always generates a Globals struct, empty if no mutable globals)
     let globals_init = if has_mut_globals {
         let mut fields = String::from("Globals { ");
         let mut first = true;
@@ -113,7 +113,7 @@ pub fn generate_constructor<B: Backend>(
         fields.push_str(" }");
         fields
     } else {
-        "()".to_string()
+        "Globals {}".to_string()
     };
 
     // Table initialization
@@ -125,7 +125,8 @@ pub fn generate_constructor<B: Backend>(
 
     if info.has_memory {
         let needs_mut = !info.data_segments.is_empty() || !info.element_segments.is_empty();
-        let globals_type = if has_mut_globals { "Globals" } else { "()" };
+        // Always use Globals type (may be empty struct)
+        let globals_type = "Globals";
         let table_size_str = if info.has_table() { "TABLE_MAX" } else { "0" };
         // Use try_init (in-place initialisation) instead of try_new to avoid
         // materialising a large Result<Module<…>, E> on the call stack. In

crates/herkos-core/src/codegen/env.rs

@@ -0,0 +1,96 @@
+//! Environment struct generation.
+//!
+//! Generates the uniform `Env<H>` context struct that bundles host + globals,
+//! along with the `ModuleHostTrait` and `Globals` struct that it contains.
+
+use crate::ir::*;
+
+/// Generate the environment block: ModuleHostTrait, Globals struct, and Env<H> struct.
+///
+/// This always generates:
+/// - `pub trait ModuleHostTrait { ... }` (empty if no imports, with methods if imports)
+/// - `impl ModuleHostTrait for herkos_runtime::NoHost {}` (only for modules with NO imports)
+/// - `pub struct Globals { ... }` (empty struct if no mutable globals, fields otherwise)
+/// - `struct Env<H: ModuleHostTrait> { pub host: H, pub globals: Globals }`
+pub fn generate_env_block(info: &ModuleInfo) -> String {
+    let mut code = String::new();
+
+    // Generate ModuleHostTrait (unified, with all imports merged)
+    code.push_str(&generate_module_host_trait(info));
+    code.push('\n');
+
+    // Generate NoHost impl only for modules with NO imports
+    let has_imports = !info.func_imports.is_empty() || !info.imported_globals.is_empty();
+    if !has_imports {
+        code.push_str("impl ModuleHostTrait for herkos_runtime::NoHost {}\n\n");
+    }
+
+    // Generate Globals struct
+    code.push_str(&generate_globals_struct(info));
+    code.push('\n');
+
+    // Generate Env<H> struct
+    code.push_str("#[allow(dead_code)]\n");
+    code.push_str("struct Env<'a, H: ModuleHostTrait + ?Sized> {\n");
+    code.push_str("    pub host: &'a mut H,\n");
+    code.push_str("    pub globals: &'a mut Globals,\n");
+    code.push_str("}\n\n");
+
+    code
+}
+
+/// Generate the unified ModuleHostTrait from both function and global imports.
+fn generate_module_host_trait(info: &ModuleInfo) -> String {
+    let mut code = String::from("pub trait ModuleHostTrait {\n");
+
+    // Add all function import methods
+    for imp in &info.func_imports {
+        let mut params = vec!["&mut self".to_string()];
+        for (i, ty) in imp.params.iter().enumerate() {
+            let rust_ty = crate::codegen::types::wasm_type_to_rust(ty);
+            params.push(format!("arg{}: {}", i, rust_ty));
+        }
+
+        let return_ty = crate::codegen::types::format_return_type(imp.return_type.as_ref());
+        code.push_str(&format!(
+            "    fn {}({}) -> {};\n",
+            imp.func_name,
+            params.join(", "),
+            return_ty
+        ));
+    }
+
+    // Add all global import accessors
+    for g in &info.imported_globals {
+        let rust_ty = crate::codegen::types::wasm_type_to_rust(&g.wasm_type);
+
+        // Getter (always)
+        code.push_str(&format!("    fn get_{}(&self) -> {};\n", g.name, rust_ty));
+
+        // Setter (only if mutable)
+        if g.mutable {
+            code.push_str(&format!(
+                "    fn set_{}(&mut self, val: {});\n",
+                g.name, rust_ty
+            ));
+        }
+    }
+
+    code.push_str("}\n");
+    code
+}
+
+/// Generate the Globals struct containing all mutable globals.
+fn generate_globals_struct(info: &ModuleInfo) -> String {
+    let mut code = String::from("pub struct Globals {\n");
+
+    for (idx, g) in info.globals.iter().enumerate() {
+        if g.mutable {
+            let rust_ty = crate::codegen::types::wasm_type_to_rust(&g.init_value.ty());
+            code.push_str(&format!("    pub g{}: {},\n", idx, rust_ty));
+        }
+    }
+
+    code.push_str("}\n");
+    code
+}

crates/herkos-core/src/codegen/export.rs

@@ -1,15 +1,15 @@
 //! Export implementation generation.
 //!
 //! Generates the `impl WasmModule { ... }` block with methods for all functions.
-//! Exported functions use their export names, internal functions use func_N names.
+//! Exported functions are thin wrappers that construct an Env<H> and forward to internal functions.
 
 use crate::backend::Backend;
 use crate::ir::*;
 
 /// Generate the `impl WasmModule { ... }` block with accessor methods for all functions.
 pub fn generate_export_impl<B: Backend>(_backend: &B, info: &ModuleInfo) -> String {
     let mut code = String::new();
-    let has_mut_globals = info.has_mutable_globals();
+    let has_imports = !info.func_imports.is_empty() || !info.imported_globals.is_empty();
 
     code.push_str("impl WasmModule {\n");
 
@@ -31,24 +31,16 @@ pub fn generate_export_impl<B: Backend>(_backend: &B, info: &ModuleInfo) -> Stri
             format!("func_{}", func_idx)
         };
 
-        // Determine trait bounds for this export
-        let trait_bounds_opt = if ir_func.needs_host {
-            crate::codegen::traits::build_trait_bounds(info)
-        } else {
-            None
-        };
-        let has_multiple_bounds = trait_bounds_opt.as_ref().is_some_and(|b| b.contains(" + "));
-
-        // Build generics: handle both H (host) and MP (imported memory size)
+        // Build generics
         let mut generics: Vec<String> = Vec::new();
         if info.has_memory_import {
             generics.push("const MP: usize".to_string());
         }
-        if has_multiple_bounds {
-            generics.push(format!("H: {}", trait_bounds_opt.as_ref().unwrap()));
+        if has_imports {
+            generics.push("H: ModuleHostTrait".to_string());
         }
 
-        // Method signature with optional generic parameter
+        // Method signature
         let mut param_parts: Vec<String> = Vec::new();
         param_parts.push("&mut self".to_string());
         for (i, (_, ty)) in ir_func.params.iter().enumerate() {
@@ -61,15 +53,9 @@ pub fn generate_export_impl<B: Backend>(_backend: &B, info: &ModuleInfo) -> Stri
             param_parts.push("memory: &mut IsolatedMemory<MP>".to_string());
         }
 
-        // Add host parameter if function needs it
-        if let Some(trait_bounds) = &trait_bounds_opt {
-            if has_multiple_bounds {
-                // Use generic parameter H
-                param_parts.push("host: &mut H".to_string());
-            } else {
-                // Single trait bound - use impl directly
-                param_parts.push(format!("host: &mut impl {trait_bounds}"));
-            }
+        // Add host parameter if module has imports
+        if has_imports {
+            param_parts.push("host: &mut H".to_string());
         }
 
         let return_type = crate::codegen::types::format_return_type(ir_func.return_type.as_ref());
@@ -88,18 +74,21 @@ pub fn generate_export_impl<B: Backend>(_backend: &B, info: &ModuleInfo) -> Stri
             return_type
         ));
 
-        // Forward call to internal function
+        // Construct Env and forward call to internal function
+        if has_imports {
+            code.push_str("        let mut env = Env { host, globals: &mut self.0.globals };\n");
+        } else {
+            code.push_str("        let mut __host = herkos_runtime::NoHost;\n");
+            code.push_str(
+                "        let mut env = Env { host: &mut __host, globals: &mut self.0.globals };\n",
+            );
+        }
+
+        // Build call arguments: wasm params + env + memory (if owned) + table
         let mut call_args: Vec<String> =
             (0..ir_func.params.len()).map(|i| format!("v{i}")).collect();
+        call_args.push("&mut env".to_string());
 
-        // Forward host parameter if needed
-        if ir_func.needs_host {
-            call_args.push("host".to_string());
-        }
-
-        if has_mut_globals {
-            call_args.push("&mut self.0.globals".to_string());
-        }
         if info.has_memory {
             call_args.push("&mut self.0.memory".to_string());
         } else if info.has_memory_import {

crates/herkos-core/src/codegen/function.rs

@@ -21,7 +21,7 @@ pub fn generate_function_with_info<B: Backend>(
     let mut output = String::new();
 
     // Suppress warnings for generated code patterns that are hard to avoid
-    output.push_str("#[allow(unused_mut, unused_variables, unused_assignments, clippy::needless_return, clippy::manual_range_contains, clippy::never_loop)]\n");
+    output.push_str("#[allow(unused_mut, unused_variables, unused_assignments, clippy::only_used_in_recursion, clippy::needless_return, clippy::manual_range_contains, clippy::never_loop)]\n");
 
     // Generate function signature
     output.push_str(&generate_signature_with_info(
@@ -172,20 +172,12 @@ pub fn generate_function_with_info<B: Backend>(
     output.push_str("    loop {\n");
     output.push_str("        match __current_block {\n");
 
-    // Compute whether this function has a host parameter in scope.
-    // Use ir_func.needs_host — the IR builder's fixpoint is authoritative.
-    let caller_has_host = ir_func.needs_host;
-
     for (idx, block) in ir_func.blocks.iter().enumerate() {
         output.push_str(&format!("            Block::B{} => {{\n", idx));
 
         for instr in &block.instructions {
-            let code = crate::codegen::instruction::generate_instruction_with_info(
-                backend,
-                instr,
-                info,
-                caller_has_host,
-            )?;
+            let code =
+                crate::codegen::instruction::generate_instruction_with_info(backend, instr, info)?;
             output.push_str(&code);
             output.push('\n');
         }
@@ -220,26 +212,13 @@ fn generate_signature_with_info<B: Backend>(
 ) -> String {
     let visibility = if is_public { "pub " } else { "" };
 
-    // Use ir_func.needs_host — the IR builder's fixpoint algorithm is authoritative.
-    // It handles: direct imports, imported globals, call_indirect with imports,
-    // and transitive calls to functions that need host.
-    let needs_host = ir_func.needs_host;
-    let trait_bounds_opt = if needs_host {
-        crate::codegen::traits::build_trait_bounds(info)
-    } else {
-        None
-    };
-
-    let has_multiple_bounds = trait_bounds_opt.as_ref().is_some_and(|b| b.contains(" + "));
-
     // Build generics: handle both H (host) and MP (imported memory size)
     let mut generics: Vec<String> = Vec::new();
     if info.has_memory_import {
         generics.push("const MP: usize".to_string());
     }
-    if has_multiple_bounds {
-        generics.push(format!("H: {}", trait_bounds_opt.as_ref().unwrap()));
-    }
+    // All internal functions have H: ModuleHostTrait generic
+    generics.push("H: ModuleHostTrait".to_string());
 
     let generic_part = if generics.is_empty() {
         String::new()
@@ -259,21 +238,8 @@ fn generate_signature_with_info<B: Backend>(
         })
         .collect();
 
-    // Add host parameter if function needs imports or global imports
-    if let Some(trait_bounds) = trait_bounds_opt {
-        if has_multiple_bounds {
-            // Use generic parameter H
-            param_parts.push("host: &mut H".to_string());
-        } else {
-            // Single trait bound - use impl directly
-            param_parts.push(format!("host: &mut impl {trait_bounds}"));
-        }
-    }
-
-    // Add globals parameter if module has mutable globals
-    if info.has_mutable_globals() {
-        param_parts.push("globals: &mut Globals".to_string());
-    }
+    // Always add env parameter
+    param_parts.push("env: &mut Env<'_, H>".to_string());
 
     // Add memory parameter — either const MAX_PAGES or generic MP
     if info.has_memory {