annoscript 0.1.0__py3-none-any.whl

annoscript/p0_grammar_tests.py

@@ -0,0 +1,147 @@
+"""Pinned decoder acceptance tests (run with XGrammar substring2, no inference)."""
+import itertools
+import json
+import random
+import unittest
+from importlib.metadata import version
+from types import SimpleNamespace
+import xgrammar as xgr
+from annoscript.backends._bounded_grammar import grammar
+
+
+def compiled(text):
+    return xgr.GrammarCompiler(xgr.TokenizerInfo([]),cache_enabled=False).compile_grammar(text)
+
+def accepts(ctx,text):
+    m=xgr.GrammarMatcher(ctx,terminate_without_stop_token=True)
+    return m.accept_string(text) and m.is_terminated()
+
+def wire(items):return json.dumps({'items':items},ensure_ascii=False,separators=(',',':'))
+
+def call(op='extract',id='a',**args):
+    source = args.pop('source', None)
+    return SimpleNamespace(id=id,operation=op,arguments=args,extraction_source=source)
+
+
+class BoundedGrammarTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):assert version('xgrammar')=='0.2.5+lev.substring2.3'
+
+    def test_example_locations_never_expand_current_source(self):
+        from annoscript._test_synthetic import REVIEW_INPUT, LAUNCH_INPUT
+        from annoscript.anno import Example
+        from annoscript.backends.base import normalize_call, InferenceContext
+        c = normalize_call("q", "extract", ("地点",),
+                           {"source_field": "text", "max_items": 2, "max_length_per_item": 80,
+                            "examples": [Example(state=REVIEW_INPUT, answer=[]),
+                                         Example(state=LAUNCH_INPUT, answer=["广州"])]},
+                           InferenceContext(REVIEW_INPUT))
+        g = compiled(grammar([c], extract_mode="substring_chars"))
+        self.assertTrue(accepts(g, wire([])))
+        self.assertTrue(accepts(g, wire(["条目"])))  # Grammar only guarantees source, not semantics.
+        self.assertFalse(accepts(g, wire(["广州"])))
+
+    def test_macro_language_consumed_length_not_sam_longest(self):
+        rng=random.Random(20261018)
+        sources=['abbb','ababa','abcab','aaaa','', 'a']+[''.join(rng.choices('abc',k=rng.randint(1,7))) for _ in range(16)]
+        for source in sources:
+            for limit in (0,1,2,4,10):
+                for empty in (False,True):
+                    args=', '.join(json.dumps(c) for c in source)
+                    g=xgr.Grammar.from_ebnf('root ::= Substring('+args+(', ' if args else '')+f'allow_empty={str(empty).lower()}, max_chunks={limit})')
+                    copies=[g,xgr.Grammar.from_ebnf(str(g)),xgr.Grammar.deserialize_json(g.serialize_json())]
+                    expected={source[i:j] for i in range(len(source)) for j in range(i+1,min(len(source),i+limit)+1)}
+                    if empty:expected.add('')
+                    for clone in copies:
+                        c=compiled(clone)
+                        for n in range(4):
+                            for parts in itertools.product('abc',repeat=n):
+                                text=''.join(parts)
+                                self.assertEqual(accepts(c,text),text in expected,(source,limit,empty,text))
+        for bad in ['-1','1.2','true','"2"','2147483648']:
+            with self.assertRaises(Exception):xgr.Grammar.from_ebnf('root ::= Substring("a", max_chunks='+bad+')')
+
+    def test_empty_chunks_consume_a_block_in_nullable_mode(self):
+        for empty in (True,False):
+            g=compiled('root ::= Substring("a", "", "b", allow_empty='+str(empty).lower()+', max_chunks=2)')
+            self.assertEqual(accepts(g,'ab'),not empty)
+            self.assertEqual(accepts(g,''),empty)
+            self.assertTrue(accepts(g,'a'));self.assertTrue(accepts(g,'b'))
+
+    def test_extract_json_unicode_count_and_boundaries(self):
+        for source in ['甲"乙\\丙\n😀','e\u0301／＠ @','\x00\x1f\b\f\r\t','\\u4e2d','aaaa','']:
+            for maximum in (1,2,4):
+                c=compiled(grammar([call(source=source,min_items=0,max_items=2,max_length_per_item=maximum)],extract_mode='substring_chars'))
+                self.assertTrue(accepts(c,wire([])))
+                self.assertFalse(accepts(c,wire([''])))
+                for i in range(len(source)):
+                    for j in range(i+1,len(source)+1):
+                        self.assertEqual(accepts(c,wire([source[i:j]])),True,(source,i,j))
+                if source:
+                    self.assertFalse(accepts(c,wire([source[0]]*3)))
+                for text in ['不存在','甲丙',source[::-1]]:
+                    if text not in source:self.assertFalse(accepts(c,wire([text])))
+                for bad in [r'{"items":["\u00"]}',r'{"items":["\"]}','{"items":["\n"]}']:
+                    self.assertFalse(accepts(c,bad))
+
+    def test_list_and_raw_generate_codepoints(self):
+        for lo,hi in ((0,0),(0,1),(1,1),(1,2),(2,3)):
+            g=grammar([call('generate_list',min_items=1,max_items=2,min_length_per_item=lo,max_length_per_item=hi)])
+            c=compiled(g)
+            for text in ('','a','甲','😀','e\u0301','"','\\','\n','甲😀中'):
+                for escaped in (False,True):
+                    value=json.dumps({'items':[text]},ensure_ascii=escaped,separators=(',',':'))
+                    self.assertEqual(accepts(c,value),bool(text) or lo==0,(lo,hi,value))
+            self.assertFalse(accepts(c,wire([])))
+            self.assertFalse(accepts(c,wire(['a','a','a'])))
+            raw=compiled(grammar([call('generate',min_length=lo,max_length=hi)]))
+            for text in ('','a','甲','😀','e\u0301','\n','甲😀中'):
+                self.assertEqual(accepts(raw,text),bool(text) or lo==0,(lo,hi,text))
+        zero=compiled(grammar([call('generate_list',max_items=0)]))
+        self.assertTrue(accepts(zero,wire([])));self.assertFalse(accepts(zero,wire([''])))
+
+    def test_mixed_same_source_different_limits_and_isolation(self):
+        cs=[call(id='a',source='甲乙丙',max_items=1,max_length_per_item=1),
+            call(id='b',source='甲乙丙',min_items=1,max_items=2,max_length_per_item=2),
+            call(id='c',source='丁戊',max_items=1,max_length_per_item=1),
+            call('generate','d',max_length=2),call('generate_list','e',max_items=1,max_length_per_item=1)]
+        g=grammar(cs,batch=True,extract_mode='substring_chars');self.assertEqual(g.count('Substring('),3)
+        c=compiled(g)
+        def wire_batch(a='甲',b='甲乙',cc='丁',d='😀',e=None):
+            return json.dumps({'answers':{'q1':[a],'q2':[b],'q3':[cc],'q4':d,'q5':e or ['甲']}},ensure_ascii=False,separators=(',',':'))
+        self.assertTrue(accepts(c,wire_batch()))
+        for kwargs in ({'a':'丁'},{'cc':'甲'},{'e':['甲','乙']},{'b':''}):
+            self.assertFalse(accepts(c,wire_batch(**kwargs)),kwargs)
+
+    def test_named_keys_preserve_source_and_value_constraints(self):
+        name = ' 地点"\\\n😀 '
+        cs = [call(id='internal-a', field_name=name, source='广州', min_items=1, max_items=1),
+              call('generate', 'internal-b', field_name='event_title', min_length=1, max_length=60)]
+        g = compiled(grammar(cs, batch=True, extract_mode='substring_chars'))
+        def encoded(place='广州', title='上市'):
+            # Key literals use ASCII escapes; substring values retain source bytes.
+            return ('{"answers":{' + json.dumps(name) + ':' + json.dumps([place], ensure_ascii=False)
+                    + ',"event_title":' + json.dumps(title, ensure_ascii=False) + '}}')
+        good = encoded()
+        self.assertTrue(accepts(g, good))
+        self.assertFalse(accepts(g, encoded('北京')))
+        self.assertFalse(accepts(g, encoded(title='')))
+        self.assertFalse(accepts(g, good.replace('event_title', 'q2')))
+
+    def test_nullable_generate_single_and_mixed_rule_cache_isolation(self):
+        optional = call('generate', 'a', min_length=1, max_length=24, nullable=True)
+        ordinary = call('generate', 'b', min_length=1, max_length=24)
+        single = compiled(grammar([optional]))
+        for raw in ('null', '"null"', '"中文"', '"\\\""'):
+            self.assertTrue(accepts(single, raw), raw)
+        for raw in ('[]', '1', 'true', '{}', '""', '', '未加引号'):
+            self.assertFalse(accepts(single, raw), raw)
+        for calls in ([optional, ordinary], [ordinary, optional]):
+            mixed = compiled(grammar(calls, batch=True))
+            for values in ((None, '标题'), ('null', '标题'), ('标题', None), ('标题', ''), ('标题', '标题')):
+                expected = all((v is None and c.arguments.get('nullable', False)) or
+                               (isinstance(v, str) and bool(v)) for c, v in zip(calls, values))
+                raw = json.dumps({'answers': dict(zip(('q1', 'q2'), values))}, ensure_ascii=False, separators=(',', ':'))
+                self.assertEqual(accepts(mixed, raw), expected, (calls, values))
+
+if __name__=='__main__':unittest.main()

