Chained and in kernels does not short-circuit — nullable array access evaluated unconditionally #1329

"""
Minimal reproducer: Warp codegen does not short-circuit chained `and`
when the first operand is a nullable array.

The array element access in the third operand is evaluated even when
the first operand (the array itself) is None, causing:
  - CUDA: error 700 (illegal memory access)
  - CPU:  segfault / undefined behavior

Expected: `if arr and idx >= 0 and arr[idx] == 0:` should skip
          arr[idx] when arr is None.
Actual:   arr[idx] is evaluated unconditionally.

Workaround: use nested `if` statements.
"""

import warp as wp

wp.init()


@wp.kernel
def chained_and_kernel(
    arr: wp.array(dtype=int),
    result: wp.array(dtype=int),
):
    tid = wp.tid()
    # BUG: arr[tid] is evaluated even when arr is None
    if arr and tid >= 0 and arr[tid] == 0:
        result[tid] = -1
        return
    result[tid] = 1


@wp.kernel
def nested_if_kernel(
    arr: wp.array(dtype=int),
    result: wp.array(dtype=int),
):
    tid = wp.tid()
    # WORKAROUND: nested if correctly guards the access
    if arr:
        if tid >= 0 and arr[tid] == 0:
            result[tid] = -1
            return
    result[tid] = 1


for device in ("cpu", "cuda:0"):
    if device.startswith("cuda") and not wp.is_cuda_available():
        print(f"Skipping {device} (no CUDA)")
        continue

    print(f"\n=== Device: {device} ===")
    result = wp.zeros(3, dtype=int, device=device)

    # Nested if with None array — should work fine
    print("Nested if (None array)... ", end="", flush=True)
    try:
        wp.launch(nested_if_kernel, dim=3, inputs=[None, result], device=device)
        wp.synchronize_device(device)
        print(f"OK  result={result.numpy().tolist()}")
    except Exception as e:
        print(f"FAIL  {e}")

    # Chained and with None array — triggers the bug
    print("Chained and (None array)... ", end="", flush=True)
    try:
        wp.launch(chained_and_kernel, dim=3, inputs=[None, result], device=device)
        wp.synchronize_device(device)
        print(f"OK  result={result.numpy().tolist()}")
    except Exception as e:
        print(f"FAIL  {e}")