Skip to main content

polydat_grammar/comprehension/
metadata.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Metadata algebra (comprehension_forms.md §10.7).
5//!
6//! Every well-formed comprehension AST node carries a four-field
7//! [`Metadata`] bundle computed bottom-up from its children's
8//! metadata and its own scalar parameters. The bundle is a
9//! monoid: propagation composes under composition, and every
10//! field is either a closed enum (capability bit) or a
11//! closed-form numeric/symbolic descriptor.
12//!
13//! This module owns:
14//!
15//! - [`Metadata`] — the four-field bundle.
16//! - [`IndexFn`] — closed-form addressing schemes (six variants
17//!   covering cartesian, zip Strict/Truncate, zip Cycle, union,
18//!   continuous, hybrid).
19//! - [`NaturalOrder`] — how a node enumerates by default.
20//! - [`Materialization`] — streaming or sized-barrier
21//!   classification (comprehension_forms.md §6.2).
22//! - [`Comprehension::metadata`] — propagation entry point.
23//!
24//! The propagation rules are total, constant-time per node, and
25//! cannot fail. Dependent-source cartesians produce
26//! `index_addressable = None`; this is the **only** place
27//! metadata propagation consults child-internal information
28//! beyond the published bundles — and it does so at the
29//! cartesian node, by walking the children's source expressions
30//! for back-references to earlier-axis names.
31
32use serde::{Deserialize, Serialize};
33
34use super::ast::Comprehension;
35use super::cardinality::{CardinalityClass, Hybrid, Interval, ProductMeasure};
36use super::source::Source;
37use super::strategy::{StrategyName, ZipMode};
38
39/// The metadata bundle carried by every well-formed AST node.
40///
41/// Computed bottom-up; never mutated after propagation. Each
42/// field is a closed enum or a closed-form descriptor — no
43/// callbacks, no fail-able analyses.
44#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
45pub struct Metadata {
46    /// Cardinality class (comprehension_forms.md §6.1).
47    pub cardinality: CardinalityClass,
48
49    /// Closed-form bijection from `0..|c|` to the node's
50    /// dispensed tuples. `None` when the node has no
51    /// addressable index space (raw filter output, dependent
52    /// cartesian, a truncated `Lex` order over either). An order
53    /// other than an untruncated `Lex` is a one-axis `Lattice` of its
54    /// selection: position `i` is the input's tuple at the `i`-th
55    /// selected position.
56    pub index_addressable: Option<IndexFn>,
57
58    /// How this node enumerates by default.
59    pub natural_order: NaturalOrder,
60
61    /// Streaming-vs-barrier classification (comprehension_forms.md §6.2).
62    pub materialization: Materialization,
63}
64
65/// Closed-form addressing schemes (comprehension_forms.md §10.7.1).
66///
67/// Six variants. Each describes the bijection from a
68/// `0..cardinality` index range to the node's tuple shape.
69/// `Continuous` and `Hybrid` carry the cardinality's
70/// interval+measure descriptors directly so the R2 push-down
71/// rule (comprehension_forms.md §10.2) dispatches on them without
72/// recomputing.
73#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
74#[serde(tag = "kind", rename_all = "snake_case")]
75pub enum IndexFn {
76    /// Discrete cartesian. `axis_sizes[i]` is the i-th axis's
77    /// element count. Multi-index `(i₀, i₁, …)` maps to the
78    /// per-axis tuple at those positions.
79    Lattice {
80        /// Element count per axis.
81        axis_sizes: Vec<u64>,
82    },
83
84    /// Zip Strict / Truncate. One index `i ∈ 0..length` maps
85    /// to the per-child tuple at position i.
86    Lockstep {
87        /// The common length.
88        length: u64,
89    },
90
91    /// Zip Cycle. Modular addressing — index `i` maps to each
92    /// child at `i mod child.cardinality`. At least one child
93    /// must be bounded (the cycling target). The index range is
94    /// [`cycle_length`] of the sizes: the longest child's, or empty
95    /// when any child is empty.
96    Modular {
97        /// Element count per child.
98        axis_sizes: Vec<u64>,
99    },
100
101    /// Union of index-addressable children. Index `i ∈
102    /// 0..Σsegment_sizes` maps to segment k where k is the
103    /// smallest such that `Σ₀^k segment_sizes > i`, position
104    /// `i - Σ₀^{k-1} segment_sizes` within that segment.
105    Concatenation {
106        /// Element count per segment, in order.
107        segment_sizes: Vec<u64>,
108    },
109
110    /// Continuous K-D box. Strategy push-down rules (Halton /
111    /// Sobol / Lhs / Extrema on Continuous) draw from this
112    /// directly; the discrete-to-continuous mapping is
113    /// strategy-specific.
114    Continuous {
115        /// The interval of each axis.
116        intervals: Vec<Interval>,
117        /// The measure drawn from.
118        measure: ProductMeasure,
119    },
120
121    /// Mixed discrete × continuous cartesian. Discrete axes get
122    /// integer indexing; continuous axes get measure-weighted
123    /// sampling. Strategy push-down dispatches per-axis.
124    Hybrid {
125        /// Element count per discrete axis.
126        discrete_axes: Vec<u64>,
127        /// The interval of each continuous axis.
128        continuous_axes: Vec<Interval>,
129        /// The measure over the continuous axes.
130        measure: ProductMeasure,
131    },
132}
133
134impl IndexFn {
135    /// `true` if this index function carries any continuous
136    /// axis. Used by per-strategy V4 checks to reject
137    /// strategies that don't accept continuous inputs.
138    pub fn has_continuous_axis(&self) -> bool {
139        matches!(self, IndexFn::Continuous { .. } | IndexFn::Hybrid { .. })
140    }
141
142    /// `true` if this index function is a multi-axis Lattice
143    /// (discrete cartesian with ≥2 axes). Required by
144    /// lattice-geometric strategies (Extrema / Shells /
145    /// Diagonal / Antidiagonal) for non-degenerate behavior.
146    pub fn is_multi_axis_lattice(&self) -> bool {
147        matches!(self, IndexFn::Lattice { axis_sizes } if axis_sizes.len() >= 2)
148    }
149}
150
151/// Natural enumeration order (comprehension_forms.md §10.7.1).
152#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
153#[serde(tag = "kind", rename_all = "snake_case")]
154pub enum NaturalOrder {
155    /// Lex order — rightmost axis varies fastest. Produced by
156    /// cartesian, single-axis clause, and `order(_, Lex, _)`.
157    Lex,
158
159    /// Lockstep — zip's natural order. One tuple per i, all
160    /// children at position i.
161    Lockstep,
162
163    /// Sequential — union's natural order. Drain child 0,
164    /// then child 1, etc.
165    Sequential,
166
167    /// Strategy-driven — produced by `order(_, non-Lex, _)`.
168    /// The wrapped strategy determines the emission order.
169    Strategy(StrategyName),
170
171    /// A continuous source that no sampling order wraps. V8
172    /// refuses to dispense it.
173    PendingSampling,
174}
175
176/// Streaming-vs-barrier classification (comprehension_forms.md §6.2).
177#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
178#[serde(tag = "kind", rename_all = "snake_case")]
179pub enum Materialization {
180    /// O(operator-local state) per pull; no input materialized.
181    Streaming,
182
183    /// Holds a finite working set; size declared at compile
184    /// time. The two natural barriers (comprehension_forms.md §6.3):
185    /// `zip(Cycle)` shorter children + non-Lex `order`.
186    BoundedBarrier {
187        /// Tuples the barrier holds at most.
188        working_set_size: u64,
189    },
190
191    /// Working set is unbounded. Always V6-rejected per spec
192    /// §5; this variant exists for representational
193    /// completeness but should never propagate through to a
194    /// valid AST's metadata.
195    UnboundedBarrier,
196}
197
198impl Comprehension {
199    /// Compute this node's metadata bundle (comprehension_forms.md
200    /// §10.7.2).
201    ///
202    /// Bottom-up: every child's metadata is computed first,
203    /// then this node's. Constant-time per node above the
204    /// child cost. Total — never fails, never partial.
205    ///
206    /// For non-leaf nodes the metadata is recomputed on every
207    /// call (no caching at this layer); consumers that need
208    /// memoization wrap it externally. The propagation cost is
209    /// O(N) in the nodes, and the optimizer re-propagates after
210    /// each rewrite.
211    pub fn metadata(&self) -> Metadata {
212        match self {
213            Comprehension::Clause { source, .. } => clause_metadata(source),
214            Comprehension::Cartesian { children } => cartesian_metadata(children),
215            Comprehension::Zip { children, mode } => zip_metadata(children, *mode),
216            Comprehension::Union { children } => union_metadata(children),
217            Comprehension::Filter { child, .. } => filter_metadata(child),
218            Comprehension::Order {
219                child,
220                strategy,
221                truncation,
222                ..
223            } => order_metadata(child, *strategy, *truncation),
224        }
225    }
226}
227
228fn clause_metadata(source: &Source) -> Metadata {
229    let cardinality = source.cardinality();
230    let (index_addressable, natural_order) = match &cardinality {
231        CardinalityClass::Bounded(n) => (
232            Some(IndexFn::Lattice {
233                axis_sizes: vec![*n],
234            }),
235            NaturalOrder::Lex,
236        ),
237        CardinalityClass::Continuous { intervals, measure } => (
238            Some(IndexFn::Continuous {
239                intervals: intervals.clone(),
240                measure: measure.clone(),
241            }),
242            NaturalOrder::PendingSampling,
243        ),
244        // BoundedAtMost / Unbounded / ContinuousAtMost — no
245        // closed-form addressing function exists.
246        _ => (None, NaturalOrder::Lex),
247    };
248    Metadata {
249        cardinality,
250        index_addressable,
251        natural_order,
252        materialization: Materialization::Streaming,
253    }
254}
255
256fn cartesian_metadata(children: &[Comprehension]) -> Metadata {
257    // First detect dependent sources: any child whose source
258    // expression references an earlier child's coordinate name.
259    // Dependent → index_addressable = None.
260    let dependent = detect_dependent_sources(children);
261
262    let child_meta: Vec<Metadata> = children.iter().map(|c| c.metadata()).collect();
263    let cardinality = combine_cartesian_cardinality(&child_meta);
264
265    let index_addressable = if dependent {
266        None
267    } else {
268        combine_cartesian_index_fn(&child_meta)
269    };
270
271    let natural_order = if matches!(
272        cardinality,
273        CardinalityClass::Continuous { .. } | CardinalityClass::Hybrid(_)
274    ) {
275        NaturalOrder::PendingSampling
276    } else {
277        NaturalOrder::Lex
278    };
279
280    Metadata {
281        cardinality,
282        index_addressable,
283        natural_order,
284        materialization: Materialization::Streaming,
285    }
286}
287
288fn zip_metadata(children: &[Comprehension], mode: ZipMode) -> Metadata {
289    let child_meta: Vec<Metadata> = children.iter().map(|c| c.metadata()).collect();
290    let cardinality = combine_zip_cardinality(&child_meta, mode);
291    let index_addressable = combine_zip_index_fn(&child_meta, mode);
292
293    let materialization = match mode {
294        ZipMode::Strict | ZipMode::Truncate => Materialization::Streaming,
295        ZipMode::Cycle => cycle_materialization(&cycle_operands(&child_meta)),
296    };
297
298    Metadata {
299        cardinality,
300        index_addressable,
301        natural_order: NaturalOrder::Lockstep,
302        materialization,
303    }
304}
305
306/// The tuple count of a `zip(Cycle)` over operands of these counts:
307/// the longest operand's, or zero when any operand is empty. Every
308/// tuple binds every operand's names, and an empty operand has no
309/// tuple to cycle.
310pub fn cycle_length(counts: &[u64]) -> u64 {
311    if counts.contains(&0) {
312        0
313    } else {
314        counts.iter().copied().max().unwrap_or(0)
315    }
316}
317
318/// `true` when metadata alone shows the operand yields no tuple.
319fn known_empty(m: &Metadata) -> bool {
320    matches!(
321        m.cardinality,
322        CardinalityClass::Bounded(0) | CardinalityClass::BoundedAtMost(0)
323    )
324}
325
326/// How a `zip(Cycle)` holds one operand while it cycles (comprehension_forms.md §6.2,
327/// §6.3).
328///
329/// Cycling re-emits an operand's earlier tuples once it is exhausted
330/// and a longer operand is not. An index-addressable operand is read
331/// at `i mod |operand|` directly and holds nothing; one operand that
332/// is not addressable streams, and is restarted when it runs out
333/// before the zip does; every other operand that is not addressable
334/// is buffered in full. An operand found empty empties the zip, so an
335/// executor checks the indexed operands' lengths and drains the
336/// buffered operands in ascending bound before it holds any tuple,
337/// and stops at the first empty one.
338#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
339#[serde(tag = "kind", rename_all = "snake_case")]
340pub enum CycleOperand {
341    /// Read at `i mod |operand|` through its index function.
342    Indexed,
343    /// Pulled once per tuple, and restarted when it runs out before
344    /// the zip does.
345    Streamed,
346    /// Held in full and replayed. A bound of zero marks an operand
347    /// known empty: the executor drains it first and holds nothing.
348    Buffered {
349        /// Tuples the buffer holds at most; `None` when the operand's
350        /// count is unknown before it is evaluated.
351        bound: Option<u64>,
352    },
353}
354
355/// The plan a `zip(Cycle)` over operands with these bundles executes:
356/// an operand known empty is buffered with bound zero; any other
357/// addressable discrete operand is [`CycleOperand::Indexed`]; of the
358/// rest, the first whose count is unknown streams, or when every count
359/// is known, the first with the largest bound; every other operand is
360/// buffered.
361pub fn cycle_operands(children: &[Metadata]) -> Vec<CycleOperand> {
362    let indexed = |m: &Metadata| {
363        !known_empty(m)
364            && m.index_addressable
365                .as_ref()
366                .is_some_and(|idx| !idx.has_continuous_axis())
367    };
368    let bound = |m: &Metadata| match &m.cardinality {
369        CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n) => Some(*n),
370        _ => None,
371    };
372    let rest: Vec<usize> = (0..children.len())
373        .filter(|&i| !indexed(&children[i]) && !known_empty(&children[i]))
374        .collect();
375    let streamed = rest
376        .iter()
377        .copied()
378        .find(|&i| bound(&children[i]).is_none())
379        .or_else(|| {
380            rest.iter()
381                .copied()
382                .rev()
383                .max_by_key(|&i| bound(&children[i]))
384        });
385    children
386        .iter()
387        .enumerate()
388        .map(|(i, m)| {
389            if indexed(m) {
390                CycleOperand::Indexed
391            } else if Some(i) == streamed {
392                CycleOperand::Streamed
393            } else {
394                CycleOperand::Buffered { bound: bound(m) }
395            }
396        })
397        .collect()
398}
399
400/// `true` when `plan` holds an operand known empty, so the zip yields
401/// no tuple and buffers nothing.
402pub fn cycle_plan_is_empty(plan: &[CycleOperand]) -> bool {
403    plan.contains(&CycleOperand::Buffered { bound: Some(0) })
404}
405
406/// A `zip(Cycle)`'s working set under `plan`: the buffered operands'
407/// bounds summed, unbounded when one of them has no bound, and
408/// streaming when nothing is buffered or an operand is known empty.
409pub fn cycle_materialization(plan: &[CycleOperand]) -> Materialization {
410    if cycle_plan_is_empty(plan) {
411        return Materialization::Streaming;
412    }
413    let mut total: u64 = 0;
414    let mut buffered = false;
415    for operand in plan {
416        if let CycleOperand::Buffered { bound } = operand {
417            buffered = true;
418            match bound {
419                Some(n) => total = total.saturating_add(*n),
420                None => return Materialization::UnboundedBarrier,
421            }
422        }
423    }
424    if buffered {
425        Materialization::BoundedBarrier {
426            working_set_size: total,
427        }
428    } else {
429        Materialization::Streaming
430    }
431}
432
433fn union_metadata(children: &[Comprehension]) -> Metadata {
434    let child_meta: Vec<Metadata> = children.iter().map(|c| c.metadata()).collect();
435    let cardinality = combine_union_cardinality(&child_meta);
436    let index_addressable = combine_union_index_fn(&child_meta);
437    Metadata {
438        cardinality,
439        index_addressable,
440        natural_order: NaturalOrder::Sequential,
441        materialization: Materialization::Streaming,
442    }
443}
444
445fn filter_metadata(child: &Comprehension) -> Metadata {
446    let child_meta = child.metadata();
447    let cardinality = match &child_meta.cardinality {
448        // Filtering nothing keeps exactly nothing.
449        CardinalityClass::Bounded(0) | CardinalityClass::BoundedAtMost(0) => {
450            CardinalityClass::Bounded(0)
451        }
452        CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n) => {
453            CardinalityClass::BoundedAtMost(*n)
454        }
455        CardinalityClass::Unbounded => CardinalityClass::Unbounded,
456        CardinalityClass::Continuous { intervals, measure }
457        | CardinalityClass::ContinuousAtMost {
458            intervals,
459            measure_at_most: measure,
460        } => CardinalityClass::ContinuousAtMost {
461            intervals: intervals.clone(),
462            measure_at_most: measure.clone(),
463        },
464        CardinalityClass::Hybrid(h) => CardinalityClass::Hybrid(h.clone()),
465    };
466    Metadata {
467        cardinality,
468        index_addressable: None, // filter destroys the bijection
469        natural_order: child_meta.natural_order,
470        materialization: child_meta.materialization,
471    }
472}
473
474fn order_metadata(
475    child: &Comprehension,
476    strategy: StrategyName,
477    truncation: Option<u64>,
478) -> Metadata {
479    let child_meta = child.metadata();
480    let cardinality = order_cardinality(child, &child_meta.cardinality, strategy, truncation);
481
482    // An order's output is addressed through its selection: position `i`
483    // is the input's tuple at the `i`-th selected position, so the
484    // output is one axis as long as the selection (comprehension_forms.md
485    // §3.6). The axis is sized by the order's count, or its bound.
486    let selected = || match &cardinality {
487        CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n) => {
488            Some(IndexFn::Lattice {
489                axis_sizes: vec![*n],
490            })
491        }
492        _ => None,
493    };
494    let (index_addressable, natural_order, materialization) = match strategy {
495        StrategyName::Lex => (
496            // An untruncated `Lex` passes its input through, positions
497            // and all. A truncated one selects a prefix of an addressable
498            // input's positions; over any other input it counts the
499            // tuples as they stream and addresses nothing.
500            match truncation {
501                None => child_meta.index_addressable,
502                Some(_) => child_meta.index_addressable.and_then(|_| selected()),
503            },
504            NaturalOrder::Lex,
505            child_meta.materialization, // counter wrapper at most
506        ),
507        non_lex => {
508            // Over an addressable input the strategy selects positions
509            // and holds only its selection (R2); over any other input
510            // the input is buffered in full first.
511            let materialization = match &child_meta.index_addressable {
512                Some(_) => Materialization::BoundedBarrier {
513                    working_set_size: strategy_working_set(
514                        non_lex,
515                        &child_meta.index_addressable,
516                        truncation,
517                    ),
518                },
519                None => match &child_meta.cardinality {
520                    CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n) => {
521                        Materialization::BoundedBarrier {
522                            working_set_size: *n,
523                        }
524                    }
525                    _ => Materialization::UnboundedBarrier,
526                },
527            };
528            (selected(), NaturalOrder::Strategy(non_lex), materialization)
529        }
530    };
531
532    Metadata {
533        cardinality,
534        index_addressable,
535        natural_order,
536        materialization,
537    }
538}
539
540// ---- cardinality combinators ----
541
542fn combine_cartesian_cardinality(children: &[Metadata]) -> CardinalityClass {
543    let mut has_continuous = false;
544    let mut has_discrete = false;
545    let mut counts: Vec<Count> = Vec::new();
546    let mut discrete_axes: Vec<u64> = Vec::new();
547    let mut continuous_intervals: Vec<Interval> = Vec::new();
548    let mut continuous_measures: Vec<ProductMeasure> = Vec::new();
549
550    for m in children {
551        match &m.cardinality {
552            CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n) => {
553                has_discrete = true;
554                discrete_axes.push(*n); // the count, or its upper bound
555                counts.extend(Count::of(&m.cardinality));
556            }
557            CardinalityClass::Unbounded => {
558                has_discrete = true;
559                discrete_axes.push(0);
560                counts.push(Count::Unknown);
561            }
562            CardinalityClass::Continuous { intervals, measure }
563            | CardinalityClass::ContinuousAtMost {
564                intervals,
565                measure_at_most: measure,
566            } => {
567                has_continuous = true;
568                continuous_intervals.extend(intervals.iter().cloned());
569                continuous_measures.push(measure.clone());
570            }
571            CardinalityClass::Hybrid(h) => {
572                has_continuous = true;
573                has_discrete = true;
574                discrete_axes.extend(h.discrete_axes.iter().copied());
575                continuous_intervals.extend(h.continuous_axes.iter().cloned());
576                continuous_measures.push(h.measure.clone());
577            }
578        }
579    }
580
581    if has_continuous && has_discrete {
582        CardinalityClass::Hybrid(Hybrid {
583            discrete_axes,
584            continuous_axes: continuous_intervals,
585            measure: simplify_measures(continuous_measures),
586        })
587    } else if has_continuous {
588        CardinalityClass::Continuous {
589            intervals: continuous_intervals,
590            measure: simplify_measures(continuous_measures),
591        }
592    } else {
593        cartesian_count(&counts).class()
594    }
595}
596
597fn combine_cartesian_index_fn(children: &[Metadata]) -> Option<IndexFn> {
598    // All children must be addressable for the cartesian to be.
599    let all_addressable = children.iter().all(|m| m.index_addressable.is_some());
600    if !all_addressable {
601        return None;
602    }
603
604    let mut all_discrete = true;
605    let mut all_continuous = true;
606    let mut discrete_axes: Vec<u64> = Vec::new();
607    let mut continuous_intervals: Vec<Interval> = Vec::new();
608    let mut continuous_measures: Vec<ProductMeasure> = Vec::new();
609
610    for m in children {
611        match m.index_addressable.as_ref().unwrap() {
612            IndexFn::Lattice { axis_sizes } => {
613                all_continuous = false;
614                discrete_axes.extend(axis_sizes.iter().copied());
615            }
616            IndexFn::Continuous { intervals, measure } => {
617                all_discrete = false;
618                continuous_intervals.extend(intervals.iter().cloned());
619                continuous_measures.push(measure.clone());
620            }
621            IndexFn::Hybrid {
622                discrete_axes: d,
623                continuous_axes: c,
624                measure,
625            } => {
626                all_discrete = false;
627                all_continuous = false;
628                discrete_axes.extend(d.iter().copied());
629                continuous_intervals.extend(c.iter().cloned());
630                continuous_measures.push(measure.clone());
631            }
632            // Lockstep / Modular / Concatenation — these don't
633            // combine as cartesian axes (they're 1-D index
634            // spaces of their own), and a cartesian of a zip or a
635            // union has no addressing scheme here, so it has none.
636            IndexFn::Lockstep { .. } | IndexFn::Modular { .. } | IndexFn::Concatenation { .. } => {
637                return None;
638            }
639        }
640    }
641
642    if all_discrete {
643        Some(IndexFn::Lattice {
644            axis_sizes: discrete_axes,
645        })
646    } else if all_continuous {
647        Some(IndexFn::Continuous {
648            intervals: continuous_intervals,
649            measure: simplify_measures(continuous_measures),
650        })
651    } else {
652        Some(IndexFn::Hybrid {
653            discrete_axes,
654            continuous_axes: continuous_intervals,
655            measure: simplify_measures(continuous_measures),
656        })
657    }
658}
659
660fn combine_zip_cardinality(children: &[Metadata], mode: ZipMode) -> CardinalityClass {
661    // A continuous child is a V7 failure; its count is unknown here.
662    let counts: Vec<Count> = children
663        .iter()
664        .map(|m| Count::of(&m.cardinality).unwrap_or(Count::Unknown))
665        .collect();
666    match mode {
667        ZipMode::Strict => strict_zip_count(&counts),
668        ZipMode::Truncate => truncate_zip_count(&counts),
669        ZipMode::Cycle => cycle_zip_count(&counts),
670    }
671    .class()
672}
673
674// ---- tuple counts ----
675
676/// A discrete tuple count as metadata knows it.
677#[derive(Debug, Clone, Copy, PartialEq, Eq)]
678enum Count {
679    /// Exactly this many.
680    Exact(u64),
681    /// Between zero and this many.
682    AtMost(u64),
683    /// No known bound.
684    Unknown,
685}
686
687impl Count {
688    /// The count a discrete class states, `None` for a continuous one.
689    /// At most zero is exactly zero.
690    fn of(class: &CardinalityClass) -> Option<Self> {
691        Some(match class {
692            CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n @ 0) => {
693                Count::Exact(*n)
694            }
695            CardinalityClass::BoundedAtMost(n) => Count::AtMost(*n),
696            CardinalityClass::Unbounded => Count::Unknown,
697            _ => return None,
698        })
699    }
700
701    fn bound(self) -> Option<u64> {
702        match self {
703            Count::Exact(n) | Count::AtMost(n) => Some(n),
704            Count::Unknown => None,
705        }
706    }
707
708    /// `n` exactly when every count it was combined from is exact, at
709    /// most `n` otherwise.
710    fn combined(n: u64, counts: &[Count]) -> Self {
711        if counts.iter().all(|c| matches!(c, Count::Exact(_))) {
712            Count::Exact(n)
713        } else {
714            Count::AtMost(n)
715        }
716    }
717
718    fn class(self) -> CardinalityClass {
719        match self {
720            Count::Exact(n) | Count::AtMost(n @ 0) => CardinalityClass::Bounded(n),
721            Count::AtMost(n) => CardinalityClass::BoundedAtMost(n),
722            Count::Unknown => CardinalityClass::Unbounded,
723        }
724    }
725}
726
727/// A cartesian's count: exactly zero when an operand is exactly empty,
728/// whatever the others; unknown when an operand's count is; otherwise
729/// the product of the operands' counts, exact when every one is.
730fn cartesian_count(counts: &[Count]) -> Count {
731    if counts.contains(&Count::Exact(0)) {
732        return Count::Exact(0);
733    }
734    let Some(bounds) = counts.iter().map(|c| c.bound()).collect::<Option<Vec<_>>>() else {
735        return Count::Unknown;
736    };
737    Count::combined(bounds.into_iter().fold(1, u64::saturating_mul), counts)
738}
739
740/// A truncating zip's count: the shortest operand's. Exactly zero when
741/// an operand is exactly empty; exact when every operand is; otherwise
742/// at most the least bound, since an operand at most `m` or of unknown
743/// count may end at any length before `m`. Unknown only when no operand
744/// has a bound.
745fn truncate_zip_count(counts: &[Count]) -> Count {
746    if counts.contains(&Count::Exact(0)) {
747        return Count::Exact(0);
748    }
749    match counts.iter().filter_map(|c| c.bound()).min() {
750        Some(m) => Count::combined(m, counts),
751        None => Count::Unknown,
752    }
753}
754
755/// A strict zip's count when it yields: every operand's, so an exact
756/// operand's count exactly; otherwise at most the least bound. Operands
757/// that end apart fail the zip instead.
758fn strict_zip_count(counts: &[Count]) -> Count {
759    if let Some(exact) = counts.iter().find(|c| matches!(c, Count::Exact(_))) {
760        return *exact;
761    }
762    match counts.iter().filter_map(|c| c.bound()).min() {
763        Some(m) => Count::AtMost(m),
764        None => Count::Unknown,
765    }
766}
767
768/// A cycle zip's count ([`cycle_length`]): exactly zero when an operand
769/// is exactly empty; unknown when an operand's count is; otherwise the
770/// longest operand's, exact when every operand is exact, and at most
771/// that when one is at most, since it may be empty at open.
772fn cycle_zip_count(counts: &[Count]) -> Count {
773    if counts.contains(&Count::Exact(0)) {
774        return Count::Exact(0);
775    }
776    let Some(bounds) = counts.iter().map(|c| c.bound()).collect::<Option<Vec<_>>>() else {
777        return Count::Unknown;
778    };
779    Count::combined(bounds.into_iter().max().unwrap_or(0), counts)
780}
781
782/// A union's count: the sum of its operands', exact when every one is,
783/// unknown when one is.
784fn union_count(counts: &[Count]) -> Count {
785    let Some(bounds) = counts.iter().map(|c| c.bound()).collect::<Option<Vec<_>>>() else {
786        return Count::Unknown;
787    };
788    Count::combined(bounds.into_iter().fold(0, u64::saturating_add), counts)
789}
790
791/// An order's count. `Lex` and the strategies that truncate by tuple
792/// count (`reverse_lex`, `diagonal`, `antidiagonal`, `halton`, `sobol`,
793/// `lhs`, `shuffle`) keep `min(count, n)` of a discrete input; `extrema`
794/// and `shells` truncate by whole strata or shells, so they keep at
795/// most the input's count. Over a continuous space a sampling strategy
796/// draws `n` points, exactly `n` unless a filter in the space or a
797/// discrete axis of inexact count may leave fewer; `extrema` takes
798/// the strata of the box, each continuous axis contributing its two
799/// ends.
800fn order_cardinality(
801    child: &Comprehension,
802    child_class: &CardinalityClass,
803    strategy: StrategyName,
804    truncation: Option<u64>,
805) -> CardinalityClass {
806    let strata = matches!(strategy, StrategyName::Extrema | StrategyName::Shells);
807    let Some(count) = Count::of(child_class) else {
808        // A continuous space, sampled by a non-`Lex` order with a count
809        // (V8); any other order over one is invalid and keeps its class.
810        let Some(n) = truncation.filter(|_| !matches!(strategy, StrategyName::Lex)) else {
811            return child_class.clone();
812        };
813        let mut space = SampledSpace::default();
814        space.collect(child);
815        if space.discrete.contains(&Count::Exact(0)) {
816            return CardinalityClass::Bounded(0);
817        }
818        return if matches!(strategy, StrategyName::Extrema) {
819            // `n` strata of the box: at most all of its corners.
820            let mut axes = space.discrete;
821            axes.extend(std::iter::repeat_n(Count::Exact(2), space.continuous));
822            match cartesian_count(&axes) {
823                Count::Exact(m) | Count::AtMost(m) => Count::AtMost(m),
824                Count::Unknown => Count::Unknown,
825            }
826        } else if !space.filtered && space.discrete.iter().all(|c| matches!(c, Count::Exact(_))) {
827            Count::Exact(n)
828        } else {
829            Count::AtMost(n)
830        }
831        .class();
832    };
833    match (count, truncation) {
834        (Count::Exact(0), _) | (_, None) => count,
835        (Count::Exact(c) | Count::AtMost(c), Some(_)) if strata => Count::AtMost(c),
836        (Count::Exact(c), Some(n)) => Count::Exact(c.min(n)),
837        (Count::AtMost(c), Some(n)) => Count::AtMost(c.min(n)),
838        (Count::Unknown, Some(_)) if strata => Count::Unknown,
839        (Count::Unknown, Some(n)) => Count::AtMost(n),
840    }
841    .class()
842}
843
844/// The axes an order over a continuous space samples, walked as the
845/// runtime walks them: clauses through cartesians and filters, any
846/// other node one discrete axis of its tuples.
847#[derive(Default)]
848struct SampledSpace {
849    /// Each discrete axis's count.
850    discrete: Vec<Count>,
851    /// How many continuous axes.
852    continuous: usize,
853    /// Whether a filter sits between the order and its clauses.
854    filtered: bool,
855}
856
857impl SampledSpace {
858    fn collect(&mut self, c: &Comprehension) {
859        match c {
860            Comprehension::Clause { source, .. } => match source.cardinality() {
861                CardinalityClass::Continuous { .. } => self.continuous += 1,
862                class => self
863                    .discrete
864                    .push(Count::of(&class).unwrap_or(Count::Unknown)),
865            },
866            Comprehension::Cartesian { children } => {
867                children.iter().for_each(|child| self.collect(child));
868            }
869            Comprehension::Filter { child, .. } => {
870                self.filtered = true;
871                self.collect(child);
872            }
873            other => self
874                .discrete
875                .push(Count::of(&other.metadata().cardinality).unwrap_or(Count::Unknown)),
876        }
877    }
878}
879
880fn combine_zip_index_fn(children: &[Metadata], mode: ZipMode) -> Option<IndexFn> {
881    let all_addressable = children.iter().all(|m| m.index_addressable.is_some());
882    if !all_addressable {
883        return None;
884    }
885    let counts: Vec<u64> = children
886        .iter()
887        .filter_map(|m| match &m.cardinality {
888            CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n) => Some(*n),
889            _ => None,
890        })
891        .collect();
892    if counts.len() != children.len() {
893        return None;
894    }
895    match mode {
896        ZipMode::Strict | ZipMode::Truncate => {
897            let length = match mode {
898                ZipMode::Strict => counts[0],
899                ZipMode::Truncate => *counts.iter().min().unwrap(),
900                ZipMode::Cycle => unreachable!(),
901            };
902            Some(IndexFn::Lockstep { length })
903        }
904        ZipMode::Cycle => Some(IndexFn::Modular { axis_sizes: counts }),
905    }
906}
907
908fn combine_union_cardinality(children: &[Metadata]) -> CardinalityClass {
909    // A continuous child is a V9 failure; its count is unknown here.
910    let counts: Vec<Count> = children
911        .iter()
912        .map(|m| Count::of(&m.cardinality).unwrap_or(Count::Unknown))
913        .collect();
914    union_count(&counts).class()
915}
916
917fn combine_union_index_fn(children: &[Metadata]) -> Option<IndexFn> {
918    let all_addressable = children.iter().all(|m| m.index_addressable.is_some());
919    if !all_addressable {
920        return None;
921    }
922    let segment_sizes: Vec<u64> = children
923        .iter()
924        .filter_map(|m| match &m.cardinality {
925            CardinalityClass::Bounded(n) | CardinalityClass::BoundedAtMost(n) => Some(*n),
926            _ => None,
927        })
928        .collect();
929    if segment_sizes.len() != children.len() {
930        return None;
931    }
932    Some(IndexFn::Concatenation { segment_sizes })
933}
934
935// ---- supporting helpers ----
936
937fn simplify_measures(measures: Vec<ProductMeasure>) -> ProductMeasure {
938    match measures.len() {
939        0 => ProductMeasure::Uniform,
940        1 => measures.into_iter().next().unwrap(),
941        _ => ProductMeasure::Product(measures),
942    }
943}
944
945/// Strategy-specific working-set size for use as
946/// `BoundedBarrier.working_set_size` over an addressable input: the
947/// selection the strategy holds, since it reads tuples only at the
948/// positions it selects (R2).
949fn strategy_working_set(
950    strategy: StrategyName,
951    input: &Option<IndexFn>,
952    truncation: Option<u64>,
953) -> u64 {
954    match (strategy, input, truncation) {
955        // Halton / Sobol / Shuffle over an index-addressable
956        // input + truncation: O(n) draws.
957        (StrategyName::Halton, Some(_), Some(n))
958        | (StrategyName::Sobol, Some(_), Some(n))
959        | (StrategyName::Shuffle, Some(_), Some(n))
960        | (StrategyName::ReverseLex, Some(_), Some(n)) => n,
961        // Lhs: O(n * dim).
962        (StrategyName::Lhs, Some(idx), Some(n)) => {
963            let dim = lattice_dim(idx).max(1);
964            n.saturating_mul(dim as u64)
965        }
966        // Extrema (comprehension_forms.md §3.6) and Shells rank every multi-index of
967        // the input's index space before keeping the first strata or
968        // shells, so they hold the whole index space.
969        (StrategyName::Extrema, Some(idx), Some(_))
970        | (StrategyName::Shells, Some(idx), Some(_)) => index_fn_cardinality(idx),
971        // Diagonal / Antidiagonal walk the diagonals in order and stop
972        // at `n`.
973        (StrategyName::Diagonal, Some(_), Some(n))
974        | (StrategyName::Antidiagonal, Some(_), Some(n)) => n,
975        // No truncation: fall back to the input's cardinality.
976        (_, Some(idx), None) => index_fn_cardinality(idx),
977        // No addressable input: we can't compute a closed form;
978        // use the naïve "input cardinality" placeholder so the
979        // metadata still has a number (consumers should treat
980        // this as a conservative upper bound).
981        (_, None, Some(n)) => n,
982        (_, None, None) => 0,
983        // Lex with truncation over addressable input — counter
984        // wrapper, working set equals output size.
985        (StrategyName::Lex, Some(_), Some(n)) => n,
986    }
987}
988
989fn lattice_dim(idx: &IndexFn) -> usize {
990    match idx {
991        IndexFn::Lattice { axis_sizes } => axis_sizes.len(),
992        IndexFn::Continuous { intervals, .. } => intervals.len(),
993        IndexFn::Hybrid {
994            discrete_axes,
995            continuous_axes,
996            ..
997        } => discrete_axes.len() + continuous_axes.len(),
998        IndexFn::Lockstep { .. } | IndexFn::Modular { .. } | IndexFn::Concatenation { .. } => 1,
999    }
1000}
1001
1002fn index_fn_cardinality(idx: &IndexFn) -> u64 {
1003    match idx {
1004        IndexFn::Lattice { axis_sizes } => axis_sizes
1005            .iter()
1006            .copied()
1007            .fold(1u64, |a, b| a.saturating_mul(b)),
1008        IndexFn::Lockstep { length } => *length,
1009        IndexFn::Modular { axis_sizes } => cycle_length(axis_sizes),
1010        IndexFn::Concatenation { segment_sizes } => segment_sizes
1011            .iter()
1012            .copied()
1013            .fold(0u64, |a, b| a.saturating_add(b)),
1014        // Continuous index has no integer cardinality.
1015        IndexFn::Continuous { .. } | IndexFn::Hybrid { .. } => 0,
1016    }
1017}
1018
1019/// Walk children's source expressions for back-references to
1020/// earlier-axis coordinate names. Used by cartesian metadata
1021/// propagation to detect dependent sources (comprehension_forms.md §3.2).
1022fn detect_dependent_sources(children: &[Comprehension]) -> bool {
1023    let mut prior_names: Vec<String> = Vec::new();
1024    for child in children {
1025        // First check if the child references any prior name in
1026        // its source(s).
1027        for name in collect_source_name_references(child) {
1028            if prior_names.contains(&name) {
1029                return true;
1030            }
1031        }
1032        // Then add this child's coordinates to the prior set.
1033        for n in child.coordinate_names() {
1034            if !prior_names.contains(&n) {
1035                prior_names.push(n);
1036            }
1037        }
1038    }
1039    false
1040}
1041
1042/// The names the sources of a comprehension subtree read when they are
1043/// evaluated ([`Source::names_read`]), a composed name's leaves
1044/// included.
1045fn collect_source_name_references(c: &Comprehension) -> std::collections::BTreeSet<String> {
1046    c.source_names_read()
1047}
1048
1049#[cfg(test)]
1050mod tests {
1051    use super::*;
1052    use crate::comprehension::source::{LiteralValue, Source};
1053
1054    fn clause(name: &str, vs: &[i64]) -> Comprehension {
1055        Comprehension::clause(
1056            name,
1057            Source::Literal {
1058                values: vs.iter().map(|n| LiteralValue::Int(*n)).collect(),
1059            },
1060        )
1061    }
1062
1063    fn continuous_clause(name: &str) -> Comprehension {
1064        Comprehension::clause(
1065            name,
1066            Source::ContinuousInterval {
1067                interval: Interval::closed(0.0, 1.0),
1068                measure: ProductMeasure::Uniform,
1069            },
1070        )
1071    }
1072
1073    #[test]
1074    fn clause_metadata_for_bounded_source() {
1075        let m = clause("k", &[1, 2, 3]).metadata();
1076        assert_eq!(m.cardinality, CardinalityClass::Bounded(3));
1077        assert_eq!(
1078            m.index_addressable,
1079            Some(IndexFn::Lattice {
1080                axis_sizes: vec![3]
1081            })
1082        );
1083        assert_eq!(m.natural_order, NaturalOrder::Lex);
1084        assert_eq!(m.materialization, Materialization::Streaming);
1085    }
1086
1087    #[test]
1088    fn clause_metadata_for_continuous_source() {
1089        let m = continuous_clause("alpha").metadata();
1090        assert!(matches!(m.cardinality, CardinalityClass::Continuous { .. }));
1091        assert!(matches!(
1092            m.index_addressable,
1093            Some(IndexFn::Continuous { .. })
1094        ));
1095        assert_eq!(m.natural_order, NaturalOrder::PendingSampling);
1096        assert_eq!(m.materialization, Materialization::Streaming);
1097    }
1098
1099    #[test]
1100    fn cartesian_metadata_combines_lattice_axes() {
1101        let c =
1102            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20, 30])]);
1103        let m = c.metadata();
1104        assert_eq!(m.cardinality, CardinalityClass::Bounded(6));
1105        assert_eq!(
1106            m.index_addressable,
1107            Some(IndexFn::Lattice {
1108                axis_sizes: vec![2, 3]
1109            })
1110        );
1111        assert_eq!(m.natural_order, NaturalOrder::Lex);
1112    }
1113
1114    #[test]
1115    fn cartesian_metadata_for_hybrid() {
1116        let c =
1117            Comprehension::cartesian(vec![clause("k", &[1, 2, 3, 4]), continuous_clause("theta")]);
1118        let m = c.metadata();
1119        match m.cardinality {
1120            CardinalityClass::Hybrid(h) => {
1121                assert_eq!(h.discrete_axes, vec![4]);
1122                assert_eq!(h.continuous_axes.len(), 1);
1123            }
1124            other => panic!("expected Hybrid, got {other:?}"),
1125        }
1126        assert!(matches!(m.index_addressable, Some(IndexFn::Hybrid { .. })));
1127        assert_eq!(m.natural_order, NaturalOrder::PendingSampling);
1128    }
1129
1130    #[test]
1131    fn dependent_cartesian_produces_none_addressable() {
1132        // clause replicas references {k} from the prior clause.
1133        let dependent = Comprehension::cartesian(vec![
1134            clause("k", &[1, 2, 3]),
1135            Comprehension::clause(
1136                "replicas",
1137                Source::Generator {
1138                    expr: "range(0, 2 * {k})".into(),
1139                    cardinality_hint: Some(6),
1140                },
1141            ),
1142        ]);
1143        let m = dependent.metadata();
1144        assert!(m.index_addressable.is_none());
1145    }
1146
1147    #[test]
1148    fn zip_strict_produces_lockstep_index_fn() {
1149        let c = Comprehension::zip(
1150            vec![clause("x", &[1, 2, 3]), clause("y", &[10, 20, 30])],
1151            ZipMode::Strict,
1152        );
1153        let m = c.metadata();
1154        assert_eq!(m.index_addressable, Some(IndexFn::Lockstep { length: 3 }));
1155        assert_eq!(m.natural_order, NaturalOrder::Lockstep);
1156        assert_eq!(m.materialization, Materialization::Streaming);
1157    }
1158
1159    #[test]
1160    fn zip_cycle_produces_modular_index_fn_and_barrier() {
1161        let c = Comprehension::zip(
1162            vec![clause("k", &[1, 2, 3, 4, 5]), clause("color", &[1, 2, 3])],
1163            ZipMode::Cycle,
1164        );
1165        let m = c.metadata();
1166        match m.index_addressable {
1167            Some(IndexFn::Modular { axis_sizes }) => {
1168                assert_eq!(axis_sizes, vec![5, 3]);
1169            }
1170            other => panic!("expected Modular, got {other:?}"),
1171        }
1172        // Both operands are addressable: cycling reads the shorter one
1173        // at `i mod 3` and buffers nothing.
1174        assert_eq!(m.materialization, Materialization::Streaming);
1175    }
1176
1177    fn unknown_count(name: &str) -> Comprehension {
1178        Comprehension::clause(
1179            name,
1180            Source::Generator {
1181                expr: "values({n})".into(),
1182                cardinality_hint: None,
1183            },
1184        )
1185    }
1186
1187    /// With an operand of unknown count, that operand streams and every
1188    /// finite operand that is not addressable is buffered in full;
1189    /// addressable ones are indexed.
1190    #[test]
1191    fn zip_cycle_with_an_unknown_count_buffers_every_finite_unaddressable_operand() {
1192        let c = Comprehension::zip(
1193            vec![
1194                unknown_count("tick"),
1195                Comprehension::filter(clause("a", &(0..1000).collect::<Vec<_>>()), "{a} > 1"),
1196                Comprehension::filter(clause("b", &[1, 2, 3]), "{b} > 0"),
1197                clause("color", &[1, 2, 3]),
1198            ],
1199            ZipMode::Cycle,
1200        );
1201        let plan = cycle_operands(
1202            &match &c {
1203                Comprehension::Zip { children, .. } => children,
1204                _ => unreachable!(),
1205            }
1206            .iter()
1207            .map(Comprehension::metadata)
1208            .collect::<Vec<_>>(),
1209        );
1210        assert_eq!(
1211            plan,
1212            vec![
1213                CycleOperand::Streamed,
1214                CycleOperand::Buffered { bound: Some(1000) },
1215                CycleOperand::Buffered { bound: Some(3) },
1216                CycleOperand::Indexed,
1217            ]
1218        );
1219        assert_eq!(
1220            c.metadata().materialization,
1221            Materialization::BoundedBarrier {
1222                working_set_size: 1003
1223            }
1224        );
1225    }
1226
1227    /// With every count known, the largest operand that is not
1228    /// addressable streams and the others are buffered.
1229    #[test]
1230    fn zip_cycle_streams_the_largest_unaddressable_operand() {
1231        let c = Comprehension::zip(
1232            vec![
1233                Comprehension::filter(clause("a", &[1, 2]), "{a} > 0"),
1234                Comprehension::filter(clause("b", &[1, 2, 3, 4]), "{b} > 0"),
1235                clause("k", &(0..100).collect::<Vec<_>>()),
1236            ],
1237            ZipMode::Cycle,
1238        );
1239        assert_eq!(
1240            c.metadata().materialization,
1241            Materialization::BoundedBarrier {
1242                working_set_size: 2
1243            }
1244        );
1245    }
1246
1247    /// An operand known empty empties the zip: no tuple, no index, and
1248    /// nothing held, in any position and beside operands of unknown
1249    /// count.
1250    #[test]
1251    fn zip_cycle_with_an_operand_known_empty_is_empty() {
1252        let operands = || {
1253            vec![
1254                unknown_count("tick"),
1255                Comprehension::filter(clause("a", &[1, 2, 3]), "{a} > 1"),
1256                clause("color", &[1, 2, 3]),
1257            ]
1258        };
1259        let empties = [
1260            clause("e", &[]),
1261            Comprehension::filter(clause("e", &[]), "{e} > 1"),
1262        ];
1263        for empty in &empties {
1264            for at in 0..=3 {
1265                let mut children = operands();
1266                children.insert(at, empty.clone());
1267                let plan = cycle_operands(
1268                    &children
1269                        .iter()
1270                        .map(Comprehension::metadata)
1271                        .collect::<Vec<_>>(),
1272                );
1273                assert_eq!(plan[at], CycleOperand::Buffered { bound: Some(0) });
1274                assert!(cycle_plan_is_empty(&plan));
1275                let m = Comprehension::zip(children, ZipMode::Cycle).metadata();
1276                assert_eq!(m.cardinality, CardinalityClass::Bounded(0));
1277                assert_eq!(m.materialization, Materialization::Streaming);
1278            }
1279        }
1280        let addressable = Comprehension::zip(
1281            vec![clause("k", &[1, 2, 3]), clause("e", &[])],
1282            ZipMode::Cycle,
1283        );
1284        let m = addressable.metadata();
1285        assert_eq!(
1286            m.index_addressable,
1287            Some(IndexFn::Modular {
1288                axis_sizes: vec![3, 0]
1289            })
1290        );
1291        assert_eq!(
1292            index_fn_cardinality(m.index_addressable.as_ref().unwrap()),
1293            0
1294        );
1295        assert_eq!(cycle_length(&[3, 0]), 0);
1296        assert_eq!(cycle_length(&[3, 5]), 5);
1297    }
1298
1299    /// An operand that may be empty at open makes the zip's count an
1300    /// upper bound.
1301    #[test]
1302    fn zip_cycle_over_an_operand_at_most_counts_at_most() {
1303        let c = Comprehension::zip(
1304            vec![
1305                clause("k", &[1, 2, 3, 4, 5]),
1306                Comprehension::filter(clause("a", &[1, 2]), "{a} > 1"),
1307            ],
1308            ZipMode::Cycle,
1309        );
1310        assert_eq!(c.metadata().cardinality, CardinalityClass::BoundedAtMost(5));
1311    }
1312
1313    /// The tuple counts an operand of count `c` may have: an exact
1314    /// count its own, an at-most count every count up to its bound, an
1315    /// unknown count small and large ones.
1316    fn witnesses(c: Count) -> Vec<u64> {
1317        match c {
1318            Count::Exact(n) => vec![n],
1319            Count::AtMost(n) => (0..=n).collect(),
1320            Count::Unknown => (0..=7).chain([1000]).collect(),
1321        }
1322    }
1323
1324    /// Assert `claim` holds for every count the operator yields over
1325    /// the operands' witnesses, and is tight: exact means every yield is
1326    /// that count, at most means none exceeds the bound and one meets
1327    /// it, and unknown means some yield exceeds any bound the operands
1328    /// state.
1329    fn assert_describes(claim: Count, yields: &[u64], what: &str) {
1330        let Some(&max) = yields.iter().max() else {
1331            return; // the operator never yields, as a strict zip of unequal operands
1332        };
1333        match claim {
1334            Count::Exact(n) => assert!(yields.iter().all(|&y| y == n), "{what}: {yields:?}"),
1335            Count::AtMost(n) => {
1336                assert!(max <= n, "{what}: {yields:?} exceed {n}");
1337                assert_eq!(max, n, "{what}: the bound is not tight");
1338                assert!(n > 0, "{what}: at most zero is exactly zero");
1339                assert!(
1340                    yields.iter().any(|&y| y != n),
1341                    "{what}: always {n}, so exact"
1342                );
1343            }
1344            Count::Unknown => assert!(max >= 1000, "{what}: bounded by {max}"),
1345        }
1346    }
1347
1348    const KINDS: [Count; 5] = [
1349        Count::Exact(0),
1350        Count::Exact(3),
1351        Count::Exact(5),
1352        Count::AtMost(4),
1353        Count::Unknown,
1354    ];
1355
1356    /// Every combination of the operands' witnesses.
1357    fn combinations(counts: &[Count]) -> Vec<Vec<u64>> {
1358        counts.iter().fold(vec![Vec::new()], |acc, c| {
1359            acc.iter()
1360                .flat_map(|prefix| {
1361                    witnesses(*c).into_iter().map(move |w| {
1362                        let mut next = prefix.clone();
1363                        next.push(w);
1364                        next
1365                    })
1366                })
1367                .collect()
1368        })
1369    }
1370
1371    /// Each operator's count over every pair and triple of kinds holds
1372    /// for, and is tight over, what the operator yields.
1373    #[test]
1374    fn every_kind_combination_counts_what_the_operator_yields() {
1375        let mut shapes: Vec<Vec<Count>> = Vec::new();
1376        for a in KINDS {
1377            for b in KINDS {
1378                shapes.push(vec![a, b]);
1379                for c in KINDS {
1380                    shapes.push(vec![a, b, c]);
1381                }
1382            }
1383        }
1384        for counts in &shapes {
1385            let combos = combinations(counts);
1386            let product: Vec<u64> = combos.iter().map(|c| c.iter().product()).collect();
1387            assert_describes(
1388                cartesian_count(counts),
1389                &product,
1390                &format!("cartesian {counts:?}"),
1391            );
1392            let sum: Vec<u64> = combos.iter().map(|c| c.iter().sum()).collect();
1393            assert_describes(union_count(counts), &sum, &format!("union {counts:?}"));
1394            let shortest: Vec<u64> = combos.iter().map(|c| *c.iter().min().unwrap()).collect();
1395            assert_describes(
1396                truncate_zip_count(counts),
1397                &shortest,
1398                &format!("truncate {counts:?}"),
1399            );
1400            let cycled: Vec<u64> = combos.iter().map(|c| cycle_length(c)).collect();
1401            assert_describes(
1402                cycle_zip_count(counts),
1403                &cycled,
1404                &format!("cycle {counts:?}"),
1405            );
1406            let strict: Vec<u64> = combos
1407                .iter()
1408                .filter(|c| c.iter().all(|&n| n == c[0]))
1409                .map(|c| c[0])
1410                .collect();
1411            assert_describes(
1412                strict_zip_count(counts),
1413                &strict,
1414                &format!("strict {counts:?}"),
1415            );
1416        }
1417    }
1418
1419    fn filtered(c: Comprehension) -> Comprehension {
1420        Comprehension::filter(c, "true")
1421    }
1422
1423    /// The combinators report a filtered operand's bound as a bound: a
1424    /// product and a truncating zip over one are at most, not exactly,
1425    /// their count.
1426    #[test]
1427    fn an_operand_at_most_makes_a_combination_at_most() {
1428        let at_most = || filtered(clause("k", &[1, 2, 3, 4, 5, 6, 7, 8, 9]));
1429        let colors = || clause("c", &[1, 2]);
1430        let product = Comprehension::cartesian(vec![at_most(), colors()]);
1431        assert_eq!(
1432            product.metadata().cardinality,
1433            CardinalityClass::BoundedAtMost(18)
1434        );
1435        let zip = Comprehension::zip(vec![at_most(), colors()], ZipMode::Truncate);
1436        assert_eq!(
1437            zip.metadata().cardinality,
1438            CardinalityClass::BoundedAtMost(2)
1439        );
1440        let zip = Comprehension::zip(vec![unknown_count("u"), colors()], ZipMode::Truncate);
1441        assert_eq!(
1442            zip.metadata().cardinality,
1443            CardinalityClass::BoundedAtMost(2)
1444        );
1445        let product = Comprehension::cartesian(vec![unknown_count("u"), clause("e", &[])]);
1446        assert_eq!(product.metadata().cardinality, CardinalityClass::Bounded(0));
1447        let empty = filtered(clause("e", &[]));
1448        assert_eq!(empty.metadata().cardinality, CardinalityClass::Bounded(0));
1449    }
1450
1451    /// An order keeps `min(count, n)` under a strategy that truncates by
1452    /// tuple count, at most its input's count under one that truncates
1453    /// by strata, and `n` samples of a continuous space unless a filter
1454    /// may leave fewer.
1455    #[test]
1456    fn an_order_counts_by_its_strategy() {
1457        let order = |c, s, t| Comprehension::order(c, s, t).metadata().cardinality;
1458        let ks = || clause("k", &[1, 2, 3, 4, 5, 6]);
1459        for s in [
1460            StrategyName::Lex,
1461            StrategyName::ReverseLex,
1462            StrategyName::Diagonal,
1463            StrategyName::Halton,
1464            StrategyName::Sobol,
1465            StrategyName::Lhs,
1466            StrategyName::Shuffle,
1467        ] {
1468            assert_eq!(
1469                order(ks(), s, Some(4)),
1470                CardinalityClass::Bounded(4),
1471                "{s:?}"
1472            );
1473            assert_eq!(
1474                order(ks(), s, Some(9)),
1475                CardinalityClass::Bounded(6),
1476                "{s:?}"
1477            );
1478            assert_eq!(order(ks(), s, None), CardinalityClass::Bounded(6), "{s:?}");
1479            assert_eq!(
1480                order(filtered(ks()), s, Some(4)),
1481                CardinalityClass::BoundedAtMost(4),
1482                "{s:?}"
1483            );
1484            assert_eq!(
1485                order(unknown_count("u"), s, Some(4)),
1486                CardinalityClass::BoundedAtMost(4),
1487                "{s:?}"
1488            );
1489        }
1490        for s in [StrategyName::Extrema, StrategyName::Shells] {
1491            assert_eq!(
1492                order(ks(), s, Some(1)),
1493                CardinalityClass::BoundedAtMost(6),
1494                "{s:?}"
1495            );
1496            assert_eq!(order(ks(), s, None), CardinalityClass::Bounded(6), "{s:?}");
1497            assert_eq!(
1498                order(clause("e", &[]), s, Some(1)),
1499                CardinalityClass::Bounded(0)
1500            );
1501        }
1502        let space = || Comprehension::cartesian(vec![clause("k", &[1, 2]), continuous_clause("u")]);
1503        assert_eq!(
1504            order(space(), StrategyName::Halton, Some(5)),
1505            CardinalityClass::Bounded(5)
1506        );
1507        assert_eq!(
1508            order(filtered(space()), StrategyName::Halton, Some(5)),
1509            CardinalityClass::BoundedAtMost(5)
1510        );
1511        assert_eq!(
1512            order(space(), StrategyName::Extrema, Some(1)),
1513            CardinalityClass::BoundedAtMost(4)
1514        );
1515        let empty_axis = Comprehension::cartesian(vec![clause("k", &[]), continuous_clause("u")]);
1516        assert_eq!(
1517            order(empty_axis, StrategyName::Sobol, Some(5)),
1518            CardinalityClass::Bounded(0)
1519        );
1520    }
1521
1522    /// Two operands of unknown count: one streams, the other has no
1523    /// bound to buffer.
1524    #[test]
1525    fn zip_cycle_with_two_unknown_counts_is_unbounded() {
1526        let c = Comprehension::zip(vec![unknown_count("x"), unknown_count("y")], ZipMode::Cycle);
1527        assert_eq!(
1528            c.metadata().materialization,
1529            Materialization::UnboundedBarrier
1530        );
1531    }
1532
1533    #[test]
1534    fn union_produces_concatenation_index_fn() {
1535        let a = Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10])]);
1536        let b = Comprehension::cartesian(vec![clause("k", &[3, 4]), clause("limit", &[20])]);
1537        let u = Comprehension::union(vec![a, b]);
1538        let m = u.metadata();
1539        assert_eq!(m.cardinality, CardinalityClass::Bounded(4));
1540        assert_eq!(
1541            m.index_addressable,
1542            Some(IndexFn::Concatenation {
1543                segment_sizes: vec![2, 2]
1544            })
1545        );
1546        assert_eq!(m.natural_order, NaturalOrder::Sequential);
1547    }
1548
1549    #[test]
1550    fn filter_destroys_addressability() {
1551        let inner =
1552            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20])]);
1553        let filtered = Comprehension::filter(inner, "{k} > 0");
1554        let m = filtered.metadata();
1555        assert_eq!(m.cardinality, CardinalityClass::BoundedAtMost(4));
1556        assert_eq!(m.index_addressable, None);
1557    }
1558
1559    /// An untruncated `Lex` order passes its input's addressing through;
1560    /// a truncated one selects a prefix of its input's positions, one
1561    /// axis as long as the prefix, and over a filter it addresses
1562    /// nothing.
1563    #[test]
1564    fn lex_order_inherits_addressability_untruncated() {
1565        let inner =
1566            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20])]);
1567        let whole = Comprehension::order(inner.clone(), StrategyName::Lex, None).metadata();
1568        assert_eq!(whole.cardinality, CardinalityClass::Bounded(4));
1569        assert_eq!(
1570            whole.index_addressable,
1571            Some(IndexFn::Lattice {
1572                axis_sizes: vec![2, 2]
1573            })
1574        );
1575        assert_eq!(whole.natural_order, NaturalOrder::Lex);
1576        let prefix = Comprehension::order(inner.clone(), StrategyName::Lex, Some(3)).metadata();
1577        assert_eq!(prefix.cardinality, CardinalityClass::Bounded(3));
1578        assert_eq!(
1579            prefix.index_addressable,
1580            Some(IndexFn::Lattice {
1581                axis_sizes: vec![3]
1582            })
1583        );
1584        assert_eq!(prefix.natural_order, NaturalOrder::Lex);
1585        let streamed = Comprehension::order(
1586            Comprehension::filter(inner, "{k} > 1"),
1587            StrategyName::Lex,
1588            Some(3),
1589        )
1590        .metadata();
1591        assert_eq!(streamed.index_addressable, None);
1592    }
1593
1594    /// Any other order addresses its output through its selection: one
1595    /// axis as long as the selection, over which the next order holds
1596    /// only its own selection.
1597    #[test]
1598    fn non_lex_order_addresses_its_selection() {
1599        let inner =
1600            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20])]);
1601        let ordered = Comprehension::order(inner, StrategyName::Halton, Some(2));
1602        let m = ordered.metadata();
1603        assert_eq!(
1604            m.index_addressable,
1605            Some(IndexFn::Lattice {
1606                axis_sizes: vec![2]
1607            })
1608        );
1609        let reordered = Comprehension::order(ordered.clone(), StrategyName::Shuffle, None);
1610        let r = reordered.metadata();
1611        assert_eq!(r.cardinality, CardinalityClass::Bounded(2));
1612        assert_eq!(
1613            r.index_addressable,
1614            Some(IndexFn::Lattice {
1615                axis_sizes: vec![2]
1616            })
1617        );
1618        assert_eq!(
1619            r.materialization,
1620            Materialization::BoundedBarrier {
1621                working_set_size: 2
1622            }
1623        );
1624        match m.natural_order {
1625            NaturalOrder::Strategy(StrategyName::Halton) => {}
1626            other => panic!("expected Strategy(Halton), got {other:?}"),
1627        }
1628        assert_eq!(
1629            m.materialization,
1630            Materialization::BoundedBarrier {
1631                working_set_size: 2
1632            }
1633        );
1634    }
1635
1636    #[test]
1637    fn continuous_sampling_yields_bounded_cardinality() {
1638        let inner =
1639            Comprehension::cartesian(vec![continuous_clause("alpha"), continuous_clause("beta")]);
1640        let ordered = Comprehension::order(inner, StrategyName::Halton, Some(100));
1641        let m = ordered.metadata();
1642        assert_eq!(m.cardinality, CardinalityClass::Bounded(100));
1643        assert_eq!(
1644            m.materialization,
1645            Materialization::BoundedBarrier {
1646                working_set_size: 100
1647            }
1648        );
1649    }
1650
1651    #[test]
1652    fn metadata_propagation_is_idempotent() {
1653        let c = Comprehension::order(
1654            Comprehension::filter(
1655                Comprehension::cartesian(vec![clause("k", &[1, 2, 3]), clause("limit", &[10, 20])]),
1656                "{k} * {limit} > 5",
1657            ),
1658            StrategyName::Halton,
1659            Some(5),
1660        );
1661        let m1 = c.metadata();
1662        let m2 = c.metadata();
1663        assert_eq!(m1, m2);
1664    }
1665
1666    #[test]
1667    fn has_continuous_axis_classifier() {
1668        let lat = IndexFn::Lattice {
1669            axis_sizes: vec![3, 4],
1670        };
1671        assert!(!lat.has_continuous_axis());
1672
1673        let cont = IndexFn::Continuous {
1674            intervals: vec![Interval::closed(0.0, 1.0)],
1675            measure: ProductMeasure::Uniform,
1676        };
1677        assert!(cont.has_continuous_axis());
1678    }
1679
1680    #[test]
1681    fn multi_axis_lattice_classifier() {
1682        assert!(
1683            IndexFn::Lattice {
1684                axis_sizes: vec![3, 4]
1685            }
1686            .is_multi_axis_lattice()
1687        );
1688        assert!(
1689            !IndexFn::Lattice {
1690                axis_sizes: vec![3]
1691            }
1692            .is_multi_axis_lattice()
1693        );
1694        assert!(
1695            !IndexFn::Continuous {
1696                intervals: vec![Interval::closed(0.0, 1.0), Interval::closed(0.0, 1.0)],
1697                measure: ProductMeasure::Uniform,
1698            }
1699            .is_multi_axis_lattice()
1700        );
1701    }
1702}