ptodsl: support Python and/or short-circuit in AST rewrite (#1332)

ptodsl/docs/user_guide/05-control-flow.md

@@ -371,6 +371,58 @@ def ast_rewrite_side_effect_kernel():
         pto.pipe_barrier(pto.Pipe.ALL)
 ```
 
+### Short-circuit `and` / `or`
+
+Python `and` / `or` between PTODSL runtime values are not bitwise
+operations: they keep Python's short-circuit semantics, so the right-hand
+side is only evaluated on the device where Python would actually need it.
+The AST rewrite turns `a and b` / `a or b` into a result-bearing `scf.if`
+whose guarded region contains exactly the RHS:
+
+<!-- ptodsl-doc-test: {"mode":"compile","symbol":"ast_rewrite_short_circuit_kernel","compile":{}} -->
+```python
+@pto.jit(target="a5")
+def ast_rewrite_short_circuit_kernel():
+    value = pto.const(10, dtype=pto.i32)
+    divisor = pto.const(2, dtype=pto.i32)
+
+    pred = (divisor != 0) and ((value // divisor) > 0)
+    if pred:
+        pto.pipe_barrier(pto.Pipe.ALL)
+```
+
+`and` returns the left operand when it is falsy and the right operand
+otherwise; `or` returns the left operand when it is truthy and the right
+operand otherwise. Integer-typed operands are tested with non-zero
+truthiness. When one merged branch yields an `i1` and the other an
+integer-like value, the integer type is kept and the `i1` side widens to its
+0/1 integer value, so the integer operand value is preserved; a merge of two
+`i1` values stays `i1`. Python `bool` literals materialize as `i1` and widen
+like any other `i1` when the opposite branch is integer-typed. A Python
+`int` literal can join a runtime merge only when the opposite branch already
+has an integer type to anchor its width (e.g. `x or 2` with `x: pto.i32`);
+against an `i1` branch there is no width to infer, so `flag and 2` /
+`flag or 2` raise an error prompting you to anchor the type explicitly,
+e.g. `flag and pto.const(2, dtype=pto.i32)`. Incompatible branch types keep
+the usual branch-merge diagnostics.
+
+Statically known `bool` / `int` / float operands short-circuit at trace time
+with native Python truthiness: `False and rhs`, `True or rhs`, and
+`0.0 or rhs` never trace the RHS at all. Runtime floating-point controls are
+not valid short-circuit predicates and raise a clear error. Python float
+literals on the RHS of a runtime short-circuit merge are also rejected: the
+merge currently accepts only `bool`/`int` literals and PTO runtime scalar
+values, because a float literal (for example, `flag and 0.5`) has no
+compatible result type when the other branch preserves the left operand. Use
+an explicit `pto.if_` merge with same-typed branch values when a floating-point
+result is required.
+Assignment expressions (`:=`) on a lazily evaluated RHS are rejected because
+the rewrite uses a helper lambda and cannot preserve the Python binding scope.
+
+`and` / `or` compose with every rewritten expression context: assignments,
+call arguments, `return`, and `if` / `while` conditions.
+
+
 ### Runtime loops
 
 Native `range(...)` loops become device-side loops:
@@ -591,10 +643,10 @@ Do not pass conflicting values through both spellings. For example,
 `@pto.jit(ast_rewrite=False, frontend_options={"ast_rewrite": True})` is
 rejected.
 
-When this mode is disabled for a function, native Python `if` / `for` executes
-while tracing that function. Runtime device-side control flow should still use
-the default rewrite mode, or explicit `pto.if_` / `pto.for_` APIs when you need
-manual control.
+When this mode is disabled for a function, native Python `if` / `for` / `and` /
+`or` executes while tracing that function. Runtime device-side control flow
+should still use the default rewrite mode, or explicit `pto.if_` / `pto.for_`
+APIs when you need manual control.
 
 The structured `frontend_options` argument is reserved for frontend rewrite
 debugging and future rewrite passes. Today it accepts the same AST rewrite
@@ -638,3 +690,4 @@ do not use Python `return` as a dynamic loop exit.
 | `pto.for_` | Device-side | Dynamic bounds, runtime loop counts |
 | `pto.for_(...).carry(...)` | Device-side | Loops with accumulated state across iterations |
 | `pto.if_` | Device-side | Runtime conditions, data-dependent branching |
+| Python `and` / `or` | Device-side short-circuit `scf.if` | Runtime predicates with Python short-circuit semantics |

ptodsl/ptodsl/_ast_rewrite.py

@@ -29,8 +29,10 @@ def rewrite_jit_function(
 ):
     """Return a function with PTODSL lexical sections lowered safely.
 
-    ``pto.section`` is a physical SSA region, not a Python ``with`` hint.  The
-    section body therefore gets a small source-level lexical rewrite even when
+    ``rewrite_control_flow`` controls both runtime statement rewriting and the
+    ``and``/``or`` short-circuit rewrite.  ``pto.section`` is a physical SSA
+    region, not a Python ``with`` hint.  Its body therefore gets a small
+    source-level lexical rewrite even when
     the optional control-flow rewrite is disabled.  This keeps Python's
     function-local assignment rules from leaking a section-local SSA value into
     a sibling physical section.
@@ -68,6 +70,8 @@ def rewrite_jit_function(
     static_env.update(static_bindings or {})
     _inject_closure_defaults(function_def, closure_vars.nonlocals)
     _sanitize_signature_for_exec(function_def)
+    if rewrite_control_flow:
+        function_def = _BoolOpRewriter().visit(function_def)
     function_def = _ConditionalExpressionNormalizer().visit(function_def)
     section_rewriter = _SectionLexicalRewriter()
     function_def = section_rewriter.visit(function_def)
@@ -381,6 +385,77 @@ def _normalize_ifexp_assignment(self, stmt, value):
         return ast.copy_location(self.generic_visit(if_stmt), stmt)
 
 
+def _zero_arg_lambda(body):
+    return ast.Lambda(
+        args=ast.arguments(
+            posonlyargs=[],
+            args=[],
+            vararg=None,
+            kwonlyargs=[],
+            kw_defaults=[],
+            kwarg=None,
+            defaults=[],
+        ),
+        body=body,
+    )
+
+
+class _BoolOpRewriter(ast.NodeTransformer):
+    """Rewrite Python ``and``/``or`` into device-side short-circuit helpers.
+
+    Python ``and``/``or`` are not overloadable operators: evaluating ``a and b``
+    forces a truth check on ``a``, which calls ``__bool__`` on a PTODSL runtime
+    value during tracing and raises.  This pass replaces every ``ast.BoolOp``
+    with a right-nested lazy helper call so the RHS is only traced inside a
+    device-side ``scf.if`` region::
+
+        a and b and c  ->  pto._short_circuit_and(a, lambda: pto._short_circuit_and(b, lambda: c))
+        a or  b or  c  ->  pto._short_circuit_or(a,  lambda: pto._short_circuit_or(b,  lambda: c))
+
+    The transformation is purely syntactic and composes with every expression
+    context: assignments, call arguments, ``return``, and ``if``/``while``
+    conditions.  Nested function bodies are rewritten as well, mirroring the
+    control-flow rewriter.  Statically-known ``bool``/``int`` operands are
+    short-circuited at trace time by the helpers themselves, so the RHS is not
+    traced at all when Python semantics would skip it.
+    """
+
+    @staticmethod
+    def _reject_rhs_assignment_expressions(node):
+        """Reject walrus expressions that would move into a generated lambda."""
+        if any(
+            isinstance(child, ast.NamedExpr)
+            for value in node.values[1:]
+            for child in ast.walk(value)
+        ):
+            raise PTODSLAstRewriteError(
+                "ast_rewrite=True cannot rewrite an and/or expression containing "
+                "an assignment expression (walrus operator) on the RHS; the "
+                "assignment would bind inside the generated helper lambda"
+            )
+
+    def visit_BoolOp(self, node):
+        self._reject_rhs_assignment_expressions(node)
+        node = self.generic_visit(node)
+        if isinstance(node.op, ast.And):
+            helper = "_short_circuit_and"
+        elif isinstance(node.op, ast.Or):
+            helper = "_short_circuit_or"
+        else:
+            return node
+        values = list(node.values)
+        if len(values) < 2:
+            return values[0] if values else node
+        result = values[-1]
+        for value in reversed(values[:-1]):
+            result = ast.Call(
+                func=_pto_attr(helper),
+                args=[value, _zero_arg_lambda(result)],
+                keywords=[],
+            )
+        return result
+
+
 @dataclass(frozen=True)
 class _NameInfo:
     loads: set[str]