polydat-core 0.5.3

// Copyright 2024-2026 Jonathan Shook
// SPDX-License-Identifier: Apache-2.0

//! `Shuffle` strategy — spec §3.6.
//!
//! Random permutation. PRNG seed captured at materialization
//! per spec §3.6 — same comprehension instance produces the
//! same shuffle on every dispense pass. Spec §9.5.2's
//! independence contract has different `CoordinateStream`
//! instances against the same comprehension get independent
//! shuffles via a per-streamer seed.
//!
//! The seed here is a module constant plus the input length;
//! per-streamer seeding is not implemented.
//!
//! ## References
//!
//! - R. A. Fisher & F. Yates, *Statistical Tables for Biological,
//!   Agricultural and Medical Research*, 3rd ed. (1948), the original
//!   shuffle. The in-place O(n) form (used via [`super::prng::Prng::shuffle`])
//!   is R. Durstenfeld, "Algorithm 235: Random permutation,"
//!   *Comm. ACM* 7(7) (1964), 420.
//!   doi:[10.1145/364520.364540](https://doi.org/10.1145/364520.364540);
//!   see also Knuth, *TAOCP* Vol. 2 §3.4.2 (Algorithm P). Correctness
//!   = the output is a *permutation* of the input (each element
//!   exactly once), verified in `tests::apply_preserves_elements`.
//!
//! Accepts any non-`None` `IndexFn` including continuous: over a
//! continuous or hybrid space it draws `n` codes, one 53-bit unit
//! fraction per continuous axis, and the runtime's sampler carries
//! each onto its axis's measure (spec §3.6, §10.2 R2).

use super::{
    EvaluatedInput, MultiIndex, Strategy, Tuple, index_fn_dim, index_fn_size,
    index_fn_supports_lookup, multi_index_to_flat, prng::Prng,
};
use crate::iteration::comprehension::metadata::IndexFn;
use crate::iteration::comprehension::strategy::StrategyName;

/// A seeded permutation.
pub struct Shuffle;

/// Seed base when none is authored; the input length is added per
/// call. Per-streamer
/// seeding is not implemented.
const DEFAULT_SEED: u64 = 0xD1CE_5EED_C0FF_EE42;

impl Strategy for Shuffle {
    fn name(&self) -> StrategyName {
        StrategyName::Shuffle
    }

    fn accepts_input(&self, idx: Option<&IndexFn>) -> bool {
        idx.is_some()
    }

    fn has_closed_form_for(&self, _idx: &IndexFn) -> bool {
        true
    }

    fn apply(&self, input: &EvaluatedInput, truncation: Option<u64>) -> Vec<Tuple> {
        self.apply_seeded(input, truncation, None)
    }

    fn apply_seeded(
        &self,
        input: &EvaluatedInput,
        truncation: Option<u64>,
        seed: Option<u64>,
    ) -> Vec<Tuple> {
        if index_fn_supports_lookup(&input.index_fn) {
            let mis = shuffle_multi_indices(&input.index_fn, truncation, seed);
            mis.into_iter()
                .filter_map(|mi| multi_index_to_flat(&input.index_fn, &mi))
                .filter_map(|flat| input.tuples.get(flat).cloned())
                .collect()
        } else {
            naive_shuffle_over_tuples(input.tuples.clone(), truncation, seed)
        }
    }
}

fn naive_shuffle_over_tuples(
    mut input: Vec<Tuple>,
    truncation: Option<u64>,
    seed: Option<u64>,
) -> Vec<Tuple> {
    let mut rng = Prng::new(
        seed.unwrap_or(DEFAULT_SEED)
            .wrapping_add(input.len() as u64),
    );
    rng.shuffle(&mut input);
    match truncation {
        Some(n) => input.into_iter().take(n as usize).collect(),
        None => input,
    }
}

/// The multi-indices of a shuffle over `idx`, `truncation` of them,
/// from the authored `seed` or the default.
pub(crate) fn shuffle_multi_indices(
    idx: &IndexFn,
    truncation: Option<u64>,
    seed: Option<u64>,
) -> Vec<MultiIndex> {
    try_shuffle_multi_indices(idx, truncation, seed).unwrap_or_else(|e| panic!("{e}"))
}

