Add support for NewWithFlag expression in AST and related components. Update parser, type checker, and IR generation to handle object allocation with GFP flags.

Author: congwang-mk

examples/simple_gfp_test.ks

@@ -0,0 +1,32 @@
+// Simple test to verify GFP flag validation 
+
+struct TestData {
+    value: u64,
+}
+
+@kfunc
+fn valid_kfunc_allocation() -> i32 {
+    // Basic allocation (valid in kernel context)
+    var basic_ptr = new TestData(GFP_ATOMIC)
+    delete basic_ptr
+    return 0
+}
+
+// This should succeed - basic allocation in eBPF context
+@xdp
+fn valid_ebpf_allocation(ctx: *xdp_md) -> xdp_action {
+    var ptr = new TestData()
+    delete ptr
+    return XDP_PASS
+}
+
+// This should succeed - basic allocation in userspace
+fn valid_userspace_allocation() -> i32 {
+    var ptr = new TestData()
+    delete ptr
+    return 0
+}
+
+fn main() -> i32 {
+    return valid_userspace_allocation()
+} 

src/ast.ml

@@ -171,6 +171,7 @@ and expr_desc =
   | StructLiteral of string * (string * expr) list
   | Match of expr * match_arm list  (* match (expr) { arms } *)
   | New of bpf_type  (* new Type() - object allocation *)
+  | NewWithFlag of bpf_type * expr  (* new Type(gfp_flag) - object allocation with flag *)
 
 (** Module function call *)
 and module_call = {
@@ -684,6 +685,8 @@ let rec string_of_expr expr =
       let arms_str = String.concat ",\n    " (List.map string_of_match_arm arms) in
       Printf.sprintf "match (%s) {\n    %s\n}" (string_of_expr expr) arms_str
   | New typ -> Printf.sprintf "new %s()" (string_of_bpf_type typ)
+  | NewWithFlag (typ, flag_expr) -> 
+      Printf.sprintf "new %s(%s)" (string_of_bpf_type typ) (string_of_expr flag_expr)
 
 and string_of_match_pattern = function
   | ConstantPattern lit -> string_of_literal lit

src/ebpf_c_codegen.ml

@@ -477,6 +477,8 @@ let rec collect_string_sizes_from_instr ir_instr =
       (collect_string_sizes_from_value instance_val) @ (collect_string_sizes_from_value struct_ops_val)
   | IRObjectNew (dest_val, _) ->
       collect_string_sizes_from_value dest_val
+  | IRObjectNewWithFlag (dest_val, _, flag_val) ->
+      (collect_string_sizes_from_value dest_val) @ (collect_string_sizes_from_value flag_val)
   | IRObjectDelete ptr_val ->
       collect_string_sizes_from_value ptr_val
 
@@ -2646,6 +2648,8 @@ let rec generate_c_instruction ctx ir_instr =
               List.iter collect_in_value args
           | IRObjectNew (dest_val, _) ->
               collect_in_value dest_val
+          | IRObjectNewWithFlag (dest_val, _, flag_val) ->
+              collect_in_value dest_val; collect_in_value flag_val
           | IRObjectDelete ptr_val ->
               collect_in_value ptr_val
           | IRJump _ | IRComment _ | IRBreak | IRContinue | IRThrow _ -> ()
@@ -2856,6 +2860,11 @@ let rec generate_c_instruction ctx ir_instr =
       (* Use proper kernel pattern: ptr = bpf_obj_new(type) *)
       emit_line ctx (sprintf "%s = bpf_obj_new(%s);" dest_str type_str)
 
+  | IRObjectNewWithFlag _ ->
+      (* GFP flags should never reach eBPF code generation - this is an internal error *)
+      failwith ("Internal error: GFP allocation flags are not supported in eBPF context. " ^
+                "This should have been caught by the type checker.")
+      
   | IRObjectDelete ptr_val ->
       let ptr_str = generate_c_value ctx ptr_val in
       (* Use the proper kernel bpf_obj_drop(ptr) macro *)
@@ -3067,6 +3076,8 @@ let collect_registers_in_function ir_func =
         collect_in_value instance_val; collect_in_value struct_ops_val
     | IRObjectNew (dest_val, _) ->
         collect_in_value dest_val
+    | IRObjectNewWithFlag (dest_val, _, flag_val) ->
+        collect_in_value dest_val; collect_in_value flag_val
     | IRObjectDelete ptr_val ->
         collect_in_value ptr_val
   in

src/evaluator.ml

@@ -754,6 +754,11 @@ and eval_expression ctx expr =
       (* For evaluator, object allocation returns a mock pointer value *)
       (* This is just for testing - real allocation happens in generated code *)
       PointerValue (Random.int 1000000)
+      
+  | NewWithFlag (_, _) ->
+      (* For evaluator, object allocation with flag returns a mock pointer value *)
+      (* This is just for testing - real allocation happens in generated code *)
+      PointerValue (Random.int 1000000)
 
 (** Evaluate statements *)
 and eval_statements ctx stmts =

src/ir.ml

@@ -297,6 +297,7 @@ and ir_instr_desc =
   | IRMapStore of ir_value * ir_value * ir_value * map_store_type
   | IRMapDelete of ir_value * ir_value
   | IRObjectNew of ir_value * ir_type  (* target_pointer, object_type *)
+  | IRObjectNewWithFlag of ir_value * ir_type * ir_value  (* target_pointer, object_type, flag_expr *)
   | IRObjectDelete of ir_value  (* pointer_to_delete *)
   | IRConfigFieldUpdate of ir_value * ir_value * string * ir_value (* map, key, field, value *)
   | IRStructFieldAssignment of ir_value * string * ir_value (* object, field, value *)
@@ -899,6 +900,8 @@ let rec string_of_ir_instruction instr =
       Printf.sprintf "delete(%s, %s)" (string_of_ir_value map) (string_of_ir_value key)
   | IRObjectNew (dest, obj_type) ->
       Printf.sprintf "%s = object_new(%s)" (string_of_ir_value dest) (string_of_ir_type obj_type)
+  | IRObjectNewWithFlag (dest, obj_type, flag_expr) ->
+      Printf.sprintf "%s = object_new(%s, %s)" (string_of_ir_value dest) (string_of_ir_type obj_type) (string_of_ir_value flag_expr)
   | IRObjectDelete ptr ->
       Printf.sprintf "object_delete(%s)" (string_of_ir_value ptr)
   | IRConfigFieldUpdate (map, key, field, value) ->

src/ir_generator.ml

@@ -937,6 +937,20 @@ let rec lower_expression ctx (expr : Ast.expr) =
       emit_instruction ctx alloc_instr;
 
       result_val
+      
+  | Ast.NewWithFlag (typ, flag_expr) ->
+      (* Object allocation with GFP flag - only valid in kernel context *)
+      let ir_type = ast_type_to_ir_type typ in
+      let result_reg = allocate_register ctx in
+      let result_val = make_ir_value (IRRegister result_reg) (IRPointer (ir_type, make_bounds_info ())) expr.expr_pos in
+      
+      (* Lower the flag expression *)
+      let flag_val = lower_expression ctx flag_expr in
+      
+      let alloc_instr = make_ir_instruction (IRObjectNewWithFlag (result_val, ir_type, flag_val)) expr.expr_pos in
+      emit_instruction ctx alloc_instr;
+      
+      result_val
 
 (** Helper function to handle register() builtin function calls *)
 and handle_register_builtin_call ctx args expr_pos ?target_register ?target_type () =
@@ -1335,7 +1349,7 @@ and lower_statement ctx stmt =
 
       (* Handle function call and new expression declarations elegantly by proper instruction ordering *)
       (match expr_opt with
-       | Some expr when (match expr.expr_desc with Ast.Call _ | Ast.New _ -> true | _ -> false) ->
+       | Some expr when (match expr.expr_desc with Ast.Call _ | Ast.New _ | Ast.NewWithFlag _ -> true | _ -> false) ->
            (* For function calls and new expressions: emit declaration first, then operation with assignment *)
            let target_type = match typ_opt with
              | Some ast_type -> resolve_type_alias ctx reg ast_type
@@ -1348,6 +1362,9 @@ and lower_statement ctx stmt =
                        | Ast.New typ -> 
                            let ir_type = ast_type_to_ir_type typ in
                            IRPointer (ir_type, make_bounds_info ())
+                       | Ast.NewWithFlag (typ, _) ->
+                           let ir_type = ast_type_to_ir_type typ in
+                           IRPointer (ir_type, make_bounds_info ())
                        | _ -> IRU32))
            in
 
@@ -1385,6 +1402,13 @@ and lower_statement ctx stmt =
                 let result_val = make_ir_value (IRRegister reg) target_type expr.Ast.expr_pos in
                 let alloc_instr = make_ir_instruction (IRObjectNew (result_val, ir_type)) expr.Ast.expr_pos in
                 emit_instruction ctx alloc_instr
