Add docstrings to core modules and helper functions

Co-authored-by: varun-r-mallya <100590632+varun-r-mallya@users.noreply.github.com>

Author: Copilot

pythonbpf/binary_ops.py

@@ -35,6 +35,17 @@ def get_operand_value(operand, builder, local_sym_tab):
 
 
 def handle_binary_op_impl(rval, builder, local_sym_tab):
+    """
+    Handle binary operations and emit corresponding LLVM IR instructions.
+    
+    Args:
+        rval: The AST BinOp node representing the binary operation
+        builder: LLVM IR builder for emitting instructions
+        local_sym_tab: Symbol table mapping variable names to their IR representations
+    
+    Returns:
+        The LLVM IR value representing the result of the binary operation
+    """
     op = rval.op
     left = get_operand_value(rval.left, builder, local_sym_tab)
     right = get_operand_value(rval.right, builder, local_sym_tab)
@@ -63,6 +74,18 @@ def handle_binary_op_impl(rval, builder, local_sym_tab):
 
 
 def handle_binary_op(rval, builder, var_name, local_sym_tab):
+    """
+    Handle binary operations and optionally store the result to a variable.
+    
+    Args:
+        rval: The AST BinOp node representing the binary operation
+        builder: LLVM IR builder for emitting instructions
+        var_name: Optional variable name to store the result
+        local_sym_tab: Symbol table mapping variable names to their IR representations
+    
+    Returns:
+        A tuple of (result_value, result_type)
+    """
     result = handle_binary_op_impl(rval, builder, local_sym_tab)
     if var_name and var_name in local_sym_tab:
         logger.info(

pythonbpf/codegen.py

@@ -37,6 +37,14 @@ def find_bpf_chunks(tree):
 
 
 def processor(source_code, filename, module):
+    """
+    Process Python source code and convert BPF-decorated functions to LLVM IR.
+    
+    Args:
+        source_code: The Python source code to process
+        filename: The name of the source file
+        module: The LLVM IR module to populate
+    """
     tree = ast.parse(source_code, filename)
     logger.debug(ast.dump(tree, indent=4))
 
@@ -56,6 +64,17 @@ def processor(source_code, filename, module):
 
 
 def compile_to_ir(filename: str, output: str, loglevel=logging.INFO):
+    """
+    Compile a Python BPF program to LLVM IR.
+    
+    Args:
+        filename: Path to the Python source file containing BPF programs
+        output: Path where the LLVM IR (.ll) file will be written
+        loglevel: Logging level for compilation messages
+    
+    Returns:
+        Path to the generated LLVM IR file
+    """
     logging.basicConfig(
         level=loglevel, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
     )
@@ -129,6 +148,18 @@ def compile_to_ir(filename: str, output: str, loglevel=logging.INFO):
 
 
 def compile(loglevel=logging.INFO) -> bool:
+    """
+    Compile the calling Python BPF program to an object file.
+    
+    This function should be called from a Python file containing BPF programs.
+    It will compile the calling file to LLVM IR and then to a BPF object file.
+    
+    Args:
+        loglevel: Logging level for compilation messages
+    
+    Returns:
+        True if compilation succeeded, False otherwise
+    """
     # Look one level up the stack to the caller of this function
     caller_frame = inspect.stack()[1]
     caller_file = Path(caller_frame.filename).resolve()
@@ -162,6 +193,18 @@ def compile(loglevel=logging.INFO) -> bool:
 
 
 def BPF(loglevel=logging.INFO) -> BpfProgram:
+    """
+    Compile the calling Python BPF program and return a BpfProgram object.
+    
+    This function compiles the calling file's BPF programs to an object file
+    and loads it into a BpfProgram object for immediate use.
+    
+    Args:
+        loglevel: Logging level for compilation messages
+    
+    Returns:
+        A BpfProgram object that can be used to load and attach BPF programs
+    """
     caller_frame = inspect.stack()[1]
     src = inspect.getsource(caller_frame.frame)
     with tempfile.NamedTemporaryFile(

pythonbpf/decorators.py

@@ -23,6 +23,15 @@ def struct(cls):
 
 
 def section(name: str):
+    """
+    Decorator to specify the ELF section name for a BPF program.
+    
+    Args:
+        name: The section name (e.g., 'xdp', 'tracepoint/syscalls/sys_enter_execve')
+    
+    Returns:
+        A decorator function that marks the function with the section name
+    """
     def wrapper(fn):
         fn._section = name
         return fn

pythonbpf/expr/expr_pass.py

@@ -332,6 +332,21 @@ def eval_expr(
     map_sym_tab,
     structs_sym_tab=None,
 ):
+    """
+    Evaluate an expression and return its LLVM IR value and type.
+    
+    Args:
+        func: The LLVM IR function being built
+        module: The LLVM IR module
+        builder: LLVM IR builder
+        expr: The AST expression node to evaluate
+        local_sym_tab: Local symbol table
+        map_sym_tab: Map symbol table
+        structs_sym_tab: Struct symbol table
+    
+    Returns:
+        A tuple of (value, type) or None if evaluation fails
+    """
     logger.info(f"Evaluating expression: {ast.dump(expr)}")
     if isinstance(expr, ast.Name):
         return _handle_name_expr(expr, local_sym_tab, builder)

pythonbpf/expr/type_normalization.py

@@ -17,7 +17,15 @@
 
 
 def _get_base_type_and_depth(ir_type):
-    """Get the base type for pointer types."""
+    """
+    Get the base type and pointer depth for an LLVM IR type.
+    
+    Args:
+        ir_type: The LLVM IR type to analyze
+    
+    Returns:
+        A tuple of (base_type, depth) where depth is the number of pointer levels
+    """
     cur_type = ir_type
     depth = 0
     while isinstance(cur_type, ir.PointerType):
@@ -27,7 +35,18 @@ def _get_base_type_and_depth(ir_type):
 
 
 def _deref_to_depth(func, builder, val, target_depth):
-    """Dereference a pointer to a certain depth."""
+    """
+    Dereference a pointer to a certain depth with null checks.
+    
+    Args:
+        func: The LLVM IR function being built
+        builder: LLVM IR builder
+        val: The pointer value to dereference
+        target_depth: Number of levels to dereference
+    
+    Returns:
+        The dereferenced value, or None if dereferencing fails
+    """
 
     cur_val = val
     cur_type = val.type
@@ -73,7 +92,18 @@ def _deref_to_depth(func, builder, val, target_depth):
 
 
 def _normalize_types(func, builder, lhs, rhs):
-    """Normalize types for comparison."""
+    """
+    Normalize types for comparison by casting or dereferencing as needed.
+    
+    Args:
+        func: The LLVM IR function being built
+        builder: LLVM IR builder
+        lhs: Left-hand side value
+        rhs: Right-hand side value
+    
+    Returns:
+        A tuple of (normalized_lhs, normalized_rhs) or (None, None) on error
+    """
 
     logger.info(f"Normalizing types: {lhs.type} vs {rhs.type}")
     if isinstance(lhs.type, ir.IntType) and isinstance(rhs.type, ir.IntType):
@@ -99,7 +129,16 @@ def _normalize_types(func, builder, lhs, rhs):
 
 
 def convert_to_bool(builder, val):
-    """Convert a value to boolean."""
+    """
+    Convert an LLVM IR value to a boolean (i1) type.
+    
+    Args:
+        builder: LLVM IR builder
+        val: The value to convert
+    
+    Returns:
+        An i1 boolean value
+    """
     if val.type == ir.IntType(1):
         return val
     if isinstance(val.type, ir.PointerType):
@@ -110,7 +149,19 @@ def convert_to_bool(builder, val):
 
 
 def handle_comparator(func, builder, op, lhs, rhs):
-    """Handle comparison operations."""
+    """
+    Handle comparison operations between two values.
+    
+    Args:
+        func: The LLVM IR function being built
+        builder: LLVM IR builder
+        op: The AST comparison operator node
+        lhs: Left-hand side value
+        rhs: Right-hand side value
+    
+    Returns:
+        A tuple of (result, ir.IntType(1)) or None on error
+    """
 
     if lhs.type != rhs.type:
         lhs, rhs = _normalize_types(func, builder, lhs, rhs)

pythonbpf/functions/functions_pass.py

@@ -243,6 +243,21 @@ def handle_assign(
 def handle_cond(
     func, module, builder, cond, local_sym_tab, map_sym_tab, structs_sym_tab=None
 ):
+    """
+    Evaluate a condition expression and convert it to a boolean value.
+    
+    Args:
+        func: The LLVM IR function being built
+        module: The LLVM IR module
+        builder: LLVM IR builder
+        cond: The AST condition node to evaluate
+        local_sym_tab: Local symbol table
+        map_sym_tab: Map symbol table
+        structs_sym_tab: Struct symbol table
+    
+    Returns:
+        LLVM IR boolean value representing the condition result
+    """
     val = eval_expr(
         func, module, builder, cond, local_sym_tab, map_sym_tab, structs_sym_tab
     )[0]
@@ -298,6 +313,18 @@ def handle_if(
 
 
 def handle_return(builder, stmt, local_sym_tab, ret_type):
+    """
+    Handle return statements in BPF functions.
+    
+    Args:
+        builder: LLVM IR builder
+        stmt: The AST Return node
+        local_sym_tab: Local symbol table
+        ret_type: Expected return type
+    
+    Returns:
+        True if a return was emitted, False otherwise
+    """
     logger.info(f"Handling return statement: {ast.dump(stmt)}")
     if stmt.value is None:
         return _handle_none_return(builder)
@@ -329,6 +356,23 @@ def process_stmt(
     did_return,
     ret_type=ir.IntType(64),
 ):
+    """
+    Process a single statement in a BPF function.
+    
+    Args:
+        func: The LLVM IR function being built
+        module: The LLVM IR module
+        builder: LLVM IR builder
+        stmt: The AST statement node to process
+        local_sym_tab: Local symbol table
+        map_sym_tab: Map symbol table
+        structs_sym_tab: Struct symbol table
+        did_return: Whether a return has been emitted
+        ret_type: Expected return type
+    
+    Returns:
+        True if a return was emitted, False otherwise
+    """
     logger.info(f"Processing statement: {ast.dump(stmt)}")
     if isinstance(stmt, ast.Expr):
         handle_expr(
@@ -363,6 +407,25 @@ def process_stmt(
 def allocate_mem(
     module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
 ):
+    """
+    Pre-allocate stack memory for local variables in a BPF function.
+    
+    This function scans the function body and creates alloca instructions
+    for all local variables before processing the function statements.
+    
+    Args:
+        module: The LLVM IR module
+        builder: LLVM IR builder
+        body: List of AST statements in the function body
+        func: The LLVM IR function being built
+        ret_type: Expected return type
+        map_sym_tab: Map symbol table
+        local_sym_tab: Local symbol table to populate
+        structs_sym_tab: Struct symbol table
+    
+    Returns:
+        Updated local symbol table
+    """
     for stmt in body:
         has_metadata = False
         if isinstance(stmt, ast.If):
@@ -556,6 +619,16 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_t
 
 
 def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab):
+    """
+    Process all BPF function chunks and generate LLVM IR.
+    
+    Args:
+        tree: The Python AST (not used in current implementation)
+        module: The LLVM IR module to add functions to
+        chunks: List of AST function nodes decorated with @bpf
+        map_sym_tab: Map symbol table
+        structs_sym_tab: Struct symbol table
+    """
     for func_node in chunks:
         is_global = False
         for decorator in func_node.decorator_list:
@@ -581,6 +654,18 @@ def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab):
 
 
 def infer_return_type(func_node: ast.FunctionDef):
+    """
+    Infer the return type of a BPF function from annotations or return statements.
+    
+    Args:
+        func_node: The AST function node
+    
+    Returns:
+        String representation of the return type (e.g., 'c_int64')
+    
+    Raises:
+        TypeError: If func_node is not a FunctionDef
+    """
     if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
         raise TypeError("Expected ast.FunctionDef")
     if func_node.returns is not None: