feat: add map[key].field op= rhs compound assignment

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

src/ast.ml

@@ -285,6 +285,8 @@ and stmt_desc =
   | Assignment of string * expr
   | CompoundAssignment of string * binary_op * expr  (* var op= expr *)
   | CompoundIndexAssignment of expr * expr * binary_op * expr  (* map[key] op= expr *)
+  | CompoundFieldIndexAssignment of expr * expr * string * binary_op * expr
+      (* map[key].field op= expr *)
   | FieldAssignment of expr * string * expr  (* object.field = value *)
   | ArrowAssignment of expr * string * expr  (* pointer->field = value *)
   | IndexAssignment of expr * expr * expr  (* map[key] = value *)
@@ -818,6 +820,10 @@ and string_of_stmt stmt =
       Printf.sprintf "%s %s= %s;" name (string_of_binary_op op) (string_of_expr expr)
   | CompoundIndexAssignment (map_expr, key_expr, op, value_expr) ->
       Printf.sprintf "%s[%s] %s= %s;" (string_of_expr map_expr) (string_of_expr key_expr) (string_of_binary_op op) (string_of_expr value_expr)
+  | CompoundFieldIndexAssignment (map_expr, key_expr, field, op, value_expr) ->
+      Printf.sprintf "%s[%s].%s %s= %s;"
+        (string_of_expr map_expr) (string_of_expr key_expr) field
+        (string_of_binary_op op) (string_of_expr value_expr)
   | FieldAssignment (obj_expr, field, value_expr) ->
       Printf.sprintf "%s.%s = %s;" (string_of_expr obj_expr) field (string_of_expr value_expr)
   | ArrowAssignment (obj_expr, field, value_expr) ->

src/ebpf_c_codegen.ml

@@ -1952,6 +1952,10 @@ and generate_c_instruction ctx ir_instr =
                   (* Other string expressions (concatenation, etc.) *)
                   let init_str = generate_c_expression ctx init_expr in
                   emit_line ctx (sprintf "%s %s = %s;" type_str var_name init_str)
+              | IRPointer _, IRValue src_val when (match src_val.value_desc with IRMapAccess _ -> true | _ -> false) ->
+                  (* Pointer-typed variable initialized from a map lookup: keep the pointer. *)
+                  let init_str = generate_c_value ~auto_deref_map_access:false ctx src_val in
+                  emit_line ctx (sprintf "%s %s = %s;" type_str var_name init_str)
               | _ ->
                   (* Regular non-string assignment *)
                   let init_str = generate_c_expression ctx init_expr in

src/evaluator.ml

@@ -804,6 +804,11 @@ and eval_statement ctx stmt =
       in
       let result = eval_binary_op current_val op value_val stmt.stmt_pos in
       Hashtbl.replace map_store key_str result
+
+  | CompoundFieldIndexAssignment (_, _, _, _, _) ->
+      (* The interpreter is used for compile-time evaluation only;
+         struct-field compound assignment on map values is a runtime construct. *)
+      eval_error "map[key].field op= rhs is not supported in the interpreter" stmt.stmt_pos
 
   | FieldAssignment (obj_expr, _field, value_expr) ->
       (* For evaluation purposes, treat config field assignment as no-op *)

src/ir_generator.ml

@@ -2044,6 +2044,49 @@ and lower_statement ctx stmt =
           stmt.stmt_pos in
         emit_instruction ctx if_instr
 
