From 6f6f101a867a379531b90bba02b0f884f81f3543 Mon Sep 17 00:00:00 2001
From: varun-r-mallya <varunrmallya@gmail.com>
Date: Thu, 11 Sep 2025 02:29:16 +0530
Subject: [PATCH] Add support for basic arithmetic operations

---
 examples/execve3.py               | 17 +++++----
 pythonbpf/functions_pass.py       | 17 ++++++++-
 pythonbpf/unary_and_binary_ops.py | 58 +++++++++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 8 deletions(-)
 create mode 100644 pythonbpf/unary_and_binary_ops.py

diff --git a/examples/execve3.py b/examples/execve3.py
index faf524a..d7998a5 100644
--- a/examples/execve3.py
+++ b/examples/execve3.py
@@ -39,20 +39,23 @@ def hello_again(ctx: c_void_p) -> c_int64:
             print("we did not prevail")
     ts = ktime()
     last().update(key, ts)
-    keys = 2
-    last().update(keys, ts)
-    key = 4
-    last().update(key, ts)
-    key = 5
+
     st = "st"
     last().update(key, ts)
-    return c_int64(0)
 
+    keena = 2 + 1
+    # below breaks
+    # keela = keena + 1
+    keema = 8 * 9
+    keesa = 10 - 11
+    keeda = 10 / 5
+    # below does not spit IR
+    keeda += 1
+    return c_int64(0)
 
 @bpf
 @bpfglobal
 def LICENSE() -> str:
     return "GPL"
 
-
 compile()
diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py
index 029ebab..1122af1 100644
--- a/pythonbpf/functions_pass.py
+++ b/pythonbpf/functions_pass.py
@@ -3,6 +3,7 @@ import ast
 
 from .bpf_helper_handler import helper_func_list, handle_helper_call
 from .type_deducer import ctypes_to_ir
+from .unary_and_binary_ops import handle_binary_op, handle_unary_op
 
 
 def get_probe_string(func_node):
@@ -21,6 +22,16 @@ def get_probe_string(func_node):
                     return arg.value
     return "helper"
 
+def handle_unary_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab):
+    """Handle unary assignment statements in the function body."""
+    SyntaxError("Unary assignment not supported")
+    target = stmt.target
+    if not isinstance(target, ast.Name):
+        SyntaxError("Unsupported assignment target")
+        return
+    else:
+        handle_unary_op(func, module, builder, stmt, map_sym_tab, local_sym_tab)
+        return
 
 def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab):
     """Handle assignment statements in the function body."""
@@ -95,6 +106,8 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab):
                 print("Unsupported assignment call structure")
         else:
             print("Unsupported assignment call function type")
+    elif isinstance(rval, ast.BinOp):
+        handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab)
     else:
         print("Unsupported assignment value type")
 
@@ -237,11 +250,13 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
 
 
 def process_stmt(func, module, builder, stmt, local_sym_tab, map_sym_tab, did_return, ret_type=ir.IntType(64)):
-    # print(f"Processing statement: {ast.dump(stmt)}")
+    print(f"Processing statement: {ast.dump(stmt)}")
     if isinstance(stmt, ast.Expr):
         handle_expr(func, module, builder, stmt, local_sym_tab, map_sym_tab)
     elif isinstance(stmt, ast.Assign):
         handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab)
+    elif isinstance(stmt, ast.AugAssign):
+        handle_unary_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab)
     elif isinstance(stmt, ast.If):
         handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab)
     elif isinstance(stmt, ast.Return):
diff --git a/pythonbpf/unary_and_binary_ops.py b/pythonbpf/unary_and_binary_ops.py
new file mode 100644
index 0000000..32ed0a7
--- /dev/null
+++ b/pythonbpf/unary_and_binary_ops.py
@@ -0,0 +1,58 @@
+import ast
+from llvmlite import ir
+
+def handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab):
+    left = rval.left
+    right = rval.right
+    op = rval.op
+
+    if isinstance(left, ast.Name):
+        left = local_sym_tab[left.id]
+    elif isinstance(left, ast.Constant):
+        left = ir.Constant(ir.IntType(64), left.value)
+    else:
+        print("Unsupported left operand type")
+
+    if isinstance(right, ast.Name):
+        right = local_sym_tab[right.id]
+    elif isinstance(right, ast.Constant):
+        right = ir.Constant(ir.IntType(64), right.value)
+    else:
+        SyntaxError("Unsupported right operand type")
+
+    if isinstance(op, ast.Add):
+        result = builder.add(left, right)
+    elif isinstance(op, ast.Sub):
+        result = builder.sub(left, right)
+    elif isinstance(op, ast.Mult):
+        result = builder.mul(left, right)
+    elif isinstance(op, ast.Div):
+        result = builder.sdiv(left, right)
+    else:
+        result = "fuck type errors"
+        SyntaxError("Unsupported binary operation")
+
+    return result
+
+def handle_unary_op(rval, module, builder, func, local_sym_tab, map_sym_tab):
+    print("UNARY ASSIGNMENT DOES NOT WORK")
+    return
+    # TODO: heavy fixxing needed
+    operand = rval.gay
+    op = rval.op
+    if isinstance(operand, ast.Name):
+        operand = local_sym_tab[operand.id]
+    elif isinstance(operand, ast.Constant):
+        operand = ir.Constant(ir.IntType(64), operand.value)
+    else:
+        SyntaxError("Unsupported operand type")
+
+    if isinstance(op, ast.UAdd):
+        result = builder.add(ir.Constant(ir.IntType(64), 0), operand)
+    elif isinstance(op, ast.USub):
+        result = builder.sub(ir.Constant(ir.IntType(64), 0), operand)
+    else:
+        result = "all my homies hate type errors"
+        SyntaxError("Unsupported unary operation")
+
+    return result