annoscript/passes/__init__.py

@@ -0,0 +1,40 @@
+"""Conservative, deterministic optimization of structured AnnoScript IR.
+
+The input is never mutated and no source, model, or user function is executed.
+Passes preserve ordinary execution and constraint bindings; they neither solve
+constraints nor treat an asserted condition as a fact. Unknown effects remain.
+Individual passes are implementation details; optimize_program is the public API.
+"""
+
+from annoscript.ir import Program
+from .control_flow import simplify_control_flow
+from .dead_assignments import eliminate_dead_assignments
+from .folding import fold_constants
+from .propagation import propagate_values
+from .temporaries import merge_temporaries
+
+__all__ = ["optimize_program"]
+
+_PASSES = (
+    propagate_values, fold_constants, simplify_control_flow,
+    eliminate_dead_assignments, merge_temporaries,
+)
+
+
+def optimize_program(program: Program) -> Program:
+    """Optimize a Program to a fixed point without changing the IR v1 schema.
+
+The passes only propagate atoms, fold bounded scalars, or remove structure;
+they do not expand loops, duplicate statements, or introduce new variables.
+JSON equality is intentional: dataclass equality conflates True, 1 and 1.0.
+"""
+    if not isinstance(program, Program):
+        raise TypeError("optimize_program expects an ir.Program")
+    previous = program.to_json(indent=None)
+    while True:
+        for transform in _PASSES:
+            program = transform(program)
+        current = program.to_json(indent=None)
+        if current == previous:
+            return program
+        previous = current

