refactor: replace usize indices with typed index wrappers for improved type safety

Author: Arnaud Riess

crates/herkos/src/codegen/constructor.rs

@@ -39,13 +39,14 @@ pub fn emit_const_globals<B: Backend>(_backend: &B, info: &ModuleInfo) -> String
 pub fn emit_element_segments(info: &ModuleInfo, table_receiver: &str) -> String {
     let mut code = String::new();
     for seg in &info.element_segments {
-        for (i, func_idx) in seg.func_indices.iter().enumerate() {
+        for (i, func_gidx) in seg.func_indices.iter().enumerate() {
             let table_idx = seg.offset + i;
-            let local_func_idx = *func_idx - info.num_imported_functions();
+            let global_func_idx_usize = func_gidx.as_usize();
+            let local_func_idx = global_func_idx_usize - info.num_imported_functions();
+            let local_func_idx_wrapped = LocalFuncIdx::new(local_func_idx);
             let type_idx = info
-                .ir_functions
-                .get(local_func_idx)
-                .map(|f| f.type_idx)
+                .ir_function(local_func_idx_wrapped)
+                .map(|f| f.type_idx.as_usize())
                 .unwrap_or(0);
             code.push_str(&format!(
                 "    {}.set({}, Some(FuncRef {{ type_index: {}, func_index: {} }})).unwrap();\n",

crates/herkos/src/codegen/export.rs

@@ -17,7 +17,7 @@ pub fn generate_export_impl<B: Backend>(_backend: &B, info: &ModuleInfo) -> Stri
     let export_names: std::collections::HashMap<usize, &str> = info
         .func_exports
         .iter()
-        .map(|e| (e.func_index, e.name.as_str()))
+        .map(|e| (e.func_index.as_usize(), e.name.as_str()))
         .collect();
 
     // Generate accessor methods for all functions

crates/herkos/src/codegen/function.rs

@@ -70,20 +70,11 @@ pub fn generate_function_with_info<B: Backend>(
                     func_idx,
                     ..
                 } => {
-                    // func_idx is in global space (imports + locals).
-                    // For Milestone 1, we error on imported functions during codegen,
-                    // so just use a fallback type here if it's an import.
-                    let ty = if *func_idx >= info.num_imported_functions() {
-                        let local_idx = func_idx - info.num_imported_functions();
-                        info.ir_functions
-                            .get(local_idx)
-                            .and_then(|f| f.return_type)
-                            .unwrap_or(WasmType::I32)
-                    } else {
-                        // Call to imported function — will error during codegen.
-                        // Use fallback type for now.
-                        WasmType::I32
-                    };
+                    // func_idx is in local space (imports already excluded)
+                    let ty = info
+                        .ir_function(func_idx.clone())
+                        .and_then(|f| f.return_type)
+                        .unwrap_or(WasmType::I32);
                     var_types.insert(*dest, ty);
                 }
                 IrInstr::CallImport {
@@ -93,8 +84,7 @@ pub fn generate_function_with_info<B: Backend>(
                 } => {
                     // Look up import signature from func_imports
                     let ty = info
-                        .func_imports
-                        .get(*import_idx)
+                        .func_import(import_idx.clone())
                         .and_then(|imp| imp.return_type)
                         .unwrap_or(WasmType::I32);
                     var_types.insert(*dest, ty);
@@ -107,17 +97,9 @@ pub fn generate_function_with_info<B: Backend>(
                     }
                 }
                 IrInstr::GlobalGet { dest, index } => {
-                    // Distinguish imported globals (lower indices) from local globals
-                    let ty = if *index < info.imported_globals.len() {
-                        // Imported global
-                        info.imported_globals[*index].wasm_type
-                    } else {
-                        // Local global — adjust index by removing imported count
-                        let local_idx = *index - info.imported_globals.len();
-                        info.globals
-                            .get(local_idx)
-                            .map(|g| g.init_value.ty())
-                            .unwrap_or(WasmType::I32)
+                    let ty = match info.resolve_global(*index) {
+                        ResolvedGlobal::Imported(_idx, g) => g.wasm_type,
+                        ResolvedGlobal::Local(_idx, g) => g.init_value.ty(),
                     };
                     var_types.insert(*dest, ty);
                 }
@@ -127,8 +109,7 @@ pub fn generate_function_with_info<B: Backend>(
                     ..
                 } => {
                     let ty = info
-                        .type_signatures
-                        .get(*type_idx)
+                        .type_signature(type_idx.clone())
                         .and_then(|s| s.return_type)
                         .unwrap_or(WasmType::I32);
                     var_types.insert(*dest, ty);
@@ -351,7 +332,7 @@ fn has_global_import_access(ir_func: &IrFunction, num_imported_globals: usize) -
     }
     ir_func.blocks.iter().any(|block| {
         block.instructions.iter().any(|instr| {
-            matches!(instr, IrInstr::GlobalGet { index, .. } | IrInstr::GlobalSet { index, .. } if *index < num_imported_globals)
+            matches!(instr, IrInstr::GlobalGet { index, .. } | IrInstr::GlobalSet { index, .. } if index.as_usize() < num_imported_globals)
         })
     })
 }

crates/herkos/src/codegen/instruction.rs

@@ -46,7 +46,14 @@ pub fn generate_instruction_with_info<B: Backend>(
             let has_globals = info.has_mutable_globals();
             let has_memory = info.has_memory;
             let has_table = info.has_table();
-            backend.emit_call(*dest, *func_idx, args, has_globals, has_memory, has_table)
+            backend.emit_call(
+                *dest,
+                func_idx.as_usize(),
+                args,
+                has_globals,
+                has_memory,
+                has_table,
+            )
         }
 
         IrInstr::CallImport {
@@ -62,38 +69,26 @@ pub fn generate_instruction_with_info<B: Backend>(
             type_idx,
             table_idx,
             args,
-        } => generate_call_indirect(*dest, *type_idx, *table_idx, args, info),
+        } => generate_call_indirect(*dest, type_idx.clone(), *table_idx, args, info),
 
         IrInstr::Assign { dest, src } => backend.emit_assign(*dest, *src),
 
-        IrInstr::GlobalGet { dest, index } => {
-            if *index < info.imported_globals.len() {
-                // Imported global — access via host trait getter
-                let g = &info.imported_globals[*index];
+        IrInstr::GlobalGet { dest, index } => match info.resolve_global(*index) {
+            ResolvedGlobal::Imported(_idx, g) => {
                 format!("                {} = host.get_{}();", dest, g.name)
-            } else {
-                // Local global — use corrected index and backend
-                let local_idx = index - info.imported_globals.len();
-                let is_mutable = info
-                    .globals
-                    .get(local_idx)
-                    .map(|g| g.mutable)
-                    .unwrap_or(true);
-                backend.emit_global_get(*dest, local_idx, is_mutable)
             }
-        }
+            ResolvedGlobal::Local(idx, g) => {
+                let is_mutable = g.mutable;
+                backend.emit_global_get(*dest, idx.as_usize(), is_mutable)
+            }
+        },
 
-        IrInstr::GlobalSet { index, value } => {
-            if *index < info.imported_globals.len() {
-                // Imported global — access via host trait setter
-                let g = &info.imported_globals[*index];
+        IrInstr::GlobalSet { index, value } => match info.resolve_global(*index) {
+            ResolvedGlobal::Imported(_idx, g) => {
                 format!("                host.set_{}({});", g.name, value)
-            } else {
-                // Local global — use corrected index and backend
-                let local_idx = index - info.imported_globals.len();
-                backend.emit_global_set(local_idx, *value)
             }
-        }
+            ResolvedGlobal::Local(idx, _g) => backend.emit_global_set(idx.as_usize(), *value),
+        },
 
         IrInstr::MemorySize { dest } => backend.emit_memory_size(*dest),
 
@@ -163,7 +158,7 @@ pub fn generate_terminator_with_mapping<B: Backend>(
 /// 3. Dispatches to the matching function via a match on func_index
 fn generate_call_indirect(
     dest: Option<VarId>,
-    type_idx: usize,
+    type_idx: TypeIdx,
     table_idx: VarId,
     args: &[VarId],
     info: &ModuleInfo,
@@ -174,11 +169,12 @@ fn generate_call_indirect(
 
     // Canonicalize the type index for structural equivalence (Wasm spec §4.4.9).
     // Two different type indices with identical (params, results) must match.
+    let type_idx_usize = type_idx.as_usize();
     let canon_idx = info
         .canonical_type
-        .get(type_idx)
+        .get(type_idx_usize)
         .copied()
-        .unwrap_or(type_idx);
+        .unwrap_or(type_idx_usize);
 
     let mut code = String::new();
 
@@ -218,7 +214,7 @@ fn generate_call_indirect(
     ));
 
     for (func_idx, ir_func) in info.ir_functions.iter().enumerate() {
-        if ir_func.type_idx == canon_idx {
+        if ir_func.type_idx.as_usize() == canon_idx {
             code.push_str(&format!(
                 "                    {} => func_{}({})?,\n",
                 func_idx, func_idx, args_str

crates/herkos/src/codegen/mod.rs

@@ -195,7 +195,7 @@ mod tests {
             }],
             entry_block: BlockId(0),
             return_type: Some(WasmType::I32),
-            type_idx: 0,
+            type_idx: TypeIdx::new(0),
             needs_host: false,
         };
 
@@ -245,7 +245,7 @@ mod tests {
             }],
             entry_block: BlockId(0),
             return_type: None,
-            type_idx: 0,
+            type_idx: TypeIdx::new(0),
             needs_host: false,
         };
 
@@ -350,7 +350,7 @@ mod tests {
             }],
             entry_block: BlockId(0),
             return_type: Some(WasmType::I64),
-            type_idx: 0,
+            type_idx: TypeIdx::new(0),
             needs_host: false,
         };
 
@@ -413,7 +413,7 @@ mod tests {
             }],
             entry_block: BlockId(0),
             return_type: Some(WasmType::I32),
-            type_idx: 0,
+            type_idx: TypeIdx::new(0),
             needs_host: false,
         };
 
@@ -457,15 +457,15 @@ mod tests {
                 id: BlockId(0),
                 instructions: vec![IrInstr::GlobalGet {
                     dest: VarId(0),
-                    index: 0,
+                    index: GlobalIdx::new(0),
                 }],
                 terminator: IrTerminator::Return {
                     value: Some(VarId(0)),
                 },
             }],
             entry_block: BlockId(0),
             return_type: Some(WasmType::I32),
-            type_idx: 0,
+            type_idx: TypeIdx::new(0),
             needs_host: false,
         };
 
@@ -484,7 +484,7 @@ mod tests {
             data_segments: Vec::new(),
             func_exports: vec![FuncExport {
                 name: "get_value".to_string(),
-                func_index: 0,
+                func_index: LocalFuncIdx::new(0),
             }],
             type_signatures: Vec::new(),
             canonical_type: Vec::new(),
@@ -530,7 +530,7 @@ mod tests {
             }],
             entry_block: BlockId(0),
             return_type: Some(WasmType::I32),
-            type_idx: 0,
+            type_idx: TypeIdx::new(0),
             needs_host: false,
         };
 
@@ -549,7 +549,7 @@ mod tests {
             }],
             func_exports: vec![FuncExport {
                 name: "load_word".to_string(),
-                func_index: 0,
+                func_index: LocalFuncIdx::new(0),
             }],
             type_signatures: Vec::new(),
             canonical_type: Vec::new(),
@@ -587,15 +587,15 @@ mod tests {
                 id: BlockId(0),
                 instructions: vec![IrInstr::GlobalGet {
                     dest: VarId(0),
-                    index: 0,
+                    index: GlobalIdx::new(0),
                 }],
                 terminator: IrTerminator::Return {
                     value: Some(VarId(0)),
                 },
             }],
             entry_block: BlockId(0),
             return_type: Some(WasmType::I32),
-            type_idx: 0,
+            type_idx: TypeIdx::new(0),
             needs_host: false,
         };
 
@@ -614,7 +614,7 @@ mod tests {
             data_segments: Vec::new(),
             func_exports: vec![FuncExport {
                 name: "get_const".to_string(),
-                func_index: 0,
+                func_index: LocalFuncIdx::new(0),
             }],
             type_signatures: Vec::new(),
             canonical_type: Vec::new(),

crates/herkos/src/ir/builder/assembly.rs

@@ -88,7 +88,11 @@ fn build_element_segments(parsed: &ParsedModule) -> Vec<ElementSegmentDef> {
         .iter()
         .map(|es| ElementSegmentDef {
             offset: es.offset as usize,
-            func_indices: es.func_indices.iter().map(|idx| *idx as usize).collect(),
+            func_indices: es
+                .func_indices
+                .iter()
+                .map(|idx| GlobalFuncIdx::new(*idx as usize))
+                .collect(),
         })
         .collect()
 }
@@ -104,7 +108,7 @@ fn build_function_exports(parsed: &ParsedModule, num_imported_functions: usize)
         .filter(|e| e.kind == ExportKind::Func && (e.index as usize) >= num_imported_functions)
         .map(|e| FuncExport {
             name: e.name.clone(),
-            func_index: (e.index as usize) - num_imported_functions,
+            func_index: LocalFuncIdx::new((e.index as usize) - num_imported_functions),
         })
         .collect()
 }
@@ -122,7 +126,7 @@ fn enrich_ir_functions(
     let num_imported_globals = imported_globals.len();
     for (func_idx, func) in parsed.functions.iter().enumerate() {
         if let Some(ir_func) = ir_functions.get_mut(func_idx) {
-            ir_func.type_idx = canonical_type[func.type_idx as usize];
+            ir_func.type_idx = TypeIdx::new(canonical_type[func.type_idx as usize]);
             ir_func.needs_host = function_calls_imports(ir_func, num_imported_globals);
         }
     }
@@ -138,7 +142,7 @@ fn function_calls_imports(ir_func: &IrFunction, num_imported_globals: usize) ->
                         instr,
                         IrInstr::GlobalGet { index, .. }
                             | IrInstr::GlobalSet { index, .. }
-                            if *index < num_imported_globals
+                            if index.as_usize() < num_imported_globals
                     ))
         })
     })
@@ -162,7 +166,7 @@ fn build_call_indirect_signatures(parsed: &ParsedModule) -> Vec<FuncSignature> {
             FuncSignature {
                 params,
                 return_type,
-                type_idx: 0,
+                type_idx: TypeIdx::new(0),
                 needs_host: false,
             }
         })

crates/herkos/src/ir/builder/core.rs

@@ -233,8 +233,8 @@ impl IrBuilder {
             blocks: self.blocks.clone(),
             entry_block: entry,
             return_type,
-            type_idx: 0,       // Set by enrich_ir_functions during assembly
-            needs_host: false, // Set by enrich_ir_functions during assembly
+            type_idx: TypeIdx::new(0), // Set by enrich_ir_functions during assembly
+            needs_host: false,         // Set by enrich_ir_functions during assembly
         })
     }