+  | Ast.CompoundFieldIndexAssignment (map_expr, key_expr, field, op, value_expr) ->
+      (* Desugar `map[k].field op= rhs` to:
+           if (var __cidx_field_N = map[k]) {
+             __cidx_field_N.field = __cidx_field_N.field op rhs
+           }
+         The IfLet lowering (above) handles the presence check; the field-store
+         lowers to a pointer-checked `ptr->field = ...` via IRStructFieldAssignment.
+         We look up the field's AST type from the map's value-struct definition so
+         that the synthesized FieldAccess/BinaryOp get correct expr_type — without
+         this the IR generator defaults to IRU32, mis-sizing wider fields. *)
+      let pos = stmt.stmt_pos in
+      let synth_name = generate_temp_variable ctx "cidx_field" in
+      let map_name = match map_expr.expr_desc with
+        | Ast.Identifier mn -> mn
+        | _ -> failwith "Compound field-index assignment requires a map identifier"
+      in
+      let map_def = Hashtbl.find ctx.maps map_name in
+      let field_ast_type =
+        let rec resolve_struct_name = function
+          | Ast.Struct n | Ast.UserType n -> Some n
+          | _ -> None
+        and resolve t =
+          match resolve_struct_name t with
+          | Some n ->
+              (match Symbol_table.lookup_symbol ctx.symbol_table n with
+               | Some { kind = Symbol_table.TypeDef (Ast.StructDef (_, fs, _)); _ } ->
+                   (try Some (List.assoc field fs) with Not_found -> None)
+               | _ -> None)
+          | None -> None
+        in
+        resolve map_def.ast_value_type
+      in
+      let mk_expr ?ty d =
+        { Ast.expr_desc = d; expr_pos = pos; expr_type = ty;
+          type_checked = false; program_context = None; map_scope = None } in
+      let access = mk_expr (Ast.ArrayAccess (map_expr, key_expr)) in
+      let tmp_id = mk_expr (Ast.Identifier synth_name) in
+      let cur_field = mk_expr ?ty:field_ast_type (Ast.FieldAccess (tmp_id, field)) in
+      let bin = mk_expr ?ty:field_ast_type (Ast.BinaryOp (cur_field, op, value_expr)) in
+      let store = { Ast.stmt_desc = Ast.FieldAssignment (tmp_id, field, bin); stmt_pos = pos } in
+      let iflet = { Ast.stmt_desc = Ast.IfLet (synth_name, access, [store], None); stmt_pos = pos } in
+      lower_statement ctx iflet
+
   | Ast.IfLet (name, expr, then_stmts, else_opt) ->
       (* Desugar `if (var name = expr) { T } else { E }` into:
            var name = expr

src/multi_program_analyzer.ml

@@ -187,6 +187,10 @@ let analyze_map_usage (programs: program_def list) (global_maps: map_declaration
         analyze_expr_for_maps prog_name map_expr;
         analyze_expr_for_maps prog_name key_expr;
         analyze_expr_for_maps prog_name value_expr
+    | CompoundFieldIndexAssignment (map_expr, key_expr, _, _, value_expr) ->
+        analyze_expr_for_maps prog_name map_expr;
+        analyze_expr_for_maps prog_name key_expr;
+        analyze_expr_for_maps prog_name value_expr
     | FieldAssignment (obj_expr, _, value_expr) ->
         analyze_expr_for_maps prog_name obj_expr;
         analyze_expr_for_maps prog_name value_expr

src/parse.ml

@@ -96,6 +96,8 @@ let validate_ast ast =
     | CompoundAssignment (_, _, expr) -> validate_expr expr
     | CompoundIndexAssignment (map_expr, key_expr, _, value_expr) ->
         validate_expr map_expr && validate_expr key_expr && validate_expr value_expr
+    | CompoundFieldIndexAssignment (map_expr, key_expr, _, _, value_expr) ->
+        validate_expr map_expr && validate_expr key_expr && validate_expr value_expr
     | FieldAssignment (obj_expr, _, value_expr) ->
         validate_expr obj_expr && validate_expr value_expr
     | ArrowAssignment (obj_expr, _, value_expr) ->

src/parser.mly

@@ -126,6 +126,7 @@
 %type <Ast.statement> assignment_or_expression_statement
 %type <Ast.statement> compound_assignment_statement
 %type <Ast.statement> compound_index_assignment_statement
+%type <Ast.statement> compound_field_index_assignment_statement
 %type <Ast.statement> field_assignment_statement
 %type <Ast.statement> arrow_assignment_statement
 %type <Ast.statement> index_assignment_statement
@@ -313,6 +314,7 @@ statement:
   | index_assignment_statement { $1 }
   | compound_assignment_statement { $1 }
   | compound_index_assignment_statement { $1 }
+  | compound_field_index_assignment_statement { $1 }
   | assignment_or_expression_statement { $1 }
   | return_statement { $1 }
   | if_statement { $1 }
@@ -380,6 +382,18 @@ compound_index_assignment_statement:
   | expression LBRACKET expression RBRACKET MODULO_ASSIGN expression
     { make_stmt (CompoundIndexAssignment ($1, $3, Mod, $6)) (make_pos ()) }
 
+compound_field_index_assignment_statement:
+  | expression LBRACKET expression RBRACKET DOT IDENTIFIER PLUS_ASSIGN expression
+    { make_stmt (CompoundFieldIndexAssignment ($1, $3, $6, Add, $8)) (make_pos ()) }
+  | expression LBRACKET expression RBRACKET DOT IDENTIFIER MINUS_ASSIGN expression
+    { make_stmt (CompoundFieldIndexAssignment ($1, $3, $6, Sub, $8)) (make_pos ()) }
+  | expression LBRACKET expression RBRACKET DOT IDENTIFIER MULTIPLY_ASSIGN expression
+    { make_stmt (CompoundFieldIndexAssignment ($1, $3, $6, Mul, $8)) (make_pos ()) }
+  | expression LBRACKET expression RBRACKET DOT IDENTIFIER DIVIDE_ASSIGN expression
+    { make_stmt (CompoundFieldIndexAssignment ($1, $3, $6, Div, $8)) (make_pos ()) }
+  | expression LBRACKET expression RBRACKET DOT IDENTIFIER MODULO_ASSIGN expression
+    { make_stmt (CompoundFieldIndexAssignment ($1, $3, $6, Mod, $8)) (make_pos ()) }
+
 return_statement:
   | RETURN { make_stmt (Return None) (make_pos ()) }
   | RETURN expression { make_stmt (Return (Some $2)) (make_pos ()) }

src/safety_checker.ml

@@ -242,6 +242,10 @@ let analyze_statement_bounds stmt =
         errors := check_array_bounds map_expr @ !errors;
         errors := check_array_bounds key_expr @ !errors;
         errors := check_array_bounds value_expr @ !errors
+    | CompoundFieldIndexAssignment (map_expr, key_expr, _, _, value_expr) ->
+        errors := check_array_bounds map_expr @ !errors;
+        errors := check_array_bounds key_expr @ !errors;
+        errors := check_array_bounds value_expr @ !errors
     | FieldAssignment (obj_expr, _, value_expr) ->
         errors := check_array_bounds obj_expr @ !errors;
         errors := check_array_bounds value_expr @ !errors

src/symbol_table.ml

@@ -688,6 +688,10 @@ and process_statement table stmt =
       process_expression table map_expr;
       process_expression table key_expr;
       process_expression table value_expr
+  | CompoundFieldIndexAssignment (map_expr, key_expr, _field, _, value_expr) ->
+      process_expression table map_expr;
+      process_expression table key_expr;
+      process_expression table value_expr
   | FieldAssignment (obj_expr, _field, value_expr) ->
       process_expression table obj_expr;
       process_expression table value_expr

src/type_checker.ml

@@ -94,6 +94,8 @@ and typed_stmt_desc =
   | TAssignment of string * typed_expr
   | TCompoundAssignment of string * binary_op * typed_expr  (* var op= expr *)
   | TCompoundIndexAssignment of typed_expr * typed_expr * binary_op * typed_expr  (* map[key] op= expr *)
+  | TCompoundFieldIndexAssignment of typed_expr * typed_expr * string * binary_op * typed_expr
+      (* map[key].field op= expr *)
   | TFieldAssignment of typed_expr * string * typed_expr  (* object, field, value *)
   | TArrowAssignment of typed_expr * string * typed_expr  (* pointer, field, value *)
   | TIndexAssignment of typed_expr * typed_expr * typed_expr
@@ -1820,7 +1822,57 @@ and type_check_statement ctx stmt =
                      type_error ("Operator " ^ string_of_binary_op op ^ 
                                 " not supported for type " ^ string_of_bpf_type element_type) stmt.stmt_pos)
             | _ -> type_error ("Compound index assignment can only be used on maps or arrays") stmt.stmt_pos))
