Core: Fix loop edge cases the lowering got wrong

- `pass` is accepted as a statement, the natural body of a loop run for its
  side effects.
- Falling off the end of a function returns 0 in the function's own return
  type. After `while True` that block is unreachable, but a c_int32 function
  still emitted `ret i64 0`, which llc rejected.
- A range() step other than 1 could wrap the counter past INT64_MAX or
  UINT64_MAX back inside the range and loop forever. The step block now ends
  the loop when the distance left to stop, exact as an unsigned number, is
  no longer than the step. That is a normal exit, so the else-branch runs.
  Steps of 1 cannot overshoot and keep the plain loop, so the helper-bounded
  loop the verifier walks is unchanged.
- A step of 2**64 or more no longer truncates to 0: it is a compile error,
  since it does not fit the counter. Steps from 2**63 up are emitted by
  their bits.
- The counter's signedness folds each bound into a signed 64-bit type rather
  than joining the bounds alone, so only a c_uint64 bound makes it unsigned:
  range(-1, n) over a c_uint32 n runs 4 times, as in C with an __s64 counter.

Author: varun-r-mallya

pythonbpf/allocation_pass.py

@@ -102,16 +102,19 @@ def handle_for_allocation(compilation_context, builder, stmt, local_sym_tab):
     start, stop, step = parse_range(stmt)
 
     # range() yields Python ints; like any undeclared local they are 64-bit,
-    # signed unless the bounds make C's arithmetic unsigned.
+    # signed unless the bounds make C's arithmetic unsigned. That arithmetic
+    # is the counter's against each bound, so it starts from a signed 64-bit
+    # counter, and only an unsigned 64-bit bound turns it unsigned: every
+    # narrower unsigned bound fits in the signed counter, as in C.
     bound_types = [
         infer_int_type(bound, local_sym_tab, compilation_context)
         for bound in (start, stop)
         if bound is not None
     ]
     signed = True
     if all(ty is not None for ty in bound_types):
-        common = bound_types[0]
-        for ty in bound_types[1:]:
+        common = IntTy(64, True)
+        for ty in bound_types:
             common = usual_arithmetic_conversions(common, ty)
         signed = signedness(common)
     if not signed and step < 0:
@@ -191,6 +194,11 @@ def parse_range(stmt):
     step = -literal.value if negate else literal.value
     if step == 0:
         raise ValueError(f"range() arg 3 must not be zero (line {stmt.lineno})")
+    if abs(step) >= 1 << 64:
+        raise ValueError(
+            f"range() step on line {stmt.lineno} does not fit the 64-bit "
+            f"loop counter: {ast.unparse(step_node)}"
+        )
     return start, stop, step
 
 

pythonbpf/functions/functions_pass.py

@@ -491,7 +491,22 @@ def bound(expr):
     )
 
     builder.position_at_end(inc_block)
-    next_idx = builder.add(builder.load(counter.var), ir.Constant(loop_ty, step))
+    idx = builder.load(counter.var)
+    if abs(step) > 1:
+        # The body only runs with the counter strictly inside the range, so
+        # the distance left to stop is exact as an unsigned number. A step no
+        # shorter than that distance ends the loop here, before the add can
+        # wrap the counter back inside the range (a step of 1 cannot
+        # overshoot stop, so it keeps the plain loop).
+        left = builder.sub(stop_val, idx) if step > 0 else builder.sub(idx, stop_val)
+        step_in_range = func.append_basic_block(name="for.step")
+        builder.cbranch(
+            builder.icmp_unsigned(">", left, ir.Constant(loop_ty, _wrap64(abs(step)))),
+            step_in_range,
+            else_block or end_block,
+        )
+        builder.position_at_end(step_in_range)
+    next_idx = builder.add(idx, ir.Constant(loop_ty, _wrap64(step)))
     builder.store(next_idx, counter.var)
     builder.branch(cond_block)
 
@@ -509,6 +524,13 @@ def bound(expr):
     )
 
 
+def _wrap64(value):
+    """`value` modulo 2**64, as the signed Python int LLVM parses for an i64:
+    a step of 2**63 or more (or its negation) is written by its bits."""
+    value %= 1 << 64
+    return value - (1 << 64) if value >= 1 << 63 else value
+
+
 def handle_loop_jump(compilation_context, builder, stmt):
     """Handle `break` and `continue`: branch to the innermost loop's exit or
     next-iteration block."""
@@ -595,6 +617,8 @@ def process_stmt(
         handle_for(func, compilation_context, builder, stmt, local_sym_tab, ret_type)
     elif isinstance(stmt, (ast.Break, ast.Continue)):
         handle_loop_jump(compilation_context, builder, stmt)
+    elif isinstance(stmt, ast.Pass):
+        pass
     elif isinstance(stmt, ast.Return):
         did_return = handle_return(
             func, builder, stmt, local_sym_tab, ret_type, compilation_context
@@ -734,7 +758,11 @@ def process_func_body(
         )
 
     if not did_return:
-        builder.ret(ir.Constant(ir.IntType(64), 0))
+        # Falling off the end returns 0, in the function's own return type:
+        # after `while True` this block is unreachable, but it still has to
+        # type-check.
+        default_ty = ret_type if isinstance(ret_type, ir.IntType) else ir.IntType(64)
+        builder.ret(ir.Constant(default_ty, 0))
 
 
 def process_bpf_chunk(func_node, compilation_context, return_type):