annoscript/passes/_analysis.py

@@ -0,0 +1,178 @@
+"""Shared structural analyses; none of these helpers evaluate user code.
+
+Pure IR expressions may still raise, read mutable data, or allocate objects.
+Safety below concerns supported IR values, excluding resource exhaustion.
+"""
+
+from collections import Counter
+from dataclasses import fields, replace
+
+from annoscript import ir
+
+
+def map_value(value, transform):
+    if isinstance(value, ir.Expr):
+        return transform(value)
+    if isinstance(value, tuple):
+        return tuple(map_value(item, transform) for item in value)
+    return value
+
+
+def rewrite_expr(expr, transform):
+    children = {f.name: map_value(getattr(expr, f.name), lambda e: rewrite_expr(e, transform))
+                for f in fields(expr)}
+    return transform(replace(expr, **children))
+
+
+def map_expressions(statement, transform):
+    """Map immediate expression fields, never bodies or constraint expressions."""
+    if isinstance(statement, ir.Constrain):
+        return statement
+    return replace(statement, **{
+        f.name: map_value(getattr(statement, f.name), transform) for f in fields(statement)
+    })
+
+
+def expression_reads(expr):
+    result = Counter()
+
+    def visit(node):
+        if isinstance(node, ir.Var):
+            result[node.name] += 1
+        return node
+
+    rewrite_expr(expr, visit)
+    return result
+
+
+def direct_reads(statement):
+    result = Counter()
+
+    def visit(expr):
+        result.update(expression_reads(expr))
+        return expr
+
+    # Constraints participate in liveness even though transformations skip them.
+    for f in fields(statement):
+        map_value(getattr(statement, f.name), visit)
+    return result
+
+
+def target_names(target):
+    if isinstance(target, ir.NameTarget):
+        return {target.name}
+    if isinstance(target, ir.TupleTarget):
+        return set().union(*(target_names(item) for item in target.items))
+    raise TypeError(f"Unknown target: {type(target).__name__}")
+
+
+def definitions(statement):
+    if isinstance(statement, (ir.Assign, ir.Infer)):
+        return {statement.target}
+    if isinstance(statement, (ir.Unpack, ir.For)):
+        return target_names(statement.target)
+    return set()
+
+
+def child_blocks(statement):
+    if isinstance(statement, ir.If):
+        return (statement.then_body, statement.else_body)
+    if isinstance(statement, ir.For):
+        return (statement.body,)
+    return ()
+
+
+def walk_statements(body):
+    for statement in body:
+        yield statement
+        for block in child_blocks(statement):
+            yield from walk_statements(block)
+
+
+def block_writes(body):
+    return set().union(*(definitions(s) for s in walk_statements(body)))
+
+
+def literal_key(value):
+    # Python equality conflates bool/int/float and the signs of floating zero.
+    return (type(value), value.hex() if type(value) is float else value)
+
+
+def atom_key(expr):
+    if isinstance(expr, ir.Literal):
+        return (ir.Literal, literal_key(expr.value))
+    if isinstance(expr, ir.Var):
+        return (ir.Var, expr.name)
+    raise TypeError("Expected a literal or variable")
+
+
+def resolve(expr, facts):
+    seen = set()
+    while isinstance(expr, ir.Var) and expr.name in facts and expr.name not in seen:
+        seen.add(expr.name)
+        expr = facts[expr.name]
+    return expr
+
+
+def invalidate(facts, names):
+    """Forget copies of old bindings as well as the overwritten bindings."""
+    resolved = {name: resolve(value, facts) for name, value in facts.items()}
+    return {name: value for name, value in resolved.items()
+            if name not in names and not (isinstance(value, ir.Var) and value.name in names)}
+
+
+def join_facts(left, right):
+    return {name: value for name, value in left.items()
+            if name in right and atom_key(value) == atom_key(right[name])}
+
+
+def join_exits(left, left_falls, right, right_falls):
+    if not left_falls:
+        return right if right_falls else {}
+    if not right_falls:
+        return left
+    return join_facts(left, right)
+
+
+def bound_before(program):
+    """Definitely bound names at each statement, including loop body entry."""
+    points = {}
+
+    def block(body, incoming):
+        bound = set(incoming)
+        falls = True
+        for statement in body:
+            points[statement.id] = frozenset(bound)
+            if isinstance(statement, ir.If):
+                a, af = block(statement.then_body, bound)
+                b, bf = block(statement.else_body, bound)
+                bound = a & b if af and bf else a if af else b if bf else set()
+                falls = af or bf
+            elif isinstance(statement, ir.For):
+                block(statement.body, bound | definitions(statement))
+                # The loop can execute zero times.
+            elif isinstance(statement, ir.Return):
+                falls = False
+            else:
+                bound |= definitions(statement)
+            if not falls:
+                break
+        return bound, falls
+
+    block(program.body, set())
+    return points
+
+
+def discardable(expr, bound):
+    """Prove that evaluating and discarding an expression cannot fail or act."""
+    if isinstance(expr, (ir.Literal, ir.State)):
+        return True
+    if isinstance(expr, ir.Var):
+        return expr.name in bound
+    if isinstance(expr, (ir.ListExpr, ir.TupleExpr)):
+        return all(discardable(item, bound) for item in expr.items)
+    if isinstance(expr, ir.DictExpr):
+        # Only literal keys are known to be hashable without a type analysis.
+        return all(isinstance(key, ir.Literal) and discardable(value, bound)
+                   for key, value in expr.entries)
+    return False

