Allow local symbols to be used within return

pythonbpf/binary_ops.py

@@ -63,4 +63,6 @@ def handle_binary_op_impl(rval, module, builder, local_sym_tab):
 
 def handle_binary_op(rval, module, builder, var_name, local_sym_tab):
     result = handle_binary_op_impl(rval, module, builder, local_sym_tab)
-    builder.store(result, local_sym_tab[var_name].var)
+    if var_name in local_sym_tab:
+        builder.store(result, local_sym_tab[var_name].var)
+    return result, result.type

pythonbpf/codegen.py

@@ -48,7 +48,7 @@ def processor(source_code, filename, module):
     globals_processing(tree, module)
 
 
-def compile_to_ir(filename: str, output: str, loglevel=logging.WARNING):
+def compile_to_ir(filename: str, output: str, loglevel=logging.INFO):
     logging.basicConfig(
         level=loglevel, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
     )
@@ -121,7 +121,7 @@ def compile_to_ir(filename: str, output: str, loglevel=logging.WARNING):
     return output
 
 
-def compile(loglevel=logging.WARNING) -> bool:
+def compile(loglevel=logging.INFO) -> bool:
     # Look one level up the stack to the caller of this function
     caller_frame = inspect.stack()[1]
     caller_file = Path(caller_frame.filename).resolve()
@@ -154,7 +154,7 @@ def compile(loglevel=logging.WARNING) -> bool:
     return success
 
 
-def BPF(loglevel=logging.WARNING) -> BpfProgram:
+def BPF(loglevel=logging.INFO) -> BpfProgram:
     caller_frame = inspect.stack()[1]
     src = inspect.getsource(caller_frame.frame)
     with tempfile.NamedTemporaryFile(

pythonbpf/functions_pass.py

@@ -385,24 +385,51 @@ def process_stmt(
         )
     elif isinstance(stmt, ast.Return):
         if stmt.value is None:
-            builder.ret(ir.Constant(ir.IntType(32), 0))
+            builder.ret(ir.Constant(ir.IntType(64), 0))
             did_return = True
         elif (
             isinstance(stmt.value, ast.Call)
             and isinstance(stmt.value.func, ast.Name)
             and len(stmt.value.args) == 1
-            and isinstance(stmt.value.args[0], ast.Constant)
-            and isinstance(stmt.value.args[0].value, int)
         ):
-            call_type = stmt.value.func.id
-            if ctypes_to_ir(call_type) != ret_type:
-                raise ValueError(
-                    "Return type mismatch: expected"
-                    f"{ctypes_to_ir(call_type)}, got {call_type}"
+            if isinstance(stmt.value.args[0], ast.Constant) and isinstance(
+                stmt.value.args[0].value, int
+            ):
+                call_type = stmt.value.func.id
+                if ctypes_to_ir(call_type) != ret_type:
+                    raise ValueError(
+                        "Return type mismatch: expected"
+                        f"{ctypes_to_ir(call_type)}, got {call_type}"
+                    )
+                else:
+                    builder.ret(ir.Constant(ret_type, stmt.value.args[0].value))
+                    did_return = True
+            elif isinstance(stmt.value.args[0], ast.BinOp):
+                # TODO: Should be routed through eval_expr
+                val = handle_binary_op(
+                    stmt.value.args[0], module, builder, None, local_sym_tab
                 )
-            else:
-                builder.ret(ir.Constant(ret_type, stmt.value.args[0].value))
+                if val is None:
+                    raise ValueError("Failed to evaluate return expression")
+                if val[1] != ret_type:
+                    raise ValueError(
+                        "Return type mismatch: expected" f"{ret_type}, got {val[1]}"
+                    )
+                builder.ret(val[0])
                 did_return = True
+            elif isinstance(stmt.value.args[0], ast.Name):
+                if stmt.value.args[0].id in local_sym_tab:
+                    var = local_sym_tab[stmt.value.args[0].id].var
+                    val = builder.load(var)
+                    if val.type != ret_type:
+                        raise ValueError(
+                            "Return type mismatch: expected"
+                            f"{ret_type}, got {val.type}"
+                        )
+                    builder.ret(val)
+                    did_return = True
+                else:
+                    raise ValueError("Failed to evaluate return expression")
         elif isinstance(stmt.value, ast.Name):
             if stmt.value.id == "XDP_PASS":
                 builder.ret(ir.Constant(ret_type, 2))
@@ -568,7 +595,7 @@ def process_func_body(
         )
 
     if not did_return:
-        builder.ret(ir.Constant(ir.IntType(32), 0))
+        builder.ret(ir.Constant(ir.IntType(64), 0))
 
 
 def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_tab):

tests/passing_tests/var_rval.py

@@ -0,0 +1,20 @@
+import logging
+
+from pythonbpf import compile, bpf, section, bpfglobal, compile_to_ir
+from ctypes import c_void_p, c_int64
+
+
+@bpf
+@section("sometag1")
+def sometag(ctx: c_void_p) -> c_int64:
+    a = 1 - 1
+    return c_int64(a)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+compile_to_ir("var_rval.py", "var_rval.ll")
+compile(loglevel=logging.INFO)