diff --git a/src/srdatalog/ir/dialects/relation/sorted_array/__init__.py b/src/srdatalog/ir/dialects/relation/sorted_array/__init__.py index e9c7646..3378da4 100644 --- a/src/srdatalog/ir/dialects/relation/sorted_array/__init__.py +++ b/src/srdatalog/ir/dialects/relation/sorted_array/__init__.py @@ -61,15 +61,15 @@ _mir_for_use_declarative.ConstantBind, _mir_for_use_declarative.InsertInto, _mir_for_use_declarative.Scan, - # B-CJ-single (per docs/phase_b_lowering_dispatcher.md §4 row - # B-CJ-single, difficulty `medium`): adds `mir.ColumnJoin` to - # the ratchet. NOTE: the new `@lowering`-dispatch path only - # owns the SINGLE-source case (`len(sources) == 1`); the - # multi-source case (`len(sources) >= 2`) stays on the legacy - # `_lower_nested_cj_multi` / `_lower_root_cj_multi` branches - # until B-CJ-multi lands. The `_should_use_declarative` helper - # in `lowerings/__init__.py` enforces the "type AND structural - # shape" gate; B-CJ-multi drops the source-count guard there. + # B-CJ-single + B-CJ-multi (per docs/phase_b_lowering_dispatcher.md + # §4 rows B-CJ-single + B-CJ-multi, orders 5 + 6): `mir.ColumnJoin` + # entry covers BOTH shapes after B-CJ-multi. The chain dispatcher + # in `lowerings/__init__.py:_lower_inner_chain` routes single- + # source via `lower_mir_cj_single_in_chain` and multi-source via + # `lower_mir_cj_multi_in_chain`; the root dispatcher in + # `lower_scan_pipeline` routes multi-source via + # `lower_mir_cj_multi_root` (root-position single-source is not + # a supported pipeline shape today). _mir_for_use_declarative.ColumnJoin, _mir_for_use_declarative.Negation, _mir_for_use_declarative.Aggregate, @@ -320,21 +320,28 @@ def _lower_mir_insert_into_registered(op: Any, ctx: Any) -> Any: def _lower_mir_scan_registered(op: Any, ctx: Any) -> Any: return lower_mir_scan(op, ctx) - # Wave 2A / B-CJ-single (per docs/phase_b_lowering_dispatcher.md §4 - # row B-CJ-single, difficulty `medium`): `mir.ColumnJoin` migrates - # to a standalone `@lowering` registration in its own file. This - # PR owns only the SINGLE-source case (`len(sources) == 1`); the - # multi-source case stays on the legacy `_lower_nested_cj_multi` - # / `_lower_root_cj_multi` branches in `lowerings/__init__.py` - # and is migrated by the next PR (B-CJ-multi). + # Wave 2A / B-CJ-single + B-CJ-multi (per docs/phase_b_lowering_dispatcher.md + # §4 rows B-CJ-single + B-CJ-multi, orders 5 + 6): `mir.ColumnJoin` + # gets ONE dialect-level `@lowering(target=iir.cf, source=mir.ColumnJoin)` + # registration that covers both source-count shapes for discipline + # ownership. The shape-aware dispatcher below picks the right per-op + # stub by `len(op.sources)`: # - # The `_should_use_declarative` helper in `lowerings/__init__.py` - # encodes the "type AND structural shape" gate so the chain - # dispatcher routes single-source through `lower_mir_cj_single_in_chain` - # and lets multi-source fall through to the legacy branch. - # Same split-with-stub pattern as B-Filter / B-Scan: the registered - # stub asserts on direct invocation, and the chain-aware variant - # is the real entry. + # - `len(sources) == 1` (B-CJ-single, PR #59): `lower_mir_cj_single` + # - `len(sources) >= 2` (B-CJ-multi, this PR): `lower_mir_cj_multi` + # + # Both stubs assert on direct registry invocation — the actual + # dispatch flows through `_lower_inner_chain` (mid-chain) and + # `lower_scan_pipeline` (root) in `lowerings/__init__.py`, which + # call the chain-aware variants (`lower_mir_cj_single_in_chain`, + # `lower_mir_cj_multi_in_chain`, `lower_mir_cj_multi_root`) with + # the surrounding pipeline `tail` / `rest` in scope. The + # `USE_DECLARATIVE` ratchet routes both shapes through the new + # path while this registration pins the dialect-ownership + # contract. + from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_cj_multi import ( + lower_mir_cj_multi, + ) from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_cj_single import ( lower_mir_cj_single, ) @@ -345,8 +352,15 @@ def _lower_mir_scan_registered(op: Any, ctx: Any) -> Any: consumes=('mir',), produces=('iir.cf',), ) - def _lower_mir_cj_single_registered(op: Any, ctx: Any) -> Any: - return lower_mir_cj_single(op, ctx) + def _lower_mir_column_join_registered(op: Any, ctx: Any) -> Any: + # Shape-aware dispatch by source count. Both stubs assert (the + # real entries are the chain-aware variants in their respective + # `lower_mir_cj_*.py` modules) — this dispatcher exists so the + # single dialect-level `@lowering` registration produces a + # consistent error message regardless of which shape was passed. + if len(op.sources) == 1: + return lower_mir_cj_single(op, ctx) + return lower_mir_cj_multi(op, ctx) # Wave 2A / B-Negation (per docs/phase_b_lowering_dispatcher.md §4 # row B-Negation, difficulty `hard`, order 9): `mir.Negation` diff --git a/src/srdatalog/ir/dialects/relation/sorted_array/lowerings/__init__.py b/src/srdatalog/ir/dialects/relation/sorted_array/lowerings/__init__.py index b08d398..adaf1c6 100644 --- a/src/srdatalog/ir/dialects/relation/sorted_array/lowerings/__init__.py +++ b/src/srdatalog/ir/dialects/relation/sorted_array/lowerings/__init__.py @@ -344,13 +344,22 @@ def lower_scan_pipeline( # module import graph linear (the sibling per-op modules import # back from this file). # - # B-CJ-single: `mir.ColumnJoin` is in `USE_DECLARATIVE` but only - # the single-source case routes through the new path — and even - # then only mid-chain (see `_lower_inner_chain`). Root-position - # single-source CJ is not a fixture today (`_supported_pipeline` - # requires `len(sources) >= 2` for root CJ). `_should_use_declarative` - # encodes the gate so root-position multi-source CJ keeps falling - # through to `_lower_root_cj_multi` below. + # B-CJ-multi (per docs/phase_b_lowering_dispatcher.md §4 row + # B-CJ-multi, order 6, difficulty `hard`): root-position multi- + # source `mir.ColumnJoin` routes through `lower_mir_cj_multi_root`, + # which delegates to `_lower_root_cj_multi`. Single-source CJ is + # not a fixture at root today (`_supported_pipeline` requires + # `len(sources) >= 2` for root CJ); the shape split here mirrors + # the chain dispatcher's split in `_lower_inner_chain`. + # + # BlockGroup coexistence: the `BlockGroupRoot` unwrap above this + # block (entered BEFORE `_should_use_declarative` is consulted) + # handles the typed-pragma path; the legacy dual-write path + # (`ep.block_group=True` + bare `ColumnJoin` head) reaches this + # branch with `ctx.bg_enabled` already True, and the + # `_lower_root_cj_multi` helper's own first-statement guard + # (`if ctx.bg_enabled: return _lower_root_cj_bg(...)`) re-routes + # to the BG variant. Both BG paths remain byte-equivalent. if _should_use_declarative(head): if isinstance(head, mir.Scan): from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_scan import ( @@ -371,6 +380,12 @@ def lower_scan_pipeline( ) return lower_mir_cart_root(head, rest, ctx) + if isinstance(head, mir.ColumnJoin) and len(head.sources) >= 2: + from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_cj_multi import ( + lower_mir_cj_multi_root, + ) + + return lower_mir_cj_multi_root(head, rest, ctx) raise AssertionError( f'lower_scan_pipeline: USE_DECLARATIVE contains {type(head).__name__!r} ' f'(and `_should_use_declarative` returned True) but no root-' @@ -1076,32 +1091,32 @@ def _should_use_declarative(head: mir.MirNode) -> bool: '''Gate: does the head op route through the new per-op `@lowering` path (instead of the legacy `if isinstance` branches below)? - For most MIR op types this is just `type(head) in USE_DECLARATIVE`. - For `mir.ColumnJoin` it additionally checks the source count: the - single-source case (`len(sources) == 1`) is owned by B-CJ-single - (this PR); the multi-source case (`len(sources) >= 2`) stays on - the legacy `_lower_nested_cj_multi` / `_lower_root_cj_multi` - branches until B-CJ-multi lands. + For every MIR op type covered by Wave 2A this is just + `type(head) in USE_DECLARATIVE`. The helper exists as a single + decision point so neither `_lower_inner_chain` nor + `lower_scan_pipeline` ends up with two isinstance-cascades on the + same op type. - B-CJ-multi extends this helper to drop the source-count guard - (the single `return True` for `mir.ColumnJoin`) — the multi- - source path then has its own `lower_mir_cj_multi_in_chain` - + the same `USE_DECLARATIVE` membership wires both shapes - through the new dispatch. + Historical note (B-CJ-single, PR #59): this helper used to gate + `mir.ColumnJoin` on `len(sources) == 1` — the single-source case + routed through the new path while multi-source fell through to + the legacy `_lower_nested_cj_multi` / `_lower_root_cj_multi` + branches. B-CJ-multi (this PR) drops that guard: ALL ColumnJoin + variants route through the new per-op path. The chain-aware / + root-aware shape split lives in the per-op module + (`lower_mir_cj_single*` vs `lower_mir_cj_multi*`); the chain + dispatcher / root dispatcher pick the right entry from + `len(head.sources)`. Per docs/phase_b_lowering_dispatcher.md §4 — the "type AND - structural shape" gate keeps the dispatch atomic so neither - `_lower_inner_chain` nor `lower_scan_pipeline` ends up with two - isinstance-cascades on the same op type. + structural shape" gate keeps the dispatch atomic. Today the + structural shape is encoded in the per-op modules' own dispatch + (not here); the helper is kept as a stable extension point if + future B-PRs need a per-op gate again. ''' from srdatalog.ir.dialects.relation.sorted_array import USE_DECLARATIVE - if type(head) not in USE_DECLARATIVE: - return False - if isinstance(head, mir.ColumnJoin): - # B-CJ-single owns single-source; B-CJ-multi (next PR) owns multi. - return len(head.sources) == 1 - return True + return type(head) in USE_DECLARATIVE def _lower_inner_chain( @@ -1170,14 +1185,24 @@ def _lower_inner_chain( return lower_mir_insert_into_in_chain(head, tail, ctx) if isinstance(head, mir.ColumnJoin): - # B-CJ-single: guard already enforced single-source by - # `_should_use_declarative`; the multi-source fall-through - # below stays the legacy path until B-CJ-multi. - from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_cj_single import ( - lower_mir_cj_single_in_chain, + # B-CJ-single + B-CJ-multi: shape-aware dispatch by source count. + # `_should_use_declarative` no longer gates by source count for + # `mir.ColumnJoin` (it was a single `return True` from B-CJ-multi + # onward) — both single-source and multi-source CJ route through + # their respective per-op entry points here. The single dialect- + # level `@lowering(mir.ColumnJoin, ...)` registration covers + # both shapes for discipline ownership. + if len(head.sources) == 1: + from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_cj_single import ( + lower_mir_cj_single_in_chain, + ) + + return lower_mir_cj_single_in_chain(head, tail, ctx) + from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_cj_multi import ( + lower_mir_cj_multi_in_chain, ) - return lower_mir_cj_single_in_chain(head, tail, ctx) + return lower_mir_cj_multi_in_chain(head, tail, ctx) if isinstance(head, mir.Negation): from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_negation import ( lower_mir_negation_in_chain, diff --git a/src/srdatalog/ir/dialects/relation/sorted_array/lowerings/lower_mir_cj_multi.py b/src/srdatalog/ir/dialects/relation/sorted_array/lowerings/lower_mir_cj_multi.py new file mode 100644 index 0000000..581e4ac --- /dev/null +++ b/src/srdatalog/ir/dialects/relation/sorted_array/lowerings/lower_mir_cj_multi.py @@ -0,0 +1,203 @@ +'''Lowering: `mir.ColumnJoin` (multi-source) -> iir.cf — sixth +Wave 2A per-op migration. + +Per `docs/phase_b_lowering_dispatcher.md` §4 (per-MIR-op work-unit +table, row B-CJ-multi, difficulty `hard`, migration order 6) and +§4.1 (per-PR template): each MIR op type gets one `@lowering(target= +iir.cf, source=mir.X)` rule in its own file under +`dialects/relation/sorted_array/lowerings/`. The dialect-level +`@lowering(...)` registration for `mir.ColumnJoin` lives alongside +the other lowerings in `__init__.py:_register_passes`; this file is +referenced from the chain dispatcher (`_lower_inner_chain`) and the +root dispatcher (`lower_scan_pipeline`) in `lowerings/__init__.py`. + +Source-shape split (B-CJ-single vs B-CJ-multi): + +`mir.ColumnJoin` carries a `sources: list[ColumnSource]` field +whose length determines the structural shape: + + - `len(sources) == 1` (B-CJ-single, PR #59): handled by + `lower_mir_cj_single_in_chain`. Always a middle-of-chain op + (`_supported_pipeline` does not accept it as a root). + + - `len(sources) >= 2` (THIS PR, B-CJ-multi): both root-position + and middle-of-chain. Root-position is the production-load- + bearing case (every TC-like fixture roots a multi-source + intersection); middle-of-chain handles nested intersections + (e.g. CJ-under-CJ in the deepest TC shape). + +The `_should_use_declarative` helper in `lowerings/__init__.py`, +after this PR, no longer gates by source count for `mir.ColumnJoin`: +ALL ColumnJoin variants route through the new per-op path. The +single dialect-level `@lowering(mir.ColumnJoin, ...)` registration +covers both shapes for ownership / discipline purposes. + +Byte-equivalence contract (§4.2): the migrated path produces the +same IIR tree as the legacy `if isinstance(head, mir.ColumnJoin) and +len(head.sources) >= 2:` branches (root + nested) by delegating to +the same `_lower_root_cj_multi` / `_lower_nested_cj_multi` helpers +the legacy branches called directly. Layer 3 cleanup deletes those +helpers once the chain-aware variants here own the emission outright. + +BlockGroup coexistence — the load-bearing constraint of this PR: + +When the EP carries `ep.block_group=True` (legacy dual-write path) +or its `pipeline[0]` is a `BlockGroupRoot` wrap op (typed-pragma +path, post C3 materialization), the BG variant of the lowering +(`_lower_root_cj_bg`) is the correct emission target — NOT the +standard `_lower_root_cj_multi`. The dispatch ordering in +`lower_scan_pipeline` keeps the BG path intact: + + 1. `lower_scan_pipeline` first checks for `BlockGroupRoot` head and + dispatches to `_lower_root_cj_bg` directly. The check sits + BEFORE the `_should_use_declarative` branch — wrap ops never + reach the new path. + 2. For non-wrap pipelines (the `ep.block_group=True` dual-write + case), `_should_use_declarative` accepts the bare multi-source + `ColumnJoin` and routes through `lower_mir_cj_multi_root`. The + latter delegates to `_lower_root_cj_multi`, whose first + statement is `if ctx.bg_enabled: return _lower_root_cj_bg(...)`. + `compile_kernel_body` threads `bg_enabled` from `ep.block_group`, + so the BG variant fires before any non-BG scaffolding runs. + +Both delegation paths preserve byte-equivalence by construction: +the new entries call the same legacy helpers the legacy `if +isinstance` branches called, with the same `ctx` and `rest`. + +Chain-aware split (mirrors B-CJ-single's `lower_mir_cj_single_in_chain` +pattern + every other Wave 2A row that needs the trailing chain): + + - `lower_mir_cj_multi_root(op, rest, ctx)`: root-position entry, + called from `lower_scan_pipeline`. Delegates to + `_lower_root_cj_multi`. + + - `lower_mir_cj_multi_in_chain(op, tail, ctx)`: middle-of-chain + entry, called from `_lower_inner_chain`. Delegates to + `_lower_nested_cj_multi`. + + - `lower_mir_cj_multi(op, ctx)`: the `@lowering`-registered stub + (sibling of `lower_mir_cj_single` in the registry; both share + the single dialect-level `mir.ColumnJoin` registration via the + shape-aware dispatcher `_dispatch_mir_cj_registered` in + `__init__.py`). Asserts on direct invocation because the chain- + aware variants need either `rest` or `tail` from the enclosing + dispatcher — same split rationale as every other chain-aware + Wave 2A entry. +''' + +from __future__ import annotations + +from typing import Any + +import srdatalog.ir.mir.types as mir +from srdatalog.ir.core import Op + + +def lower_mir_cj_multi_root( + op: mir.ColumnJoin, + rest: list[Any], + ctx: Any, +) -> Op: + '''Emit the IIR for a root-position multi-source `mir.ColumnJoin` + with trailing `rest`. + + Mirrors the legacy `if isinstance(head, mir.ColumnJoin):` branch + in `lower_scan_pipeline` byte-for-byte by delegating to + `_lower_root_cj_multi`. The helper internally re-dispatches to + `_lower_root_cj_bg` when `ctx.bg_enabled` is set — that BG entry + is the load-bearing dual-write anchor (see module docstring). + + Caller-side guard: `_should_use_declarative` in `lower_scan_pipeline` + already filtered through `USE_DECLARATIVE` membership. We re-assert + `len(op.sources) >= 2` below as the structural contract: a future + refactor that bypassed the helper and called this entry with a + single-source CJ would silently misroute (single-source CJ has its + own root-emission scaffolding — though no fixture exercises it). + ''' + from srdatalog.ir.dialects.relation.sorted_array.lowerings import ( + _lower_root_cj_multi, + ) + + if len(op.sources) < 2: + raise AssertionError( + f'lower_mir_cj_multi_root: expected at least two ColumnSources, ' + f'got {len(op.sources)}. Single-source ColumnJoin is owned by ' + f'B-CJ-single (lower_mir_cj_single) — `_should_use_declarative` ' + f'should have routed it through the single-source entry.' + ) + + return _lower_root_cj_multi(op, rest, ctx) + + +def lower_mir_cj_multi_in_chain( + op: mir.ColumnJoin, + tail: list[Any], + ctx: Any, +) -> Op: + '''Emit the IIR for a middle-of-chain multi-source `mir.ColumnJoin` + with trailing `tail`. + + Mirrors the legacy `if isinstance(head, mir.ColumnJoin) and + len(head.sources) >= 2:` branch in `_lower_inner_chain` byte-for- + byte by delegating to `_lower_nested_cj_multi`. The legacy helper + stays LIVE until Layer 3 cleanup deletes both the legacy branch + AND `USE_DECLARATIVE`. + + Caller-side guard: `_should_use_declarative` in `_lower_inner_chain` + already filtered through `USE_DECLARATIVE` membership. We re-assert + `len(op.sources) >= 2` below as a structural contract — see + `lower_mir_cj_multi_root` for the same rationale. + ''' + from srdatalog.ir.dialects.relation.sorted_array.lowerings import ( + _lower_nested_cj_multi, + ) + + if len(op.sources) < 2: + raise AssertionError( + f'lower_mir_cj_multi_in_chain: expected at least two ColumnSources, ' + f'got {len(op.sources)}. Single-source ColumnJoin is owned by ' + f'B-CJ-single (lower_mir_cj_single_in_chain) — ' + f'`_should_use_declarative` should have routed it through the ' + f'single-source entry.' + ) + + return _lower_nested_cj_multi(op, tail, ctx) + + +def lower_mir_cj_multi(op: mir.ColumnJoin, ctx: Any) -> Op: + '''Framework-registry stub for the multi-source half of + `@lowering(target=iir.cf, source=mir.ColumnJoin)`. The actual + dispatch lives in `lowerings._lower_inner_chain` (mid-chain) and + `lowerings.lower_scan_pipeline` (root) — both call the + corresponding chain-aware variant with the surrounding pipeline + slice in scope. + + Note: the dialect-level `@lowering(mir.ColumnJoin, ...)` + registration in `relation.sorted_array.__init__._register_passes` + is shared between this file and `lower_mir_cj_single.py`. The + registered dispatcher inspects `len(op.sources)` and forwards to + this stub OR `lower_mir_cj_single` accordingly. Both stubs assert + because direct registry invocation has no access to the trailing + pipeline `rest` / `tail` that the chain-aware variants need. + + Calling this stub directly raises a structural assertion — the + framework path is reserved for future readers who can plumb in + the missing `tail` / `rest` (e.g. a refactored MIR-IIR walker that + no longer routes through `_lower_inner_chain`). + ''' + raise AssertionError( + 'lower_mir_cj_multi: dispatch goes through ' + '`lowerings._lower_inner_chain` -> `lower_mir_cj_multi_in_chain` ' + '(mid-chain) or `lowerings.lower_scan_pipeline` -> ' + '`lower_mir_cj_multi_root` (root) so the trailing pipeline is ' + 'in scope. Direct invocation indicates a refactor that bypassed ' + 'the dispatcher — plumb the `tail` / `rest` through and call the ' + 'chain-aware variant instead.' + ) + + +__all__ = [ + 'lower_mir_cj_multi', + 'lower_mir_cj_multi_in_chain', + 'lower_mir_cj_multi_root', +] diff --git a/tests/test_lower_mir_cj_multi_byte_equivalent.py b/tests/test_lower_mir_cj_multi_byte_equivalent.py new file mode 100644 index 0000000..37a0107 --- /dev/null +++ b/tests/test_lower_mir_cj_multi_byte_equivalent.py @@ -0,0 +1,673 @@ +'''Byte-equivalence test for Wave 2A B-CJ-multi migration. + +Per `docs/phase_b_lowering_dispatcher.md` §4.2 (per-PR acceptance +gate): each per-MIR-op migration ships a +`test_lower_mir__byte_equivalent` test that runs the migrated +path on every relevant fixture and asserts byte equality with the +legacy `if isinstance(head, mir.X):` branch. + +For multi-source ColumnJoin (`len(sources) >= 2`) the relevant +fixtures are: + + - root-position multi-source CJ (the production-load-bearing + case — every TC-like fixture uses this shape), + - mid-chain (nested) multi-source CJ — CJ under CJ, + - 2-source and 3-source variants, + - prefix-narrowed sources (parent handle aliased from the + enclosing CJ scope) versus fresh sources, + - BG-eligible shapes: bare multi-source root CJ with + `ctx.bg_enabled=True` (the legacy `block_group: bool` dual- + write path) AND `BlockGroupRoot(inner=multi-source CJ)` + (the typed-pragma path post C3 materialization), + - full-pipeline shapes via `compile_pipeline`, + - Filter / ConstantBind interleaved before / after the CJ. + +The test compares two compilation paths: + + - LEGACY: `USE_DECLARATIVE` patched to NOT contain + `mir.ColumnJoin`, so `_lower_inner_chain` and + `lower_scan_pipeline` fall into the legacy + `_lower_nested_cj_multi` / `_lower_root_cj_multi` branches + directly. + + - NEW: `USE_DECLARATIVE` left alone (`ColumnJoin` IS in the set; + `_should_use_declarative` no longer gates by source count + post B-CJ-multi), so the dispatchers route through + `lower_mir_cj_multi_in_chain` / `lower_mir_cj_multi_root`. + +Byte-equivalence is asserted on the rendered IIR text. The chain- +aware variants delegate to the same legacy helpers the legacy +branches called, so equivalence holds by construction; the test +pins it explicitly as the per-PR acceptance gate. + +CRITICAL — BlockGroup coexistence: when `ctx.bg_enabled` is True +or `pipeline[0]` is a `BlockGroupRoot` wrap op, the BG variant +(`_lower_root_cj_bg`) is the correct emission target — NOT the +standard `_lower_root_cj_multi`. The dispatch ordering in +`lower_scan_pipeline` keeps both BG paths intact (the +`BlockGroupRoot` unwrap sits BEFORE `_should_use_declarative`; +the bare multi-source CJ path with `ctx.bg_enabled=True` delegates +into `_lower_root_cj_multi`, whose first statement re-dispatches +to `_lower_root_cj_bg`). The BG fixtures here pin both routes. +''' + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any + +import pytest + +import srdatalog.ir.mir.types as mir +from srdatalog.ir.codegen.cuda.emit import EmitCtx, emit +from srdatalog.ir.dialects.iir.cf import Block as IirBlock +from srdatalog.ir.dialects.relation.sorted_array.lowerings import ( + LoweringCtx, + _lower_inner_chain, + _lower_nested_cj_multi, + _lower_root_cj_bg, + _lower_root_cj_multi, + _should_use_declarative, + _state_key, + lower_scan_pipeline, +) +from srdatalog.ir.dialects.relation.sorted_array.lowerings.lower_mir_cj_multi import ( + lower_mir_cj_multi, + lower_mir_cj_multi_in_chain, + lower_mir_cj_multi_root, +) +from srdatalog.ir.hir.types import Version + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + + +def _render(op: Any) -> str: + '''Render an IIR op tree through the CUDA emitter — the same + surface the byte-equivalence harnesses exercise.''' + return emit(op, EmitCtx(indent_level=2)) + + +@contextmanager +def _force_legacy_branch(): + '''Temporarily strip `mir.ColumnJoin` from `USE_DECLARATIVE` so + the dispatchers fall into the legacy multi-source CJ branches + (`_lower_nested_cj_multi` / `_lower_root_cj_multi`) directly. + + Save / restore as a context manager — `_should_use_declarative` + re-reads the (mutable) dialect-module attribute on each call, so + swapping the frozenset for the duration of one block is enough + to force the legacy path. + + Mirrors the same fixture in `test_lower_mir_cj_single_byte_equivalent.py`. + ''' + import srdatalog.ir.dialects.relation.sorted_array as sa_dialect + + saved = sa_dialect.USE_DECLARATIVE + sa_dialect.USE_DECLARATIVE = frozenset(saved - {mir.ColumnJoin}) + try: + yield + finally: + sa_dialect.USE_DECLARATIVE = saved + + +def _new_ctx(**kwargs: Any) -> LoweringCtx: + '''Fresh LoweringCtx — counters reset so both paths bump identically.''' + view_var_names = kwargs.pop('view_var_names', None) or { + '0': 'v_A_full', + '1': 'v_B_full', + '2': 'v_C_full', + '3': 'v_D_full', + } + return LoweringCtx(output_var='ctx0', view_var_names=view_var_names, **kwargs) + + +def _src( + rel: str, + handle_start: int, + *, + prefix_vars: tuple[str, ...] = (), + index: tuple[int, ...] = (0, 1), + version: Version = Version.FULL, +) -> mir.ColumnSource: + return mir.ColumnSource( + rel_name=rel, + version=version, + index=list(index), + prefix_vars=list(prefix_vars), + handle_start=handle_start, + ) + + +def _multi_src_cj( + *sources: mir.ColumnSource, + var_name: str = 'y', + handle_start: int = -1, +) -> mir.ColumnJoin: + return mir.ColumnJoin( + var_name=var_name, + sources=list(sources), + handle_start=handle_start, + ) + + +def _insert_into(vars_: tuple[str, ...] = ('y',)) -> mir.InsertInto: + return mir.InsertInto( + rel_name='Dst', + version=Version.NEW, + vars=list(vars_), + index=list(range(len(vars_))), + ) + + +# ----------------------------------------------------------------------------- +# 1. Mid-chain (nested) multi-source CJ — 2 sources, fresh +# ----------------------------------------------------------------------------- + + +def test_nested_multi_cj_2_sources_fresh_byte_equivalent(): + '''Nested multi-source CJ with two fresh-root sources (no prefix + narrowing). Both paths emit identical intersect_handles + + IntersectIter scaffold. + ''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + with _force_legacy_branch(): + legacy_text = _render(_lower_inner_chain([cj, ins], _new_ctx())) + new_text = _render(_lower_inner_chain([cj, ins], _new_ctx())) + + assert legacy_text == new_text + assert 'intersect_handles' in new_text + + +def test_nested_multi_cj_3_sources_byte_equivalent(): + '''Nested multi-source CJ with three sources stress-tests the + N-way intersect ordering. Both paths must match. + ''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), _src('C', 2), var_name='y') + ins = _insert_into(('y',)) + + with _force_legacy_branch(): + legacy_text = _render(_lower_inner_chain([cj, ins], _new_ctx())) + new_text = _render(_lower_inner_chain([cj, ins], _new_ctx())) + + assert legacy_text == new_text + + +# ----------------------------------------------------------------------------- +# 2. Mid-chain multi-source CJ with prefix-narrowed sources +# ----------------------------------------------------------------------------- + + +def test_nested_multi_cj_prefix_narrowed_byte_equivalent(): + '''Nested multi-source CJ where the sources carry `prefix_vars` + — each one aliases a parent handle from the enclosing scope. + Both paths emit the same `auto h = parent;` bindings. + ''' + cj = _multi_src_cj( + _src('A', 0, prefix_vars=('x',)), + _src('B', 1, prefix_vars=('x',)), + var_name='y', + ) + ins = _insert_into(('x', 'y')) + + parent_key_a = _state_key('A', [0, 1], ['x'], Version.FULL) + parent_key_b = _state_key('B', [0, 1], ['x'], Version.FULL) + + with _force_legacy_branch(): + ctx = _new_ctx() + ctx.handle_vars[parent_key_a] = 'parent_h_A' + ctx.handle_vars[parent_key_b] = 'parent_h_B' + ctx.bound_vars.append('x') + legacy_text = _render(_lower_inner_chain([cj, ins], ctx)) + + ctx = _new_ctx() + ctx.handle_vars[parent_key_a] = 'parent_h_A' + ctx.handle_vars[parent_key_b] = 'parent_h_B' + ctx.bound_vars.append('x') + new_text = _render(_lower_inner_chain([cj, ins], ctx)) + + assert legacy_text == new_text + assert 'parent_h_A' in new_text + assert 'parent_h_B' in new_text + + +# ----------------------------------------------------------------------------- +# 3. Interleaved Filter / ConstantBind around a nested multi-source CJ +# ----------------------------------------------------------------------------- + + +def test_nested_multi_cj_with_trailing_filter_byte_equivalent(): + '''Chain: nested multi-source CJ + Filter + InsertInto. Filter + body sits inside the CJ's iter loop.''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + filt = mir.Filter(vars=['y'], code='return y > 0;') + ins = _insert_into(('y',)) + + with _force_legacy_branch(): + legacy_text = _render(_lower_inner_chain([cj, filt, ins], _new_ctx())) + new_text = _render(_lower_inner_chain([cj, filt, ins], _new_ctx())) + + assert legacy_text == new_text + assert 'if (y > 0)' in new_text + + +def test_nested_multi_cj_with_trailing_constant_bind_byte_equivalent(): + '''Chain: nested multi-source CJ + ConstantBind + InsertInto.''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + cbind = mir.ConstantBind(var_name='yy', code='y + 1', deps=['y']) + ins = _insert_into(('y', 'yy')) + + with _force_legacy_branch(): + legacy_text = _render(_lower_inner_chain([cj, cbind, ins], _new_ctx())) + new_text = _render(_lower_inner_chain([cj, cbind, ins], _new_ctx())) + + assert legacy_text == new_text + assert 'auto yy = y + 1;' in new_text + + +# ----------------------------------------------------------------------------- +# 4. Root-position multi-source CJ (the production-load-bearing shape) +# ----------------------------------------------------------------------------- + + +def test_root_multi_cj_2_sources_byte_equivalent(): + '''Root-position multi-source CJ — every TC-like fixture uses + this shape. The new path routes via `lower_mir_cj_multi_root`, + which delegates to `_lower_root_cj_multi`; the legacy path + invokes `_lower_root_cj_multi` directly through + `lower_scan_pipeline`'s isinstance-cascade. + ''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + with _force_legacy_branch(): + legacy_text = _render(lower_scan_pipeline([cj, ins], _new_ctx())) + new_text = _render(lower_scan_pipeline([cj, ins], _new_ctx())) + + assert legacy_text == new_text + assert 'root_unique_values' in new_text + + +def test_root_multi_cj_3_sources_byte_equivalent(): + '''Root-position three-source CJ.''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), _src('C', 2), var_name='y') + ins = _insert_into(('y',)) + + with _force_legacy_branch(): + legacy_text = _render(lower_scan_pipeline([cj, ins], _new_ctx())) + new_text = _render(lower_scan_pipeline([cj, ins], _new_ctx())) + + assert legacy_text == new_text + + +# ----------------------------------------------------------------------------- +# 5. CJ under CJ: root-multi + nested-multi (deepest TC shape) +# ----------------------------------------------------------------------------- + + +def test_root_multi_cj_with_nested_multi_cj_byte_equivalent(): + '''Root multi-source CJ binding `y`, then a nested multi-source + CJ binding `z` with two fresh sources (no prefix narrowing). The + deepest shape `_lower_root_cj_multi` -> `_lower_inner_chain` -> + `_lower_nested_cj_multi` is exercised end-to-end. + + Fresh sources are used because root CJ only registers state keys + for its own sources (A, B), not for the deeper fresh ones (C, D). + Prefix-narrowed nested CJ would need parent handles registered by + matching root sources — the dedicated mid-chain prefix-narrowed + test exercises that path with pre-populated `ctx.handle_vars`. + ''' + root_cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + nested_cj = _multi_src_cj( + _src('C', 2), + _src('D', 3), + var_name='z', + ) + ins = _insert_into(('y', 'z')) + + with _force_legacy_branch(): + legacy_text = _render(lower_scan_pipeline([root_cj, nested_cj, ins], _new_ctx())) + new_text = _render(lower_scan_pipeline([root_cj, nested_cj, ins], _new_ctx())) + + assert legacy_text == new_text + assert 'root_unique_values' in new_text + assert 'intersect_handles' in new_text + + +# ----------------------------------------------------------------------------- +# 6. BlockGroup coexistence — bare multi CJ with ctx.bg_enabled +# ----------------------------------------------------------------------------- + + +def test_root_multi_cj_with_bg_enabled_routes_to_bg_variant_byte_equivalent(): + '''C3 dual-write path: a bare multi-source `mir.ColumnJoin` with + `ctx.bg_enabled=True` (the `ep.block_group=True` flag carried by + `compile_kernel_body`). Both paths reach `_lower_root_cj_multi` + whose first statement is `if ctx.bg_enabled: return _lower_root_cj_bg(...)`, + so both emit the BG scaffold (BgRootCjMulti). + + This is the load-bearing dual-write contract — the BG variant + fires before any non-BG scaffolding runs. + ''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + with _force_legacy_branch(): + ctx = _new_ctx() + ctx.bg_enabled = True + legacy_text = _render(lower_scan_pipeline([cj, ins], ctx)) + + ctx = _new_ctx() + ctx.bg_enabled = True + new_text = _render(lower_scan_pipeline([cj, ins], ctx)) + + assert legacy_text == new_text + # BG path emits the `BgRootCjMulti` scaffold (binary search + + # warp redistribution) — a marker that the non-BG `_lower_root_cj_multi` + # would NOT produce. + assert 'bg_key_idx' in new_text or 'block_group' in new_text.lower() + + +# ----------------------------------------------------------------------------- +# 7. BlockGroup coexistence — BlockGroupRoot wrap op (typed-pragma path) +# ----------------------------------------------------------------------------- + + +def test_block_group_root_wrap_routes_via_lower_scan_pipeline_byte_equivalent(): + '''C3 typed-pragma path: `pipeline[0]` is `BlockGroupRoot(inner= + multi-source ColumnJoin)`. `lower_scan_pipeline` recognizes the + wrap head BEFORE the `_should_use_declarative` check fires, slices + off `rest`, and dispatches into `_lower_root_cj_bg` directly with + `ctx.bg_enabled=True`. + + Both LEGACY and NEW must produce the same BG scaffold — the + BlockGroupRoot unwrap path is invariant under B-CJ-multi (we did + NOT touch the wrap dispatch in `lower_scan_pipeline`; it sits + above the new `_should_use_declarative` branch). + + Pins the load-bearing constraint that B-CJ-multi did NOT perturb + the typed-pragma BG dispatch. + ''' + inner_cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + wrap = mir.BlockGroupRoot(inner=inner_cj) + ins = _insert_into(('y',)) + + with _force_legacy_branch(): + legacy_text = _render(lower_scan_pipeline([wrap, ins], _new_ctx())) + new_text = _render(lower_scan_pipeline([wrap, ins], _new_ctx())) + + assert legacy_text == new_text + # BG scaffold marker (same as the dual-write test above). + assert 'bg_key_idx' in new_text or 'block_group' in new_text.lower() + + +# ----------------------------------------------------------------------------- +# 8. Full pipeline via compile_pipeline — root multi CJ + InsertInto +# ----------------------------------------------------------------------------- + + +def test_compile_pipeline_root_multi_cj_byte_equivalent(): + '''End-to-end smoke test: a multi-source root CJ + InsertInto + via `compile_pipeline`. Both paths produce identical CUDA text. + ''' + from srdatalog.compile import compile_pipeline + + s1 = _src('A', 0) + s2 = _src('B', 1) + cj = _multi_src_cj(s1, s2, var_name='y') + ins = mir.InsertInto( + rel_name='Dst', + version=Version.NEW, + vars=['y'], + index=[0], + ) + ep = mir.ExecutePipeline( + pipeline=[cj, ins], + source_specs=[s1, s2], + dest_specs=[ins], + rule_name='CjMultiRoot', + ) + + with _force_legacy_branch(): + legacy_out = compile_pipeline(ep) + new_out = compile_pipeline(ep) + + assert legacy_out == new_out + assert new_out + + +def test_compile_pipeline_root_multi_cj_with_filter_byte_equivalent(): + '''End-to-end: root multi-source CJ with a trailing Filter before + the InsertInto.''' + from srdatalog.compile import compile_pipeline + + s1 = _src('A', 0) + s2 = _src('B', 1) + cj = _multi_src_cj(s1, s2, var_name='y') + filt = mir.Filter(vars=['y'], code='return y > 0;') + ins = mir.InsertInto( + rel_name='Dst', + version=Version.NEW, + vars=['y'], + index=[0], + ) + ep = mir.ExecutePipeline( + pipeline=[cj, filt, ins], + source_specs=[s1, s2], + dest_specs=[ins], + rule_name='CjMultiRootFilter', + ) + + with _force_legacy_branch(): + legacy_out = compile_pipeline(ep) + new_out = compile_pipeline(ep) + + assert legacy_out == new_out + + +# ----------------------------------------------------------------------------- +# 9. Direct lowering call returns the expected IIR shape +# ----------------------------------------------------------------------------- + + +def test_lower_mir_cj_multi_in_chain_emits_block(): + '''Pin the IIR tree shape: nested multi-source CJ -> a `Block` + containing alias binds + IntersectIter (or a wrapping + D2lSegmentLoop for D2L FULL_VER sources). + ''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + out = lower_mir_cj_multi_in_chain(cj, [ins], _new_ctx()) + assert isinstance(out, IirBlock) + + +def test_lower_mir_cj_multi_root_emits_block(): + '''Pin the IIR tree shape: root multi-source CJ -> a `Block` + containing the ParallelFor + GridStrideLoop scaffold.''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + out = lower_mir_cj_multi_root(cj, [ins], _new_ctx()) + assert isinstance(out, IirBlock) + + +def test_lower_mir_cj_multi_in_chain_delegates_to_helper(): + '''The chain entry delegates to `_lower_nested_cj_multi` directly + — the rendered text matches calling the helper.''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + via_entry = _render(lower_mir_cj_multi_in_chain(cj, [ins], _new_ctx())) + via_helper = _render(_lower_nested_cj_multi(cj, [ins], _new_ctx())) + assert via_entry == via_helper + + +def test_lower_mir_cj_multi_root_delegates_to_helper(): + '''The root entry delegates to `_lower_root_cj_multi` directly.''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + via_entry = _render(lower_mir_cj_multi_root(cj, [ins], _new_ctx())) + via_helper = _render(_lower_root_cj_multi(cj, [ins], _new_ctx())) + assert via_entry == via_helper + + +def test_lower_mir_cj_multi_root_delegates_to_bg_when_enabled(): + '''When `ctx.bg_enabled=True`, the root entry's delegation flows + through `_lower_root_cj_multi`'s internal `if ctx.bg_enabled:` + dispatch into `_lower_root_cj_bg`. Pins the dual-write contract + from the entry surface. + ''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ins = _insert_into(('y',)) + + ctx_a = _new_ctx() + ctx_a.bg_enabled = True + via_entry = _render(lower_mir_cj_multi_root(cj, [ins], ctx_a)) + + ctx_b = _new_ctx() + ctx_b.bg_enabled = True + via_helper = _render(_lower_root_cj_bg(cj, [ins], ctx_b)) + + assert via_entry == via_helper + + +def test_lower_mir_cj_multi_in_chain_rejects_single_source(): + '''The chain entry's structural guard fires when a single-source + CJ slips through — pins that the source-count gate isn't only + enforced by the caller. Single-source CJ is owned by B-CJ-single + and must NOT route through this entry. + ''' + single_cj = _multi_src_cj(_src('A', 0), var_name='y') + ins = _insert_into(('y',)) + with pytest.raises(AssertionError, match=r'expected at least two ColumnSources'): + lower_mir_cj_multi_in_chain(single_cj, [ins], _new_ctx()) + + +def test_lower_mir_cj_multi_root_rejects_single_source(): + '''Mirror of `*_in_chain_rejects_single_source` for the root entry.''' + single_cj = _multi_src_cj(_src('A', 0), var_name='y') + ins = _insert_into(('y',)) + with pytest.raises(AssertionError, match=r'expected at least two ColumnSources'): + lower_mir_cj_multi_root(single_cj, [ins], _new_ctx()) + + +# ----------------------------------------------------------------------------- +# 10. Registry contract — stub asserts on direct call +# ----------------------------------------------------------------------------- + + +def test_lower_mir_cj_multi_registry_stub_asserts(): + '''The multi-source half of the `@lowering(target=iir.cf, source= + mir.ColumnJoin)` registry entry is a stub that asserts on direct + invocation — dispatch is expected to flow through `_lower_inner_chain` + / `lower_scan_pipeline` -> the chain-aware variant. + + Mirrors the B-CJ-single stub assertion. + ''' + cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + ctx = _new_ctx() + with pytest.raises(AssertionError, match=r'lower_mir_cj_multi_in_chain'): + lower_mir_cj_multi(cj, ctx) + + +def test_lower_mir_column_join_remains_a_single_registration(): + '''Pin that B-CJ-multi did NOT add a SECOND dialect-level + `@lowering(mir.ColumnJoin, ...)` registration. The single + registration covers both single-source and multi-source shapes + via a shape-aware dispatcher in `_register_passes`. + ''' + from srdatalog.ir.dialects.relation.sorted_array import DIALECT as SA_DIALECT + + matched = [low for low in SA_DIALECT.lowerings if low.matches is mir.ColumnJoin] + assert len(matched) == 1 + assert matched[0].consumes == ('mir',) + assert 'iir.cf' in matched[0].produces + + +# ----------------------------------------------------------------------------- +# 11. _should_use_declarative now accepts both shapes +# ----------------------------------------------------------------------------- + + +def test_should_use_declarative_multi_source_cj_returns_true(): + '''Post-B-CJ-multi the helper no longer gates `mir.ColumnJoin` + on source count — both shapes return True.''' + multi_cj = _multi_src_cj(_src('A', 0), _src('B', 1), var_name='y') + assert _should_use_declarative(multi_cj) is True + + +def test_should_use_declarative_single_source_cj_still_returns_true(): + '''Regression guard: single-source CJ stays in the new path.''' + single_cj = _multi_src_cj(_src('A', 0), var_name='z') + assert _should_use_declarative(single_cj) is True + + +# ----------------------------------------------------------------------------- +# 12. BG end-to-end via compile_kernel_body — load-bearing dual-write +# ----------------------------------------------------------------------------- + + +def test_compile_kernel_body_bg_enabled_byte_equivalent(): + '''The most load-bearing path for the C3 dual-write transition: + `compile_kernel_body(ep, bg_enabled=True)` threads `bg_enabled` + through to `LoweringCtx`, which the legacy `_lower_root_cj_multi` + consults to re-dispatch into `_lower_root_cj_bg`. Pins that the + new dispatch ordering preserves the legacy BG emission. + ''' + from srdatalog.compile import compile_kernel_body + + s1 = _src('R', 0) + s2 = _src('S', 1) + cj = _multi_src_cj(s1, s2, var_name='x') + ins = mir.InsertInto( + rel_name='Out', + version=Version.NEW, + vars=['x'], + index=[0], + ) + ep = mir.ExecutePipeline( + pipeline=[cj, ins], + source_specs=[s1, s2], + dest_specs=[ins], + rule_name='BgDualWrite', + block_group=True, + ) + + with _force_legacy_branch(): + legacy_text = compile_kernel_body( + ep, is_counting=False, bg_enabled=True, output_var_name='output_ctx_0' + ) + new_text = compile_kernel_body( + ep, is_counting=False, bg_enabled=True, output_var_name='output_ctx_0' + ) + + assert legacy_text == new_text + # BG-specific scaffolding marker — the BG path emits the + # `bg_*` work-balanced scaffold via `BgRootCjMulti`. + assert 'bg_key_idx' in new_text or 'cumulative_work' in new_text + + +# ----------------------------------------------------------------------------- +# 13. USE_DECLARATIVE invariants +# ----------------------------------------------------------------------------- + + +def test_use_declarative_contains_column_join_after_b_cj_multi(): + '''Pin the ratchet: `mir.ColumnJoin` remains in `USE_DECLARATIVE` + after B-CJ-multi (the dispatch surface widened, the ratchet + membership did not change). All prior Wave 2A migrations remain. + ''' + from srdatalog.ir.dialects.relation.sorted_array import USE_DECLARATIVE + + assert mir.ColumnJoin in USE_DECLARATIVE + assert mir.Filter in USE_DECLARATIVE + assert mir.ConstantBind in USE_DECLARATIVE + assert mir.InsertInto in USE_DECLARATIVE + assert mir.Scan in USE_DECLARATIVE + assert mir.Negation in USE_DECLARATIVE + assert mir.Aggregate in USE_DECLARATIVE diff --git a/tests/test_lower_mir_cj_single_byte_equivalent.py b/tests/test_lower_mir_cj_single_byte_equivalent.py index 83551c5..ac5565e 100644 --- a/tests/test_lower_mir_cj_single_byte_equivalent.py +++ b/tests/test_lower_mir_cj_single_byte_equivalent.py @@ -419,10 +419,18 @@ def test_should_use_declarative_single_source_cj_returns_true(): assert _should_use_declarative(cj) is True -def test_should_use_declarative_multi_source_cj_returns_false(): - '''The shape-gate rejects multi-source ColumnJoin — the legacy - `_lower_nested_cj_multi` branch retains ownership until - B-CJ-multi lands.''' +def test_should_use_declarative_multi_source_cj_returns_true_after_b_cj_multi(): + '''B-CJ-multi dropped the source-count guard in `_should_use_declarative`. + Multi-source ColumnJoin now routes through the new path + (`lower_mir_cj_multi_in_chain` mid-chain, `lower_mir_cj_multi_root` + at root) instead of falling through to the legacy + `_lower_nested_cj_multi` / `_lower_root_cj_multi` branches. + + Historical note: pre-B-CJ-multi this returned False — the helper + encoded a `len(sources) == 1` gate so only single-source CJ used + the new path. The new entries delegate to the same legacy helpers, + so the rendered IIR remains byte-equivalent. + ''' src_a = mir.ColumnSource( rel_name='A', version=Version.FULL, index=[0, 1], prefix_vars=[], handle_start=0 ) @@ -430,7 +438,7 @@ def test_should_use_declarative_multi_source_cj_returns_false(): rel_name='B', version=Version.FULL, index=[0, 1], prefix_vars=[], handle_start=1 ) multi_cj = mir.ColumnJoin(var_name='y', sources=[src_a, src_b]) - assert _should_use_declarative(multi_cj) is False + assert _should_use_declarative(multi_cj) is True def test_should_use_declarative_filter_returns_true(): @@ -469,13 +477,18 @@ def test_use_declarative_contains_column_join(): # ----------------------------------------------------------------------------- -def test_multi_source_cj_still_uses_legacy_branch(): - '''Regression guard: with `mir.ColumnJoin` in `USE_DECLARATIVE`, - the dispatch helper must still route multi-source CJ through the - legacy branch. The smoke test is the end-to-end pipeline - `_lower_inner_chain` call — if the new path were - incorrectly hit, the multi-source CJ would route to - `lower_mir_cj_single_in_chain` and fail the source-count guard. +def test_multi_source_cj_still_emits_intersect_handles(): + '''Regression guard: with `mir.ColumnJoin` in `USE_DECLARATIVE` AND + B-CJ-multi's `lower_mir_cj_multi_in_chain` wiring, multi-source CJ + still emits the `intersect_handles + iter` scaffold by delegating + to `_lower_nested_cj_multi`. The rendered text is byte-equivalent + to the legacy path. + + Pre-B-CJ-multi: this test exercised the "multi-source falls + through to the legacy branch" contract. Post-B-CJ-multi: multi- + source CJ routes through `lower_mir_cj_multi_in_chain`, which + delegates to `_lower_nested_cj_multi` — same emission, new + dispatch surface. ''' src_a = mir.ColumnSource( rel_name='A', version=Version.FULL, index=[0, 1], prefix_vars=[], handle_start=0 @@ -487,8 +500,8 @@ def test_multi_source_cj_still_uses_legacy_branch(): ins = _insert_into(('y',)) ctx = _new_ctx() - # No exception: dispatch goes through `_lower_nested_cj_multi`, - # which is the legacy path (B-CJ-multi territory). + # No exception: dispatch goes through `_lower_nested_cj_multi` + # (now via `lower_mir_cj_multi_in_chain`, post-B-CJ-multi). out = _lower_inner_chain([multi_cj, ins], ctx) rendered = _render(out) # The multi-source CJ emits an intersect_handles + iter scaffold —