annoscript/passes/control_flow.py

@@ -0,0 +1,68 @@
+"""Remove proven dead paths and redundant boolean completion guards."""
+
+from dataclasses import replace
+
+from annoscript import ir
+from ._analysis import block_writes, definitions, invalidate, join_exits, resolve
+
+
+def _completion_pair(first, second, facts):
+    if not (isinstance(first, ir.If) and not first.else_body and first.then_body
+            and isinstance(second, ir.If) and not second.else_body):
+        return False
+    guard = second.condition
+    if not (isinstance(guard, ir.Unary) and guard.op == "not" and isinstance(guard.value, ir.Var)):
+        return False
+    name = guard.value.name
+    initial = resolve(ir.Var(name), facts)
+    last = first.then_body[-1]
+    return (isinstance(initial, ir.Literal) and initial.value is False
+            and isinstance(last, ir.Assign) and last.target == name
+            and isinstance(last.value, ir.Literal) and last.value.value is True)
+
+
+def simplify_control_flow(program: ir.Program) -> ir.Program:
+    def block(body, incoming):
+        facts = dict(incoming)
+        output = []
+        index = 0
+        falls = True
+        while index < len(body):
+            statement = body[index]
+            if index + 1 < len(body) and _completion_pair(statement, body[index + 1], facts):
+                statement = replace(statement, else_body=body[index + 1].then_body)
+                index += 1
+            index += 1
+            if isinstance(statement, ir.If):
+                condition = resolve(statement.condition, facts)
+                if isinstance(condition, ir.Literal):
+                    selected = statement.then_body if condition.value else statement.else_body
+                    nested, facts, falls = block(selected, facts)
+                    output.extend(nested)
+                else:
+                    yes, a, af = block(statement.then_body, facts)
+                    no, b, bf = block(statement.else_body, facts)
+                    output.append(replace(statement, then_body=yes, else_body=no))
+                    facts = join_exits(a, af, b, bf)
+                    falls = af or bf
+            elif isinstance(statement, ir.For):
+                facts = invalidate(facts, definitions(statement) | block_writes(statement.body))
+                nested, _, _ = block(statement.body, facts)
+                output.append(replace(statement, body=nested))
+            else:
+                output.append(statement)
+                if isinstance(statement, ir.Assign):
+                    value = resolve(statement.value, facts)
+                    facts = invalidate(facts, {statement.target})
+                    if isinstance(value, ir.Literal):
+                        facts[statement.target] = value
+                elif isinstance(statement, ir.Constrain):
+                    facts.clear()
+                else:
+                    facts = invalidate(facts, definitions(statement))
+                falls = not isinstance(statement, ir.Return)
+            if not falls:
+                break
+        return tuple(output), facts, falls
+
+    return replace(program, body=block(program.body, {})[0])

