Enhance type resolution in IR generation by adding support for type aliases and enums in the context of AST to IR type conversion. This includes preserving alias names and properly resolving struct and enum definitions from the symbol table.

Author: congwang-mk

src/ir.ml

@@ -686,7 +686,30 @@ let rec ast_type_to_ir_type = function
 (* Helper function that preserves type aliases when converting AST types to IR types *)
 let rec ast_type_to_ir_type_with_context symbol_table ast_type =
   match ast_type with
-  | UserType name | Struct name ->
+  | UserType name ->
+      (* Check if this is a type alias or struct by looking up the symbol *)
+      (match Symbol_table.lookup_symbol symbol_table name with
+         | Some symbol ->
+             (match symbol.kind with
+              | Symbol_table.TypeDef (Ast.TypeAlias (_, underlying_type)) -> 
+                  (* Create IRTypeAlias to preserve the alias name *)
+                  IRTypeAlias (name, ast_type_to_ir_type underlying_type)
+              | Symbol_table.TypeDef (Ast.StructDef (_, fields, kernel_defined)) ->
+                  (* Resolve struct fields properly with type aliases preserved *)
+                  let ir_fields = List.map (fun (field_name, field_type) ->
+                    (field_name, ast_type_to_ir_type_with_context symbol_table field_type)
+                  ) fields in
+                  IRStruct (name, ir_fields, kernel_defined)
+              | Symbol_table.TypeDef (Ast.EnumDef (_, values, kernel_defined)) -> 
+                  let ir_values = List.map (fun (enum_name, opt_value) ->
+                    (enum_name, Option.value ~default:0 opt_value)
+                  ) values in
+                  IREnum (name, ir_values, kernel_defined)
+              | _ -> ast_type_to_ir_type ast_type)
+         | None ->
+             (* Fallback to regular conversion *)
+             ast_type_to_ir_type ast_type)
+  | Struct name ->
       (* Check if this is a type alias or struct by looking up the symbol *)
       (match Symbol_table.lookup_symbol symbol_table name with
          | Some symbol ->
@@ -717,6 +740,18 @@ let rec ast_type_to_ir_type_with_context symbol_table ast_type =
       (* Recursively handle array element types with context *)
       let bounds = make_bounds_info ~min_size:size ~max_size:size () in
       IRArray (ast_type_to_ir_type_with_context symbol_table elem_type, size, bounds)
+  | Enum name ->
+      (* Check if this enum is defined in the symbol table *)
+      (match Symbol_table.lookup_symbol symbol_table name with
+         | Some symbol ->
+             (match symbol.kind with
+              | Symbol_table.TypeDef (Ast.EnumDef (_, values, kernel_defined)) -> 
+                  let ir_values = List.map (fun (enum_name, opt_value) ->
+                    (enum_name, Option.value ~default:0 opt_value)
+                  ) values in
+                  IREnum (name, ir_values, kernel_defined)
+              | _ -> ast_type_to_ir_type ast_type)
+         | None -> ast_type_to_ir_type ast_type)
   | _ -> ast_type_to_ir_type ast_type
 
 let ast_map_type_to_ir_map_type = function

src/ir_generator.ml

@@ -990,9 +990,9 @@ and resolve_type_alias ctx reg ast_type =
                 Hashtbl.replace ctx.register_aliases reg (alias_name, underlying_ir_type);
                 (* Create IRTypeAlias to preserve the alias name *)
                 IRTypeAlias (alias_name, underlying_ir_type)
-            | _ -> ast_type_to_ir_type ast_type)
-        | None -> ast_type_to_ir_type ast_type)
-  | _ -> ast_type_to_ir_type ast_type
+            | _ -> ast_type_to_ir_type_with_context ctx.symbol_table ast_type)
+        | None -> ast_type_to_ir_type_with_context ctx.symbol_table ast_type)
+  | _ -> ast_type_to_ir_type_with_context ctx.symbol_table ast_type
 
 (** Helper function to calculate stack usage for a type *)
 and calculate_stack_usage = function

tests/test_type_checker.ml

@@ -1221,6 +1221,75 @@ let test_tail_call_cross_program_type_restriction _ =
               true (contains_substr msg "incompatible program type")
    | _ -> failwith "Expected TypeError for cross-program-type tail call")
 
+(** Test map index type resolution bug fix - structs, enums, and type aliases as map keys *)
+let test_map_index_type_resolution_bug_fix _ =
+  let source_code = {|
+    // Type alias 
+    type IpAddress = u32
+    type Counter = u64
+    
+    // Enum type
+    enum Protocol {
+      TCP = 6,
+      UDP = 17,
+      ICMP = 1
+    }
+    
+    // Struct type
+    struct PacketInfo {
+      src_ip: IpAddress,
+      dst_ip: IpAddress,
+      protocol: u8
+    }
+    
+    // Maps using different key types
+    map<IpAddress, Counter> connection_count : HashMap(1024)      // Type alias key
+    map<Protocol, Counter> protocol_stats : PercpuArray(32)       // Enum key  
+    map<PacketInfo, u32> packet_filter : LruHash(512)             // Struct key
+    
+    @helper
+    fn test_indexing() -> u32 {
+      // Create test values
+      var ip: IpAddress = 0xC0A80001
+      var proto = TCP
+      var info = PacketInfo { src_ip: ip, dst_ip: ip, protocol: 6 }
+      
+      // These should all work without "Array index must be integer type" error
+      var count1 = connection_count[ip]        // Type alias as key
+      var count2 = protocol_stats[proto]       // Enum as key  
+      var result = packet_filter[info]         // Struct as key
+      
+      if (count1 != none && count2 != none && result != none) {
+        return count1 + count2 + result
+      } else {
+        return 0
+      }
+    }
+    
+    @xdp fn packet_handler(ctx: *xdp_md) -> xdp_action {
+      return XDP_PASS
+    }
+  |} in
+  
+  try
+    let ast = parse_string source_code in
+    let symbol_table = Test_utils.Helpers.create_test_symbol_table ast in
+    let _typed_ast = type_check_and_annotate_ast ~symbol_table:(Some symbol_table) ast in
+    
+    (* If we reach here, type checking succeeded *)
+    check bool "map index type resolution works for structs, enums, and type aliases" true true
+  with
+  | Type_error (msg, _) when String.contains msg 'A' && String.contains msg 'r' && String.contains msg 'i' ->
+      (* If we get "Array index must be integer type" error, the test fails *)
+      fail ("Bug regression - map indexing should work with user types: " ^ msg)
+  | Type_error (msg, _) ->
+      (* Other type errors might be valid (e.g., map key type mismatches) *)
+      fail ("Unexpected type error: " ^ msg)
+  | Parse_error (msg, _) ->
+      fail ("Parse error: " ^ msg)
+  | e ->
+      fail ("Unexpected error: " ^ Printexc.to_string e)
+
 let type_checker_tests = [
   "type_unification", `Quick, test_type_unification;
   "basic_type_inference", `Quick, test_basic_type_inference;
@@ -1258,6 +1327,7 @@ let type_checker_tests = [
   "kernel_to_kernel_function_calls", `Quick, test_kernel_to_kernel_function_calls;
   "function_call_user_type_resolution", `Quick, test_function_call_user_type_resolution;
   "tail_call_cross_program_type_restriction", `Quick, test_tail_call_cross_program_type_restriction;
+  "map_index_type_resolution_bug_fix", `Quick, test_map_index_type_resolution_bug_fix;
 ]
 
 let () =