Implement include system in KernelScript.

This update introduces a new `include` keyword for including header files (.kh) alongside the existing `import` functionality for modules.

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

Author: congwang-mk

SPEC.md

@@ -89,9 +89,9 @@ fn main(args: Args) -> i32 {
 }
 ```
 
-### 1.4 Unified Import System
+### 1.4 Unified Import and Include System
 
-KernelScript supports importing both KernelScript modules and external language modules using a unified syntax. Import behavior is automatically determined by file extension:
+KernelScript supports both importing modules and including headers using distinct keywords for different use cases:
 
 ```kernelscript
 // Import KernelScript modules (.ks files)
@@ -132,6 +132,30 @@ fn main() -> i32 {
 }
 ```
 
+#### Include System for Headers (.kh files)
+
+```kernelscript
+// Include KernelScript headers (.kh files) - declarations only, flattened into global namespace
+include "generated/common_kfuncs.kh"           // extern kfunc declarations
+include "generated/xdp_kfuncs.kh"              // XDP-specific kfuncs  
+include "types/networking.kh"                  // Type definitions
+
+@xdp
+fn packet_processor(ctx: *xdp_md) -> xdp_action {
+    // Direct access to included extern kfuncs (no namespace)
+    var timestamp = bpf_ktime_get_ns()          // From common_kfuncs.kh
+    bpf_xdp_adjust_head(ctx, -14)               // From xdp_kfuncs.kh
+    
+    return XDP_PASS
+}
+```
+
+**Key Distinctions:**
+- **`import name from "file"`**: Creates namespace, works with full implementations (.ks/.py files)
+- **`include "file"`**: Flattens into global namespace, works with headers only (.kh files)
+- **Use cases**: Import for libraries/modules, include for extern declarations and types
+- **Validation**: Include validates that .kh files contain only declarations (no function bodies)
+
 ## 2. Lexical Structure
 
 ### 2.1 Keywords
@@ -141,7 +165,7 @@ pin         type        struct      enum        if          else
 while       loop        break       continue    return      import
 pub         priv        impl        true        false       null
 try         catch       throw       defer       delete      match
-extern
+extern      include
 ```
 
 **Note**: The `pin` keyword is used for both maps and global variables to enable filesystem persistence.
@@ -4824,16 +4848,23 @@ import_declaration = "import" identifier "from" string_literal
 (* External kernel function declarations - for importing existing kernel kfuncs *)
 extern_declaration = "extern" identifier "(" parameter_list ")" [ "->" type_annotation ]
 
+(* Include declarations - for KernelScript headers (.kh files) *)
+include_declaration = "include" string_literal
+
 (* Examples:
    import utils from "./common/utils.ks"          // KernelScript import
    import ml_analysis from "./ml/threat.py"       // Python import (userspace only)
    
    extern bpf_ktime_get_ns() -> u64               // Import existing kernel kfunc
    extern bpf_trace_printk(fmt: *u8, fmt_size: u32) -> i32  // Import with parameters
    
+   include "common_kfuncs.kh"                    // Include header with extern declarations
+   include "types/networking.kh"                 // Include header with type definitions
+   
    Import behavior is determined by file extension:
    - .ks files: Import KernelScript symbols (functions, types, maps, configs)  
    - .py files: Import Python functions with automatic FFI bridging (userspace only)
+   - .kh files: Include headers with declarations only (flattened into global namespace)
 *) 
 
 (* Identifiers and basic tokens *)

examples/common_kfuncs.kh

@@ -0,0 +1,8 @@
+// Common kernel function declarations
+extern bpf_ktime_get_ns() -> u64
+extern bpf_trace_printk(fmt: *u8, fmt_size: u32) -> i32
+extern bpf_get_current_pid_tgid() -> u64
+
+// Common type definitions
+type Timestamp = u64
+type ProcessId = u32

examples/include_demo.ks

@@ -0,0 +1,27 @@
+// Example demonstrating include functionality
+// This shows how to include KernelScript headers (.kh files)
+
+// Include declarations from header files
+include "common_kfuncs.kh"
+include "xdp_kfuncs.kh"
+
+// XDP program that uses included kfunc declarations
+@xdp
+fn packet_processor(ctx: *xdp_md) -> xdp_action {
+    // These functions are available from the included headers
+    var timestamp = bpf_ktime_get_ns()
+    var pid_tgid = bpf_get_current_pid_tgid()
+    var result = bpf_xdp_adjust_head(ctx, -14)
+    
+    // Use the timestamp and pid to suppress unused variable warnings
+    var action: XdpAction = 2  // XDP_PASS
+    if (timestamp > 0 && pid_tgid > 0 && result >= 0) {
+        return action
+    }
+    
+    return 2  // XDP_PASS
+}
+
+fn main() -> i32 {
+    return 0
+}

examples/xdp_kfuncs.kh

@@ -0,0 +1,6 @@
+// XDP-specific kernel function declarations
+extern bpf_xdp_adjust_head(ctx: *xdp_md, delta: i32) -> i32
+extern bpf_xdp_adjust_tail(ctx: *xdp_md, delta: i32) -> i32
+
+// XDP-specific types
+type XdpAction = u32

