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/// Extract `{name}` interpolation references from source
1043/// expressions in a comprehension subtree. Sources that carry
1044/// raw strings (`Generator`, `WorkloadParamList`) are walked;
1045/// `Literal`, `IntRange`, `ContinuousInterval`, `Distribution`
1046/// contain no string references.
1047fn collect_source_name_references(c: &Comprehension) -> Vec<String> {
1048    let mut out = Vec::new();
1049    walk_source_refs(c, &mut out);
1050    out
1051}
1052
1053fn walk_source_refs(c: &Comprehension, out: &mut Vec<String>) {
1054    match c {
1055        Comprehension::Clause { source, .. } => {
1056            extract_source_refs(source, out);
1057        }
1058        Comprehension::Cartesian { children }
1059        | Comprehension::Zip { children, .. }
1060        | Comprehension::Union { children } => {
1061            for c in children {
1062                walk_source_refs(c, out);
1063            }
1064        }
1065        Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
1066            walk_source_refs(child, out);
1067        }
1068    }
1069}
1070
1071fn extract_source_refs(source: &Source, out: &mut Vec<String>) {
1072    let s = match source {
1073        Source::Generator { expr, .. } => expr.as_str(),
1074        Source::WorkloadParamList { name, .. } => name.as_str(),
1075        _ => return,
1076    };
1077    let bytes = s.as_bytes();
1078    let mut i = 0;
1079    while i < bytes.len() {
1080        if bytes[i] == b'{'
1081            && let Some(close) = s[i + 1..].find('}')
1082        {
1083            let name = s[i + 1..i + 1 + close].trim();
1084            if !name.is_empty()
1085                && name.chars().all(|c| c.is_alphanumeric() || c == '_')
1086                && !out.contains(&name.to_string())
1087            {
1088                out.push(name.to_string());
1089            }
1090            i += close + 2;
1091            continue;
1092        }
1093        i += 1;
1094    }
1095}
1096
1097#[cfg(test)]
1098mod tests {
1099    use super::*;
1100    use crate::comprehension::source::{LiteralValue, Source};
1101
1102    fn clause(name: &str, vs: &[i64]) -> Comprehension {
1103        Comprehension::clause(
1104            name,
1105            Source::Literal {
1106                values: vs.iter().map(|n| LiteralValue::Int(*n)).collect(),
1107            },
1108        )
1109    }
1110
1111    fn continuous_clause(name: &str) -> Comprehension {
1112        Comprehension::clause(
1113            name,
1114            Source::ContinuousInterval {
1115                interval: Interval::closed(0.0, 1.0),
1116                measure: ProductMeasure::Uniform,
1117            },
1118        )
1119    }
1120
1121    #[test]
1122    fn clause_metadata_for_bounded_source() {
1123        let m = clause("k", &[1, 2, 3]).metadata();
1124        assert_eq!(m.cardinality, CardinalityClass::Bounded(3));
1125        assert_eq!(
1126            m.index_addressable,
1127            Some(IndexFn::Lattice {
1128                axis_sizes: vec![3]
1129            })
1130        );
1131        assert_eq!(m.natural_order, NaturalOrder::Lex);
1132        assert_eq!(m.materialization, Materialization::Streaming);
1133    }
1134
1135    #[test]
1136    fn clause_metadata_for_continuous_source() {
1137        let m = continuous_clause("alpha").metadata();
1138        assert!(matches!(m.cardinality, CardinalityClass::Continuous { .. }));
1139        assert!(matches!(
1140            m.index_addressable,
1141            Some(IndexFn::Continuous { .. })
1142        ));
1143        assert_eq!(m.natural_order, NaturalOrder::PendingSampling);
1144        assert_eq!(m.materialization, Materialization::Streaming);
1145    }
1146
1147    #[test]
1148    fn cartesian_metadata_combines_lattice_axes() {
1149        let c =
1150            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20, 30])]);
1151        let m = c.metadata();
1152        assert_eq!(m.cardinality, CardinalityClass::Bounded(6));
1153        assert_eq!(
1154            m.index_addressable,
1155            Some(IndexFn::Lattice {
1156                axis_sizes: vec![2, 3]
1157            })
1158        );
1159        assert_eq!(m.natural_order, NaturalOrder::Lex);
1160    }
1161
1162    #[test]
1163    fn cartesian_metadata_for_hybrid() {
1164        let c =
1165            Comprehension::cartesian(vec![clause("k", &[1, 2, 3, 4]), continuous_clause("theta")]);
1166        let m = c.metadata();
1167        match m.cardinality {
1168            CardinalityClass::Hybrid(h) => {
1169                assert_eq!(h.discrete_axes, vec![4]);
1170                assert_eq!(h.continuous_axes.len(), 1);
1171            }
1172            other => panic!("expected Hybrid, got {other:?}"),
1173        }
1174        assert!(matches!(m.index_addressable, Some(IndexFn::Hybrid { .. })));
1175        assert_eq!(m.natural_order, NaturalOrder::PendingSampling);
1176    }
1177
1178    #[test]
1179    fn dependent_cartesian_produces_none_addressable() {
1180        // clause replicas references {k} from the prior clause.
1181        let dependent = Comprehension::cartesian(vec![
1182            clause("k", &[1, 2, 3]),
1183            Comprehension::clause(
1184                "replicas",
1185                Source::Generator {
1186                    expr: "range(0, 2 * {k})".into(),
1187                    cardinality_hint: Some(6),
1188                },
1189            ),
1190        ]);
1191        let m = dependent.metadata();
1192        assert!(m.index_addressable.is_none());
1193    }
1194
1195    #[test]
1196    fn zip_strict_produces_lockstep_index_fn() {
1197        let c = Comprehension::zip(
1198            vec![clause("x", &[1, 2, 3]), clause("y", &[10, 20, 30])],
1199            ZipMode::Strict,
1200        );
1201        let m = c.metadata();
1202        assert_eq!(m.index_addressable, Some(IndexFn::Lockstep { length: 3 }));
1203        assert_eq!(m.natural_order, NaturalOrder::Lockstep);
1204        assert_eq!(m.materialization, Materialization::Streaming);
1205    }
1206
1207    #[test]
1208    fn zip_cycle_produces_modular_index_fn_and_barrier() {
1209        let c = Comprehension::zip(
1210            vec![clause("k", &[1, 2, 3, 4, 5]), clause("color", &[1, 2, 3])],
1211            ZipMode::Cycle,
1212        );
1213        let m = c.metadata();
1214        match m.index_addressable {
1215            Some(IndexFn::Modular { axis_sizes }) => {
1216                assert_eq!(axis_sizes, vec![5, 3]);
1217            }
1218            other => panic!("expected Modular, got {other:?}"),
1219        }
1220        // Both operands are addressable: cycling reads the shorter one
1221        // at `i mod 3` and buffers nothing.
1222        assert_eq!(m.materialization, Materialization::Streaming);
1223    }
1224
1225    fn unknown_count(name: &str) -> Comprehension {
1226        Comprehension::clause(
1227            name,
1228            Source::Generator {
1229                expr: "values({n})".into(),
1230                cardinality_hint: None,
1231            },
1232        )
1233    }
1234
1235    /// With an operand of unknown count, that operand streams and every
1236    /// finite operand that is not addressable is buffered in full;
1237    /// addressable ones are indexed.
1238    #[test]
1239    fn zip_cycle_with_an_unknown_count_buffers_every_finite_unaddressable_operand() {
1240        let c = Comprehension::zip(
1241            vec![
1242                unknown_count("tick"),
1243                Comprehension::filter(clause("a", &(0..1000).collect::<Vec<_>>()), "{a} > 1"),
1244                Comprehension::filter(clause("b", &[1, 2, 3]), "{b} > 0"),
1245                clause("color", &[1, 2, 3]),
1246            ],
1247            ZipMode::Cycle,
1248        );
1249        let plan = cycle_operands(
1250            &match &c {
1251                Comprehension::Zip { children, .. } => children,
1252                _ => unreachable!(),
1253            }
1254            .iter()
1255            .map(Comprehension::metadata)
1256            .collect::<Vec<_>>(),
1257        );
1258        assert_eq!(
1259            plan,
1260            vec![
1261                CycleOperand::Streamed,
1262                CycleOperand::Buffered { bound: Some(1000) },
1263                CycleOperand::Buffered { bound: Some(3) },
1264                CycleOperand::Indexed,
1265            ]
1266        );
1267        assert_eq!(
1268            c.metadata().materialization,
1269            Materialization::BoundedBarrier {
1270                working_set_size: 1003
1271            }
1272        );
1273    }
1274
1275    /// With every count known, the largest operand that is not
1276    /// addressable streams and the others are buffered.
1277    #[test]
1278    fn zip_cycle_streams_the_largest_unaddressable_operand() {
1279        let c = Comprehension::zip(
1280            vec![
1281                Comprehension::filter(clause("a", &[1, 2]), "{a} > 0"),
1282                Comprehension::filter(clause("b", &[1, 2, 3, 4]), "{b} > 0"),
1283                clause("k", &(0..100).collect::<Vec<_>>()),
1284            ],
1285            ZipMode::Cycle,
1286        );
1287        assert_eq!(
1288            c.metadata().materialization,
1289            Materialization::BoundedBarrier {
1290                working_set_size: 2
1291            }
1292        );
1293    }
1294
1295    /// An operand known empty empties the zip: no tuple, no index, and
1296    /// nothing held, in any position and beside operands of unknown
1297    /// count.
1298    #[test]
1299    fn zip_cycle_with_an_operand_known_empty_is_empty() {
1300        let operands = || {
1301            vec![
1302                unknown_count("tick"),
1303                Comprehension::filter(clause("a", &[1, 2, 3]), "{a} > 1"),
1304                clause("color", &[1, 2, 3]),
1305            ]
1306        };
1307        let empties = [
1308            clause("e", &[]),
1309            Comprehension::filter(clause("e", &[]), "{e} > 1"),
1310        ];
1311        for empty in &empties {
1312            for at in 0..=3 {
1313                let mut children = operands();
1314                children.insert(at, empty.clone());
1315                let plan = cycle_operands(
1316                    &children
1317                        .iter()
1318                        .map(Comprehension::metadata)
1319                        .collect::<Vec<_>>(),
1320                );
1321                assert_eq!(plan[at], CycleOperand::Buffered { bound: Some(0) });
1322                assert!(cycle_plan_is_empty(&plan));
1323                let m = Comprehension::zip(children, ZipMode::Cycle).metadata();
1324                assert_eq!(m.cardinality, CardinalityClass::Bounded(0));
1325                assert_eq!(m.materialization, Materialization::Streaming);
1326            }
1327        }
1328        let addressable = Comprehension::zip(
1329            vec![clause("k", &[1, 2, 3]), clause("e", &[])],
1330            ZipMode::Cycle,
1331        );
1332        let m = addressable.metadata();
1333        assert_eq!(
1334            m.index_addressable,
1335            Some(IndexFn::Modular {
1336                axis_sizes: vec![3, 0]
1337            })
1338        );
1339        assert_eq!(
1340            index_fn_cardinality(m.index_addressable.as_ref().unwrap()),
1341            0
1342        );
1343        assert_eq!(cycle_length(&[3, 0]), 0);
1344        assert_eq!(cycle_length(&[3, 5]), 5);
1345    }
1346
1347    /// An operand that may be empty at open makes the zip's count an
1348    /// upper bound.
1349    #[test]
1350    fn zip_cycle_over_an_operand_at_most_counts_at_most() {
1351        let c = Comprehension::zip(
1352            vec![
1353                clause("k", &[1, 2, 3, 4, 5]),
1354                Comprehension::filter(clause("a", &[1, 2]), "{a} > 1"),
1355            ],
1356            ZipMode::Cycle,
1357        );
1358        assert_eq!(c.metadata().cardinality, CardinalityClass::BoundedAtMost(5));
1359    }
1360
1361    /// The tuple counts an operand of count `c` may have: an exact
1362    /// count its own, an at-most count every count up to its bound, an
1363    /// unknown count small and large ones.
1364    fn witnesses(c: Count) -> Vec<u64> {
1365        match c {
1366            Count::Exact(n) => vec![n],
1367            Count::AtMost(n) => (0..=n).collect(),
1368            Count::Unknown => (0..=7).chain([1000]).collect(),
1369        }
1370    }
1371
1372    /// Assert `claim` holds for every count the operator yields over
1373    /// the operands' witnesses, and is tight: exact means every yield is
1374    /// that count, at most means none exceeds the bound and one meets
1375    /// it, and unknown means some yield exceeds any bound the operands
1376    /// state.
1377    fn assert_describes(claim: Count, yields: &[u64], what: &str) {
1378        let Some(&max) = yields.iter().max() else {
1379            return; // the operator never yields, as a strict zip of unequal operands
1380        };
1381        match claim {
1382            Count::Exact(n) => assert!(yields.iter().all(|&y| y == n), "{what}: {yields:?}"),
1383            Count::AtMost(n) => {
1384                assert!(max <= n, "{what}: {yields:?} exceed {n}");
1385                assert_eq!(max, n, "{what}: the bound is not tight");
1386                assert!(n > 0, "{what}: at most zero is exactly zero");
1387                assert!(
1388                    yields.iter().any(|&y| y != n),
1389                    "{what}: always {n}, so exact"
1390                );
1391            }
1392            Count::Unknown => assert!(max >= 1000, "{what}: bounded by {max}"),
1393        }
1394    }
1395
1396    const KINDS: [Count; 5] = [
1397        Count::Exact(0),
1398        Count::Exact(3),
1399        Count::Exact(5),
1400        Count::AtMost(4),
1401        Count::Unknown,
1402    ];
1403
1404    /// Every combination of the operands' witnesses.
1405    fn combinations(counts: &[Count]) -> Vec<Vec<u64>> {
1406        counts.iter().fold(vec![Vec::new()], |acc, c| {
1407            acc.iter()
1408                .flat_map(|prefix| {
1409                    witnesses(*c).into_iter().map(move |w| {
1410                        let mut next = prefix.clone();
1411                        next.push(w);
1412                        next
1413                    })
1414                })
1415                .collect()
1416        })
1417    }
1418
1419    /// Each operator's count over every pair and triple of kinds holds
1420    /// for, and is tight over, what the operator yields.
1421    #[test]
1422    fn every_kind_combination_counts_what_the_operator_yields() {
1423        let mut shapes: Vec<Vec<Count>> = Vec::new();
1424        for a in KINDS {
1425            for b in KINDS {
1426                shapes.push(vec![a, b]);
1427                for c in KINDS {
1428                    shapes.push(vec![a, b, c]);
1429                }
1430            }
1431        }
1432        for counts in &shapes {
1433            let combos = combinations(counts);
1434            let product: Vec<u64> = combos.iter().map(|c| c.iter().product()).collect();
1435            assert_describes(
1436                cartesian_count(counts),
1437                &product,
1438                &format!("cartesian {counts:?}"),
1439            );
1440            let sum: Vec<u64> = combos.iter().map(|c| c.iter().sum()).collect();
1441            assert_describes(union_count(counts), &sum, &format!("union {counts:?}"));
1442            let shortest: Vec<u64> = combos.iter().map(|c| *c.iter().min().unwrap()).collect();
1443            assert_describes(
1444                truncate_zip_count(counts),
1445                &shortest,
1446                &format!("truncate {counts:?}"),
1447            );
1448            let cycled: Vec<u64> = combos.iter().map(|c| cycle_length(c)).collect();
1449            assert_describes(
1450                cycle_zip_count(counts),
1451                &cycled,
1452                &format!("cycle {counts:?}"),
1453            );
1454            let strict: Vec<u64> = combos
1455                .iter()
1456                .filter(|c| c.iter().all(|&n| n == c[0]))
1457                .map(|c| c[0])
1458                .collect();
1459            assert_describes(
1460                strict_zip_count(counts),
1461                &strict,
1462                &format!("strict {counts:?}"),
1463            );
1464        }
1465    }
1466
1467    fn filtered(c: Comprehension) -> Comprehension {
1468        Comprehension::filter(c, "true")
1469    }
1470
1471    /// The combinators report a filtered operand's bound as a bound: a
1472    /// product and a truncating zip over one are at most, not exactly,
1473    /// their count.
1474    #[test]
1475    fn an_operand_at_most_makes_a_combination_at_most() {
1476        let at_most = || filtered(clause("k", &[1, 2, 3, 4, 5, 6, 7, 8, 9]));
1477        let colors = || clause("c", &[1, 2]);
1478        let product = Comprehension::cartesian(vec![at_most(), colors()]);
1479        assert_eq!(
1480            product.metadata().cardinality,
1481            CardinalityClass::BoundedAtMost(18)
1482        );
1483        let zip = Comprehension::zip(vec![at_most(), colors()], ZipMode::Truncate);
1484        assert_eq!(
1485            zip.metadata().cardinality,
1486            CardinalityClass::BoundedAtMost(2)
1487        );
1488        let zip = Comprehension::zip(vec![unknown_count("u"), colors()], ZipMode::Truncate);
1489        assert_eq!(
1490            zip.metadata().cardinality,
1491            CardinalityClass::BoundedAtMost(2)
1492        );
1493        let product = Comprehension::cartesian(vec![unknown_count("u"), clause("e", &[])]);
1494        assert_eq!(product.metadata().cardinality, CardinalityClass::Bounded(0));
1495        let empty = filtered(clause("e", &[]));
1496        assert_eq!(empty.metadata().cardinality, CardinalityClass::Bounded(0));
1497    }
1498
1499    /// An order keeps `min(count, n)` under a strategy that truncates by
1500    /// tuple count, at most its input's count under one that truncates
1501    /// by strata, and `n` samples of a continuous space unless a filter
1502    /// may leave fewer.
1503    #[test]
1504    fn an_order_counts_by_its_strategy() {
1505        let order = |c, s, t| Comprehension::order(c, s, t).metadata().cardinality;
1506        let ks = || clause("k", &[1, 2, 3, 4, 5, 6]);
1507        for s in [
1508            StrategyName::Lex,
1509            StrategyName::ReverseLex,
1510            StrategyName::Diagonal,
1511            StrategyName::Halton,
1512            StrategyName::Sobol,
1513            StrategyName::Lhs,
1514            StrategyName::Shuffle,
1515        ] {
1516            assert_eq!(
1517                order(ks(), s, Some(4)),
1518                CardinalityClass::Bounded(4),
1519                "{s:?}"
1520            );
1521            assert_eq!(
1522                order(ks(), s, Some(9)),
1523                CardinalityClass::Bounded(6),
1524                "{s:?}"
1525            );
1526            assert_eq!(order(ks(), s, None), CardinalityClass::Bounded(6), "{s:?}");
1527            assert_eq!(
1528                order(filtered(ks()), s, Some(4)),
1529                CardinalityClass::BoundedAtMost(4),
1530                "{s:?}"
1531            );
1532            assert_eq!(
1533                order(unknown_count("u"), s, Some(4)),
1534                CardinalityClass::BoundedAtMost(4),
1535                "{s:?}"
1536            );
1537        }
1538        for s in [StrategyName::Extrema, StrategyName::Shells] {
1539            assert_eq!(
1540                order(ks(), s, Some(1)),
1541                CardinalityClass::BoundedAtMost(6),
1542                "{s:?}"
1543            );
1544            assert_eq!(order(ks(), s, None), CardinalityClass::Bounded(6), "{s:?}");
1545            assert_eq!(
1546                order(clause("e", &[]), s, Some(1)),
1547                CardinalityClass::Bounded(0)
1548            );
1549        }
1550        let space = || Comprehension::cartesian(vec![clause("k", &[1, 2]), continuous_clause("u")]);
1551        assert_eq!(
1552            order(space(), StrategyName::Halton, Some(5)),
1553            CardinalityClass::Bounded(5)
1554        );
1555        assert_eq!(
1556            order(filtered(space()), StrategyName::Halton, Some(5)),
1557            CardinalityClass::BoundedAtMost(5)
1558        );
1559        assert_eq!(
1560            order(space(), StrategyName::Extrema, Some(1)),
1561            CardinalityClass::BoundedAtMost(4)
1562        );
1563        let empty_axis = Comprehension::cartesian(vec![clause("k", &[]), continuous_clause("u")]);
1564        assert_eq!(
1565            order(empty_axis, StrategyName::Sobol, Some(5)),
1566            CardinalityClass::Bounded(0)
1567        );
1568    }
1569
1570    /// Two operands of unknown count: one streams, the other has no
1571    /// bound to buffer.
1572    #[test]
1573    fn zip_cycle_with_two_unknown_counts_is_unbounded() {
1574        let c = Comprehension::zip(vec![unknown_count("x"), unknown_count("y")], ZipMode::Cycle);
1575        assert_eq!(
1576            c.metadata().materialization,
1577            Materialization::UnboundedBarrier
1578        );
1579    }
1580
1581    #[test]
1582    fn union_produces_concatenation_index_fn() {
1583        let a = Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10])]);
1584        let b = Comprehension::cartesian(vec![clause("k", &[3, 4]), clause("limit", &[20])]);
1585        let u = Comprehension::union(vec![a, b]);
1586        let m = u.metadata();
1587        assert_eq!(m.cardinality, CardinalityClass::Bounded(4));
1588        assert_eq!(
1589            m.index_addressable,
1590            Some(IndexFn::Concatenation {
1591                segment_sizes: vec![2, 2]
1592            })
1593        );
1594        assert_eq!(m.natural_order, NaturalOrder::Sequential);
1595    }
1596
1597    #[test]
1598    fn filter_destroys_addressability() {
1599        let inner =
1600            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20])]);
1601        let filtered = Comprehension::filter(inner, "{k} > 0");
1602        let m = filtered.metadata();
1603        assert_eq!(m.cardinality, CardinalityClass::BoundedAtMost(4));
1604        assert_eq!(m.index_addressable, None);
1605    }
1606
1607    /// An untruncated `Lex` order passes its input's addressing through;
1608    /// a truncated one selects a prefix of its input's positions, one
1609    /// axis as long as the prefix, and over a filter it addresses
1610    /// nothing.
1611    #[test]
1612    fn lex_order_inherits_addressability_untruncated() {
1613        let inner =
1614            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20])]);
1615        let whole = Comprehension::order(inner.clone(), StrategyName::Lex, None).metadata();
1616        assert_eq!(whole.cardinality, CardinalityClass::Bounded(4));
1617        assert_eq!(
1618            whole.index_addressable,
1619            Some(IndexFn::Lattice {
1620                axis_sizes: vec![2, 2]
1621            })
1622        );
1623        assert_eq!(whole.natural_order, NaturalOrder::Lex);
1624        let prefix = Comprehension::order(inner.clone(), StrategyName::Lex, Some(3)).metadata();
1625        assert_eq!(prefix.cardinality, CardinalityClass::Bounded(3));
1626        assert_eq!(
1627            prefix.index_addressable,
1628            Some(IndexFn::Lattice {
1629                axis_sizes: vec![3]
1630            })
1631        );
1632        assert_eq!(prefix.natural_order, NaturalOrder::Lex);
1633        let streamed = Comprehension::order(
1634            Comprehension::filter(inner, "{k} > 1"),
1635            StrategyName::Lex,
1636            Some(3),
1637        )
1638        .metadata();
1639        assert_eq!(streamed.index_addressable, None);
1640    }
1641
1642    /// Any other order addresses its output through its selection: one
1643    /// axis as long as the selection, over which the next order holds
1644    /// only its own selection.
1645    #[test]
1646    fn non_lex_order_addresses_its_selection() {
1647        let inner =
1648            Comprehension::cartesian(vec![clause("k", &[1, 2]), clause("limit", &[10, 20])]);
1649        let ordered = Comprehension::order(inner, StrategyName::Halton, Some(2));
1650        let m = ordered.metadata();
1651        assert_eq!(
1652            m.index_addressable,
1653            Some(IndexFn::Lattice {
1654                axis_sizes: vec![2]
1655            })
1656        );
1657        let reordered = Comprehension::order(ordered.clone(), StrategyName::Shuffle, None);
1658        let r = reordered.metadata();
1659        assert_eq!(r.cardinality, CardinalityClass::Bounded(2));
1660        assert_eq!(
1661            r.index_addressable,
1662            Some(IndexFn::Lattice {
1663                axis_sizes: vec![2]
1664            })
1665        );
1666        assert_eq!(
1667            r.materialization,
1668            Materialization::BoundedBarrier {
1669                working_set_size: 2
1670            }
1671        );
1672        match m.natural_order {
1673            NaturalOrder::Strategy(StrategyName::Halton) => {}
1674            other => panic!("expected Strategy(Halton), got {other:?}"),
1675        }
1676        assert_eq!(
1677            m.materialization,
1678            Materialization::BoundedBarrier {
1679                working_set_size: 2
1680            }
1681        );
1682    }
1683
1684    #[test]
1685    fn continuous_sampling_yields_bounded_cardinality() {
1686        let inner =
1687            Comprehension::cartesian(vec![continuous_clause("alpha"), continuous_clause("beta")]);
1688        let ordered = Comprehension::order(inner, StrategyName::Halton, Some(100));
1689        let m = ordered.metadata();
1690        assert_eq!(m.cardinality, CardinalityClass::Bounded(100));
1691        assert_eq!(
1692            m.materialization,
1693            Materialization::BoundedBarrier {
1694                working_set_size: 100
1695            }
1696        );
1697    }
1698
1699    #[test]
1700    fn metadata_propagation_is_idempotent() {
1701        let c = Comprehension::order(
1702            Comprehension::filter(
1703                Comprehension::cartesian(vec![clause("k", &[1, 2, 3]), clause("limit", &[10, 20])]),
1704                "{k} * {limit} > 5",
1705            ),
1706            StrategyName::Halton,
1707            Some(5),
1708        );
1709        let m1 = c.metadata();
1710        let m2 = c.metadata();
1711        assert_eq!(m1, m2);
1712    }
1713
1714    #[test]
1715    fn has_continuous_axis_classifier() {
1716        let lat = IndexFn::Lattice {
1717            axis_sizes: vec![3, 4],
1718        };
1719        assert!(!lat.has_continuous_axis());
1720
1721        let cont = IndexFn::Continuous {
1722            intervals: vec![Interval::closed(0.0, 1.0)],
1723            measure: ProductMeasure::Uniform,
1724        };
1725        assert!(cont.has_continuous_axis());
1726    }
1727
1728    #[test]
1729    fn multi_axis_lattice_classifier() {
1730        assert!(
1731            IndexFn::Lattice {
1732                axis_sizes: vec![3, 4]
1733            }
1734            .is_multi_axis_lattice()
1735        );
1736        assert!(
1737            !IndexFn::Lattice {
1738                axis_sizes: vec![3]
1739            }
1740            .is_multi_axis_lattice()
1741        );
1742        assert!(
1743            !IndexFn::Continuous {
1744                intervals: vec![Interval::closed(0.0, 1.0), Interval::closed(0.0, 1.0)],
1745                measure: ProductMeasure::Uniform,
1746            }
1747            .is_multi_axis_lattice()
1748        );
1749    }
1750}