Skip to main content

earthsci_ast/
relational.rs

1//! Build-time relational engine — the five primitives the unified-IR
2//! value-invention pass (RFC `semiring-faq-unified-ir` §5.5, §6.1) runs **once at
3//! setup**, off the per-timestep hot path, to materialise the data-derived index
4//! sets and dense IDs that the numeric stencil then consumes:
5//!
6//! 1. [`distinct`]        — deduplicate tuples (unique mesh edges from face→vertex lists)
7//! 2. [`equijoin`]        — value-equality equi-join (connectivity inversion, *edges of cell i*;
8//!    also the `join.on` equality gate's candidate set, §5.5.8)
9//! 3. [`skolem`] / [`skolem_edge`] — deterministic content-addressed key from a tuple
10//! 4. [`rank`]            — dense integer renumbering of a distinct set
11//! 5. [`group_aggregate`] — group-by + associative/commutative semiring `⊕` (sum/min/max/…)
12//!
13//! # Determinism (the reason this module exists)
14//!
15//! `earthsci-ast` is **parallel native implementations** (Julia, Rust, Python)
16//! verified by a conformance suite, not one core behind FFI. So the hard problem is
17//! **bit-for-bit determinism across the bindings**: identical deduped sets,
18//! identical dense IDs, identical skolem keys. The governing principle
19//! (`CONFORMANCE_SPEC.md` §5.5 = RFC §5.7) is that *every emitted set, key, and
20//! dense ID is a pure function of a defined total order over tuples* — **no
21//! observable output may depend on hash-table iteration order or a language-native
22//! hash value** (Rust `HashMap`/`HashSet` use a per-instance random SipHash seed,
23//! so their iteration order is not portable).
24//!
25//! This is the Rust side of bead `ess-my4.3.4`; the Julia
26//! (`EarthSciAST.Relational`, `ess-my4.3.3`) and Python (`ess-my4.3.5`)
27//! bindings implement the same §5.5 contract, so all three produce byte-identical
28//! index sets and identical (base-normalised) dense-ID arrays. Concretely, per
29//! `CONFORMANCE_SPEC.md` §5.5.1:
30//!
31//! - **Total order** — lexicographic over tuple fields; integers by value; strings
32//!   by Unicode code-point (UTF-8 byte) order. Rust's derived [`Ord`] on [`Key`]
33//!   gives exactly this: `i64` by value, `String` by UTF-8 bytes (= code-point
34//!   order), and `Vec<Key>` lexicographically. **Floats are forbidden in keys**
35//!   (rule 1) — enforced *by construction*: [`Key`] has no float variant, and the
36//!   JSON boundary constructor [`Key::try_from_json`] rejects float components with
37//!   [`FloatKeyError`]. (Float *values* in a [`group_aggregate`] are allowed; see
38//!   [`Num`].)
39//! - **`distinct`** — sort by the total order, drop *adjacent* duplicates; output
40//!   order *is* sorted order, never first-seen / insertion order (rule 2).
41//! - **`rank`** — dense IDs by position in the sorted distinct sequence. Rust emits
42//!   **0-based** (`CONFORMANCE_SPEC.md` §5.5.1 rule 3 pins Rust = base 0 = the
43//!   canonical numbering); [`rank_with_base`] supports an arbitrary base for the
44//!   conformance base-pin round-trip (`canonical = reported − base`).
45//! - **`skolem`** — a canonical *tuple*, never a hash (rule 4): sort components for
46//!   a symmetric relation (undirected edge `(min,max)`), preserve order for a
47//!   directed one.
48//! - **`join` / group-by** — hashing may bucket only (we use [`indexmap::IndexMap`],
49//!   whose iteration order is insertion order, independent of the hasher); the
50//!   result is emitted **sorted by the canonical key** (rule 5). The semiring `⊕`
51//!   must be associative + commutative; for a floating-point `⊕` the per-bucket
52//!   reduction is done sequentially in canonical value order to avoid last-ULP
53//!   drift.
54//!
55//! # Implementation notes (RFC Appendix A.4, Rust)
56//!
57//! Built on `indexmap` (already a dep) `IndexMap` for bucketing + `sort_unstable`
58//! on the full tuple. The float-normalisation helper [`format_canonical_float`] is
59//! reused from [`crate::canonicalize`] for the skolem/distinct float-value
60//! serialisation base, and the sort-then-enumerate `rank` pattern mirrors
61//! `src/performance.rs`. `polars`/`datafusion`/`arrow` are rejected (heavy, out of
62//! proportion — RFC A.3/A.4); a non-portable fast hasher (`ahash`,
63//! `rustc-hash`/FxHash) MUST NEVER drive emitted order or keys (the join
64//! buckets below hash with it; their output order comes from a sort).
65
66use crate::canonicalize::format_canonical_float;
67use indexmap::IndexMap;
68use std::cmp::Ordering;
69
70/// An `IndexMap` on the fast hasher, for bucketing whose output order comes
71/// from a sort, never from the map.
72type FxIndexMap<K, V> = IndexMap<K, V, rustc_hash::FxBuildHasher>;
73
74/// Raised when a relational key contains a floating-point (or otherwise
75/// non-integer / out-of-domain) component, violating `CONFORMANCE_SPEC.md`
76/// §5.5.1 rule 1 ("floats are forbidden in keys"). Normalise the value to an
77/// integer / categorical ID before the build-time pre-pass.
78///
79/// In Rust rule 1 is enforced *by construction* — [`Key`] simply has no float
80/// variant — so this error only arises at the untyped JSON boundary
81/// ([`Key::try_from_json`]), the analogue of the Julia/Python reference's
82/// per-primitive `assert_key` runtime check.
83#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
84#[error("FloatKeyError: {0}")]
85pub struct FloatKeyError(pub String);
86
87/// A relational **key** component: an integer / categorical ID, or a tuple
88/// thereof. This is the orderable, float-free domain over which the five
89/// primitives operate.
90///
91/// The derived [`Ord`] *is* the `CONFORMANCE_SPEC.md` §5.5.1 rule-1 total order
92/// **within a type**: [`Key::Int`] by value, [`Key::Str`] by UTF-8 byte order
93/// (= Unicode code-point order, *not* locale collation), and [`Key::Tuple`]
94/// lexicographically. Cross-type comparison (e.g. an `Int` against a `Str`) falls
95/// back to the variant declaration order; it is deterministic but never exercised
96/// by a conformant index set, whose rows are homogeneous.
97///
98/// There is deliberately **no float variant** — that is how rule 1 ("floats are
99/// forbidden in keys") is enforced in a statically-typed binding. Float *values*
100/// for aggregation live in [`Num`], not here.
101#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
102pub enum Key {
103    /// An integer ID (vertex / cell index, scale 10⁴–10⁷).
104    Int(i64),
105    /// A categorical string ID. Ordered by UTF-8 byte / Unicode code-point order.
106    Str(String),
107    /// A boolean ID (boolean-or keys). Ordered `false < true`.
108    Bool(bool),
109    /// A tuple of components (e.g. a mesh edge `(v_lo, v_hi)`). Ordered
110    /// lexicographically.
111    Tuple(Vec<Key>),
112}
113
114impl Key {
115    /// Construct a [`Key`] from an untyped JSON value, **rejecting floats**
116    /// (rule 1) at the boundary — the analogue of the Julia/Python reference's
117    /// `assert_key`. A JSON number with a fractional part or exponent (e.g.
118    /// `1.5`, `1.0`) is a float and raises [`FloatKeyError`]; a bare integer
119    /// (`1`, `-3`) becomes [`Key::Int`]. `null` and objects are not valid keys.
120    pub fn try_from_json(value: &serde_json::Value) -> Result<Key, FloatKeyError> {
121        match value {
122            serde_json::Value::Bool(b) => Ok(Key::Bool(*b)),
123            serde_json::Value::String(s) => Ok(Key::Str(s.clone())),
124            serde_json::Value::Number(n) => {
125                if let Some(i) = n.as_i64() {
126                    Ok(Key::Int(i))
127                } else if let Some(u) = n.as_u64() {
128                    Err(FloatKeyError(format!(
129                        "integer {u} out of range for a relational key (i64 domain)"
130                    )))
131                } else {
132                    Err(FloatKeyError(format!(
133                        "float {n} forbidden in a relational key; keys must be \
134                         integer / categorical IDs (CONFORMANCE_SPEC.md §5.5.1 \
135                         rule 1) — normalise to an ID before the build-time \
136                         relational pre-pass"
137                    )))
138                }
139            }
140            serde_json::Value::Array(items) => {
141                let mut out = Vec::with_capacity(items.len());
142                for item in items {
143                    out.push(Key::try_from_json(item)?);
144                }
145                Ok(Key::Tuple(out))
146            }
147            serde_json::Value::Null | serde_json::Value::Object(_) => Err(FloatKeyError(format!(
148                "unsupported relational key component {value}; keys must be \
149                 integer / categorical IDs or tuples thereof"
150            ))),
151        }
152    }
153
154    /// Append this key's canonical-JSON token to `out` (`CONFORMANCE_SPEC.md`
155    /// §5.5.3): integers as bare digits, booleans as `true`/`false`, strings as
156    /// JSON-escaped quoted strings (UTF-8, no `\uXXXX` escaping of non-ASCII),
157    /// tuples as nested arrays. Compact — no spaces.
158    fn write_token(&self, out: &mut String) {
159        match self {
160            Key::Int(i) => out.push_str(&i.to_string()),
161            Key::Bool(b) => out.push_str(if *b { "true" } else { "false" }),
162            // serde_json gives exactly the canonical JSON string escaping
163            // (escape `"`, `\`, and control chars; raw UTF-8 otherwise), matching
164            // the Julia `JSON3.write` and Python `json.dumps(ensure_ascii=False)`
165            // reference serialisers byte-for-byte.
166            Key::Str(s) => {
167                out.push_str(&serde_json::to_string(s).expect("a string always serialises to JSON"))
168            }
169            Key::Tuple(items) => write_array(out, items, |o, k| k.write_token(o)),
170        }
171    }
172}
173
174/// A numeric **value** for [`group_aggregate`]. Unlike [`Key`], values may be
175/// floating-point (only *keys* are float-forbidden, rule 1). Integer aggregates
176/// stay integer (so a sum serialises as `5`, not `5.0`); a float anywhere in a
177/// bucket promotes the reduction to `f64`.
178#[derive(Debug, Clone, Copy, PartialEq)]
179pub enum Num {
180    /// An exact integer value.
181    Int(i64),
182    /// A floating-point value.
183    Float(f64),
184}
185
186impl Num {
187    fn as_f64(self) -> f64 {
188        match self {
189            Num::Int(i) => i as f64,
190            Num::Float(f) => f,
191        }
192    }
193
194    /// Total order over values, used only to put each bucket's values into a
195    /// canonical order before the sequential reduction (rule 5). `Int` sorts
196    /// before `Float` (the two never mix in a conformant bucket); floats use
197    /// [`f64::total_cmp`] so `-0.0`/`NaN` are still totally ordered.
198    fn canonical_cmp(&self, other: &Num) -> Ordering {
199        match (self, other) {
200            (Num::Int(a), Num::Int(b)) => a.cmp(b),
201            (Num::Float(a), Num::Float(b)) => a.total_cmp(b),
202            (Num::Int(_), Num::Float(_)) => Ordering::Less,
203            (Num::Float(_), Num::Int(_)) => Ordering::Greater,
204        }
205    }
206
207    /// Append this value's canonical-JSON token to `out`. Integers are bare
208    /// digits; floats use [`format_canonical_float`] (the RFC §5.4.6 form, reused
209    /// from [`crate::canonicalize`]).
210    fn write_token(&self, out: &mut String) {
211        match self {
212            Num::Int(i) => out.push_str(&i.to_string()),
213            Num::Float(f) => out.push_str(&format_canonical_float(*f)),
214        }
215    }
216}
217
218/// An associative + commutative semiring combiner for [`group_aggregate`]
219/// (rule 5). Every registry `⊕` is one of these, so input and bucket order can
220/// never change the result.
221#[derive(Debug, Clone, Copy, PartialEq, Eq)]
222pub enum SemiringOp {
223    /// Sum (`+`).
224    Sum,
225    /// Product (`*`).
226    Prod,
227    /// Minimum.
228    Min,
229    /// Maximum.
230    Max,
231}
232
233impl SemiringOp {
234    /// Combine two values. Integer inputs stay integer (with wrapping `i64`
235    /// arithmetic, matching Julia's native `Int64`); if either input is a float
236    /// the reduction promotes to `f64`.
237    fn apply(self, a: Num, b: Num) -> Num {
238        match (a, b) {
239            (Num::Int(x), Num::Int(y)) => Num::Int(match self {
240                SemiringOp::Sum => x.wrapping_add(y),
241                SemiringOp::Prod => x.wrapping_mul(y),
242                SemiringOp::Min => x.min(y),
243                SemiringOp::Max => x.max(y),
244            }),
245            _ => {
246                // The float arm is a reduction's ⊕ and therefore one of the
247                // evaluator's operations: under `element_type: "Float32"`
248                // (esm-spec §11.3.1) it rounds per accumulation step, so an
249                // N-row group sums with N roundings rather than in a binary64
250                // accumulator that would be more accurate than the binary32
251                // reference. Identity under Float64, and the INTEGER arm above
252                // is exact and untouched. `round(x op y)` on binary32 operands
253                // is the binary32 operation (binary64 carries 53 ≥ 2·24+2 bits).
254                let prec = crate::precision::active();
255                let (x, y) = (prec.round(a.as_f64()), prec.round(b.as_f64()));
256                Num::Float(match self {
257                    SemiringOp::Sum => prec.round(x + y),
258                    SemiringOp::Prod => prec.round(x * y),
259                    SemiringOp::Min => x.min(y),
260                    SemiringOp::Max => x.max(y),
261                })
262            }
263        }
264    }
265}
266
267// ── Primitive 3: skolem (canonical-tuple content-addressed key) ─────────────
268
269/// Canonical key for an **undirected** pair (a symmetric relation):
270/// `(min(a,b), max(a,b))`. The deterministic, content-addressed identity of a
271/// mesh edge (RFC §5.5 generalises ESI `pack`). It is **not** a hash (rule 4) —
272/// the tuple itself is the key, so the dense ID later assigned by [`rank`] is
273/// reproducible across bindings. Equal components are preserved (`(7,7)`).
274pub fn skolem_edge(a: Key, b: Key) -> Key {
275    if a <= b {
276        Key::Tuple(vec![a, b])
277    } else {
278        Key::Tuple(vec![b, a])
279    }
280}
281
282/// Canonical-tuple Skolem key (rule 4). For a `symmetric` relation the
283/// components are sorted (generalising [`skolem_edge`] to arity > 2); for a
284/// directed relation the order is preserved, so `(1, 2)` and `(2, 1)` stay
285/// distinct. Never a native hash — the tuple *is* the content-addressed key; the
286/// dense ID then comes from [`rank`].
287pub fn skolem(mut components: Vec<Key>, symmetric: bool) -> Key {
288    if symmetric {
289        components.sort_unstable();
290    }
291    Key::Tuple(components)
292}
293
294// ── Primitive 1: distinct (sort + drop adjacent duplicates) ─────────────────
295
296/// Set semantics over `rows`: sort by the §5.5.1 total order, then drop
297/// **adjacent** duplicates (rule 2). The returned order **is** the sorted order
298/// — never first-seen / insertion order. A pure function of the input multiset,
299/// so duplicate, reversed, and permuted inputs all collapse to the identical
300/// output.
301///
302/// Mirrors the DuckDB oracle `SELECT DISTINCT … ORDER BY …`.
303pub fn distinct(rows: &[Key]) -> Vec<Key> {
304    let mut v = rows.to_vec();
305    v.sort_unstable();
306    v.dedup(); // removes *consecutive* equal elements ⇒ adjacent dedup after sort
307    v
308}
309
310// ── Primitive 2: equijoin (inner value-equality join) ───────────────────────
311
312/// Inner equi-join of two key columns (rule 5): every `(l, r)` position pair
313/// whose keys compare equal under the §5.5.1 total order.
314///
315/// This is the kernel behind the **`join.on` equality gate** — the value-equality
316/// mirror of the `join.overlap` broad phase (`CONFORMANCE_SPEC.md` §5.5.6 /
317/// §5.5.8). It is deliberately the *only* hashing path for value-equality
318/// matching in this crate: the dense evaluator's gate
319/// (`crate::simulate_array::eval`) and the value-invention producer both build
320/// their candidate pair set here rather than rolling a second one.
321///
322/// **Semantics (RFC `semiring-faq-unified-ir` §5.3).**
323/// - **Inner only.** A left or right position whose key matches nothing appears
324///   in no pair; downstream it therefore contributes the additive identity `0̄`.
325/// - **Many-to-many is defined.** A key occurring `m` times left and `n` times
326///   right yields all `m·n` pairs — categorical disaggregation, not an error.
327/// - **Exact equality only.** [`Key`] has no float variant, so rule 1 ("floats
328///   are forbidden in keys") holds by construction; composite (multi-column)
329///   keys are [`Key::Tuple`]s built with [`skolem`].
330///
331/// **Determinism (rule 5).** Hashing buckets ONLY: the emitted pair list is
332/// ordered by the **canonical key** first (§5.5.1 total order over the matched
333/// keys), then by left position, then by right position — never by
334/// hash-iteration or first-seen order. Duplicate, reversed and permuted inputs
335/// therefore give byte-identical output. (A consumer that needs
336/// position-ascending drive order — the dense contraction odometer does, to stay
337/// an order-preserving subsequence of the full product — sorts the pairs itself;
338/// both orders are pure functions of the input.)
339///
340/// Cost is `O(|left| + |right| + |matches|)` plus one `O(K log K)` sort over the
341/// `K` matched distinct keys — never the `O(|left|·|right|)` product.
342///
343/// Mirrors the DuckDB oracle
344/// `SELECT l.pos, r.pos FROM l JOIN r USING (key) ORDER BY key, l.pos, r.pos`.
345pub fn equijoin(left: &[Key], right: &[Key]) -> Vec<(usize, usize)> {
346    // Bucket the RIGHT side by key (IndexMap: iteration order is insertion
347    // order, independent of the hasher — and we never iterate it anyway; the
348    // emitted order comes from the sorted key list below, so the fast hasher
349    // decides nothing about the output).
350    let mut buckets: FxIndexMap<&Key, Vec<usize>> =
351        FxIndexMap::with_capacity_and_hasher(right.len(), Default::default());
352    for (j, k) in right.iter().enumerate() {
353        buckets.entry(k).or_default().push(j);
354    }
355    // Probe with the LEFT side, collecting the matched left positions per key.
356    let mut matched: FxIndexMap<&Key, Vec<usize>> = FxIndexMap::default();
357    for (i, k) in left.iter().enumerate() {
358        if buckets.contains_key(k) {
359            matched.entry(k).or_default().push(i);
360        }
361    }
362    // Emit in canonical key order (rule 5), then left position, then right.
363    let mut keys: Vec<&Key> = matched.keys().copied().collect();
364    keys.sort_unstable();
365    let total: usize = keys
366        .iter()
367        .map(|k| matched[*k].len() * buckets[*k].len())
368        .sum();
369    let mut out = Vec::with_capacity(total);
370    for k in keys {
371        // Both position lists were collected by ascending index, so the nested
372        // walk is already (l asc, r asc) within the key.
373        for &i in &matched[k] {
374            for &j in &buckets[k] {
375                out.push((i, j));
376            }
377        }
378    }
379    out
380}
381
382/// How many pairs [`equijoin`] WOULD emit, without emitting them.
383///
384/// `O(|left| + |right|)` in both time and memory — it buckets the right side by
385/// key exactly as [`equijoin`] does, then sums each left key's bucket size
386/// instead of writing out the cross product within the key. The count is the
387/// same number [`equijoin`] reserves its output with.
388///
389/// This exists so a planner can PRICE a gate before paying for it. A match set
390/// costs memory proportional to its size and a join whose key is
391/// low-cardinality can match hundreds of millions of pairs from tables of a few
392/// hundred thousand rows; knowing that before materialising is the difference
393/// between declining the gate and spending gigabytes on one that another clause
394/// has already made redundant.
395#[must_use]
396pub(crate) fn equijoin_match_count(left: &[Key], right: &[Key]) -> usize {
397    let mut buckets: FxIndexMap<&Key, usize> =
398        FxIndexMap::with_capacity_and_hasher(right.len(), Default::default());
399    for k in right {
400        *buckets.entry(k).or_default() += 1;
401    }
402    left.iter()
403        .filter_map(|k| buckets.get(k))
404        .fold(0usize, |acc, n| acc.saturating_add(*n))
405}
406
407// ── Primitive 4: rank (dense integer renumbering) ───────────────────────────
408
409/// Result of [`rank`].
410///
411/// - `order` — the distinct tuples in §5.5.1 total order.
412/// - `id` — `id[t]` is the dense integer assigned to tuple `t` (insertion order
413///   equals sorted order, since it is filled in `order` sequence).
414/// - `base` — the emission base. Rust emits **0-based** (`CONFORMANCE_SPEC.md`
415///   §5.5.1 rule 3, the canonical numbering); the conformance adapter recovers
416///   the canonical 0-based ID via `canonical = reported − base`.
417#[derive(Debug, Clone)]
418pub struct Ranking {
419    /// The distinct tuples in total order.
420    pub order: Vec<Key>,
421    /// Map from tuple to its dense integer ID.
422    pub id: IndexMap<Key, i64>,
423    /// The emission base (Rust native = 0).
424    pub base: i64,
425}
426
427impl Ranking {
428    /// The emitted dense IDs in `order` sequence: `[base, base+1, …]`.
429    pub fn dense_ids(&self) -> Vec<i64> {
430        (0..self.order.len() as i64)
431            .map(|i| i + self.base)
432            .collect()
433    }
434
435    /// The canonical 0-based dense IDs (`reported − base`): always `[0, 1, …]`.
436    /// This is what the cross-binding conformance suite asserts on.
437    pub fn canonical_dense_ids(&self) -> Vec<i64> {
438        (0..self.order.len() as i64).collect()
439    }
440}
441
442/// Dense integer renumbering (rule 3) at Rust's native **0-based** emission:
443/// assign IDs by position in the sorted [`distinct`] sequence. Equivalent to SQL
444/// `dense_rank() OVER (ORDER BY …)` over the deduplicated rows.
445pub fn rank(rows: &[Key]) -> Ranking {
446    rank_with_base(rows, 0)
447}
448
449/// [`rank`] with an explicit emission `base`. `base = 0` is the canonical
450/// numbering the conformance suite asserts on; `base = 1` mirrors Julia's native
451/// 1-based emission for the base-pin round-trip test.
452pub fn rank_with_base(rows: &[Key], base: i64) -> Ranking {
453    let order = distinct(rows);
454    let mut id = IndexMap::with_capacity(order.len());
455    for (i, t) in order.iter().enumerate() {
456        id.insert(t.clone(), i as i64 + base);
457    }
458    Ranking { order, id, base }
459}
460
461// ── Primitive 5: group-by + semiring aggregate ──────────────────────────────
462
463/// Group-by + semiring aggregate (rule 5). Bucket `rows` by their [`Key`]
464/// (hashing only to bucket, via [`IndexMap`]), combine each group's [`Num`]
465/// values with the semiring `op` (`⊕`), and emit `(key, aggregate)` pairs
466/// **sorted by the canonical key**.
467///
468/// `op` is associative + commutative (every registry `⊕`), so the result is
469/// independent of input / bucket order. The per-bucket reduction is a **left
470/// fold over the values in canonical (sorted) order** so that a **floating-point**
471/// `op` produces a reproducible last-ULP result (rule 5); the integer path uses
472/// the same canonical order (immaterial there, but one code path).
473///
474/// Mirrors the DuckDB oracle `SELECT key, ⊕(value) … GROUP BY key ORDER BY key`.
475pub fn group_aggregate(rows: &[(Key, Num)], op: SemiringOp) -> Vec<(Key, Num)> {
476    let mut buckets: IndexMap<Key, Vec<Num>> = IndexMap::new();
477    for (k, v) in rows {
478        buckets.entry(k.clone()).or_default().push(*v);
479    }
480    let mut keys: Vec<Key> = buckets.keys().cloned().collect();
481    keys.sort_unstable(); // canonical key order — never IndexMap bucket order
482    let mut out = Vec::with_capacity(keys.len());
483    for key in keys {
484        let mut vals = buckets.swap_remove(&key).expect("key came from the map");
485        vals.sort_by(Num::canonical_cmp); // canonical value order ⇒ reproducible float ⊕
486        let agg = vals
487            .into_iter()
488            .reduce(|acc, x| op.apply(acc, x))
489            .expect("each bucket has ≥ 1 value");
490        out.push((key, agg));
491    }
492    out
493}
494
495// ── Canonical serialization (CONFORMANCE_SPEC.md §5.5.3) ─────────────────────
496
497/// Canonical byte form of an index set (`CONFORMANCE_SPEC.md` §5.5.3): the
498/// [`distinct`] rows, each tuple serialised as a JSON array, in §5.5.1 sorted
499/// order, as **compact JSON** (`,` / `:` separators, no spaces, UTF-8, no
500/// `\uXXXX` escaping). Two conforming bindings MUST produce byte-for-byte
501/// identical output for the same input multiset.
502///
503/// This is the artifact the adversarial conformance harness (§5.5.4) compares
504/// byte-for-byte across duplicate / reversed / permuted inputs.
505pub fn canonical_index_set_json(rows: &[Key]) -> String {
506    serialize_keys(&distinct(rows))
507}
508
509/// Serialise an already-distinct, already-sorted slice of keys to the canonical
510/// compact-JSON index-set form (no further dedup/sort).
511pub fn serialize_keys(rows: &[Key]) -> String {
512    let mut out = String::new();
513    write_array(&mut out, rows, |o, k| k.write_token(o));
514    out
515}
516
517/// Serialise the `(key, aggregate)` output of [`group_aggregate`] to the
518/// canonical compact-JSON form, each pair as a two-element array `[key, value]`.
519pub fn serialize_pairs(rows: &[(Key, Num)]) -> String {
520    let mut out = String::new();
521    write_array(&mut out, rows, |o, (k, v)| {
522        o.push('[');
523        k.write_token(o);
524        o.push(',');
525        v.write_token(o);
526        o.push(']');
527    });
528    out
529}
530
531/// Write `items` as a compact JSON array `[e0,e1,…]` using `emit` per element.
532fn write_array<T>(out: &mut String, items: &[T], emit: impl Fn(&mut String, &T)) {
533    out.push('[');
534    for (i, item) in items.iter().enumerate() {
535        if i > 0 {
536            out.push(',');
537        }
538        emit(out, item);
539    }
540    out.push(']');
541}
542
543#[cfg(test)]
544mod tests {
545    use super::*;
546
547    // ── concise constructors for tuples/keys in tests ───────────────────────
548    fn ki(i: i64) -> Key {
549        Key::Int(i)
550    }
551    fn ks(s: &str) -> Key {
552        Key::Str(s.to_string())
553    }
554    fn t2(a: i64, b: i64) -> Key {
555        Key::Tuple(vec![ki(a), ki(b)])
556    }
557    fn tup(items: Vec<Key>) -> Key {
558        Key::Tuple(items)
559    }
560    /// Directed edges of each face (consecutive vertices, with wraparound),
561    /// canonicalised to undirected `(min,max)` skolem keys — the mesh-edge
562    /// enumeration producer (matches the harness `directed_edges_from_faces` +
563    /// undirected `skolem`).
564    fn undirected_edges(faces: &[Vec<i64>]) -> Vec<Key> {
565        let mut edges = Vec::new();
566        for face in faces {
567            let n = face.len();
568            for i in 0..n {
569                edges.push(skolem_edge(ki(face[i]), ki(face[(i + 1) % n])));
570            }
571        }
572        edges
573    }
574
575    // ── Primitive 1: distinct ───────────────────────────────────────────────
576    #[test]
577    fn distinct_sorts_and_dedups() {
578        // Integers: output is SORTED, not first-seen order [3,1,2] (rule 2 /
579        // the `first_seen_order` negative control).
580        assert_eq!(
581            distinct(&[ki(3), ki(1), ki(2), ki(1)]),
582            vec![ki(1), ki(2), ki(3)]
583        );
584        // Tuples, lexicographic.
585        assert_eq!(
586            distinct(&[t2(2, 1), t2(1, 2), t2(2, 1)]),
587            vec![t2(1, 2), t2(2, 1)]
588        );
589        // Strings: code-point order ("B" < "Z" < "a"), capitals before lowercase.
590        assert_eq!(
591            distinct(&[ks("a"), ks("Z"), ks("B"), ks("a")]),
592            vec![ks("B"), ks("Z"), ks("a")]
593        );
594        // Empty input → empty.
595        assert_eq!(distinct(&[]), Vec::<Key>::new());
596    }
597
598    // ── Primitive 2: equijoin ───────────────────────────────────────────────
599    #[test]
600    fn match_count_is_what_equijoin_would_emit() {
601        // The planner declines a gate on this number without ever building the
602        // pair list, so the two have to agree exactly — including the cases
603        // that make them easy to get wrong: a many-to-many key (the count is a
604        // product within the key, not a sum), keys matching on one side only,
605        // and either side empty.
606        let cases: Vec<(Vec<Key>, Vec<Key>)> = vec![
607            // Many-to-many: 3 left and 2 right on key 1 is six pairs, not five.
608            (
609                vec![ki(1), ki(1), ki(1), ki(2)],
610                vec![ki(1), ki(1), ki(2), ki(3)],
611            ),
612            // Nothing in common.
613            (vec![ki(1), ki(2)], vec![ki(3), ki(4)]),
614            // One low-cardinality key against many rows — the shape the
615            // planner exists for.
616            (
617                (0..50).map(|i| ki(i % 2)).collect(),
618                (0..200).map(|i| ki(i % 2)).collect(),
619            ),
620            // Mixed key variants, including composites.
621            (
622                vec![ks("a"), t2(1, 2), ki(7), Key::Bool(true)],
623                vec![t2(1, 2), t2(1, 2), ks("a"), Key::Bool(false)],
624            ),
625            // Either side empty.
626            (vec![], vec![ki(1)]),
627            (vec![ki(1)], vec![]),
628            (vec![], vec![]),
629        ];
630        for (left, right) in cases {
631            assert_eq!(
632                equijoin_match_count(&left, &right),
633                equijoin(&left, &right).len(),
634                "count disagrees with the pair list it is standing in for: \
635                 left={left:?} right={right:?}"
636            );
637        }
638    }
639
640    // ── Primitive 3: skolem ─────────────────────────────────────────────────
641    #[test]
642    fn skolem_edge_canonicalises_undirected() {
643        assert_eq!(skolem_edge(ki(2), ki(5)), t2(2, 5));
644        assert_eq!(skolem_edge(ki(5), ki(2)), t2(2, 5)); // reversed → same
645        assert_eq!(skolem_edge(ki(7), ki(7)), t2(7, 7)); // equal preserved
646    }
647
648    #[test]
649    fn skolem_directed_preserves_symmetric_sorts() {
650        // directed: order preserved, (1,2) ≠ (2,1)
651        assert_eq!(skolem(vec![ki(1), ki(2)], false), t2(1, 2));
652        assert_eq!(skolem(vec![ki(2), ki(1)], false), t2(2, 1));
653        // symmetric: components sorted (generalises skolem_edge to arity > 2)
654        assert_eq!(
655            skolem(vec![ki(3), ki(1), ki(2)], true),
656            tup(vec![ki(1), ki(2), ki(3)])
657        );
658    }
659
660    // ── Primitive 4: rank ───────────────────────────────────────────────────
661    #[test]
662    fn rank_dense_ids_zero_based_and_base_pin_roundtrip() {
663        let rows = [ki(30), ki(10), ki(20), ki(10)];
664        let rk = rank(&rows); // Rust native base 0
665        assert_eq!(rk.order, vec![ki(10), ki(20), ki(30)]);
666        assert_eq!(rk.base, 0);
667        assert_eq!(rk.id[&ki(10)], 0);
668        assert_eq!(rk.id[&ki(20)], 1);
669        assert_eq!(rk.id[&ki(30)], 2);
670        assert_eq!(rk.dense_ids(), vec![0, 1, 2]);
671        assert_eq!(rk.canonical_dense_ids(), vec![0, 1, 2]);
672
673        // base-pin round-trip: Julia emits 1-based; `reported − base` recovers
674        // the canonical 0-based numbering (CONFORMANCE_SPEC.md §5.5.1 rule 3).
675        let rk1 = rank_with_base(&rows, 1);
676        assert_eq!(rk1.id[&ki(10)], 1);
677        assert_eq!(rk1.id[&ki(20)], 2);
678        assert_eq!(rk1.id[&ki(30)], 3);
679        for t in &rk.order {
680            assert_eq!(rk.id[t] - rk.base, rk1.id[t] - rk1.base);
681        }
682        assert_eq!(rk1.canonical_dense_ids(), vec![0, 1, 2]);
683    }
684
685    // ── Primitive 5: group_aggregate ────────────────────────────────────────
686    #[test]
687    fn group_aggregate_semirings_and_permutation_invariance() {
688        let rows = [
689            (ks("b"), Num::Int(3)),
690            (ks("a"), Num::Int(1)),
691            (ks("b"), Num::Int(4)),
692            (ks("a"), Num::Int(10)),
693            (ks("c"), Num::Int(5)),
694        ];
695        let sum = group_aggregate(&rows, SemiringOp::Sum);
696        assert_eq!(
697            sum,
698            vec![
699                (ks("a"), Num::Int(11)),
700                (ks("b"), Num::Int(7)),
701                (ks("c"), Num::Int(5))
702            ]
703        );
704        let max = group_aggregate(&rows, SemiringOp::Max);
705        assert_eq!(
706            max,
707            vec![
708                (ks("a"), Num::Int(10)),
709                (ks("b"), Num::Int(4)),
710                (ks("c"), Num::Int(5))
711            ]
712        );
713        let min = group_aggregate(&rows, SemiringOp::Min);
714        assert_eq!(
715            min,
716            vec![
717                (ks("a"), Num::Int(1)),
718                (ks("b"), Num::Int(3)),
719                (ks("c"), Num::Int(5))
720            ]
721        );
722
723        // Permuting the input rows cannot change any bucket result (assoc+comm).
724        let rows_r: Vec<(Key, Num)> = rows.iter().rev().cloned().collect();
725        assert_eq!(group_aggregate(&rows_r, SemiringOp::Sum), sum);
726
727        // Product semiring.
728        assert_eq!(
729            group_aggregate(
730                &[
731                    (ki(1), Num::Int(2)),
732                    (ki(1), Num::Int(3)),
733                    (ki(1), Num::Int(4))
734                ],
735                SemiringOp::Prod
736            ),
737            vec![(ki(1), Num::Int(24))]
738        );
739    }
740
741    #[test]
742    fn group_aggregate_bool_keys_allowed() {
743        // Bool is an allowed categorical key (boolean-or keys); false < true.
744        let rows = [
745            (Key::Bool(true), Num::Int(1)),
746            (Key::Bool(false), Num::Int(2)),
747            (Key::Bool(true), Num::Int(3)),
748        ];
749        assert_eq!(
750            group_aggregate(&rows, SemiringOp::Sum),
751            vec![
752                (Key::Bool(false), Num::Int(2)),
753                (Key::Bool(true), Num::Int(4))
754            ]
755        );
756    }
757
758    #[test]
759    fn group_aggregate_float_values_reduced_in_canonical_order() {
760        // Float *values* are allowed (only keys are float-forbidden). The
761        // per-bucket reduction is a left fold over ascending-sorted values, so
762        // the last-ULP result is reproducible: ((0.1 + 0.2) + 0.3).
763        let rows = [
764            (ki(1), Num::Float(0.3)),
765            (ki(1), Num::Float(0.1)),
766            (ki(1), Num::Float(0.2)),
767        ];
768        let got = group_aggregate(&rows, SemiringOp::Sum);
769        let want = (0.1_f64 + 0.2_f64) + 0.3_f64;
770        assert_eq!(got, vec![(ki(1), Num::Float(want))]);
771        // And it serialises through the canonical float formatter.
772        assert_eq!(
773            serialize_pairs(&got),
774            format!("[[1,{}]]", format_canonical_float(want))
775        );
776    }
777
778    // ── Canonical serialization (§5.5.3) ────────────────────────────────────
779    #[test]
780    fn canonical_index_set_json_matches_reference() {
781        assert_eq!(
782            canonical_index_set_json(&[ki(3), ki(1), ki(2), ki(1)]),
783            "[1,2,3]"
784        );
785        assert_eq!(
786            canonical_index_set_json(&[t2(2, 1), t2(1, 2)]),
787            "[[1,2],[2,1]]"
788        );
789        // code-point order + JSON-escaped quoted strings
790        assert_eq!(
791            canonical_index_set_json(&[ks("a"), ks("B")]),
792            "[\"B\",\"a\"]"
793        );
794        assert_eq!(canonical_index_set_json(&[]), "[]");
795    }
796
797    // ── §5.5.4 golden + adversarial collapse ────────────────────────────────
798    // Golden values are the cross-binding determinism manifest
799    // (tests/conformance/determinism/manifest.json) = the DuckDB throwaway oracle
800    // output (SELECT DISTINCT … ORDER BY …; dense_rank() OVER (ORDER BY …)),
801    // 0-based dense IDs. Every adversarial variant MUST collapse to it.
802
803    #[test]
804    fn fixture_edge_enumeration_undirected() {
805        let golden_set = vec![t2(1, 2), t2(1, 3), t2(2, 3), t2(2, 4), t2(3, 4)];
806        let golden_json = "[[1,2],[1,3],[2,3],[2,4],[3,4]]";
807        let golden_ids = vec![0, 1, 2, 3, 4];
808
809        let canonical = vec![vec![1, 2, 3], vec![3, 2, 4]];
810        let variants = [
811            ("permuted_faces", vec![vec![3, 2, 4], vec![1, 2, 3]]),
812            ("reversed_winding", vec![vec![3, 2, 1], vec![4, 2, 3]]),
813            (
814                "duplicate_face",
815                vec![vec![1, 2, 3], vec![3, 2, 4], vec![1, 2, 3]],
816            ),
817        ];
818        for faces in std::iter::once(&("canonical", canonical.clone()))
819            .chain(variants.iter())
820            .map(|(_, f)| f)
821        {
822            let edges = undirected_edges(faces);
823            assert_eq!(distinct(&edges), golden_set);
824            assert_eq!(canonical_index_set_json(&edges), golden_json);
825            assert_eq!(rank(&edges).canonical_dense_ids(), golden_ids);
826        }
827    }
828
829    #[test]
830    fn fixture_directed_arcs() {
831        let golden_set = vec![t2(1, 2), t2(2, 1), t2(2, 3)];
832        let golden_json = "[[1,2],[2,1],[2,3]]";
833        let golden_ids = vec![0, 1, 2];
834
835        let canonical = vec![t2(1, 2), t2(2, 1), t2(2, 3), t2(1, 2)];
836        let variants: Vec<Vec<Key>> = vec![
837            vec![t2(2, 3), t2(1, 2), t2(2, 1), t2(1, 2)], // permuted_input
838            vec![t2(1, 2), t2(2, 1), t2(2, 3), t2(1, 2), t2(2, 3)], // duplicate_arc
839        ];
840        for arcs in std::iter::once(&canonical).chain(variants.iter()) {
841            // skolem directed preserves order; distinct sorts + dedups.
842            let keyed: Vec<Key> = arcs.iter().map(|k| skolem(field_pair(k), false)).collect();
843            assert_eq!(distinct(&keyed), golden_set);
844            assert_eq!(canonical_index_set_json(&keyed), golden_json);
845            assert_eq!(rank(&keyed).canonical_dense_ids(), golden_ids);
846        }
847    }
848
849    /// Split a 2-tuple key back into its components (for re-skolemising arcs).
850    fn field_pair(k: &Key) -> Vec<Key> {
851        match k {
852            Key::Tuple(v) => v.clone(),
853            _ => panic!("expected a tuple arc"),
854        }
855    }
856
857    #[test]
858    fn fixture_group_by_sum_string_keys() {
859        let golden = vec![
860            (ks("B"), Num::Int(5)),
861            (ks("Z"), Num::Int(1)),
862            (ks("a"), Num::Int(9)),
863        ];
864        let golden_json = "[[\"B\",5],[\"Z\",1],[\"a\",9]]";
865        let golden_ids = vec![0, 1, 2];
866
867        let canonical = vec![
868            (ks("B"), Num::Int(2)),
869            (ks("a"), Num::Int(5)),
870            (ks("B"), Num::Int(3)),
871            (ks("Z"), Num::Int(1)),
872            (ks("a"), Num::Int(4)),
873        ];
874        let permuted = vec![
875            (ks("a"), Num::Int(4)),
876            (ks("Z"), Num::Int(1)),
877            (ks("B"), Num::Int(2)),
878            (ks("a"), Num::Int(5)),
879            (ks("B"), Num::Int(3)),
880        ];
881        for rows in [&canonical, &permuted] {
882            let agg = group_aggregate(rows, SemiringOp::Sum);
883            assert_eq!(agg, golden);
884            assert_eq!(serialize_pairs(&agg), golden_json);
885            // dense IDs rank the emitted (sorted) distinct keys.
886            let keys: Vec<Key> = agg.iter().map(|(k, _)| k.clone()).collect();
887            assert_eq!(rank(&keys).canonical_dense_ids(), golden_ids);
888        }
889    }
890
891    // ── Negative controls (manifest negative_controls) ──────────────────────
892    #[test]
893    fn float_in_key_rejected_at_boundary() {
894        // `float_in_key`: a float key component must raise, not silently bucket
895        // on a platform-dependent float repr (rule 1).
896        assert!(Key::try_from_json(&serde_json::json!(1.5)).is_err());
897        assert!(Key::try_from_json(&serde_json::json!(1.0)).is_err()); // 1.0 is a float in JSON
898        assert!(Key::try_from_json(&serde_json::json!([1, 2.5])).is_err());
899        // Integers and strings are fine.
900        assert_eq!(Key::try_from_json(&serde_json::json!(7)).unwrap(), ki(7));
901        assert_eq!(
902            Key::try_from_json(&serde_json::json!("B")).unwrap(),
903            ks("B")
904        );
905        assert_eq!(
906            Key::try_from_json(&serde_json::json!([1, 2])).unwrap(),
907            t2(1, 2)
908        );
909        // Bool is an allowed categorical key.
910        assert_eq!(
911            Key::try_from_json(&serde_json::json!(true)).unwrap(),
912            Key::Bool(true)
913        );
914    }
915}