/// [`shuffle_multi_indices`], refusing a draw count that cannot be
/// held. Over a continuous space the count is the order's own, from
/// the spec text; over a discrete one it is bounded by the tuples the
/// input already holds.
pub(crate) fn try_shuffle_multi_indices(
    idx: &IndexFn,
    truncation: Option<u64>,
    seed: Option<u64>,
) -> Result<Vec<MultiIndex>, String> {
    let total = index_fn_size(idx);
    let continuous = matches!(idx, IndexFn::Continuous { .. } | IndexFn::Hybrid { .. });
    // A continuous space has no tuple count: the truncation is the
    // number of draws.
    let n = match (truncation, continuous) {
        (Some(t), true) => t,
        (Some(t), false) => t.min(total),
        (None, true) => return Ok(Vec::new()),
        (None, false) => total,
    };
    if n == 0 {
        return Ok(Vec::new());
    }

    let dim = index_fn_dim(idx);
    let axis_sizes = axis_sizes_for(idx);
    // The seed follows the draw count over a continuous space, which
    // has no tuple count of its own.
    let base = seed.unwrap_or(DEFAULT_SEED);
    let mut rng = Prng::new(base.wrapping_add(if continuous { n } else { total }));

    Ok(match idx {
        IndexFn::Continuous { intervals, .. } => {
            let _ = intervals;
            let mut out = crate::derive_support::try_buffer_for(n, "order shuffle")?;
            out.extend((0..n).map(|_| (0..dim).map(|_| rng.next_u64() >> 11).collect()));
            out
        }
        IndexFn::Hybrid {
            discrete_axes,
            continuous_axes,
            ..
        } => {
            let _ = continuous_axes;
            let mut out = crate::derive_support::try_buffer_for(n, "order shuffle")?;
            out.extend((0..n).map(|_| {
                let mut mi = Vec::with_capacity(dim);
                for size in discrete_axes {
                    mi.push(rng.next_bounded(*size));
                }
                for _ in 0..continuous_axes.len() {
                    mi.push(rng.next_u64() >> 11);
                }
                mi
            }));
            out
        }
        _ => {
            if n == total {
                let mut indices: Vec<u64> = (0..total).collect();
                rng.shuffle(&mut indices);
                indices
                    .into_iter()
                    .map(|i| linear_to_multi(i, &axis_sizes))
                    .collect()
            } else {
                let mut pool: Vec<u64> = (0..total).collect();
                let mut out = Vec::with_capacity(n as usize);
                for i in 0..n {
                    let j = rng.next_bounded(total - i);
                    let pick = pool[j as usize];
                    out.push(linear_to_multi(pick, &axis_sizes));
                    let last = pool.len() - 1;
                    pool.swap(j as usize, last);
                    pool.pop();
                }
                out
            }
        }
    })
}

fn axis_sizes_for(idx: &IndexFn) -> Vec<u64> {
    match idx {
        IndexFn::Lattice { axis_sizes } | IndexFn::Modular { axis_sizes } => axis_sizes.clone(),
        IndexFn::Lockstep { length } => vec![*length],
        IndexFn::Concatenation { segment_sizes } => vec![segment_sizes.iter().sum()],
        IndexFn::Continuous { .. } | IndexFn::Hybrid { .. } => Vec::new(),
    }
}

