vitaminc-permutation 0.5.0

//! Constant-time oblivious generation of random permutations.
//!
//! A uniform random permutation is generated by sorting random keys through a
//! fixed sorting network (the construction used by djbsort and NTRU Prime):
//!
//! 1. **Pack**: build `[u64; N]` where `w[i] = (rng.next_u64() << 8) | i`.
//!    The high 56 bits are the random sort key; the low 8 bits carry the
//!    index as payload, so the sort routes the payload for free.
//! 2. **Sort**: run a Batcher odd-even mergesort network — a fixed,
//!    data-independent schedule of compare-exchange gates — with a branchless
//!    constant-time gate.
//! 3. **Collision check**: if any two random keys collide, the whole batch
//!    is rejected and generation **fails** — the seed must be discarded and a
//!    fresh one generated (see [`random_permutation`] for the lifecycle
//!    argument). A tie broken by the packed index would bias the output
//!    toward identity order, and patching only the colliding keys would leak
//!    *which* positions collided.
//! 4. **Strip**: the low bytes of the sorted array *are* the permutation.
//!
//! Instruction trace, memory trace, and per-instruction latency are functions
//! of `N` only, with one deliberate exception: the final accept/reject of the
//! whole batch in step 3, which is a single branch on the collision predicate
//! and is the documented failure mode. Gate indices come from the public
//! compile-time schedule, and the comparison outcome inside every gate is
//! absorbed into a `subtle` select mask. No memory address is ever derived
//! from a secret value.

use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, ConstantTimeGreater};
use vitaminc_random::{RandomError, Rng, SafeRand};
use zeroize::Zeroizing;

use crate::private::IsPermutable;

/// Emits the Batcher odd-even mergesort network for `n` inputs (Knuth 5.3.4,
/// iterative form) into `out`, returning the total gate count. `n` must be a
/// power of two: the network is only used at those sizes, and for them every
/// merge block is full, so no partial-block guard is needed on the inner
/// loop. The division guard keeps comparisons within the pair of runs being
/// merged. Correctness rests on the zero-one principle — a network sorts
/// everything iff it sorts all binary inputs — so the tests target this
/// transcription, not the theorem.
///
/// This is the single source of truth for the network: the counting pass
/// calls it with an empty slice (gates beyond `out.len()` are counted but not
/// stored), and the schedule pass calls it with the full-size array.
const fn batcher_fill(n: usize, out: &mut [(u8, u8)]) -> usize {
    assert!(n.is_power_of_two(), "network size must be a power of two");
    let mut gates = 0;
    let mut p = 1;
    while p < n {
        let mut k = p;
        while k != 0 {
            let mut j = k % p;
            while j + k < n {
                let mut i = 0;
                while i < k {
                    // Every candidate pair lies inside the network. This is
                    // implied by the loop bounds for a power of two, and
                    // asserting it makes an off-by-one in those bounds fail
                    // at compile time instead of being absorbed by the
                    // division guard below.
                    assert!(i + j + k < n, "gate index out of bounds");
                    if (i + j) / (2 * p) == (i + j + k) / (2 * p) {
                        if gates < out.len() {
                            out[gates] = ((i + j) as u8, (i + j + k) as u8);
                        }
                        gates += 1;
                    }
                    i += 1;
                }
                j += 2 * k;
            }
            k /= 2;
        }
        p *= 2;
    }
    gates
}

/// Number of compare-exchange gates in the Batcher network for `n` inputs;
/// for `n = 2^k` the closed form is `(k² − k + 4)·2^(k−2) − 1`.
pub(crate) const fn batcher_gate_count(n: usize) -> usize {
    batcher_fill(n, &mut [])
}

/// The Batcher network for `n` inputs as a fixed compile-time schedule of
/// compare-exchange gates. This is what each [`IsPermutable`] impl stores as
/// its `SCHEDULE`, so a length without a network is a missing associated
/// const — a compile error in this crate — rather than a runtime lookup.
pub(crate) const fn batcher_schedule<const G: usize>(n: usize) -> [(u8, u8); G] {
    assert!(n <= 256, "gate indices must fit in u8");
    let mut out = [(0u8, 0u8); G];
    let gates = batcher_fill(n, &mut out);
    assert!(gates == G, "schedule length must match the gate count");
    out
}

/// Branchless compare-exchange for `a < b`: both locations are read and
/// written unconditionally on every gate, and the comparison outcome only
/// ever feeds a select mask (`cmov`/`csel`), never a branch or an address.
/// The swap is `subtle`'s xor-mask form, so the only transient it creates is
/// the masked difference of the two words, never a copy of either.
#[inline(always)]
fn compare_exchange(w: &mut [u64], a: usize, b: usize) {
    let swap = w[a].ct_gt(&w[b]);
    let (lo, hi) = w.split_at_mut(b);
    u64::conditional_swap(&mut lo[a], &mut hi[0], swap);
}