-  
+
+  | CompoundFieldIndexAssignment (map_expr, key_expr, field, op, value_expr) ->
+      let typed_key = type_check_expression ctx key_expr in
+      let typed_value = type_check_expression ctx value_expr in
+      let map_name = match map_expr.expr_desc with
+        | Identifier name when Hashtbl.mem ctx.maps name -> name
+        | _ -> type_error "Compound field-index assignment requires a map identifier" stmt.stmt_pos
+      in
+      let map_decl = Hashtbl.find ctx.maps map_name in
+      (* Key type *)
+      let resolved_key_type = resolve_user_type ctx map_decl.ast_key_type in
+      let resolved_typed_key_type = resolve_user_type ctx typed_key.texpr_type in
+      (match unify_types resolved_key_type resolved_typed_key_type with
+       | Some _ -> ()
+       | None -> type_error "Map key type mismatch" stmt.stmt_pos);
+      (* Resolve the map's value type to a struct *)
+      let resolved_value_type = resolve_user_type ctx map_decl.ast_value_type in
+      let struct_name = match resolved_value_type with
+        | Struct n | UserType n -> n
+        | _ -> type_error "map[k].field op= rhs requires the map's value type to be a struct" stmt.stmt_pos
+      in
+      let fields =
+        try
+          (match Hashtbl.find ctx.types struct_name with
+           | StructDef (_, fs, _) -> fs
+           | _ -> type_error (struct_name ^ " is not a struct") stmt.stmt_pos)
+        with Not_found -> type_error ("Undefined struct: " ^ struct_name) stmt.stmt_pos
+      in
+      let field_type =
+        try List.assoc field fields
+        with Not_found ->
+          type_error ("Field not found: " ^ field ^ " in struct " ^ struct_name) stmt.stmt_pos
+      in
+      (* rhs must match field type *)
+      let resolved_field_type = resolve_user_type ctx field_type in
+      let resolved_typed_value_type = resolve_user_type ctx typed_value.texpr_type in
+      (match unify_types resolved_field_type resolved_typed_value_type with
+       | Some _ -> ()
+       | None -> type_error ("Field value type mismatch for " ^ field) stmt.stmt_pos);
+      (* op must be valid for the field type *)
+      (match op, resolved_field_type with
+       | (Add | Sub | Mul | Div | Mod), (U8 | U16 | U32 | U64 | I8 | I16 | I32 | I64) ->
+           let typed_map = { texpr_desc = TIdentifier map_name;
+                             texpr_type = Map (map_decl.ast_key_type, map_decl.ast_value_type, map_decl.ast_map_type, map_decl.max_entries);
+                             texpr_pos = map_expr.expr_pos } in
+           { tstmt_desc = TCompoundFieldIndexAssignment (typed_map, typed_key, field, op, typed_value);
+             tstmt_pos = stmt.stmt_pos }
+       | _, _ ->
+           type_error ("Operator " ^ string_of_binary_op op ^
+                      " not supported for field type " ^ string_of_bpf_type resolved_field_type) stmt.stmt_pos)
+
   | Declaration (name, type_opt, expr_opt) ->
       let typed_expr_opt = Option.map (type_check_expression ctx) expr_opt in
 
@@ -2600,6 +2652,8 @@ and typed_stmt_to_stmt tstmt =
     | TCompoundAssignment (name, op, expr) -> CompoundAssignment (name, op, typed_expr_to_expr expr)
     | TCompoundIndexAssignment (map_expr, key_expr, op, value_expr) ->
         CompoundIndexAssignment (typed_expr_to_expr map_expr, typed_expr_to_expr key_expr, op, typed_expr_to_expr value_expr)
+    | TCompoundFieldIndexAssignment (map_expr, key_expr, field, op, value_expr) ->
+        CompoundFieldIndexAssignment (typed_expr_to_expr map_expr, typed_expr_to_expr key_expr, field, op, typed_expr_to_expr value_expr)
     | TFieldAssignment (obj_expr, field, value_expr) ->
         FieldAssignment (typed_expr_to_expr obj_expr, field, typed_expr_to_expr value_expr)
     | TArrowAssignment (obj_expr, field, value_expr) ->
@@ -3272,6 +3326,13 @@ and populate_multi_program_context ast multi_prog_analysis =
         (match map_expr.program_context with
          | Some ctx -> map_expr.program_context <- Some { ctx with data_flow_direction = Some Write }
          | None -> ())
+    | CompoundFieldIndexAssignment (map_expr, key_expr, _, _, value_expr) ->
+        enhance_expr prog_type map_expr;
+        enhance_expr prog_type key_expr;
+        enhance_expr prog_type value_expr;
+        (match map_expr.program_context with
+         | Some ctx -> map_expr.program_context <- Some { ctx with data_flow_direction = Some Write }
+         | None -> ())
     | FieldAssignment (obj_expr, _, value_expr) ->
         enhance_expr prog_type obj_expr;
         enhance_expr prog_type value_expr
@@ -3377,6 +3438,10 @@ and populate_multi_program_context ast multi_prog_analysis =
           enhance_userspace_expr map_expr;
           enhance_userspace_expr key_expr;
           enhance_userspace_expr value_expr
+      | CompoundFieldIndexAssignment (map_expr, key_expr, _, _, value_expr) ->
+          enhance_userspace_expr map_expr;
+          enhance_userspace_expr key_expr;
+          enhance_userspace_expr value_expr
       | FieldAssignment (obj_expr, _, value_expr) ->
           enhance_userspace_expr obj_expr;
           enhance_userspace_expr value_expr