annoscript/passes/dead_assignments.py

@@ -0,0 +1,44 @@
+"""Structured backward liveness, preserving evaluations that may fail or act."""
+
+from dataclasses import replace
+
+from annoscript import ir
+from ._analysis import bound_before, definitions, direct_reads, discardable
+
+
+def eliminate_dead_assignments(program: ir.Program) -> ir.Program:
+    points = bound_before(program)
+
+    def block(body, outgoing):
+        live = set(outgoing)
+        output = []
+        for statement in reversed(body):
+            reads = set(direct_reads(statement))
+            if isinstance(statement, ir.Assign):
+                if statement.target not in live and discardable(statement.value, points.get(statement.id, ())):
+                    continue
+                live = (live - {statement.target}) | reads
+            elif isinstance(statement, ir.If):
+                yes, a = block(statement.then_body, live)
+                no, b = block(statement.else_body, live)
+                statement = replace(statement, then_body=yes, else_body=no)
+                live = a | b | reads
+            elif isinstance(statement, ir.For):
+                # A back edge reaches another target binding or exits the loop.
+                head = set(live)
+                while True:
+                    nested, before_body = block(statement.body, head)
+                    updated = live | (before_body - definitions(statement))
+                    if updated == head:
+                        break
+                    head = updated
+                statement = replace(statement, body=nested)
+                live = head | reads
+            elif isinstance(statement, ir.Return):
+                live = reads
+            else:
+                live = (live - definitions(statement)) | reads
+            output.append(statement)
+        return tuple(reversed(output)), live
+
+    return replace(program, body=block(program.body, set())[0])