src/ast.ml

@@ -351,6 +351,12 @@ type extern_kfunc_declaration = {
   extern_pos: position;
 }
 
+(** Include declaration - for KernelScript headers (.kh files) *)
+type include_declaration = {
+  include_path: string;        (* Path to .kh file *)
+  include_pos: position;
+}
+
 (** Top-level declarations *)
 type declaration =
   | AttributedFunction of attributed_function
@@ -363,6 +369,7 @@ type declaration =
   | ImplBlock of impl_block
   | ImportDecl of import_declaration
   | ExternKfuncDecl of extern_kfunc_declaration
+  | IncludeDecl of include_declaration
 
 (** Complete AST *)
 type ast = declaration list
@@ -456,6 +463,11 @@ let make_extern_kfunc_declaration name params return_type pos = {
   extern_pos = pos;
 }
 
+let make_include_declaration path pos = {
+  include_path = path;
+  include_pos = pos;
+}
+
 let make_type_def def = def
 
 let make_enum_def name values = EnumDef (name, values)
@@ -904,6 +916,8 @@ let string_of_declaration = function
         | None -> ""
       in
       Printf.sprintf "extern %s(%s)%s;" extern_decl.extern_name params_str return_str
+  | IncludeDecl include_decl ->
+      Printf.sprintf "include \"%s\"" include_decl.include_path
 
 let string_of_ast ast =
   String.concat "\n\n" (List.map string_of_declaration ast)

src/dune

@@ -6,7 +6,7 @@
           userspace_codegen evaluator safety_checker stdlib test_codegen
 tail_call_analyzer kernel_module_codegen dynptr_bridge
           btf_parser btf_binary_parser kernel_types struct_ops_registry
