Add create_targets_and_rvals, use it in handle_assign to enable tuple assignment

Author: r41k0u

pythonbpf/allocation_pass.py

@@ -22,24 +22,27 @@ def __iter__(self):
         yield self.metadata
 
 
-def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
-    """Handle memory allocation for assignment statements."""
-
-    logger.info(f"Handling assignment for allocation: {ast.dump(stmt)}")
-
-    # NOTE: Support multi-target assignments (e.g.: a, b = 1, 2)
+def create_targets_and_rvals(stmt):
+    """Create lists of targets and right-hand values from an assignment statement."""
     if isinstance(stmt.targets[0], ast.Tuple):
         if not isinstance(stmt.value, ast.Tuple):
             logger.warning("Mismatched multi-target assignment, skipping allocation")
             return
-        targets = stmt.targets[0].elts
-        rvals = stmt.value.elts
+        targets, rvals = stmt.targets[0].elts, stmt.value.elts
         if len(targets) != len(rvals):
-            logger.warning("Mismatched multi-target assignment, skipping allocation")
+            logger.warning("length of LHS != length of RHS, skipping allocation")
             return
-    else:
-        targets = stmt.targets
-        rvals = [stmt.value]
+        return targets, rvals
+    return stmt.targets, [stmt.value]
+
+
+def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
+    """Handle memory allocation for assignment statements."""
+
+    logger.info(f"Handling assignment for allocation: {ast.dump(stmt)}")
+
+    # NOTE: Support multi-target assignments (e.g.: a, b = 1, 2)
+    targets, rvals = create_targets_and_rvals(stmt)
 
     for target, rval in zip(targets, rvals):
         # Skip non-name targets (e.g., struct field assignments)

pythonbpf/functions/functions_pass.py

@@ -12,7 +12,11 @@
     handle_variable_assignment,
     handle_struct_field_assignment,
 )
-from pythonbpf.allocation_pass import handle_assign_allocation, allocate_temp_pool
+from pythonbpf.allocation_pass import (
+    handle_assign_allocation,
+    allocate_temp_pool,
+    create_targets_and_rvals,
+)
 
 from .return_utils import handle_none_return, handle_xdp_return, is_xdp_name
 from .function_metadata import get_probe_string, is_global_function, infer_return_type
@@ -140,48 +144,43 @@ def handle_assign(
 ):
     """Handle assignment statements in the function body."""
 
-    # TODO: Support this later
-    # GH #37
-    if len(stmt.targets) != 1:
-        logger.error("Multi-target assignment is not supported for now")
-        return
-
-    target = stmt.targets[0]
-    rval = stmt.value
+    # NOTE: Support multi-target assignments (e.g.: a, b = 1, 2)
+    targets, rvals = create_targets_and_rvals(stmt)
 
-    if isinstance(target, ast.Name):
-        # NOTE: Simple variable assignment case: x = 5
-        var_name = target.id
-        result = handle_variable_assignment(
-            func,
-            module,
-            builder,
-            var_name,
-            rval,
-            local_sym_tab,
-            map_sym_tab,
-            structs_sym_tab,
-        )
-        if not result:
-            logger.error(f"Failed to handle assignment to {var_name}")
-        return
+    for target, rval in zip(targets, rvals):
+        if isinstance(target, ast.Name):
+            # NOTE: Simple variable assignment case: x = 5
+            var_name = target.id
+            result = handle_variable_assignment(
+                func,
+                module,
+                builder,
+                var_name,
+                rval,
+                local_sym_tab,
+                map_sym_tab,
+                structs_sym_tab,
+            )
+            if not result:
+                logger.error(f"Failed to handle assignment to {var_name}")
+            continue
 
-    if isinstance(target, ast.Attribute):
-        # NOTE: Struct field assignment case: pkt.field = value
-        handle_struct_field_assignment(
-            func,
-            module,
-            builder,
-            target,
-            rval,
-            local_sym_tab,
-            map_sym_tab,
-            structs_sym_tab,
-        )
-        return
+        if isinstance(target, ast.Attribute):
+            # NOTE: Struct field assignment case: pkt.field = value
+            handle_struct_field_assignment(
+                func,
+                module,
+                builder,
+                target,
+                rval,
+                local_sym_tab,
+                map_sym_tab,
+                structs_sym_tab,
+            )
+            continue
 
-    # Unsupported target type
-    logger.error(f"Unsupported assignment target: {ast.dump(target)}")
+        # Unsupported target type
+        logger.error(f"Unsupported assignment target: {ast.dump(target)}")
 
 
 def handle_cond(