Add var assigning for helpers

Author: r41k0u

examples/execve2.py

@@ -25,7 +25,7 @@ def hello_again(ctx: c_void_p) -> c_int64:
     key = 0
     tsp = last().lookup(key)
     print(tsp)
-    ktime()
+    ts = ktime()
     return c_int64(0)
 
 

pythonbpf/bpf_helper_handler.py

@@ -107,7 +107,7 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
         func_name = call.func.id
         if func_name in helper_func_list:
             # it is not a map method call
-            helper_func_list[func_name](call, module, builder, func)
+            return helper_func_list[func_name](call, module, builder, func)
         else:
             raise NotImplementedError(
                 f"Function {func_name} is not implemented as a helper function.")
@@ -119,7 +119,7 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
             if map_sym_tab and map_name in map_sym_tab:
                 map_ptr = map_sym_tab[map_name]
                 if method_name in helper_func_list:
-                    helper_func_list[method_name](
+                    return helper_func_list[method_name](
                         call, map_ptr, module, builder, local_sym_tab)
                 else:
                     raise NotImplementedError(

pythonbpf/functions_pass.py

@@ -39,7 +39,6 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
     if isinstance(rval, ast.Constant):
         if isinstance(rval.value, int):
             # Assume c_int64 for now
-            # TODO: make symtab for this
             var = builder.alloca(ir.IntType(64), name=var_name)
             var.align = 8
             builder.store(ir.Constant(ir.IntType(64), rval.value), var)
@@ -57,6 +56,14 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
                 print(f"Assigned {call_type} constant "
                       f"{rval.args[0].value} to {var_name}")
                 local_sym_tab[var_name] = var
+            elif call_type in helper_func_list:
+                var = builder.alloca(ir.IntType(64), name=var_name)
+                var.align = 8
+                val = handle_helper_call(
+                    rval, module, builder, None, local_sym_tab, map_sym_tab)
+                builder.store(val, var)
+                local_sym_tab[var_name] = var
+                print(f"Assigned constant {rval.func.id} to {var_name}")
             else:
                 print(f"Unsupported assignment call type: {call_type}")
         elif isinstance(rval.func, ast.Attribute):
@@ -70,6 +77,8 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
                             rval, module, builder, None, local_sym_tab, map_sym_tab)
             else:
                 print("Unsupported assignment call structure")
+        else:
+            print("Unsupported assignment call function type")
 
 
 def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab):