Feat: chunked eval-breaker polling for range loops (#206)

docs/SPEC.md

@@ -666,7 +666,14 @@ All in `PM/commonTest`, so they run wherever the interpreter loads.
   (direct and mutual recursion) are compiled: every compiled function counts its frame (`tp_enter_call` /
   `tp_leave_call`, added to the Python frame depth and compared with `sys.getrecursionlimit()`), so recursion
   past the limit raises `RecursionError('maximum recursion depth exceeded')` where CPython does, and a C-stack
-  guard keeps the process alive at any limit (#57). `Status: implemented` for desktop host builds
+  guard keeps the process alive at any limit (#57). Compiled loops check the eval breaker so signals
+  (`KeyboardInterrupt`), pending calls and GIL hand-off reach Python during long compiled work (#141): a
+  `while` loop counts down once per iteration; a `for ... in range(...)` loop, whose trip count is known,
+  runs in chunks of `TP_POLL_INTERVAL` iterations and polls once per chunk (#206), with `break`,
+  `continue` and `return` meaning exactly what they mean without chunking (trip counts just below, at and
+  above a multiple of the interval, negative steps, empty ranges and ranges spanning the i64 limits all
+  behave as CPython's `range`). Functions holding an array copy-in or entry-globals snapshot do not poll.
+  `Status: implemented` for desktop host builds
   (`test_pipeline.py`, `test_cgen*.py`, `test_recursion.py`, `test_eval_breaker.py`, `test_fast_entry.py`,
   `test_free_threaded.py`; free-threaded 3.15t, #159). Not yet: shipping compiled modules in an app through
   pypackpack (#25), methods of fixed-layout classes (#155), native ints in impure functions (#232).

python-multiplatform-ksp/src/main/python/typedpython/cgen.py

@@ -865,6 +865,9 @@ def __init__(self, mg: _ModGen, f: ir.Function, fast: bool = False):
         self.temps: list[tuple[str, str]] = []      # (C type, name)
         self.arrays = {p.name: p for p in f.params if isinstance(p, ir.ArrayParam)}
         self.local_arrays: dict[str, Type] = {}       # name -> F64_ARRAY / I64_ARRAY
+        self.brk: list = []                    # per open loop: goto label for `break`, or None
+        self.label_n = 0
+        self.used_labels: set = set()
         self.types: dict[str, Type] = {}
         for p in f.params:
             if not isinstance(p, ir.ArrayParam):
@@ -1659,7 +1662,9 @@ def st_ForObj(self, s: ir.ForObj) -> None:
                 local = _ident("l", name)
                 self.emit(f"  {{ PyObject *tp_old = {local}; {local} = tp_u[{i}]; tp_release(&tp_old); }}")
             self.emit("}")
+        self.brk.append(None)
         self.stmts(s.body)
+        self.brk.pop()
         self.depth -= 1
         self.emit("}")
         self._own_release(it)               # after the loop and after a `break`; a return/error exits via tp_exit
@@ -1733,7 +1738,9 @@ def st_While(self, s: ir.While) -> None:
         self.poll()
         c = self._cond(s.cond)
         self.emit(f"if (!({c})) break;")
+        self.brk.append(None)
         self.stmts(s.body)
+        self.brk.pop()
         self.depth -= 1
         self.emit("}")
 
@@ -1765,13 +1772,33 @@ def st_ForRange(self, s: ir.ForRange) -> None:
             self.emit(f"    {tn} = ((uint64_t){ts} - (uint64_t){te} - 1u) / {mag} + 1u;")
             sign = "-"
         self.depth += 1
-        self.emit(f"for ({tk} = 0; {tk} < {tn}; {tk}++) {{")
-        self.depth += 1
-        self.poll()
+        if self.snapshot:
+            # no polling in a snapshot holder: a plain loop
+            self.emit(f"for ({tk} = 0; {tk} < {tn}; {tk}++) {{")
+            self.depth += 1
+            self.brk.append(None)
+        else:
+            # Chunked polling (#206): the trip count is known, so poll once per TP_POLL_INTERVAL
+            # iterations. `continue` stays inside the chunk loop; `break` must leave both loops.
+            te2, label = self.tmp("uint64_t"), f"tp_brk{self.label_n}"
+            self.label_n += 1
+            self.emit(f"for ({tk} = 0; {tk} < {tn}; ) {{")
+            self.depth += 1
+            self.poll()
+            self.emit(f"{te2} = {tn} - {tk} > (uint64_t)TP_POLL_INTERVAL ? {tk} + (uint64_t)TP_POLL_INTERVAL : {tn};")
+            self.emit(f"for (; {tk} < {te2}; {tk}++) {{")
+            self.depth += 1
+            self.brk.append(label)
         self.emit(f"{var} = (int64_t)((uint64_t){ts} {sign} {tk} * {mag});")
         self.stmts(s.body)
+        label = self.brk.pop()
         self.depth -= 1
         self.emit("}")
+        if label is not None:
+            self.depth -= 1
+            self.emit("}")
+            if label in self.used_labels:
+                self.emit(f"{label}: ;")
         self.depth -= 1
         self.emit("}")
 
@@ -1792,7 +1819,10 @@ def st_Return(self, s: ir.Return) -> None:
         self.emit("tp_rc = 0; goto tp_exit;")
 
     def st_Break(self, s) -> None:
-        self.emit("break;")
+        target = self.brk[-1] if self.brk else None
+        if target:
+            self.used_labels.add(target)
+        self.emit(f"goto {target};" if target else "break;")
 
     def st_Continue(self, s) -> None:
         self.emit("continue;")

python-multiplatform-ksp/src/test/python/test_eval_breaker.py

@@ -157,3 +157,218 @@ def handler(signum, frame):
         # come a few microseconds BEFORE the child reads t1 (3.15t checks right after a call): the
         # tolerance only has to separate that from "ran inside a call of 0.3 s or more".
         assert r["handler"] >= r["t1"] - 0.005, f"handler ran inside the call but its write is not in the result: {r}"
+
+
+# ---- chunked polling for range loops (#206, SPEC N-8) ---------------------------------------------
+
+CHUNK_SOURCE = """\
+# typedpython: compiled
+
+
+def brk(n: int, k: int) -> int:
+    s: int = 0
+    last: int = -1
+    for i in range(n):
+        if i == k:
+            break
+        s += i
+        last = i
+    return s * 3 + last
+
+
+def cont(n: int, m: int) -> int:
+    s: int = 0
+    for i in range(n):
+        if i % m == 0:
+            continue
+        s += 1
+    return s
+
+
+def ret(n: int, k: int) -> int:
+    for i in range(n):
+        if i == k:
+            return i * 2
+    return -1
+
+
+def down(n: int) -> int:
+    s: int = 0
+    for i in range(n, 0, -3):
+        s += i
+    return s
+
+
+def down1(n: int) -> int:
+    s: int = 0
+    for i in range(n, -1, -1):
+        s += 1
+    return s
+
+
+def up3(a: int, b: int) -> int:
+    s: int = 0
+    for i in range(a, b, 3):
+        s += i % 7
+    return s
+
+
+def nested(n: int, m: int) -> int:
+    s: int = 0
+    for i in range(n):
+        for j in range(m):
+            if j == i:
+                break
+            s += 1
+    return s
+
+
+def brk_in_while(n: int) -> int:
+    s: int = 0
+    for i in range(n):
+        while True:
+            s += 1
+            break
+        if i % 2 == 0:
+            continue
+        s += 1
+    return s
+
+
+def huge(a: int, b: int, cap: int) -> int:
+    c: int = 0
+    for i in range(a, b):
+        c += 1
+        if c >= cap:
+            break
+    return c
+
+
+def huge_last(a: int, b: int, cap: int) -> int:
+    c: int = 0
+    last: int = 0
+    for i in range(a, b):
+        c += 1
+        last = i
+        if c >= cap:
+            break
+    return last
+"""
+
+CHUNK_CHILD = """
+import json
+src = json.loads(sys.argv[2])
+ns = {}
+exec(src, ns)
+ref = ns
+I = 4096
+M = 2**63
+trips = [0, 1, 2, I - 1, I, I + 1, 2 * I - 1, 2 * I, 2 * I + 1, 3 * I, 3 * I + 5, 10 * I + 7]
+cases = []
+for n in trips:
+    for k in (0, 1, I - 1, I, I + 1, 2 * I, n - 1, n, n + 1):
+        cases += [("brk", (n, k)), ("ret", (n, k))]
+    for mm in (1, 2, 3, I):
+        cases.append(("cont", (n, mm)))
+    cases += [("down", (n,)), ("down1", (n,)), ("brk_in_while", (n,))]
+    cases.append(("up3", (-7, n)))
+    cases.append(("up3", (n, -7)))
+    cases.append(("nested", (n, 5)))
+cases += [("nested", (40, 4097)), ("nested", (4100, 4100)), ("nested", (3, 20000))]
+cases += [("up3", (5, 5)), ("up3", (6, 5)), ("down", (0,)), ("down", (-5,))]
+cases += [("huge", (-M, M - 1, cap)) for cap in (1, I, I + 1, 3 * I)]
+cases += [("huge", (M - 3, M - 1, 100)), ("huge", (M - 1, M - 1, 100)), ("huge", (-M, -M + 2 * I + 3, 10**9)),
+          ("huge", (-M, -M + 3 * I, 10**9)), ("huge_last", (-M, M - 1, 2 * I + 1)),
+          ("huge_last", (M - 2 * I - 3, M - 1, 10**9)), ("huge", (M - 1, -M, 10))]
+bad = []
+for name, args in cases:
+    want = getattr(ref["__builtins__"], "None", None)
+    want = ns[name](*args)
+    got = getattr(m, name)(*args)
+    if want != got or type(want) is not type(got):
+        bad.append([name, [str(a) for a in args], str(want), str(got)])
+print(json.dumps({"n": len(cases), "bad": bad[:10]}))
+"""
+
+
+def run_child_args(extension: Path, body: str, *args: str) -> dict:
+    script = LOADER + textwrap.dedent(body)
+    proc = subprocess.run([sys.executable, "-c", script, str(extension), *args],
+                          capture_output=True, text=True, timeout=300)
+    assert proc.returncode == 0, f"child failed ({proc.returncode}):\n{proc.stdout}\n{proc.stderr}"
+    return json.loads(proc.stdout.strip().splitlines()[-1])
+
+
+@pytest.fixture(scope="module")
+def chunk_module(tmp_path_factory):
+    d = tmp_path_factory.mktemp("ebchunk")
+    path = d / "loops.py"
+    path.write_text(CHUNK_SOURCE)
+    result = pipeline.compile_module(path, d / "out")
+    assert result.extension is not None, result.skipped
+    names = {f.name for f in result.module.functions}
+    assert names >= {"brk", "cont", "ret", "down", "down1", "up3", "nested", "brk_in_while", "huge",
+                     "huge_last"}, result.skipped
+    return path, result
+
+
+def _generated_c(path: Path) -> str:
+    from typedpython import cgen, frontend, verify
+    module, _ = verify.verify(frontend.lower(path))
+    return cgen.generate(module, path)
+
+
+def _function_c(c: str, name: str) -> str:
+    import re
+    start = re.search(rf"tp_impl_{name}\([^)]*\)\s*\{{", c).end()
+    return c[start:c.index("\ntp_exit:", start)]
+
+
+def test_range_loop_polls_once_per_chunk(chunk_module):
+    path, _ = chunk_module
+    body = _function_c(_generated_c(path), "down1")
+    # one poll site, outside the innermost (per-iteration) loop
+    assert body.count("tp_pc_tick") == 1, body
+    assert "TP_POLL_INTERVAL" in body, body
+    lines = body.splitlines()
+    poll = next(i for i, l in enumerate(lines) if "tp_pc_tick" in l)
+    fors = [i for i, l in enumerate(lines) if l.lstrip().startswith("for (") or l.lstrip().startswith("while (")]
+    assert len([i for i in fors if i < poll]) >= 1, body            # the chunk loop encloses the poll
+    assert len([i for i in fors if i > poll]) >= 1, body            # the per-iteration loop follows it
+    # the per-iteration loop has no poll inside it
+    inner = lines[max(i for i in fors if i > poll):]
+    assert not any("tp_pc_tick" in l for l in inner), body
+
+
+def test_while_loop_keeps_its_per_iteration_countdown(tmp_path):
+    path = tmp_path / "w.py"
+    path.write_text("# typedpython: compiled\n\n\ndef w(n: int) -> int:\n    i: int = 0\n    while i < n:\n"
+                    "        i += 1\n    return i\n")
+    body = _function_c(_generated_c(path), "w")
+    lines = body.splitlines()
+    loop = next(i for i, l in enumerate(lines) if "for (;;)" in l)
+    poll = next(i for i, l in enumerate(lines) if "tp_pc_tick" in l)
+    assert poll == loop + 1, body
+
+
+def test_range_loop_results_match_cpython_across_chunk_boundaries(chunk_module):
+    path, result = chunk_module
+    r = run_child_args(result.extension, CHUNK_CHILD, json.dumps(CHUNK_SOURCE))
+    assert r["n"] > 200
+    assert r["bad"] == [], r
+
+
+def test_keyboard_interrupt_reaches_a_long_range_loop_in_chunks(extension):
+    # spin() is a pure range loop: with chunked polling the interrupt must still arrive promptly.
+    r = run_child(extension, """
+        threading.Timer(0.05, lambda: __import__("_thread").interrupt_main()).start()
+        t0 = time.perf_counter()
+        try:
+            m.spin(9_000_000_000_000)
+            out = "completed"
+        except KeyboardInterrupt:
+            out = "interrupted"
+        print(json.dumps({"out": out, "elapsed": time.perf_counter() - t0}))
+    """)
+    assert r["out"] == "interrupted"
+    assert r["elapsed"] < 1.0, r