polydat_nodes/sampling/

metashift.rs

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

//! Galois LFSR-based deterministic permutation (MetaShift / Shuffle).
//!
//! Provides bijective, deterministic, O(1)-space permutations of integer
//! ranges. Given a range [0, N), the LFSR visits every value exactly once
//! before cycling, in a pseudo-random order determined by the feedback
//! polynomial.
//!
//! This is useful for:
//! - Generating all values in a range without repetition or memory
//! - Shuffling sequences without materializing them
//! - Deterministic reordering across distributed workers (via bank selection)
//!
//! The core algorithm is a Galois-configuration LFSR. The `Shuffle` Polydat
//! node wraps it with range normalization and rejection sampling.
//!
//! The `feedback` polynomial is exposed explicitly as a Const arg so the
//! macro can auto-emit the JIT-eligible `compiled_u64` / `jit_constants`
//! hooks (Setup-derived state disables the macro's auto-emitted u64
//! kit; a `compiled_u64 = ...` / `jit_constants = ...` override, as
//! `cycle_walk` uses, could capture it instead). Callers compute
//! `feedback` via
//! [`feedback_for_width_and_bank`] or [`feedback_for_size`].

// -----------------------------------------------------------------
// The algorithm
// -----------------------------------------------------------------
//
// The LFSR step, the bounded permutation over it, and the feedback
// polynomials all live in `polydat::numeric::permute`, which is where
// a body shared by a node and its native lowering belongs: `shuffle`
// below and `jit_shuffle` call the same function, so the interpreter
// and native code agree by construction rather than by review. The
// selection helpers keep their paths here.
pub use polydat::numeric::permute::{
    feedback_for_size, feedback_for_width, feedback_for_width_and_bank, width_for_period,
};
// Imported under the old name: `lfsr_step` is taken here by the
// macro-consumed node-authoring function below.
use polydat::numeric::permute::{lfsr_step as step, shuffle_bounded};

// -----------------------------------------------------------------
// Shuffle: bounded bijective permutation
// -----------------------------------------------------------------

/// Deterministic, bijective permutation of a bounded integer range.
///
/// Signature: `shuffle(input: u64, feedback: u64, size: u64, min: u64) -> (u64)`
///
/// Maps every value in [min, min+size) to itself in a pseudo-random
/// order, visiting each value exactly once per cycle. Uses a Galois
/// LFSR with rejection sampling to handle ranges that are not exact
/// powers of 2.
///
/// Use when you need every key in a range visited exactly once without
/// repetition and without materializing the full sequence in memory.
/// Common patterns: generating unique primary keys for bulk inserts,
/// distributing work across partitions without collision, or simulating
/// a deck-of-cards draw. Pick different `feedback` polynomial values
/// (via [`feedback_for_width_and_bank`]) for independent permutation
/// orderings across distributed workers.
///
/// JIT level: P3 — every arg + return is `u64`, so the macro auto-emits
/// `compiled_u64` with `feedback`/`size`/`min` captured by `Copy` and
/// `jit_constants` returning `[feedback, size, min]` (the layout
/// `JitOp::ShuffleConst` consumes).
#[polydat::polydat_node(category = Permutation, validate = shuffle_validate)]
fn shuffle(
    input: u64,
    #[poly_default(0u64)] feedback: Const<u64>,
    #[poly_default(0u64)] size: Const<u64>,
    #[poly_default(0u64)] min: Const<u64>,
) -> u64 {
    shuffle_bounded(input, *feedback, *size, *min)
}

/// `shuffle`'s range, `[min, min + size)`, must be a range of `u64`:
/// its last value `min + size - 1` has to fit. The node answers inside
/// the range, so for one past the top there is no right answer, and the
/// denormalizing add overflowed (a panic in debug, a wrapped value in
/// release). Node-level because no per-parameter constraint can relate
/// two parameters.
fn shuffle_validate(_name: &str, consts: &[polydat::dsl::factory::ConstArg]) -> Result<(), String> {
    let size = consts.get(1).map(|c| c.as_u64()).unwrap_or(0);
    let min = consts.get(2).map(|c| c.as_u64()).unwrap_or(0);
    if size > 0 && min.checked_add(size - 1).is_none() {
        return Err(format!(
            "the range [min, min + size) must fit in u64; min {min} + size {size} runs past {}",
            u64::MAX
        ));
    }
    Ok(())
}