/// Sorts `w` in place through the fixed network for `N`.
pub(crate) fn sort<const N: usize>(w: &mut [u64; N])
where
    [u8; N]: IsPermutable,
{
    for &(a, b) in <[u8; N] as IsPermutable>::SCHEDULE {
        compare_exchange(w, a as usize, b as usize);
    }
}

/// Sorts a batch of packed words and writes the permutation payload into
/// `out`, or returns `false` without touching `out` if any two random sort
/// keys collide. Only the random bits matter for collisions: the packed
/// indices make the full words distinct, and equal keys end up adjacent
/// after sorting. The scan accumulates into a mask so it is itself
/// branch-free; the single branch is the final accept/reject.
///
/// The extracted permutation is itself secret key material, which is why it
/// is written straight into the caller's buffer rather than returned by
/// value: the caller passes the wiped-on-drop slot the key will live in, so
/// no plain `[u8; N]` copy of the permutation is ever made.
fn permutation_from_words<const N: usize>(w: &mut [u64; N], out: &mut [u8; N]) -> bool
where
    [u8; N]: IsPermutable,
{
    sort(w);
    let mut collision = Choice::from(0u8);
    for pair in w.windows(2) {
        collision |= (pair[0] >> 8).ct_eq(&(pair[1] >> 8));
    }
    if bool::from(collision) {
        return false;
    }
    for (o, x) in out.iter_mut().zip(w.iter()) {
        *o = (x & 0xFF) as u8;
    }
    true
}

