polydat_core/iteration/comprehension/strategies/
mod.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Strategy implementations — spec §3.6 + §10.2 R2 + §10.7.8.
5//!
6//! ## Single invocation surface
7//!
8//! Every named strategy exposes one public entry point —
9//! [`Strategy::apply`]. The caller passes an [`EvaluatedInput`]
10//! carrying the materialized tuples, their cardinality, and the
11//! `IndexFn` they actually satisfy. The strategy decides
12//! internally whether to dispatch its closed-form indexed
13//! algorithm (when `has_closed_form_for(&input.index_fn)`) or
14//! its fallback reorder over the materialized tuples.
15//!
16//! Per spec §10.7.8 this is the **strategy invocation
17//! contract**: V4 fires at `apply` time against the
18//! `EvaluatedInput`'s `index_fn` — definitively, regardless of
19//! how the input source was authored (literal, range,
20//! registry-recognized generator, or workload-param).
21//!
22//! ## Internal split
23//!
24//! Per-strategy modules organise the implementation into two
25//! private helpers (`apply_indexed` for the R2 closed-form path
26//! when applicable, `apply_naive` for the generic fallback);
27//! [`Strategy::apply`] is the dispatcher. The trait surface
28//! exposes only the dispatcher plus the V4/R2 introspection
29//! predicates ([`Strategy::accepts_input`],
30//! [`Strategy::has_closed_form_for`]).
31//!
32//! Strategies are selected by [`StrategyName`]; [`for_name`]
33//! dispatches a strategy name to its boxed [`Strategy`] impl.
34
35use super::metadata::IndexFn;
36use super::strategy::StrategyName;
37
38pub mod antidiagonal;
39pub mod diagonal;
40pub mod extrema;
41pub mod halton;
42pub mod lex;
43pub mod lhs;
44pub mod prng;
45pub mod reverse_lex;
46pub mod shells;
47pub mod shuffle;
48pub mod sobol;
49
50/// A multi-coordinate index. Each component is the per-axis
51/// position in the input's index space. Length equals the
52/// input's dimensionality (1 for `Lockstep` / `Modular` /
53/// `Concatenation`; N for `Lattice` / `Continuous` /
54/// `Hybrid`).
55///
56/// `MultiIndex` is the indexed-form output type. The R2 IR
57/// opcode emitted by the IR compiler consumes these and resolves
58/// each through the input's `IndexFn` to dispense the actual
59/// tuple.
60pub type MultiIndex = Vec<u64>;
61
62/// A named-tuple value. Subset of the polydat `Value` set that
63/// is the strategy layer's currency; the runtime walker
64/// converts `Value`s to it before `apply` and maps results
65/// back. For the strategy module in isolation, this
66/// lightweight type lets tests run without pulling in the
67/// broader runtime.
68#[derive(Debug, Clone, PartialEq)]
69pub struct Tuple {
70    /// The tuple's `(name, value)` pairs, in shape order.
71    pub bindings: Vec<(String, TupleValue)>,
72}
73
74/// Subset of polydat's `Value` enum. `TupleValue` is the
75/// strategy layer's currency; the runtime walker converts
76/// `Value`s to it before `apply` and maps results back.
77#[derive(Debug, Clone, PartialEq)]
78pub enum TupleValue {
79    /// An unsigned integer.
80    U64(u64),
81    /// A signed integer.
82    I64(i64),
83    /// A float.
84    F64(f64),
85    /// A string.
86    Str(String),
87    /// A boolean.
88    Bool(bool),
89}
90
91impl Tuple {
92    /// An empty tuple.
93    pub fn new() -> Self {
94        Self {
95            bindings: Vec::new(),
96        }
97    }
98
99    /// The tuple with one more binding.
100    pub fn with<K: Into<String>>(mut self, key: K, value: TupleValue) -> Self {
101        self.bindings.push((key.into(), value));
102        self
103    }
104}
105
106impl Default for Tuple {
107    fn default() -> Self {
108        Self::new()
109    }
110}
111
112/// The materialized input to a strategy at invocation time
113/// (spec §10.7.8).
114///
115/// `tuples` are the input stream's tuples in source order (the
116/// natural enumeration of the upstream comprehension subtree).
117/// `cardinality` matches `tuples.len() as u64`. `index_fn` is
118/// the addressing scheme the input actually satisfies —
119/// derived from observed shape for Generator /
120/// WorkloadParamList leaves via the [`crate::iteration::comprehension::eval_source`]
121/// layer, combined upward by the runtime walker per spec
122/// §10.7.2 propagation rules.
123pub struct EvaluatedInput {
124    /// The input's tuples, in source order.
125    pub tuples: Vec<Tuple>,
126    /// How many tuples: `tuples.len()`.
127    pub cardinality: u64,
128    /// The addressing scheme the input satisfies.
129    pub index_fn: IndexFn,
130}
131
132/// The strategy invocation surface per spec §10.7.8.
133///
134/// Implementations are stateless — every call to [`apply`](Strategy::apply)
135/// produces the same output given the same inputs
136/// (deterministic). PRNG-based strategies (`Shuffle`, `Lhs`)
137/// derive their state from the authored seed of
138/// [`apply_seeded`](Strategy::apply_seeded), or a module constant
139/// when none is authored, plus the input length; no per-streamer
140/// seed is threaded.
141pub trait Strategy {
142    /// The strategy's name. Mirrors [`StrategyName`].
143    fn name(&self) -> StrategyName;
144
145    /// V4 input-shape check (spec §3.6). `None` represents an
146    /// input with no closed-form index function; only `Lex`
147    /// accepts that. Concrete `IndexFn` variants are accepted
148    /// per the per-strategy rules in spec §3.6's table.
149    fn accepts_input(&self, idx: Option<&IndexFn>) -> bool;
150
151    /// R2 push-down eligibility (spec §10.2 R2). `true` if
152    /// this strategy has a closed-form indexed lookup over the
153    /// given input. If `false`, [`apply`](Strategy::apply) uses the strategy's
154    /// fallback reorder over the materialized tuples.
155    fn has_closed_form_for(&self, idx: &IndexFn) -> bool;
156
157    /// Apply this strategy to the given input.
158    ///
159    /// Internally dispatches: when the strategy has a
160    /// closed-form rule for `input.index_fn`, it uses the
161    /// indexed-form algorithm (compute multi-indices over the
162    /// index space, look up against `input.tuples` via
163    /// [`multi_index_to_flat`]). Otherwise it falls back to a
164    /// per-strategy reorder over `input.tuples` directly.
165    ///
166    /// V4 is the caller's responsibility — call
167    /// `accepts_input(Some(&input.index_fn))` before `apply`
168    /// to fire V4 at strategy-invocation time per spec §10.7.8.
169    fn apply(&self, input: &EvaluatedInput, truncation: Option<u64>) -> Vec<Tuple>;
170
171    /// [`apply`](Strategy::apply) under an authored seed
172    /// (comprehension_forms.md §3.6): a seeded strategy (`Shuffle`,
173    /// `Lhs`) derives its state from `seed` and the input's
174    /// structural identity, and from its fixed default when `seed`
175    /// is `None`. Every other strategy ignores the seed.
176    fn apply_seeded(
177        &self,
178        input: &EvaluatedInput,
179        truncation: Option<u64>,
180        seed: Option<u64>,
181    ) -> Vec<Tuple> {
182        let _ = seed;
183        self.apply(input, truncation)
184    }
185}
186
187/// Dispatch a [`StrategyName`] to its concrete [`Strategy`]
188/// implementation. The returned trait object is stateless;
189/// callers can hold a single instance per strategy name for
190/// the life of the process if desired.
191pub fn for_name(name: StrategyName) -> Box<dyn Strategy + Send + Sync> {
192    match name {
193        StrategyName::Lex => Box::new(lex::Lex),
194        StrategyName::ReverseLex => Box::new(reverse_lex::ReverseLex),
195        StrategyName::Shuffle => Box::new(shuffle::Shuffle),
196        StrategyName::Halton => Box::new(halton::Halton),
197        StrategyName::Sobol => Box::new(sobol::Sobol),
198        StrategyName::Lhs => Box::new(lhs::Lhs),
199        StrategyName::Extrema => Box::new(extrema::Extrema),
200        StrategyName::Shells => Box::new(shells::Shells),
201        StrategyName::Diagonal => Box::new(diagonal::Diagonal),
202        StrategyName::Antidiagonal => Box::new(antidiagonal::Antidiagonal),
203    }
204}
205
206/// Resolve a [`MultiIndex`] to a flat position in the
207/// input's tuple list, given the input's [`IndexFn`].
208///
209/// The flat position matches the natural enumeration order
210/// the runtime walker produces:
211///
212/// - `Lattice { axis_sizes: [s0, s1, …, sN-1] }` — row-major
213///   over the axes: `flat = i0 * s1 * s2 * … + i1 * s2 * … + … + iN-1`.
214///   This matches the runtime walker's cartesian enumeration
215///   (head axis varies slowest, tail nested).
216/// - `Lockstep { length }` — one-axis identity:
217///   `flat = mi[0]`.
218/// - `Modular { axis_sizes }` — one-axis identity over `max(axis_sizes)`:
219///   `flat = mi[0]`.
220/// - `Concatenation { segment_sizes }` — one-axis identity
221///   over `Σ segment_sizes`: `flat = mi[0]`.
222/// - `Continuous` / `Hybrid` — `None`; these inputs have no
223///   pre-materialized tuple list (the strategy's multi-indices
224///   are quantiles, not lookups).
225///
226/// Returns `None` for out-of-range positions or dimension
227/// mismatches.
228pub fn multi_index_to_flat(idx: &IndexFn, mi: &MultiIndex) -> Option<usize> {
229    match idx {
230        IndexFn::Lattice { axis_sizes } => {
231            if mi.len() != axis_sizes.len() {
232                return None;
233            }
234            let mut flat: u64 = 0;
235            let mut stride: u64 = 1;
236            for i in (0..axis_sizes.len()).rev() {
237                let pos = mi[i];
238                let size = axis_sizes[i];
239                if pos >= size {
240                    return None;
241                }
242                flat = flat.checked_add(pos.checked_mul(stride)?)?;
243                stride = stride.checked_mul(size)?;
244            }
245            Some(flat as usize)
246        }
247        IndexFn::Lockstep { length } => {
248            if mi.len() != 1 || mi[0] >= *length {
249                return None;
250            }
251            Some(mi[0] as usize)
252        }
253        IndexFn::Modular { axis_sizes } => {
254            let max = axis_sizes.iter().copied().max().unwrap_or(0);
255            if mi.len() != 1 || mi[0] >= max {
256                return None;
257            }
258            Some(mi[0] as usize)
259        }
260        IndexFn::Concatenation { segment_sizes } => {
261            let total: u64 = segment_sizes.iter().copied().sum();
262            if mi.len() != 1 || mi[0] >= total {
263                return None;
264            }
265            Some(mi[0] as usize)
266        }
267        IndexFn::Continuous { .. } | IndexFn::Hybrid { .. } => None,
268    }
269}
270
271/// `true` when [`multi_index_to_flat`] returns a usable
272/// position for in-range multi-indices over this `IndexFn`.
273/// `false` for `Continuous` / `Hybrid` where the indexed
274/// strategy emits quantiles, not lookups.
275pub fn index_fn_supports_lookup(idx: &IndexFn) -> bool {
276    !matches!(idx, IndexFn::Continuous { .. } | IndexFn::Hybrid { .. })
277}
278
279/// Cardinality of an `IndexFn`. Used by strategies to size
280/// their output when no truncation is specified. Mirrors the
281/// helper in `metadata.rs` but lives here to avoid a circular
282/// dependency.
283pub(crate) fn index_fn_size(idx: &IndexFn) -> u64 {
284    match idx {
285        IndexFn::Lattice { axis_sizes } => axis_sizes
286            .iter()
287            .copied()
288            .fold(1u64, |a, b| a.saturating_mul(b)),
289        IndexFn::Lockstep { length } => *length,
290        IndexFn::Modular { axis_sizes } => axis_sizes.iter().copied().max().unwrap_or(0),
291        IndexFn::Concatenation { segment_sizes } => segment_sizes
292            .iter()
293            .copied()
294            .fold(0u64, |a, b| a.saturating_add(b)),
295        IndexFn::Continuous { .. } | IndexFn::Hybrid { .. } => 0,
296    }
297}
298
299/// Lattice dimensionality of an `IndexFn`. Used by strategies
300/// that branch on dimensionality (Extrema's corner count,
301/// Lhs's per-axis stratification).
302pub(crate) fn index_fn_dim(idx: &IndexFn) -> usize {
303    match idx {
304        IndexFn::Lattice { axis_sizes } => axis_sizes.len(),
305        IndexFn::Continuous { intervals, .. } => intervals.len(),
306        IndexFn::Hybrid {
307            discrete_axes,
308            continuous_axes,
309            ..
310        } => discrete_axes.len() + continuous_axes.len(),
311        IndexFn::Lockstep { .. } | IndexFn::Modular { .. } => 1,
312        IndexFn::Concatenation { segment_sizes } => segment_sizes.len(),
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319
320    #[test]
321    fn for_name_dispatches_to_correct_strategy() {
322        assert_eq!(for_name(StrategyName::Lex).name(), StrategyName::Lex);
323        assert_eq!(for_name(StrategyName::Halton).name(), StrategyName::Halton);
324        assert_eq!(
325            for_name(StrategyName::Extrema).name(),
326            StrategyName::Extrema
327        );
328    }
329
330    #[test]
331    fn index_fn_size_lattice() {
332        let idx = IndexFn::Lattice {
333            axis_sizes: vec![3, 4, 5],
334        };
335        assert_eq!(index_fn_size(&idx), 60);
336    }
337
338    #[test]
339    fn index_fn_size_concatenation() {
340        let idx = IndexFn::Concatenation {
341            segment_sizes: vec![10, 20, 30],
342        };
343        assert_eq!(index_fn_size(&idx), 60);
344    }
345
346    #[test]
347    fn index_fn_dim_classifies_correctly() {
348        assert_eq!(
349            index_fn_dim(&IndexFn::Lattice {
350                axis_sizes: vec![3, 4]
351            }),
352            2
353        );
354        assert_eq!(index_fn_dim(&IndexFn::Lockstep { length: 10 }), 1);
355        assert_eq!(
356            index_fn_dim(&IndexFn::Concatenation {
357                segment_sizes: vec![1, 2, 3]
358            }),
359            3
360        );
361    }
362}
polydat_core/iteration/comprehension/strategies/mod.rs

polydat_core/iteration/comprehension/strategies/
mod.rs