annoscript/passes/folding.py

@@ -0,0 +1,104 @@
+"""Bounded folding of a closed set of scalar operations, never Python source."""
+
+from dataclasses import replace
+import math
+import operator
+
+from annoscript import ir
+from ._analysis import map_expressions, rewrite_expr
+
+
+_COMPARE = {
+    "eq": operator.eq, "ne": operator.ne, "lt": operator.lt, "le": operator.le,
+    "gt": operator.gt, "ge": operator.ge, "in": lambda a, b: a in b,
+    "not_in": lambda a, b: a not in b, "is": operator.is_, "is_not": operator.is_not,
+}
+_MAX_STRING = 65536
+_MAX_BITS = 4096
+
+
+def _bounded(value):
+    if type(value) is str:
+        return len(value) <= _MAX_STRING
+    if type(value) is int:
+        return value.bit_length() <= _MAX_BITS
+    if type(value) is float:
+        return math.isfinite(value)
+    return value is None or type(value) is bool
+
+
+def _fold(expr):
+    if isinstance(expr, ir.BoolOp):
+        values = list(expr.values)
+        while len(values) > 1 and isinstance(values[0], ir.Literal):
+            first = values.pop(0)
+            if (expr.op == "and" and not first.value) or (expr.op == "or" and first.value):
+                return first
+        return values[0] if len(values) == 1 else replace(expr, values=tuple(values))
+
+    def literal(node):
+        if not isinstance(node, ir.Literal) or not _bounded(node.value):
+            raise _NotConstant
+        return node.value
+
+    try:
+        if isinstance(expr, ir.Unary):
+            value = {"not": operator.not_, "positive": operator.pos, "negative": operator.neg}[expr.op](literal(expr.value))
+        elif isinstance(expr, ir.Binary):
+            a, b = literal(expr.left), literal(expr.right)
+            if isinstance(a, str) and isinstance(b, str) and len(a) + len(b) > _MAX_STRING:
+                return expr
+            value = operator.add(a, b)
+        elif isinstance(expr, ir.Compare):
+            left = literal(expr.left)
+            value = True
+            for op, item in zip(expr.ops, expr.comparators):
+                right = literal(item)
+                # Identity is meaningful only for the IR's singleton values.
+                if op in ("is", "is_not") and not any(v is None or type(v) is bool for v in (left, right)):
+                    return expr
+                if not _COMPARE[op](left, right):
+                    value = False
+                    break
+                left = right
+        elif isinstance(expr, ir.FormatValue):
+            if expr.spec is not None and literal(expr.spec) != "":
+                return expr
+            value = literal(expr.value)
+            if expr.conversion != "none":
+                value = {"s": str, "r": repr, "a": ascii}[expr.conversion](value)
+            value = format(value, "")
+        elif isinstance(expr, ir.FormatString):
+            parts = [literal(part) for part in expr.parts]
+            if not all(type(part) is str for part in parts) or sum(map(len, parts)) > _MAX_STRING:
+                return expr
+            value = "".join(parts)
+        elif isinstance(expr, ir.Intrinsic) and expr.name in ("len", "int", "str.strip"):
+            args = [literal(arg) for arg in expr.args]
+            kwargs = {name: literal(arg) for name, arg in expr.kwargs}
+            value = {"len": len, "int": int, "str.strip": str.strip}[expr.name](*args, **kwargs)
+        else:
+            return expr
+        return ir.Literal(value) if _bounded(value) else expr
+    except (_NotConstant, TypeError, ValueError, OverflowError):
+        # A failing expression must still fail at its original execution point.
+        return expr
+
+
+class _NotConstant(Exception):
+    pass
+
+
+def fold_constants(program: ir.Program) -> ir.Program:
+    def block(body):
+        output = []
+        for statement in body:
+            statement = map_expressions(statement, lambda e: rewrite_expr(e, _fold))
+            if isinstance(statement, ir.If):
+                statement = replace(statement, then_body=block(statement.then_body), else_body=block(statement.else_body))
+            elif isinstance(statement, ir.For):
+                statement = replace(statement, body=block(statement.body))
+            output.append(statement)
+        return tuple(output)
+
+    return replace(program, body=block(program.body))

