Core: Lower if-branches in the function's return type

handle_if processed its branches without passing ret_type down, so a
return inside an if fell back to process_stmt's i64 default: a -> c_int32
function got `ret i64 7`, which llc rejects. The branches now go through
process_block, which threads ret_type and stops at a statement that ends
the block instead of emitting after its terminator.

Author: varun-r-mallya

pythonbpf/functions/functions_pass.py

@@ -318,7 +318,7 @@ def handle_cond(func, compilation_context, builder, cond, local_sym_tab):
     return convert_to_bool(builder, val)
 
 
-def handle_if(func, compilation_context, builder, stmt, local_sym_tab):
+def handle_if(func, compilation_context, builder, stmt, local_sym_tab, ret_type):
     """Handle if statements in the function body."""
     logger.info("Handling if statement")
     # start = builder.block.parent
@@ -336,22 +336,17 @@ def handle_if(func, compilation_context, builder, stmt, local_sym_tab):
         builder.cbranch(cond, then_block, merge_block)
 
     builder.position_at_end(then_block)
-    for s in stmt.body:
-        process_stmt(func, compilation_context, builder, s, local_sym_tab, False)
+    process_block(
+        func, compilation_context, builder, stmt.body, local_sym_tab, ret_type
+    )
     if not builder.block.is_terminated:
         builder.branch(merge_block)
 
     if else_block:
         builder.position_at_end(else_block)
-        for s in stmt.orelse:
-            process_stmt(
-                func,
-                compilation_context,
-                builder,
-                s,
-                local_sym_tab,
-                False,
-            )
+        process_block(
+            func, compilation_context, builder, stmt.orelse, local_sym_tab, ret_type
+        )
         if not builder.block.is_terminated:
             builder.branch(merge_block)
 
@@ -413,7 +408,7 @@ def process_stmt(
     elif isinstance(stmt, ast.Global):
         logger.debug(f"global declaration of {', '.join(stmt.names)} already bound")
     elif isinstance(stmt, ast.If):
-        handle_if(func, compilation_context, builder, stmt, local_sym_tab)
+        handle_if(func, compilation_context, builder, stmt, local_sym_tab, ret_type)
     elif isinstance(stmt, ast.Return):
         did_return = handle_return(
             func, builder, stmt, local_sym_tab, ret_type, compilation_context
@@ -428,6 +423,19 @@ def process_stmt(
     return did_return
 
 
+def process_block(func, compilation_context, builder, stmts, local_sym_tab, ret_type):
+    """Process a nested statement list, such as an if-branch, in the
+    enclosing function's return type. Stops at the first statement that ends
+    the block (a return), because whatever follows it in the same list can
+    never run, and would otherwise be emitted after a terminator."""
+    for s in stmts:
+        if builder.block.is_terminated:
+            break
+        process_stmt(
+            func, compilation_context, builder, s, local_sym_tab, False, ret_type
+        )
+
+
 # ============================================================================
 # SECTION 3: Function Body Processing
 # ============================================================================

tests/passing_tests/return/if_narrow.py

@@ -0,0 +1,22 @@
+# A return inside an if-branch is lowered in the function's declared return
+# type (i32 here), not a default i64 that llc rejects.
+from pythonbpf import bpf, section, bpfglobal, compile
+from ctypes import c_void_p, c_int32
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def hello(ctx: c_void_p) -> c_int32:
+    x = 1
+    if x == 1:
+        return 7
+    return 0
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()