+            | Ast.NewWithFlag (typ, flag_expr) ->
+                (* Handle new expression with flag: emit allocation instruction with flag *)
+                let ir_type = ast_type_to_ir_type typ in
+                let result_val = make_ir_value (IRRegister reg) target_type expr.Ast.expr_pos in
+                let flag_val = lower_expression ctx flag_expr in
+                let alloc_instr = make_ir_instruction (IRObjectNewWithFlag (result_val, ir_type, flag_val)) expr.Ast.expr_pos in
+                emit_instruction ctx alloc_instr
             | _ -> ()) (* Shouldn't happen due to our guard *)
        | _ ->
            (* Non-function call declarations: use existing logic *)

src/kernel_module_codegen.ml

@@ -179,6 +179,14 @@ let rec generate_statement_translation stmt =
       sprintf "    %s;" (generate_expression_translation expr)
   | Break -> "    break;"
   | Continue -> "    continue;"
+  | Delete (DeletePointer ptr_expr) ->
+      (* Translate pointer deletion to kfree() *)
+      let ptr_str = generate_expression_translation ptr_expr in
+      sprintf "    kfree(%s);" ptr_str
+  | Delete (DeleteMapEntry (map_expr, key_expr)) ->
+      (* Map deletion not supported in kernel modules - this should be caught earlier *)
+      sprintf "    /* Map deletion not supported in kernel modules: delete %s[%s] */" 
+        (generate_expression_translation map_expr) (generate_expression_translation key_expr)
   | _ -> "    /* TODO: Implement statement translation */"
 
 (** Generate expression translation *)
@@ -262,6 +270,15 @@ and generate_expression_translation expr =
       sprintf "%s->%s" (generate_expression_translation obj) field
   | ArrayAccess (array, index) ->
       sprintf "%s[%s]" (generate_expression_translation array) (generate_expression_translation index)
+  | New typ ->
+      (* Basic allocation with GFP_KERNEL (default for kernel context) *)
+      let c_type = kernelscript_type_to_c_type typ in
+      sprintf "kmalloc(sizeof(%s), GFP_KERNEL)" c_type
+  | NewWithFlag (typ, flag_expr) ->
+      (* Allocation with specific GFP flag *)
+      let c_type = kernelscript_type_to_c_type typ in
+      let flag_str = generate_expression_translation flag_expr in
+      sprintf "kmalloc(sizeof(%s), %s)" c_type flag_str
   | _ -> "/* TODO: Implement expression translation */"
 
 (** Generate function implementation for regular kernel module functions *)

src/main.ml

@@ -740,7 +740,10 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
         (List.length filtered_btf_declarations) 
         (List.length btf_declarations - List.length filtered_btf_declarations);
 
-      let symbol_table = Symbol_table.build_symbol_table ~project_name:base_name ~builtin_asts:[filtered_btf_declarations] compilation_ast in
+    (* Add stdlib builtin types to the symbol table *)
+    let stdlib_builtin_declarations = Stdlib.get_builtin_types () in
+    let all_builtin_declarations = stdlib_builtin_declarations @ filtered_btf_declarations in
+    let symbol_table = Symbol_table.build_symbol_table ~project_name:base_name ~builtin_asts:[all_builtin_declarations] compilation_ast in
 
     Printf.printf "✅ Symbol table created successfully with BTF types\n\n";
 

src/parse.ml

@@ -86,7 +86,8 @@ let validate_ast ast =
           | SingleExpr expr -> validate_expr expr
           | Block stmts -> List.for_all validate_stmt stmts
         ) arms
-    | New _ -> true  (* New expressions are always syntactically valid *)
+    | New _ -> true
+    | NewWithFlag (_, flag_expr) -> validate_expr flag_expr  (* New expressions are always syntactically valid *)
 
   and validate_stmt stmt =
     match stmt.stmt_desc with

src/parser.mly

@@ -433,6 +433,7 @@ primary_expression:
   | primary_expression DOT IDENTIFIER { make_expr (FieldAccess ($1, $3)) (make_pos ()) }
   | primary_expression ARROW IDENTIFIER { make_expr (ArrowAccess ($1, $3)) (make_pos ()) }
   | NEW bpf_type LPAREN RPAREN { make_expr (New $2) (make_pos ()) }
+  | NEW bpf_type LPAREN expression RPAREN { make_expr (NewWithFlag ($2, $4)) (make_pos ()) }
 
 function_call:
   | IDENTIFIER LPAREN argument_list RPAREN