/// Writes a uniform random permutation of `0..N` in gather form
/// (`out[j] = data[p[j]]` applies it) into `out`, generated obliviously:
/// timing and memory access patterns are independent of the result. `out`
/// should be the wiped-on-drop slot the key will live in; on failure it is
/// left untouched.
///
/// Exactly **one** batch is attempted. On a key collision the whole batch is
/// rejected and generation fails with [`RandomError::SeedRejected`] — the
/// caller must discard the seed and start over with a fresh one, never retry
/// from the same RNG stream. Two reasons:
///
/// - The seed is typically a long-term secret that clients retain to
///   re-derive the permutation. A seed whose first batch collides can never
///   produce a stable permutation, so it is unusable by definition.
/// - Retrying in-stream would make runtime reveal whether the *retained*
///   seed collided — a predicate of the secret seed. Failing after one batch
///   moves the retry to the seed-generation layer, where timing reveals only
///   how many independent, discarded seeds preceded the accepted one.
///
/// A caller whose generator was seeded from the OS rather than from a
/// retained seed has nothing to discard: it builds a fresh generator with
/// [`SafeRand::from_entropy`] and calls again.
///
/// With 56 random bits the collision probability is ≈ N²/2⁵⁷ (≈ 2⁻⁴³ at
/// N = 128, ≈ 2⁻⁵¹ at N = 8), so honest generation essentially never fails.
pub(crate) fn random_permutation<const N: usize>(
    rng: &mut SafeRand,
    out: &mut [u8; N],
) -> Result<(), RandomError>
where
    [u8; N]: IsPermutable,
{
    let mut w: Zeroizing<[u64; N]> = Zeroizing::new([0; N]);
    for (i, slot) in w.iter_mut().enumerate() {
        // The shift clears the low byte, so adding the index is the same as
        // or-ing it in. `+` is used because `|` and `^` are indistinguishable
        // here, which left the packing untestable by mutation.
        *slot = (rng.next_u64() << 8) + i as u64;
    }
    if permutation_from_words(&mut w, out) {
        Ok(())
    } else {
        Err(RandomError::SeedRejected)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use vitaminc_random::SeedableRng;

    fn assert_is_permutation<const N: usize>(p: &[u8; N]) {
        let mut seen = [false; N];
        for &v in p {
            assert!((v as usize) < N, "value {v} out of range");
            assert!(!seen[v as usize], "value {v} repeated");
            seen[v as usize] = true;
        }
    }

    #[test]
    fn gate_counts_match_closed_form() {
        // (k² − k + 4)·2^(k−2) − 1 for n = 2^k.
        for (n, expected) in [
            (2, 1),
            (4, 5),
            (8, 19),
            (16, 63),
            (32, 191),
            (64, 543),
            (128, 1471),
        ] {
            assert_eq!(batcher_gate_count(n), expected, "gate count for n = {n}");
        }
    }

    #[test]
    #[should_panic(expected = "network size must be a power of two")]
    fn network_rejects_a_size_that_is_not_a_power_of_two() {
        batcher_gate_count(6);
    }

    #[test]
    fn schedule_gates_are_in_bounds_and_ordered() {
        fn check<const N: usize>()
        where
            [u8; N]: IsPermutable,
        {
            let gates = <[u8; N] as IsPermutable>::SCHEDULE;
            assert_eq!(gates.len(), batcher_gate_count(N));
            for &(a, b) in gates {
                assert!(a < b, "gate ({a}, {b}) not ordered");
                assert!((b as usize) < N, "gate ({a}, {b}) out of bounds for {N}");
            }
        }
        check::<8>();
        check::<16>();
        check::<32>();
        check::<64>();
        check::<128>();
    }

    #[test]
    fn network_sorts_all_binary_inputs_for_every_length_to_16() {
        // Zero-one principle: a comparison network sorts every input iff it
        // sorts all binary inputs. Build the schedule at runtime for every
        // power of two up to 16 (exhaustive over 2^n inputs), check the gate
        // bounds, and at the shipped sizes check that the runtime build is
        // byte-for-byte the compile-time `SCHEDULE`.
        for n in [2usize, 4, 8, 16] {
            let count = batcher_gate_count(n);
            let mut gates = vec![(0u8, 0u8); count];
            assert_eq!(batcher_fill(n, &mut gates), count);
            for &(a, b) in &gates {
                assert!(a < b, "n = {n}: gate ({a}, {b}) not ordered");
                assert!((b as usize) < n, "n = {n}: gate ({a}, {b}) out of bounds");
            }
            match n {
                8 => assert_eq!(<[u8; 8] as IsPermutable>::SCHEDULE, &gates[..]),
                16 => assert_eq!(<[u8; 16] as IsPermutable>::SCHEDULE, &gates[..]),
                _ => {}
            }
            for bits in 0u32..(1 << n) {
                let mut w: Vec<u64> = (0..n).map(|i| u64::from(bits >> i) & 1).collect();
                for &(a, b) in &gates {
                    compare_exchange(&mut w, a as usize, b as usize);
                }
                assert!(w.is_sorted(), "n = {n}: failed on binary input {bits:#b}");
            }
        }
    }

    #[test]
    fn network_sorts_random_words() {
        // Every supported length: the zero-one test only covers up to 16,
        // and a length-specific transcription bug would otherwise go
        // unnoticed.
        fn check<const N: usize>(rng: &mut SafeRand)
        where
            [u8; N]: IsPermutable,
        {
            for _ in 0..100 {
                let mut w: [u64; N] = core::array::from_fn(|_| rng.next_u64());
                let mut expected = w;
                expected.sort_unstable();
                sort(&mut w);
                assert_eq!(w, expected);
            }
        }
        let mut rng = SafeRand::from_seed([42u8; 32]);
        check::<8>(&mut rng);
        check::<16>(&mut rng);
        check::<32>(&mut rng);
        check::<64>(&mut rng);
        check::<128>(&mut rng);
    }

    #[test]
    fn output_is_a_valid_permutation() {
        fn check<const N: usize>(rng: &mut SafeRand)
        where
            [u8; N]: IsPermutable,
        {
            for _ in 0..50 {
                let mut out = [0u8; N];
                random_permutation::<N>(rng, &mut out).unwrap();
                assert_is_permutation(&out);
            }
        }
        let mut rng = SafeRand::from_seed([9u8; 32]);
        check::<8>(&mut rng);
        check::<16>(&mut rng);
        check::<32>(&mut rng);
        check::<64>(&mut rng);
        check::<128>(&mut rng);
    }

    #[test]
    fn colliding_keys_reject_the_batch() {
        // Two equal random keys (high 56 bits) with different payloads must
        // reject the whole batch, even though the packed words are distinct,
        // and must leave the output untouched.
        let mut w: [u64; 8] = core::array::from_fn(|i| ((i as u64) << 8) | i as u64);
        w[3] = (7 << 8) | 3; // same sort key as w[7], different payload
        let mut out = [0xAAu8; 8];
        assert!(!permutation_from_words(&mut w, &mut out));
        assert_eq!(out, [0xAA; 8]);

        // Distinct keys must produce the payload permutation in sorted-key
        // order: descending keys reverse the payloads.
        let mut w: [u64; 8] = core::array::from_fn(|i| ((7 - i as u64) << 8) | i as u64);
        assert!(permutation_from_words(&mut w, &mut out));
        assert_eq!(out, [7, 6, 5, 4, 3, 2, 1, 0]);
    }

    #[test]
    fn output_is_deterministic_for_a_seed() {
        let mut a = [0u8; 64];
        let mut b = [0u8; 64];
        random_permutation(&mut SafeRand::from_seed([1u8; 32]), &mut a).unwrap();
        random_permutation(&mut SafeRand::from_seed([1u8; 32]), &mut b).unwrap();
        assert_eq!(a, b);
    }
}