// -----------------------------------------------------------------
// Raw LFSR step as a Polydat node (for advanced use)
// -----------------------------------------------------------------

/// Single Galois LFSR step as a Polydat node.
///
/// Signature: `lfsr_step(input: u64, feedback: u64) -> (u64)`
///
/// This is the raw bijective LFSR operation without range bounding.
/// The period is 2^width - 1. Useful for building custom permutation
/// patterns. The `feedback` polynomial selects which permutation
/// ordering is produced — use [`feedback_for_width`] or
/// [`feedback_for_width_and_bank`] to compute one for a given
/// register width.
///
/// JIT level: P3 — auto-emitted because both args + return are `u64`.
#[polydat::polydat_node(category = Permutation)]
fn lfsr_step(input: u64, feedback: Const<u64>) -> u64 {
    step(input, *feedback)
}

#[cfg(test)]
mod tests {
    use super::*;
    use polydat::ast::{PolydatNode, Value};

    /// `size` defaults to zero, so `shuffle(x)` and `shuffle(x, fb)`
    /// reach the empty range without the author asking for it. It is
    /// the range's floor, and above all it is not a panic: a division
    /// by zero here would surface as a raw arithmetic message from the
    /// build's constant fold.
    #[test]
    fn an_empty_range_answers_its_floor() {
        for min in [0u64, 7, u64::MAX] {
            let node = Shuffle::new(0, 0, min);
            let mut out = [Value::None];
            for input in [0u64, 1, 42, u64::MAX] {
                node.eval(&[Value::U64(input)], &mut out);
                assert_eq!(
                    out[0].as_u64(),
                    min,
                    "shuffle over an empty range at min={min}, input={input}"
                );
            }
        }
    }

    #[test]
    fn lfsr_step_nonzero() {
        // LFSR should never produce 0 from a nonzero input
        let feedback = feedback_for_width(8);
        let mut reg = 1u64;
        for _ in 0..255 {
            reg = step(reg, feedback);
            assert_ne!(reg, 0, "LFSR must never produce 0");
        }
    }

    #[test]
    fn lfsr_full_cycle() {
        // An 8-bit LFSR should visit all 255 nonzero values exactly once
        let feedback = feedback_for_width(8);
        let mut seen = vec![false; 256];
        let mut reg = 1u64;
        for _ in 0..255 {
            reg = step(reg, feedback);
            assert!(!seen[reg as usize], "duplicate value {reg}");
            seen[reg as usize] = true;
        }
        // Verify all nonzero values visited
        for i in 1..=255u64 {
            assert!(seen[i as usize], "value {i} not visited");
        }
    }

    #[test]
    fn lfsr_period_returns_to_start() {
        let feedback = feedback_for_width(8);
        let start = 42u64;
        let mut reg = start;
        for _ in 0..255 {
            reg = step(reg, feedback);
        }
        assert_eq!(reg, start, "LFSR should return to start after 2^N-1 steps");
    }

    /// Test-only helper: build a `Shuffle` over `[min, min+size)` using
    /// bank 0. Mirrors the historical `Shuffle::new(min, size)` shape so
    /// the in-file tests stay readable.
    fn shuf(min: u64, size: u64) -> Shuffle {
        Shuffle::new(feedback_for_size(size), size, min)
    }

    /// Test-only helper: build a `Shuffle` over `[0, size)` using bank 0.
    fn shuf0(size: u64) -> Shuffle {
        shuf(0, size)
    }

    fn apply(node: &Shuffle, input: u64) -> u64 {
        let mut out = [Value::None];
        node.eval(&[Value::U64(input)], &mut out);
        out[0].as_u64()
    }

    #[test]
    fn shuffle_bijective_small() {
        // Shuffle over [0, 31) should produce a permutation
        let node = shuf0(31);
        let mut seen = [false; 31];
        for i in 0..31u64 {
            let out = apply(&node, i);
            assert!(out < 31, "out of range: {out}");
            assert!(!seen[out as usize], "duplicate at input {i}: {out}");
            seen[out as usize] = true;
        }
        assert!(seen.iter().all(|&s| s), "not all values produced");
    }