annoscript/passes/propagation.py

@@ -0,0 +1,53 @@
+"""Forward propagation of scalar constants and copies of existing values."""
+
+from dataclasses import replace
+
+from annoscript import ir
+from ._analysis import (
+    block_writes, bound_before, definitions, invalidate, join_exits,
+    map_expressions, resolve, rewrite_expr,
+)
+
+
+def propagate_values(program: ir.Program) -> ir.Program:
+    points = bound_before(program)
+
+    def block(body, incoming):
+        facts = dict(incoming)
+        output = []
+        falls = True
+        for original in body:
+            statement = map_expressions(original, lambda e: rewrite_expr(e, lambda v: resolve(v, facts)))
+            if isinstance(statement, ir.If):
+                yes, a, af = block(statement.then_body, facts)
+                no, b, bf = block(statement.else_body, facts)
+                facts = join_exits(a, af, b, bf)
+                statement = replace(statement, then_body=yes, else_body=no)
+                falls = af or bf
+            elif isinstance(statement, ir.For):
+                writes = definitions(statement) | block_writes(statement.body)
+                facts = invalidate(facts, writes)
+                nested, _, _ = block(statement.body, facts)
+                statement = replace(statement, body=nested)
+            elif isinstance(statement, ir.Assign):
+                value = statement.value
+                if not (isinstance(value, ir.Var) and value.name == statement.target):
+                    facts = invalidate(facts, {statement.target})
+                    if isinstance(value, ir.Literal) or (
+                        isinstance(value, ir.Var) and value.name in points.get(statement.id, ())
+                    ):
+                        facts[statement.target] = value
+            elif isinstance(statement, ir.Constrain):
+                facts.clear()
+            elif isinstance(statement, ir.Return):
+                falls = False
+            else:
+                facts = invalidate(facts, definitions(statement))
+            output.append(statement)
+            if not falls:
+                # Unreachable code is removed by the control-flow pass.
+                output.extend(body[len(output):])
+                break
+        return tuple(output), facts, falls
+
+    return replace(program, body=block(program.body, {})[0])