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 seed from a module constant plus the input
138/// length; no per-streamer seed is threaded.
139pub trait Strategy {
140    /// The strategy's name. Mirrors [`StrategyName`].
141    fn name(&self) -> StrategyName;
142
143    /// V4 input-shape check (spec §3.6). `None` represents an
144    /// input with no closed-form index function; only `Lex`
145    /// accepts that. Concrete `IndexFn` variants are accepted
146    /// per the per-strategy rules in spec §3.6's table.
147    fn accepts_input(&self, idx: Option<&IndexFn>) -> bool;
148
149    /// R2 push-down eligibility (spec §10.2 R2). `true` if
150    /// this strategy has a closed-form indexed lookup over the
151    /// given input. If `false`, [`apply`](Strategy::apply) uses the strategy's
152    /// fallback reorder over the materialized tuples.
153    fn has_closed_form_for(&self, idx: &IndexFn) -> bool;
154
155    /// Apply this strategy to the given input.
156    ///
157    /// Internally dispatches: when the strategy has a
158    /// closed-form rule for `input.index_fn`, it uses the
159    /// indexed-form algorithm (compute multi-indices over the
160    /// index space, look up against `input.tuples` via
161    /// [`multi_index_to_flat`]). Otherwise it falls back to a
162    /// per-strategy reorder over `input.tuples` directly.
163    ///
164    /// V4 is the caller's responsibility — call
165    /// `accepts_input(Some(&input.index_fn))` before `apply`
166    /// to fire V4 at strategy-invocation time per spec §10.7.8.
167    fn apply(&self, input: &EvaluatedInput, truncation: Option<u64>) -> Vec<Tuple>;
168}
169
170/// Dispatch a [`StrategyName`] to its concrete [`Strategy`]
171/// implementation. The returned trait object is stateless;
172/// callers can hold a single instance per strategy name for
173/// the life of the process if desired.
174pub fn for_name(name: StrategyName) -> Box<dyn Strategy + Send + Sync> {
175    match name {
176        StrategyName::Lex => Box::new(lex::Lex),
177        StrategyName::ReverseLex => Box::new(reverse_lex::ReverseLex),
178        StrategyName::Shuffle => Box::new(shuffle::Shuffle),
179        StrategyName::Halton => Box::new(halton::Halton),
180        StrategyName::Sobol => Box::new(sobol::Sobol),
181        StrategyName::Lhs => Box::new(lhs::Lhs),
182        StrategyName::Extrema => Box::new(extrema::Extrema),
183        StrategyName::Shells => Box::new(shells::Shells),
184        StrategyName::Diagonal => Box::new(diagonal::Diagonal),
185        StrategyName::Antidiagonal => Box::new(antidiagonal::Antidiagonal),
186    }
187}
188
189/// Resolve a [`MultiIndex`] to a flat position in the
190/// input's tuple list, given the input's [`IndexFn`].
191///
192/// The flat position matches the natural enumeration order
193/// the runtime walker produces:
194///
195/// - `Lattice { axis_sizes: [s0, s1, …, sN-1] }` — row-major
196///   over the axes: `flat = i0 * s1 * s2 * … + i1 * s2 * … + … + iN-1`.
197///   This matches the runtime walker's cartesian enumeration
198///   (head axis varies slowest, tail nested).
199/// - `Lockstep { length }` — one-axis identity:
200///   `flat = mi[0]`.
201/// - `Modular { axis_sizes }` — one-axis identity over `max(axis_sizes)`:
202///   `flat = mi[0]`.
203/// - `Concatenation { segment_sizes }` — one-axis identity
204///   over `Σ segment_sizes`: `flat = mi[0]`.
205/// - `Continuous` / `Hybrid` — `None`; these inputs have no
206///   pre-materialized tuple list (the strategy's multi-indices
207///   are quantiles, not lookups).
208///
209/// Returns `None` for out-of-range positions or dimension
210/// mismatches.
211pub fn multi_index_to_flat(idx: &IndexFn, mi: &MultiIndex) -> Option<usize> {
212    match idx {
213        IndexFn::Lattice { axis_sizes } => {
214            if mi.len() != axis_sizes.len() {
215                return None;
216            }
217            let mut flat: u64 = 0;
218            let mut stride: u64 = 1;
219            for i in (0..axis_sizes.len()).rev() {
220                let pos = mi[i];
221                let size = axis_sizes[i];
222                if pos >= size {
223                    return None;
224                }
225                flat = flat.checked_add(pos.checked_mul(stride)?)?;
226                stride = stride.checked_mul(size)?;
227            }
228            Some(flat as usize)
229        }
230        IndexFn::Lockstep { length } => {
231            if mi.len() != 1 || mi[0] >= *length {
232                return None;
233            }
234            Some(mi[0] as usize)
235        }
236        IndexFn::Modular { axis_sizes } => {
237            let max = axis_sizes.iter().copied().max().unwrap_or(0);
238            if mi.len() != 1 || mi[0] >= max {
239                return None;
240            }
241            Some(mi[0] as usize)
242        }
243        IndexFn::Concatenation { segment_sizes } => {
244            let total: u64 = segment_sizes.iter().copied().sum();
245            if mi.len() != 1 || mi[0] >= total {
246                return None;
247            }
248            Some(mi[0] as usize)
249        }
250        IndexFn::Continuous { .. } | IndexFn::Hybrid { .. } => None,
251    }
252}
253
254/// `true` when [`multi_index_to_flat`] returns a usable
255/// position for in-range multi-indices over this `IndexFn`.
256/// `false` for `Continuous` / `Hybrid` where the indexed
257/// strategy emits quantiles, not lookups.
258pub fn index_fn_supports_lookup(idx: &IndexFn) -> bool {
259    !matches!(idx, IndexFn::Continuous { .. } | IndexFn::Hybrid { .. })
260}
261
262/// Cardinality of an `IndexFn`. Used by strategies to size
263/// their output when no truncation is specified. Mirrors the
264/// helper in `metadata.rs` but lives here to avoid a circular
265/// dependency.
266pub(crate) fn index_fn_size(idx: &IndexFn) -> u64 {
267    match idx {
268        IndexFn::Lattice { axis_sizes } => axis_sizes
269            .iter()
270            .copied()
271            .fold(1u64, |a, b| a.saturating_mul(b)),
272        IndexFn::Lockstep { length } => *length,
273        IndexFn::Modular { axis_sizes } => axis_sizes.iter().copied().max().unwrap_or(0),
274        IndexFn::Concatenation { segment_sizes } => segment_sizes
275            .iter()
276            .copied()
277            .fold(0u64, |a, b| a.saturating_add(b)),
278        IndexFn::Continuous { .. } | IndexFn::Hybrid { .. } => 0,
279    }
280}
281
282/// Lattice dimensionality of an `IndexFn`. Used by strategies
283/// that branch on dimensionality (Extrema's corner count,
284/// Lhs's per-axis stratification).
285pub(crate) fn index_fn_dim(idx: &IndexFn) -> usize {
286    match idx {
287        IndexFn::Lattice { axis_sizes } => axis_sizes.len(),
288        IndexFn::Continuous { intervals, .. } => intervals.len(),
289        IndexFn::Hybrid {
290            discrete_axes,
291            continuous_axes,
292            ..
293        } => discrete_axes.len() + continuous_axes.len(),
294        IndexFn::Lockstep { .. } | IndexFn::Modular { .. } => 1,
295        IndexFn::Concatenation { segment_sizes } => segment_sizes.len(),
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn for_name_dispatches_to_correct_strategy() {
305        assert_eq!(for_name(StrategyName::Lex).name(), StrategyName::Lex);
306        assert_eq!(for_name(StrategyName::Halton).name(), StrategyName::Halton);
307        assert_eq!(
308            for_name(StrategyName::Extrema).name(),
309            StrategyName::Extrema
310        );
311    }
312
313    #[test]
314    fn index_fn_size_lattice() {
315        let idx = IndexFn::Lattice {
316            axis_sizes: vec![3, 4, 5],
317        };
318        assert_eq!(index_fn_size(&idx), 60);
319    }
320
321    #[test]
322    fn index_fn_size_concatenation() {
323        let idx = IndexFn::Concatenation {
324            segment_sizes: vec![10, 20, 30],
325        };
326        assert_eq!(index_fn_size(&idx), 60);
327    }
328
329    #[test]
330    fn index_fn_dim_classifies_correctly() {
331        assert_eq!(
332            index_fn_dim(&IndexFn::Lattice {
333                axis_sizes: vec![3, 4]
334            }),
335            2
336        );
337        assert_eq!(index_fn_dim(&IndexFn::Lockstep { length: 10 }), 1);
338        assert_eq!(
339            index_fn_dim(&IndexFn::Concatenation {
340                segment_sizes: vec![1, 2, 3]
341            }),
342            3
343        );
344    }
345}
polydat_core/iteration/comprehension/strategies/mod.rs

polydat_core/iteration/comprehension/strategies/
mod.rs