fn linear_to_multi(mut linear: u64, axis_sizes: &[u64]) -> MultiIndex {
    let mut out = vec![0u64; axis_sizes.len()];
    for i in (0..axis_sizes.len()).rev() {
        out[i] = linear % axis_sizes[i];
        linear /= axis_sizes[i];
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::iteration::comprehension::strategies::TupleValue;

    fn tup(k: i64) -> Tuple {
        Tuple::new().with("k", TupleValue::I64(k))
    }

    fn input_with(tuples: Vec<Tuple>) -> EvaluatedInput {
        let n = tuples.len() as u64;
        EvaluatedInput {
            tuples,
            cardinality: n,
            index_fn: IndexFn::Lattice {
                axis_sizes: vec![n],
            },
        }
    }

    #[test]
    fn apply_preserves_elements() {
        let inp = input_with(vec![tup(1), tup(2), tup(3), tup(4), tup(5)]);
        let mut out = Shuffle.apply(&inp, None);
        let mut sorted_in = inp.tuples.clone();
        out.sort_by_key(|t| match t.bindings[0].1 {
            TupleValue::I64(v) => v,
            _ => panic!(),
        });
        sorted_in.sort_by_key(|t| match t.bindings[0].1 {
            TupleValue::I64(v) => v,
            _ => panic!(),
        });
        assert_eq!(out, sorted_in);
    }

    #[test]
    fn apply_deterministic() {
        let inp = input_with(vec![tup(1), tup(2), tup(3), tup(4), tup(5)]);
        let a = Shuffle.apply(&inp, None);
        let b = Shuffle.apply(&inp, None);
        assert_eq!(a, b);
    }

    #[test]
    fn shuffle_multi_indices_produces_unique_discrete() {
        let idx = IndexFn::Lattice {
            axis_sizes: vec![3, 4],
        };
        let out = shuffle_multi_indices(&idx, Some(10), None);
        assert_eq!(out.len(), 10);
        let mut seen = std::collections::HashSet::new();
        for mi in &out {
            assert!(seen.insert(mi.clone()), "duplicate: {mi:?}");
        }
        for mi in &out {
            assert!(mi[0] < 3);
            assert!(mi[1] < 4);
        }
    }

    #[test]
    fn shuffle_multi_indices_full_lattice() {
        let idx = IndexFn::Lattice {
            axis_sizes: vec![2, 2],
        };
        let out = shuffle_multi_indices(&idx, None, None);
        assert_eq!(out.len(), 4);
        let mut sorted = out.clone();
        sorted.sort();
        assert_eq!(sorted, vec![vec![0, 0], vec![0, 1], vec![1, 0], vec![1, 1]]);
    }

    #[test]
    fn linear_to_multi_round_trip() {
        let sizes = vec![3u64, 4, 5];
        for linear in 0..60u64 {
            let mi = linear_to_multi(linear, &sizes);
            let mut back = 0u64;
            for (s, m) in sizes.iter().zip(mi.iter()) {
                back = back * s + m;
            }
            assert_eq!(back, linear);
        }
    }

    #[test]
    fn accepts_any_non_none() {
        assert!(Shuffle.accepts_input(Some(&IndexFn::Lattice {
            axis_sizes: vec![3]
        })));
        assert!(!Shuffle.accepts_input(None));
    }

    /// Over a continuous or hybrid space the truncation is the number
    /// of draws (spec §3.6: "n PRNG draws from the measure"), and a
    /// continuous code is a 53-bit fraction of the unit interval.
    #[test]
    fn continuous_draws_are_counted_by_the_truncation() {
        use crate::iteration::comprehension::cardinality::{Interval, ProductMeasure};
        let idx = IndexFn::Continuous {
            intervals: vec![Interval::closed(2.0, 4.0)],
            measure: ProductMeasure::Uniform,
        };
        let out = shuffle_multi_indices(&idx, Some(16), None);
        assert_eq!(out.len(), 16);
        assert!(out.iter().all(|mi| mi[0] < (1u64 << 53)), "{out:?}");
        assert!(shuffle_multi_indices(&idx, None, None).is_empty());
        let idx = IndexFn::Hybrid {
            discrete_axes: vec![3],
            continuous_axes: vec![Interval::closed(0.0, 1.0)],
            measure: ProductMeasure::Uniform,
        };
        let out = shuffle_multi_indices(&idx, Some(5), None);
        assert_eq!(out.len(), 5);
        assert!(
            out.iter().all(|mi| mi[0] < 3 && mi[1] < (1u64 << 53)),
            "{out:?}"
        );
    }

    /// An authored seed selects a different permutation from the
    /// default and the same one on every call (comprehension_forms.md
    /// §3.6: state from the authored seed and structural identity).
    #[test]
    fn an_authored_seed_is_deterministic_and_distinct() {
        let idx = IndexFn::Lattice {
            axis_sizes: vec![6, 6],
        };
        let default = shuffle_multi_indices(&idx, Some(12), None);
        let seeded = shuffle_multi_indices(&idx, Some(12), Some(42));
        let again = shuffle_multi_indices(&idx, Some(12), Some(42));
        assert_eq!(seeded, again);
        assert_ne!(seeded, default);
        assert_ne!(seeded, shuffle_multi_indices(&idx, Some(12), Some(43)));
    }
}