    #[test]
    fn shuffle_bijective_non_power_of_two() {
        // Shuffle over [0, 50) — not a power of 2, requires rejection sampling
        let node = shuf0(50);
        let mut seen = [false; 50];
        for i in 0..50u64 {
            let out = apply(&node, i);
            assert!(out < 50, "out of range: {out}");
            assert!(!seen[out as usize], "duplicate at input {i}: {out}");
            seen[out as usize] = true;
        }
        assert!(seen.iter().all(|&s| s), "not all values produced");
    }

    #[test]
    fn shuffle_with_min_offset() {
        let node = shuf(100, 20);
        let mut seen = [false; 20];
        for i in 0..20u64 {
            let out = apply(&node, i);
            assert!((100..120).contains(&out), "out of range: {out}");
            seen[(out - 100) as usize] = true;
        }
        assert!(seen.iter().all(|&s| s), "not all values produced");
    }

    #[test]
    fn shuffle_deterministic() {
        let node = shuf0(100);
        let a = apply(&node, 42);
        let b = apply(&node, 42);
        assert_eq!(a, b);
    }

    #[test]
    fn shuffle_not_identity() {
        // The shuffle should reorder, not pass through
        let node = shuf0(100);
        let mut identity_count = 0;
        for i in 0..100u64 {
            if apply(&node, i) == i {
                identity_count += 1;
            }
        }
        // Some fixed points are expected, but not all
        assert!(identity_count < 50, "shuffle should reorder most values");
    }

    #[test]
    fn shuffle_polydat_node() {
        let node = shuf0(100);
        let mut out = [Value::None];
        node.eval(&[Value::U64(7)], &mut out);
        assert!(out[0].as_u64() < 100);
    }

    #[test]
    fn shuffle_compiled() {
        let node = shuf0(100);
        let op = node.compiled_u64().expect("should compile");
        let mut out = [0u64];
        op(&[7], &mut out);
        assert!(out[0] < 100);

        // Matches eval path
        let mut eval_out = [Value::None];
        node.eval(&[Value::U64(7)], &mut eval_out);
        assert_eq!(out[0], eval_out[0].as_u64());
    }

    #[test]
    fn lfsr_step_node() {
        let node = LfsrStep::new(feedback_for_width(8));
        let mut out = [Value::None];
        node.eval(&[Value::U64(1)], &mut out);
        let v = out[0].as_u64();
        assert_ne!(v, 0);
        assert_ne!(v, 1);
    }

    #[test]
    fn shuffle_large_range() {
        // Verify shuffle works for a larger range (1000)
        let node = shuf0(1000);
        let mut seen = vec![false; 1000];
        for i in 0..1000u64 {
            let out = apply(&node, i);
            assert!(out < 1000, "out of range: {out}");
            seen[out as usize] = true;
        }
        assert!(seen.iter().all(|&s| s), "not all values produced");
    }

    #[test]
    fn different_banks_different_orderings() {
        let size = 100;
        let fb0 = feedback_for_width_and_bank(width_for_period(size), 0);
        let fb1 = feedback_for_width_and_bank(width_for_period(size), 1);
        let n0 = Shuffle::new(fb0, size, 0);
        let n1 = Shuffle::new(fb1, size, 0);
        // Both should be bijective permutations
        let mut seen0 = [false; 100];
        let mut seen1 = [false; 100];
        let mut differ = false;
        for i in 0..100u64 {
            let a = apply(&n0, i);
            let b = apply(&n1, i);
            assert!(a < 100);
            assert!(b < 100);
            seen0[a as usize] = true;
            seen1[b as usize] = true;
            if a != b {
                differ = true;
            }
        }
        assert!(seen0.iter().all(|&s| s), "bank 0 not bijective");
        assert!(seen1.iter().all(|&s| s), "bank 1 not bijective");
        assert!(differ, "different banks should produce different orderings");
    }

    #[test]
    fn width_for_period_table() {
        assert_eq!(width_for_period(1), 4); // minimum is 4
        assert_eq!(width_for_period(15), 4); // 15 < 2^4
        assert_eq!(width_for_period(16), 5); // 16 = 2^4, needs 5 bits
        assert_eq!(width_for_period(31), 5);
        assert_eq!(width_for_period(32), 6);
        assert_eq!(width_for_period(255), 8);
        assert_eq!(width_for_period(256), 9);
        assert_eq!(width_for_period(1000), 10);
    }
}