-          import_resolver python_bridge kernelscript_bridge)
+          import_resolver include_resolver python_bridge kernelscript_bridge)
  (libraries unix str kernelscript_context)
  (foreign_stubs
   (language c)

src/include_resolver.ml

@@ -0,0 +1,158 @@
+(*
+ * Copyright 2025 Multikernel Technologies, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *)
+
+(** Include Resolution for KernelScript Headers (.kh files)
+    
+    This module handles including KernelScript header files (.kh files)
+    that contain only declarations (extern, type, struct, enum, config).
+    It validates that included files contain no function implementations.
+*)
+
+open Ast
+
+(** Include validation errors *)
+type include_validation_error = 
+  | FunctionBodyFound of string (* function name with body *)
+  | InvalidExtension of string (* non-.ksh extension *)
+  | InvalidDeclaration of string (* unsupported declaration type *)
+
+(** Include resolution errors *)
+exception Include_error of string * position
+exception Include_validation_error of include_validation_error * string * position
+
+let include_error msg pos = raise (Include_error (msg, pos))
+
+(** Validate that a declaration is allowed in header files *)
+let validate_header_declaration decl =
+  match decl with
+  | TypeDef _ -> true
+  | StructDecl _ -> true  
+  | ConfigDecl _ -> true
+  | ExternKfuncDecl _ -> true
+  | IncludeDecl _ -> true  (* Allow nested includes *)
+  | GlobalVarDecl _ -> true  (* Allow global variable declarations *)
+  | AttributedFunction attr_func ->
+      (* Check if this is just a declaration (no body) *)
+      (match attr_func.attr_function.func_body with
+       | [] -> true  (* Empty body = declaration only *)
+       | _ -> false) (* Has body = implementation *)
+  | GlobalFunction func ->
+      (* Check if this is just a declaration (no body) *)
+      (match func.func_body with
+       | [] -> true  (* Empty body = declaration only *)
+       | _ -> false) (* Has body = implementation *)
+  | MapDecl _ -> true  (* Allow map declarations *)
+  | ImplBlock _ -> false  (* Impl blocks not allowed in headers *)
+  | ImportDecl _ -> true  (* Allow imports in headers *)
+
+(** Validate that included file contains only valid header declarations *)
+let validate_header_file file_path ast =
+  let validate_decl decl =
+    if not (validate_header_declaration decl) then
+      let error_msg = match decl with
+      | AttributedFunction attr_func ->
+          FunctionBodyFound attr_func.attr_function.func_name
+      | GlobalFunction func ->
+          FunctionBodyFound func.func_name
+      | ImplBlock impl_block ->
+          InvalidDeclaration ("impl block '" ^ impl_block.impl_name ^ "' not allowed in headers")
+      | _ ->
+          InvalidDeclaration "unknown invalid declaration type"
+      in
+      let pos = { line = 0; column = 0; filename = file_path } in
+      raise (Include_validation_error (error_msg, file_path, pos))
+  in
+  List.iter validate_decl ast
+
+(** Validate file extension is .kh *)
+let validate_file_extension file_path =
+  if not (Filename.check_suffix file_path ".kh") then
+    let pos = { line = 0; column = 0; filename = file_path } in
+    raise (Include_validation_error (InvalidExtension file_path, file_path, pos))
+
+(** Resolve a single include declaration *)
+let resolve_include include_decl base_path =
+  (* Resolve relative paths *)
+  let file_path = if Filename.is_relative include_decl.include_path then
+    Filename.concat (Filename.dirname base_path) include_decl.include_path
+  else
+    include_decl.include_path
+  in
+  
+  (* Validate file extension *)
+  validate_file_extension file_path;
+  
+  (* Check if file exists *)
+  if not (Sys.file_exists file_path) then
+    include_error ("Include file not found: " ^ file_path) include_decl.include_pos;
+  
+  try
+    (* Parse the included file *)
+    let ic = open_in file_path in
+    let content = really_input_string ic (in_channel_length ic) in
+    close_in ic;
+    
+    let lexbuf = Lexing.from_string content in
+    let ast = Parser.program Lexer.token lexbuf in
+    
+    (* Validate that it's a proper header file *)
+    validate_header_file file_path ast;
+    
+    (* Return the parsed declarations *)
+    ast
+  with
+  | Include_validation_error (err, file_path, pos) ->
+      let error_msg = match err with
+        | FunctionBodyFound func_name ->
+            Printf.sprintf "Header file '%s' contains function implementation '%s'. Header files (.kh) should only contain declarations. Move implementations to .ks files." file_path func_name
+        | InvalidExtension file_path ->
+            Printf.sprintf "Include directive can only include .kh header files, but found: %s" file_path
+        | InvalidDeclaration desc ->
+            Printf.sprintf "Header file '%s' contains invalid declaration: %s" file_path desc
+      in
+      include_error error_msg pos
+  | Sys_error msg -> 
+      let pos = { line = 0; column = 0; filename = file_path } in
+      include_error ("Cannot read header file: " ^ msg) pos
+  | Parsing.Parse_error ->
+      let pos = { line = 0; column = 0; filename = file_path } in
+      include_error ("Parse error in header file: " ^ file_path) pos
+
+(** Process all includes in an AST and return expanded AST with included declarations *)
+let process_includes ast base_path =
+  let rec process_decls decls =
+    List.fold_left (fun acc decl ->
+      match decl with
+      | IncludeDecl include_decl ->
+          (* Resolve the include and get its declarations *)
+          let included_ast = resolve_include include_decl base_path in
+          (* Recursively process includes in the included file *)
+          let processed_included = process_decls included_ast in
+          (* Add the included declarations to our AST (flatten) *)
+          acc @ processed_included
+      | _ ->
+          (* Keep non-include declarations as-is *)
+          acc @ [decl]
+    ) [] decls
+  in
+  process_decls ast
+
+(** Get all include declarations from an AST *)
+let get_includes ast =
+  List.filter_map (function
+    | IncludeDecl include_decl -> Some include_decl
+    | _ -> None
+  ) ast

src/lexer.mll

@@ -73,6 +73,7 @@
   let lookup_keyword = function
     | "fn" -> FN
     | "extern" -> EXTERN
+    | "include" -> INCLUDE
     | "pin" -> PIN
     | "type" -> TYPE
     | "struct" -> STRUCT

src/main.ml

@@ -571,6 +571,16 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
     ) resolved_imports;
     Printf.printf "\n";
 
+    (* Phase 1.6: Include Processing *)
+    Printf.printf "Phase 1.6: Include Processing\n";
+    let includes = Include_resolver.get_includes ast in
+    let expanded_ast = Include_resolver.process_includes ast input_file in
+    Printf.printf "✅ Processed %d includes, expanded to %d declarations\n" (List.length includes) (List.length expanded_ast);
+    List.iter (fun include_decl -> 
+      Printf.printf "   📄 Header: %s\n" include_decl.Ast.include_path
+    ) includes;
+    Printf.printf "\n";
+    
     (* Determine output directory early *)
     let actual_output_dir = match output_dir with
       | Some dir -> dir
@@ -606,16 +616,17 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
     in
     copy_python_files resolved_imports actual_output_dir;
 
-    (* Store original AST before any filtering *)
-    let original_ast = ast in
+    (* Store original AST before any filtering and use expanded AST with includes *)
+    let _original_ast = ast in
+    let ast_with_includes = expanded_ast in  (* Use expanded AST with included declarations *)
 
     (* Test mode: Filter AST for @test functions *)
       let filtered_ast = if test_mode then
-    Test_codegen.filter_ast_for_testing ast input_file
-    else original_ast in
+    Test_codegen.filter_ast_for_testing ast_with_includes input_file
+    else ast_with_includes in
 
-    (* For regular eBPF compilation, always use original AST *)
-    let compilation_ast = original_ast in
+    (* Use the filtered AST (which includes expanded includes) for compilation *)
+    let compilation_ast = filtered_ast in
 
     (* Extract base name for project name *)
     let base_name = Filename.remove_extension (Filename.basename input_file) in

src/parse.ml

@@ -152,6 +152,7 @@ let validate_ast ast =
         ) impl_block.impl_items
     | ImportDecl _ -> true (* Import declarations are always valid once parsed *)
     | ExternKfuncDecl _ -> true (* Extern kfunc declarations are always valid once parsed *)
+    | IncludeDecl _ -> true (* Include declarations are always valid once parsed *)
   in
 
   List.for_all validate_declaration ast