From 6fea580693c3b0d1e7bf2e5e1d206d2db1a2871f Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyanshchaturvedi18@gmail.com>
Date: Fri, 3 Oct 2025 17:56:21 +0530
Subject: [PATCH] Fix t/f/return.py, tweak handle_binary_ops

---
 pythonbpf/binary_ops.py                       |  4 ++-
 pythonbpf/codegen.py                          |  6 ++--
 pythonbpf/functions_pass.py                   | 32 +++++++++++++------
 tests/failing_tests/var_rval.py               |  0
 .../return.py                                 |  0
 5 files changed, 29 insertions(+), 13 deletions(-)
 create mode 100644 tests/failing_tests/var_rval.py
 rename tests/{failing_tests => passing_tests}/return.py (100%)

diff --git a/pythonbpf/binary_ops.py b/pythonbpf/binary_ops.py
index 5ab393f..1b25c9c 100644
--- a/pythonbpf/binary_ops.py
+++ b/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
diff --git a/pythonbpf/codegen.py b/pythonbpf/codegen.py
index 5de23a5..cf20f06 100644
--- a/pythonbpf/codegen.py
+++ b/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(
diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py
index 430b55f..e343aa3 100644
--- a/pythonbpf/functions_pass.py
+++ b/pythonbpf/functions_pass.py
@@ -391,17 +391,31 @@ def process_stmt(
             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, ast.Name):
             if stmt.value.id == "XDP_PASS":
diff --git a/tests/failing_tests/var_rval.py b/tests/failing_tests/var_rval.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/failing_tests/return.py b/tests/passing_tests/return.py
similarity index 100%
rename from tests/failing_tests/return.py
rename to tests/passing_tests/return.py