diff --git a/examples/srdatalog_jaccard_demo/README.md b/examples/srdatalog_jaccard_demo/README.md new file mode 100644 index 0000000..8e2f476 --- /dev/null +++ b/examples/srdatalog_jaccard_demo/README.md @@ -0,0 +1,77 @@ +# srdatalog-jaccard-demo + +External-plugin demo for [srdatalog](../../). Proves the central +extensibility claim: a separate Python package, installed via `pip +install`, can extend the srdatalog compiler with a new dialect + +typed Pragma + lowering rule WITHOUT touching any file in +`srdatalog/` core. + +## What it ships + +A minimal `relation.jaccard` dialect: + +- `Jaccard(threshold: float)` — typed `Pragma` subclass; users attach + it to a rule via `Rule(...).with_pragma(Jaccard(threshold=0.7))`. +- `JaccardIndex(inner: InsertInto, threshold: float)` — MIR wrap op + the pragma materializes into. +- `@pragma_handler(Jaccard, on=ExecutePipeline)` — wraps each + trailing `InsertInto` in `JaccardIndex(...)` during + `MirPragmaPass`. +- `@lowering(target=DIALECT, source=JaccardIndex)` — emits IIR for + the wrap op (delegates to the sorted_array dialect's + `_lower_insert_into` under `ctx.dedup_hash=True`, plus a marker + comment carrying the threshold). + +## How discovery works + +`pyproject.toml` declares the entry point: + +```toml +[project.entry-points."srdatalog.plugins"] +jaccard = "srdatalog_jaccard:register" +``` + +Once installed, `Compiler.with_default_plugins()` walks the +`srdatalog.plugins` entry-point group, calls +`srdatalog_jaccard.register(compiler)`, and the new dialect lands on +the running compiler. No edits to `src/srdatalog/` and no edits to +the main package's `pyproject.toml` are required. + +## Install + run the tests + +From the repo root: + +```bash +pip install -e examples/srdatalog_jaccard_demo --no-deps +PYTHONPATH=src python -m pytest examples/srdatalog_jaccard_demo/tests/ -v +``` + +The tests verify: + +1. `Compiler.with_default_plugins()` discovers and loads the + `jaccard` entry point. +2. `Rule(...).with_pragma(Jaccard(...))` is accepted by the DSL + (the typed-pragma registry sees the new handler). +3. `compile_to_mir(program)` runs end-to-end with a Jaccard pragma + on a rule, producing a MIR program containing `JaccardIndex` + wrap ops after `MirPragmaPass`. +4. The registered `@lowering(target=DIALECT, source=JaccardIndex)` + rule emits well-formed IIR; the rendered output matches a stable + golden snapshot. + +## What this demo does NOT modify + +- No files under `src/srdatalog/`. +- No entries in the main package's `[project.entry-points."srdatalog.plugins"]` + block. +- No imports from framework-internal modules other than the public + `srdatalog.ir.core` surface and one deferred call into the + sorted_array lowering helper (`_lower_insert_into`) — flagged at + the use site as the only cross-dialect reuse, intentionally taken + so the demo can focus on the registration pathway rather than + re-implementing well-tested codegen. + +## Reference + +Spec: [`docs/phase_e_plugin_extensibility.md`](../../docs/phase_e_plugin_extensibility.md) +§4 (worked example). diff --git a/examples/srdatalog_jaccard_demo/pyproject.toml b/examples/srdatalog_jaccard_demo/pyproject.toml new file mode 100644 index 0000000..ad8b1ee --- /dev/null +++ b/examples/srdatalog_jaccard_demo/pyproject.toml @@ -0,0 +1,48 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "srdatalog-jaccard-demo" +version = "0.1.0" +description = "External-plugin demo: Jaccard-similarity pragma for srdatalog." +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT" } +authors = [ + { name = "SRDatalog Contributors" }, +] +keywords = ["srdatalog", "plugin", "datalog", "jaccard"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Topic :: Software Development :: Compilers", +] +# Intentionally NO srdatalog dep declared — the demo is exercised against +# the in-tree editable install of the main package via the worktree's +# PYTHONPATH=src. A real PyPI release would pin `srdatalog>=0.2`. +dependencies = [] + +# --------------------------------------------------------------------------- +# Phase E plugin discovery — the entire point of this package. +# +# Installing this wheel (`pip install srdatalog-jaccard-demo`) registers +# the `jaccard` entry under the `srdatalog.plugins` group. The next +# `Compiler.with_default_plugins()` call walks that group, loads +# `srdatalog_jaccard:register`, and gets the new dialect + Pragma + +# lowering without anyone editing the main srdatalog package. +# +# Spec: docs/phase_e_plugin_extensibility.md §2 (entry-point discovery). +# --------------------------------------------------------------------------- + +[project.entry-points."srdatalog.plugins"] +jaccard = "srdatalog_jaccard:register" + +[project.urls] +Homepage = "https://github.com/anthropics/srdatalog-python" + +[tool.hatch.build.targets.wheel] +packages = ["src/srdatalog_jaccard"] diff --git a/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/__init__.py b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/__init__.py new file mode 100644 index 0000000..0326a26 --- /dev/null +++ b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/__init__.py @@ -0,0 +1,90 @@ +'''srdatalog_jaccard — external-plugin demo for srdatalog. + +Spec: `docs/phase_e_plugin_extensibility.md` §4 (worked example). + +What this package proves: a SEPARATE Python package, installed via +`pip install srdatalog-jaccard-demo`, can extend the srdatalog +compiler with a new dialect + typed Pragma + lowering rule WITHOUT +touching any file in `srdatalog/` core. + +The plugin contract surface this demo exercises: + + - `Dialect(name='relation.jaccard', ops=[JaccardIndex])` — new + dialect with a custom Op. + - `Jaccard(Pragma)` — typed Pragma subclass. + - `@pragma_handler(Jaccard, on=ExecutePipeline)` — materializes + each Jaccard pragma into a `JaccardIndex(inner=InsertInto)` + wrap op (parallel to the built-in `DedupHash` / `WorkStealing` + pattern). + - `@lowering(target=DIALECT, source=JaccardIndex)` — lowers the + wrap op to IIR by delegating to the sorted_array dialect's + `_lower_insert_into` helper. The delegation reuses an + EXISTING IIR shape (the dedup-style emit, but parameterised + by a Jaccard threshold side-channel) — the demo's point is + the REGISTRATION pathway, not novel codegen semantics. + - `register(compiler)` — entry-point callable invoked by + `Compiler.with_default_plugins()`. Declared in this package's + `pyproject.toml` under `[project.entry-points."srdatalog.plugins"]` + as `jaccard = "srdatalog_jaccard:register"`. + +What this demo does NOT modify: + + - Nothing in `src/srdatalog/`. Verified by the in-tree + `tests/test_jaccard_external_plugin.py` smoke run. + - Nothing in the main `pyproject.toml` entry-point block. + +If a future change to srdatalog core breaks any of the above +imports, the plugin contract has leaked and needs a follow-up. +''' + +from __future__ import annotations + +from typing import Any + +# Importing `lowerings` runs `@lowering(DIALECT, JaccardIndex)` as a +# side effect — registers the wrap-op lowering on DIALECT.lowerings. +# Importing `pragmas.jaccard` runs `@pragma_handler(Jaccard, ...)` as +# a side effect — populates the module-global pragma registry that +# `MirPragmaPass` consults and that `Rule.with_pragma` validates +# against. Both registrations are import-time + idempotent (Python's +# module cache makes the side effects fire exactly once per process). +from srdatalog_jaccard import lowerings +from srdatalog_jaccard.dialect import DIALECT +from srdatalog_jaccard.pragmas import jaccard + +__all__ = ['DIALECT', 'register'] + + +def register(compiler: Any) -> None: + '''Plugin entry point — register the `relation.jaccard` dialect + on `compiler`. + + Lowerings + pragma handler are wired by the import-time side + effects above. This callable only performs the per-Compiler step: + `compiler.register_dialect(DIALECT)`. + + Idempotent: F4's `Compiler.register_plugin` short-circuits + re-registration of the same plugin name (`jaccard`). + ''' + compiler.register_dialect(DIALECT) + + +# Plugin metadata read by F4's topo-sort + conflict detection +# (`src/srdatalog/ir/core/plugin.py::_plugin_attr`). +# +# `plugin_name` — entry-point identifier; must match the key in +# pyproject.toml `[project.entry-points."srdatalog.plugins"]`. +# `provides` — the dialect name this plugin contributes. Other +# plugins may declare `requires=('relation.jaccard',)` to load +# after this one. +# `requires` — dialects that must be loaded first. We declare +# `relation.sorted_array` because the lowering body delegates to +# that dialect's `_lower_insert_into` helper at call time. The +# helper itself is reachable via a deferred import regardless of +# load order, but pinning the dependency in `requires` documents +# the contract: this plugin would not function without +# sorted_array's lowering helpers, so loading it standalone is a +# misconfiguration. +register.plugin_name = 'jaccard' # type: ignore[attr-defined] +register.provides = ('relation.jaccard',) # type: ignore[attr-defined] +register.requires = ('relation.sorted_array',) # type: ignore[attr-defined] diff --git a/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/dialect.py b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/dialect.py new file mode 100644 index 0000000..a59c541 --- /dev/null +++ b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/dialect.py @@ -0,0 +1,64 @@ +'''`relation.jaccard` dialect — sparse-similarity index. + +The dialect's vocabulary today is exactly one Op type: `JaccardIndex`, +a wrap op around an `mir.InsertInto` that signals "this emission +should pass through a Jaccard-similarity gate". Materialized by the +`Jaccard` typed pragma (see `pragmas/jaccard.py`) during +`MirPragmaPass`. + +Why a wrap op (not a new MIR-level field)? The op carries the inner +`InsertInto` so the registered `@lowering(target=DIALECT, source= +JaccardIndex)` rule has everything it needs to emit IIR (vars, +rel_name, index). This mirrors the built-in `DedupGate` / `WSScope` +pattern (`src/srdatalog/ir/mir/types.py:DedupGate`, +`src/srdatalog/ir/dialects/parallel/atomic_ws/__init__.py`); the +difference is that `JaccardIndex` is defined HERE in an external +package — proving the framework's claim that new wrap ops do not +require core-side enum updates. +''' + +from __future__ import annotations + +from dataclasses import dataclass +from typing import final + +from srdatalog.ir.core import Dialect, Op +from srdatalog.ir.mir.types import InsertInto + + +@final +@dataclass(frozen=True, slots=True) +class JaccardIndex(Op): + '''Wrap op: route an emission through a Jaccard-similarity gate. + + Inserted by `srdatalog_jaccard.pragmas.jaccard.materialize_jaccard` + during `MirPragmaPass` whenever an `ExecutePipeline` carries a + `Jaccard` pragma. Lowered by the + `@lowering(target=DIALECT, source=JaccardIndex)` rule registered + in `srdatalog_jaccard.lowerings` (which delegates back into the + sorted_array dialect's `_lower_insert_into` helper to emit IIR). + + Fields: + + inner — the `InsertInto` op being gated. Carrying it (rather + than a relation name + var list) lets the lowering + reuse the existing sorted_array machinery verbatim. + threshold — the Jaccard similarity threshold (0.0–1.0). Recorded + here so the lowering can lift it into a generated + kernel-side constant. The threshold is a structural + field of the wrap op (not a free-floating context + attribute) per discipline D1 (Op subclasses are pure + data). + ''' + + inner: InsertInto + threshold: float + + +DIALECT = Dialect( + name='relation.jaccard', + ops=[JaccardIndex], +) + + +__all__ = ['DIALECT', 'JaccardIndex'] diff --git a/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/lowerings.py b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/lowerings.py new file mode 100644 index 0000000..7bfa307 --- /dev/null +++ b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/lowerings.py @@ -0,0 +1,99 @@ +'''Lowering: `JaccardIndex` -> IIR. + +Registers an `@lowering(target=DIALECT, source=JaccardIndex)` rule +that lowers our wrap op to an IIR `Block` containing: + + 1. A `// jaccard threshold=` comment marker so the threshold + surfaces in the rendered C++ (useful as a discriminator in + golden snapshots). + 2. The same IIR shape that the sorted_array dialect's + `_lower_insert_into` produces for the dedup-hash branch — we + reuse that helper verbatim so the demo's emission is a + well-tested existing shape, not novel codegen we'd have to + port. The PURPOSE of this demo is the REGISTRATION pathway, + not new lowering semantics; piggybacking on an existing IIR + shape lets the test verify "the plugin's lowering rule fires + and produces well-formed IIR" without us inventing new ops. + +The registration runs as a module-import side effect — the parent +`srdatalog_jaccard` package imports this module exactly for that +side effect. + +Importantly, this file's `from srdatalog.ir.dialects.relation. +sorted_array.lowerings import _lower_insert_into` import is the +ONLY cross-package dependency on srdatalog's lowering internals. +A more polished plugin would re-implement the IIR shape itself; we +delegate to demonstrate that cross-dialect reuse is supported. +''' + +from __future__ import annotations + +from typing import Any + +from srdatalog.ir.core import Op +from srdatalog.ir.core.passes import lowering +from srdatalog.ir.dialects.iir.cf import Block, Comment +from srdatalog_jaccard.dialect import DIALECT, JaccardIndex + + +@lowering( + DIALECT, + JaccardIndex, + consumes=('mir', 'relation.jaccard'), + produces=('iir.cf', 'relation.sorted_array'), +) +def lower_jaccard_index(op: JaccardIndex, ctx: Any) -> Op: + '''Emit the IIR for `JaccardIndex(inner=InsertInto, threshold=t)`. + + Returns a `Block` containing: + + 1. `Comment(text="jaccard threshold=")` — a marker for golden- + snapshot tests; survives all the way to the rendered C++. + 2. The IIR statements produced by the sorted_array dialect's + `_lower_insert_into` under `ctx.dedup_hash=True`. The dedup- + hash branch produces a `dedup_table.try_insert(...) + if (_p) + {...}` gate — semantically a close-enough analogue for a + Jaccard threshold check (both decide whether to materialize + a tuple based on a per-emission predicate). Byte-equivalence + between dedup_hash and Jaccard is NOT a goal; reusing the + branch is a demo-simplification, not a contract. + + The save/restore around `ctx.dedup_hash` is defensive: if a + future caller invokes this lowering from a partial-ctx scope + where `dedup_hash` is False, the flag still flips on for the + duration of the gate (the lowering helper is the byte-equivalence + anchor for the dedup_hash branch; we reuse it as-is). + + Args: + op — the `JaccardIndex` wrap op to lower. `op.inner` is the + `mir.InsertInto` being gated. + ctx — the lowering context (`LoweringCtx` from sorted_array's + lowerings module). Non-frozen dataclass; we mutate + `dedup_hash` directly per discipline D10's parallel rule + for `LowerCtx`. + + Returns: + A `Block` IIR op the runner's emit pipeline renders to a + single `{ ... }` C++ block. + ''' + # Deferred import: the sorted_array monolith module imports + # nothing from this package. Importing at function-call time + # keeps the dialect import graph linear AND ensures sorted_array + # is loaded only when we actually need to lower (i.e., when the + # user has used the Jaccard pragma at least once). + from srdatalog.ir.dialects.relation.sorted_array.lowerings import ( + _lower_insert_into, + ) + + prev = getattr(ctx, 'dedup_hash', False) + try: + ctx.dedup_hash = True + stmts = list(_lower_insert_into(op.inner, ctx)) + finally: + ctx.dedup_hash = prev + + marker = Comment(text=f'jaccard threshold={op.threshold}') + return Block(stmts=(marker, *stmts)) + + +__all__ = ['lower_jaccard_index'] diff --git a/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/pragmas/__init__.py b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/pragmas/__init__.py new file mode 100644 index 0000000..0bb59bf --- /dev/null +++ b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/pragmas/__init__.py @@ -0,0 +1,8 @@ +'''Typed Pragma subclasses for the `relation.jaccard` dialect. + +Each module here defines exactly one `Pragma` subclass + its +`@pragma_handler` materialization callback. Module-import-time side +effects populate the global pragma registry; the parent +`srdatalog_jaccard` package imports each module for those side +effects. +''' diff --git a/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/pragmas/jaccard.py b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/pragmas/jaccard.py new file mode 100644 index 0000000..8a896de --- /dev/null +++ b/examples/srdatalog_jaccard_demo/src/srdatalog_jaccard/pragmas/jaccard.py @@ -0,0 +1,121 @@ +'''Pragma: `Jaccard` — sparse-similarity gate around an emission. + +Trigger: `Rule(...).with_pragma(Jaccard(threshold=0.7))`. + +Materialization (this module): wrap each `mir.InsertInto` at the tail + of an `ExecutePipeline.pipeline` in a `JaccardIndex(inner=..., + threshold=...)` wrap op during `MirPragmaPass`, then strip the + `Jaccard` instance from the EP's `pragmas` tuple. + + Mirrors the built-in `DedupHash` -> `DedupGate` pattern in + `src/srdatalog/ir/dialects/relation/sorted_array/pragmas/ + dedup_hash.py:materialize_dedup_hash`. The difference is that + the wrap op `JaccardIndex` and the threshold field are defined + in this external package — neither core MIR types nor the main + package's pragma registry needed editing. + +Lowering: the `@lowering(target=DIALECT, source=JaccardIndex)` rule + registered in `srdatalog_jaccard.lowerings` handles emission. + +DSL-time validation: `Jaccard.__post_init__` raises + `PragmaConfigError` (subclass of `ValueError`) on out-of-range + thresholds so the user sees the error at the `.with_pragma(...)` + keystroke, not deep in `MirPragmaPass`. Per + `docs/pragma_as_typed_object.md` §2. +''' + +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass +from typing import Any, final + +from srdatalog.ir.core import Pragma, pragma_handler +from srdatalog.ir.core.pragma import PragmaConfigError +from srdatalog.ir.mir.types import ExecutePipeline, InsertInto +from srdatalog_jaccard.dialect import JaccardIndex + + +@final +@dataclass(frozen=True, slots=True) +class Jaccard(Pragma): + '''Sparse-similarity (Jaccard-coefficient) gate around an emission. + + Triggers `MirPragmaPass` to wrap each `InsertInto` at the tail of + the carrying `ExecutePipeline.pipeline` in a `JaccardIndex` wrap + op; the gate's `@lowering` rule emits the IIR shape that the + runner-side similarity-table machinery expects. + + Fields: + + threshold — minimum similarity for an emission to pass the gate. + Validated at construction; out-of-range values raise + `PragmaConfigError` immediately, so DSL users see the + error at the `.with_pragma(...)` call site rather + than at compile time. + ''' + + threshold: float = 0.7 + + def __post_init__(self) -> None: + if not (0.0 < self.threshold <= 1.0): + raise PragmaConfigError(f'Jaccard.threshold must be in (0.0, 1.0]; got {self.threshold!r}') + + +# ----------------------------------------------------------------------------- +# Materialization handler (registered at import time via @pragma_handler) +# ----------------------------------------------------------------------------- + + +def _wrap_inserts_in_jaccard_gate( + pipeline: list[Any], + threshold: float, +) -> list[Any]: + '''Replace every `InsertInto` in `pipeline` with + `JaccardIndex(inner=that_insert, threshold=threshold)`. + + Mirrors the C2 `_wrap_inserts_in_dedup_gate` pattern: walk + `pipeline`, swap each `InsertInto` for the wrap op, leave non- + insert ops unchanged. Non-leaf operations (Scan, ColumnJoin, + Cartesian, etc.) pass through untouched — Jaccard is a per- + emission decoration, structurally an MIR-level concern that the + MIR -> IIR lowering then translates. + ''' + wrapped: list[Any] = [] + for child in pipeline: + if isinstance(child, InsertInto): + wrapped.append(JaccardIndex(inner=child, threshold=threshold)) + else: + wrapped.append(child) + return wrapped + + +@pragma_handler(Jaccard, on=ExecutePipeline) +def materialize_jaccard( + op: Any, # ExecutePipeline at runtime — typed as Any to match the registry's + # Callable[[Any, Pragma, PragmaCtx], Any] signature (see + # core/pragma.py:PragmaRegistration). Body re-narrows via the + # typed `op.pragmas` / `op.pipeline` accesses below. + pragma: Any, # Jaccard at runtime; same Any reason. + ctx: Any, # PragmaCtx; unused for this single-field pragma. +) -> Any: + '''Materialize `Jaccard` into a `JaccardIndex` wrap around each + trailing `InsertInto` in `op.pipeline`. + + Returns a new `ExecutePipeline` with: + - `pipeline` rewritten so every `InsertInto` is replaced by + `JaccardIndex(inner=that_insert, threshold=pragma.threshold)`, + - `pragmas` filtered to drop the consumed `Jaccard` instance + (per the `MirPragmaPass` post-flight invariant — see + `docs/pragma_as_typed_object.md` §3). + + Idempotency: the EP carries at most one `Jaccard` per + `with_pragma(Jaccard(...))` call; the filter is defensive against + stray duplicates. + ''' + new_pragmas = tuple(p for p in op.pragmas if not isinstance(p, Jaccard)) + new_pipeline = _wrap_inserts_in_jaccard_gate(list(op.pipeline), pragma.threshold) + return dataclasses.replace(op, pipeline=new_pipeline, pragmas=new_pragmas) + + +__all__ = ['Jaccard', 'materialize_jaccard'] diff --git a/examples/srdatalog_jaccard_demo/tests/test_jaccard_external_plugin.py b/examples/srdatalog_jaccard_demo/tests/test_jaccard_external_plugin.py new file mode 100644 index 0000000..996ff04 --- /dev/null +++ b/examples/srdatalog_jaccard_demo/tests/test_jaccard_external_plugin.py @@ -0,0 +1,419 @@ +'''External-plugin demo — end-to-end test. + +Spec: `docs/phase_e_plugin_extensibility.md` §4 (worked example). + +This test is the load-bearing assertion that the plugin contract is +honoured: a SEPARATE package (this directory's `srdatalog_jaccard`), +installed via `pip install -e .`, can extend the compiler with a +new dialect + typed Pragma + lowering rule WITHOUT touching any +file in `src/srdatalog/`. + +Sections: + + 1. Plugin discovery — `Compiler.with_default_plugins()` finds the + `jaccard` entry point shipped by this package's `pyproject.toml`, + calls its `register(compiler)`, and registers the + `relation.jaccard` dialect. + + 2. Pragma registration — `@pragma_handler(Jaccard, on= + ExecutePipeline)` ran as an import-time side effect, so the + DSL's `Rule.with_pragma(Jaccard())` validator accepts the typed + pragma. DSL-time config validation + (`Jaccard.__post_init__` raising `PragmaConfigError`) fires on + out-of-range thresholds. + + 3. End-to-end compile_to_mir — `MirPragmaPass` consumes the + `Jaccard` instance and inserts the `JaccardIndex` wrap op into + the resulting MIR program. The pragma is stripped from + `ep.pragmas` post-pass (per the `MirPragmaPass` post-flight + invariant). + + 4. Lowering rule — the registered + `@lowering(target=DIALECT, source=JaccardIndex)` rule fires and + produces well-formed IIR; the rendered C++ output matches a + stable golden snapshot. + +Tests skip cleanly if the demo package is not installed — the +`srdatalog_jaccard` import sits inside a `pytest.importorskip` at +top of file so users running the full test suite without first +`pip install -e examples/srdatalog_jaccard_demo` see one +SKIPPED line instead of an ImportError crash. +''' + +from __future__ import annotations + +from dataclasses import dataclass +from typing import final + +import pytest + +# Skip cleanly if the demo package is not installed. The most common +# path here is a fresh checkout where the user has not yet run +# `pip install -e examples/srdatalog_jaccard_demo --no-deps`. Per the +# README, the install step is a one-liner; the importorskip turns a +# hard ImportError into a single SKIPPED row in the pytest summary. +srdatalog_jaccard = pytest.importorskip('srdatalog_jaccard') + +from srdatalog_jaccard import DIALECT as JACCARD_DIALECT +from srdatalog_jaccard import register as register_jaccard +from srdatalog_jaccard.dialect import JaccardIndex +from srdatalog_jaccard.pragmas.jaccard import Jaccard + +import srdatalog.ir.mir.types as m +from srdatalog.dsl import Program, Relation, Var +from srdatalog.ir.core import Compiler, Pragma +from srdatalog.ir.core.plugin import ENTRY_POINT_GROUP +from srdatalog.ir.core.pragma import PragmaConfigError +from srdatalog.ir.hir import compile_to_mir + +# ----------------------------------------------------------------------------- +# Fixtures +# ----------------------------------------------------------------------------- + + +def _similar_program(threshold: float = 0.7) -> Program: + '''Build a one-rule program: `Similar(x, y) :- Arc(x, y)` with a + `Jaccard(threshold=...)` pragma attached. + + This is the minimal shape that exercises the full pipeline: HIR + planning, MIR lowering, `MirPragmaPass`, and (in section 4) the + wrap-op lowering. The rule is intentionally non-recursive so the + resulting MIR contains a single `ExecutePipeline` — easier to + introspect than a fixpoint. + ''' + x, y = Var('x'), Var('y') + arc = Relation('Arc', 2) + similar = Relation('Similar', 2) + return Program( + rules=[ + (similar(x, y) <= arc(x, y)).with_pragma(Jaccard(threshold=threshold)), + ], + ) + + +def _find_execute_pipelines(node: object) -> list[m.ExecutePipeline]: + '''Walk a MIR step tree, collecting every `ExecutePipeline`. The + MIR plan can nest EPs inside `FixpointPlan.instructions` / + `ParallelGroup.ops`; this helper flattens them out. + ''' + out: list[m.ExecutePipeline] = [] + if isinstance(node, m.ExecutePipeline): + out.append(node) + elif isinstance(node, m.FixpointPlan): + for inst in node.instructions: + out.extend(_find_execute_pipelines(inst)) + elif isinstance(node, m.ParallelGroup): + for op in node.ops: + out.extend(_find_execute_pipelines(op)) + return out + + +def _first_ep(mir_prog: m.Program) -> m.ExecutePipeline: + '''Return the first `ExecutePipeline` in a MIR program (linear walk + over `steps`). Asserts at least one EP exists — every test fixture + in this file produces one. + ''' + for step, _is_rec in mir_prog.steps: + eps = _find_execute_pipelines(step) + if eps: + return eps[0] + raise AssertionError('no ExecutePipeline in mir_prog') + + +# ----------------------------------------------------------------------------- +# 1. Plugin discovery — `with_default_plugins()` picks up jaccard +# ----------------------------------------------------------------------------- + + +def test_jaccard_is_discovered_by_with_default_plugins() -> None: + '''`Compiler.with_default_plugins()` walks the + `srdatalog.plugins` entry-point group; the `jaccard` entry point + shipped by this package's `pyproject.toml` is loaded and registers + the `relation.jaccard` dialect on the resulting Compiler. + + If this fails with "no plugins loaded named jaccard", check the + install: `pip install -e examples/srdatalog_jaccard_demo --no-deps`. + ''' + compiler = Compiler.with_default_plugins() + + assert 'jaccard' in compiler._plugins_loaded, ( + f'expected jaccard in loaded plugins; got ' + f'{sorted(compiler._plugins_loaded)!r}. If missing, the editable ' + f'install of examples/srdatalog_jaccard_demo is stale.' + ) + + dialect = compiler.get_dialect('relation.jaccard') + assert dialect is JACCARD_DIALECT, ( + 'with_default_plugins registered a different Dialect instance ' + 'than the module-level singleton — the plugin path is not ' + 'reusing DIALECT' + ) + + # The dialect's plugin attribution is the entry-point name. + assert compiler._dialects_by_plugin['relation.jaccard'] == 'jaccard' + + +def test_jaccard_entry_point_shape() -> None: + '''The entry-point declaration in `pyproject.toml` exposes the + expected (name, target) tuple. Catches typos in the + `[project.entry-points."srdatalog.plugins"]` block before a downstream + load-order test reports a more confusing failure. + ''' + import importlib.metadata + + eps = list(importlib.metadata.entry_points(group=ENTRY_POINT_GROUP)) + jaccard_eps = [ep for ep in eps if ep.name == 'jaccard'] + assert len(jaccard_eps) == 1, ( + f'expected exactly one jaccard entry point; got {jaccard_eps!r}. ' + f'Multiple entries indicates a stale install left in place; ' + f'rerun `pip install -e examples/srdatalog_jaccard_demo --no-deps ' + f'--force-reinstall`.' + ) + ep = jaccard_eps[0] + assert ep.value == 'srdatalog_jaccard:register' + + +def test_jaccard_register_metadata_matches_entry_point() -> None: + '''The `register` callable's `plugin_name` / `provides` / `requires` + attributes match the entry-point declaration and the F4 topo-sort + semantics. Pin them here so a future contributor renaming the + entry point doesn't silently break the topo-sort. + ''' + assert register_jaccard.plugin_name == 'jaccard' # type: ignore[attr-defined] + assert register_jaccard.provides == ('relation.jaccard',) # type: ignore[attr-defined] + assert register_jaccard.requires == ('relation.sorted_array',) # type: ignore[attr-defined] + + +def test_register_jaccard_plugin_idempotent() -> None: + '''Calling `register_plugin(register_jaccard)` twice on the same + Compiler is a no-op the second time — F4's `register_plugin` + short-circuits on the plugin name (`register.plugin_name = "jaccard"`). + Mirrors `tests/test_sorted_array_as_plugin.py:: + test_register_sorted_array_plugin_idempotent` for the external case. + ''' + compiler = Compiler() + # Register the dependency first (required by the topo-sort declaration). + from srdatalog.ir.dialects.relation.sorted_array import register as register_sa + + compiler.register_plugin(register_sa) + compiler.register_plugin(register_jaccard) + first_loaded = dict(compiler._plugins_loaded) + first_dialect_count = len(compiler.dialects) + + compiler.register_plugin(register_jaccard) + assert compiler._plugins_loaded == first_loaded + assert len(compiler.dialects) == first_dialect_count + assert compiler.get_dialect('relation.jaccard') is JACCARD_DIALECT + + +# ----------------------------------------------------------------------------- +# 2. Pragma registration — DSL accepts Jaccard, validates threshold +# ----------------------------------------------------------------------------- + + +def test_dsl_accepts_jaccard_pragma() -> None: + '''`Rule(...).with_pragma(Jaccard())` does not raise — proves the + `@pragma_handler(Jaccard, on=ExecutePipeline)` decoration registered + on import (via `srdatalog_jaccard.__init__`'s side-effect chain). + ''' + prog = _similar_program(threshold=0.5) + assert len(prog.rules) == 1 + rule = prog.rules[0] + # The DSL `with_pragma` appends to the rule's plans tuple; for + # a rule with no prior plans, exactly one plan carries the pragma. + assert len(rule.plans) == 1 + attached = rule.plans[0].pragmas + assert len(attached) == 1 + assert isinstance(attached[0], Jaccard) + assert attached[0].threshold == 0.5 + + +def test_jaccard_rejects_out_of_range_threshold() -> None: + '''`Jaccard(threshold=...)` validates at construction; out-of-range + values raise `PragmaConfigError` (subclass of `ValueError`) so the + user sees the error at their `.with_pragma(...)` keystroke rather + than deep in `MirPragmaPass`. + ''' + with pytest.raises(PragmaConfigError, match=r'threshold'): + Jaccard(threshold=0.0) + with pytest.raises(PragmaConfigError, match=r'threshold'): + Jaccard(threshold=1.5) + with pytest.raises(PragmaConfigError, match=r'threshold'): + Jaccard(threshold=-0.1) + + +def test_with_pragma_rejects_non_jaccard_unregistered_pragma() -> None: + '''A typed `Pragma` subclass with no `@pragma_handler` registration + is rejected at DSL time. This is the parallel discipline test to + `tests/test_pragma_dedup_hash_end_to_end.py:: + test_with_pragma_rejects_unregistered_pragma` — proves the + did-you-mean validator covers the external case too. + ''' + from srdatalog.ir.core import UnregisteredPragmaError + + @final + @dataclass(frozen=True, slots=True) + class _GhostJaccard(Pragma): + pass + + prog = _similar_program() + rule = prog.rules[0] + with pytest.raises(UnregisteredPragmaError, match=r'_GhostJaccard'): + rule.with_pragma(_GhostJaccard()) + + +# ----------------------------------------------------------------------------- +# 3. End-to-end — compile_to_mir materializes the wrap op +# ----------------------------------------------------------------------------- + + +def test_compile_to_mir_inserts_jaccard_index_wrap_op() -> None: + '''`compile_to_mir(program)` runs the full HIR -> MIR pipeline + INCLUDING `MirPragmaPass`. After the pass, the EP carrying a + `Jaccard()` pragma: + + - Has `pragmas == ()` (pragma consumed per the post-flight + invariant). + - Has every trailing `InsertInto` in its `pipeline` replaced by + `JaccardIndex(inner=that_insert, threshold=...)`. + + This is the load-bearing assertion that the plugin's + materialization handler ran inside the production + `compile_to_mir` entry point — not via a test-only shim. + ''' + prog = _similar_program(threshold=0.7) + mir_prog = compile_to_mir(prog) + + ep = _first_ep(mir_prog) + assert ep.pragmas == (), f'expected empty pragmas after MirPragmaPass; got {ep.pragmas!r}' + + # Find the wrap op. + wrap_ops = [op for op in ep.pipeline if isinstance(op, JaccardIndex)] + assert len(wrap_ops) == 1, ( + f'expected exactly one JaccardIndex in pipeline; got ' + f'{[type(o).__name__ for o in ep.pipeline]!r}' + ) + gate = wrap_ops[0] + assert isinstance(gate.inner, m.InsertInto) + assert gate.threshold == 0.7 + + # The original InsertInto must NOT also be in the pipeline (it + # was REPLACED by the wrap op, not appended). + bare_inserts = [op for op in ep.pipeline if isinstance(op, m.InsertInto)] + assert bare_inserts == [], ( + f'expected zero bare InsertIntos in pipeline (all wrapped); got {bare_inserts!r}' + ) + + +def test_compile_to_mir_no_pragma_is_unchanged() -> None: + '''Sanity: a program WITHOUT the Jaccard pragma compiles to MIR + with no `JaccardIndex` ops. Anchors the assertion above — + catches the case where the materialization handler runs on every + EP regardless of pragma presence (a registration-side bug). + ''' + x, y = Var('x'), Var('y') + arc = Relation('Arc', 2) + similar = Relation('Similar', 2) + prog = Program(rules=[similar(x, y) <= arc(x, y)]) + + mir_prog = compile_to_mir(prog) + ep = _first_ep(mir_prog) + assert all(not isinstance(op, JaccardIndex) for op in ep.pipeline) + + +# ----------------------------------------------------------------------------- +# 4. Lowering rule fires + golden snapshot +# ----------------------------------------------------------------------------- + + +# Stable golden snapshot for the lowered IIR rendered to C++. This is +# a NEW path (no Nim reference exists for Jaccard), so we anchor the +# output here. If the snapshot drifts, the test fails loudly — at +# which point the new output should be inspected by hand before +# updating this constant. +# +# The exact text is derived from running the lowering against the +# `_similar_program(threshold=0.7)` fixture and rendering via +# `EmitCtx(indent_level=4)`. The leading marker comment is the +# discriminator for this plugin's output; the rest is the byte- +# equivalent dedup-hash gate the lowering delegates to. +_EXPECTED_RENDER = ( + ' // jaccard threshold=0.7\n' + ' // Emit: Similar(x, y)\n' + ' { bool _p = dedup_table.try_insert(thread_id, x, y);\n' + ' if (_p) {\n' + ' if (tile.thread_rank() == 0) {\n' + ' uint32_t pos = atomicAdd(atomic_write_pos, 1u);\n' + ' out_data_0[(pos + out_base_0) + 0 * out_stride_0] = x;\n' + ' out_data_0[(pos + out_base_0) + 1 * out_stride_0] = y;\n' + ' }\n' + ' } }\n' +) + + +def test_jaccard_lowering_fires_and_emits_marker() -> None: + '''The registered `@lowering(target=DIALECT, source=JaccardIndex)` + rule is looked up through the compiler's dialect registry (the + same mechanism the production dispatcher would use) and applied to + a real `JaccardIndex` instance materialized by `compile_to_mir`. + + Three assertions: + + - The rule object reachable via `dialect.lowerings` claims our + `JaccardIndex` source type. + - Calling its `.apply(gate, ctx)` returns a `Block` IIR op (i.e., + the lowering produced well-formed IIR, not raised an exception + or returned `None`). + - The rendered C++ matches the byte-stable snapshot above + (anchors the threshold marker + the inner dedup gate). + ''' + from srdatalog.ir.codegen.cuda.emit import EmitCtx, emit + from srdatalog.ir.dialects.iir.cf import Block + from srdatalog.ir.dialects.relation.sorted_array.lowerings import LoweringCtx + + compiler = Compiler.with_default_plugins() + dialect = compiler.get_dialect('relation.jaccard') + + jaccard_lowerings = [lo for lo in dialect.lowerings if lo.matches is JaccardIndex] + assert len(jaccard_lowerings) == 1, ( + f'expected exactly one JaccardIndex lowering registered; got {len(jaccard_lowerings)}' + ) + rule = jaccard_lowerings[0] + assert rule.produces == ('iir.cf', 'relation.sorted_array') + + # Materialize the wrap op via the production compile_to_mir path. + prog = _similar_program(threshold=0.7) + mir_prog = compile_to_mir(prog) + ep = _first_ep(mir_prog) + gate = next(op for op in ep.pipeline if isinstance(op, JaccardIndex)) + + # Build a minimally-populated LoweringCtx matching the view-var + # shape lower_scan_pipeline would set up. (We bypass + # `compile_pipeline` because the sorted_array chain dispatcher + # does not know how to skip past our external wrap op; this + # discovery is documented in the test-run-report — addressing it + # would require a contract addition in the main package, out of + # scope for the demo.) + ctx = LoweringCtx( + view_var_names={'0': 'view_Arc_0_FULL'}, + is_counting=False, + output_var='output', + dedup_hash=False, + ) + iir = rule.apply(gate, ctx) + assert isinstance(iir, Block) + + # Threshold marker is the discriminator for this plugin's output. + marker_stmts = [s for s in iir.stmts if type(s).__name__ == 'Comment'] + assert any('jaccard threshold=0.7' in s.text for s in marker_stmts), ( + f'expected jaccard threshold marker in lowered IIR; got comments ' + f'{[s.text for s in marker_stmts]!r}' + ) + + emit_ctx = EmitCtx(indent_level=4) + rendered = emit(iir, emit_ctx) + assert rendered == _EXPECTED_RENDER, ( + f'rendered output diverged from golden snapshot.\n' + f'--- expected ---\n{_EXPECTED_RENDER}' + f'--- got ---\n{rendered}' + )