polydat_nodes/sampling/

alias.rs

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

//! Vose's alias method for O(1) sampling from discrete distributions.
//!
//! Given N outcomes with associated weights, an alias table pre-computes
//! a structure that allows selecting an outcome in constant time from a
//! single uniform u64 input.
//!
//! Two variants:
//! - [`AliasTable<T>`]: generic, works with any `Clone` outcome type.
//! - [`AliasTableU64`]: flat parallel arrays, optimized for the compiled
//!   kernel paths (closure tier and the `jit_weighted_pick` extern)
//!   where outcomes are indices 0..N.

use std::collections::VecDeque;

// -----------------------------------------------------------------
// Generic alias table
// -----------------------------------------------------------------

struct AliasSlot<T> {
    bias: f64,
    primary: T,
    alias: T,
}

/// Generic alias table for O(1) weighted sampling.
///
/// Construct with [`AliasTable::from_weights`] or [`AliasTable::uniform`],
/// then sample with [`AliasTable::sample`].
pub struct AliasTable<T> {
    slots: Vec<AliasSlot<T>>,
}

impl<T: Clone> AliasTable<T> {
    /// Build an alias table from outcomes and their weights.
    ///
    /// Weights do not need to be normalized — they are scaled
    /// internally. All weights must be non-negative; at least one
    /// must be positive.
    pub fn from_weights(outcomes: &[T], weights: &[f64]) -> Self {
        assert_eq!(
            outcomes.len(),
            weights.len(),
            "outcomes and weights must have equal length"
        );
        let n = outcomes.len();
        assert!(n > 0, "must have at least one outcome");

        let sum: f64 = weights.iter().sum();
        assert!(sum > 0.0, "total weight must be positive");

        // Normalize so weights sum to N
        let scale = n as f64 / sum;
        let mut scaled: Vec<f64> = weights.iter().map(|w| w * scale).collect();

        // Partition into small and large queues
        let mut small: VecDeque<usize> = VecDeque::new();
        let mut large: VecDeque<usize> = VecDeque::new();
        for (i, &w) in scaled.iter().enumerate() {
            if w < 1.0 {
                small.push_back(i);
            } else {
                large.push_back(i);
            }
        }

        // Build slots
        let mut slots: Vec<AliasSlot<T>> = (0..n)
            .map(|i| AliasSlot {
                bias: 1.0,
                primary: outcomes[i].clone(),
                alias: outcomes[i].clone(),
            })
            .collect();

        while let (Some(s), Some(l)) = (small.pop_front(), large.pop_front()) {
            slots[s].bias = scaled[s];
            slots[s].alias = outcomes[l].clone();

            scaled[l] -= 1.0 - scaled[s];
            if scaled[l] < 1.0 {
                small.push_back(l);
            } else {
                large.push_back(l);
            }
        }

        // Remaining items (due to floating-point drift) are their own alias
        for &i in small.iter().chain(large.iter()) {
            slots[i].bias = 1.0;
        }

        Self { slots }
    }

    /// Build a uniform alias table (all outcomes equally weighted).
    pub fn uniform(outcomes: &[T]) -> Self {
        let weights = vec![1.0; outcomes.len()];
        Self::from_weights(outcomes, &weights)
    }

    /// Sample an outcome from a uniform u64 input.
    ///
    /// The input is split into two independent parts: the low bits
    /// select a slot, the high bits determine the bias test. This
    /// avoids correlation between slot selection and the bias coin
    /// flip.
    #[inline]
    pub fn sample(&self, input: u64) -> &T {
        let n = self.slots.len();
        let slot_idx = (input as usize) % n;
        // Use upper bits for the bias test (independent of slot selection)
        let frac = (input >> 32) as f64 / u32::MAX as f64;
        let slot = &self.slots[slot_idx];
        if frac < slot.bias {
            &slot.primary
        } else {
            &slot.alias
        }
    }

    /// Number of outcomes in the table.
    pub fn len(&self) -> usize {
        self.slots.len()
    }

    /// Whether the table has no outcomes.
    pub fn is_empty(&self) -> bool {
        self.slots.is_empty()
    }
}

// -----------------------------------------------------------------
// u64-specialized alias table (flat parallel arrays)
// -----------------------------------------------------------------

/// Alias table optimized for u64 outcomes (indices 0..N).
///
/// Uses three parallel arrays for cache-friendly access. Suitable
/// for the Phase 2 compiled kernel path.
pub struct AliasTableU64 {
    biases: Vec<f64>,
    primaries: Vec<u64>,
    aliases: Vec<u64>,
}

impl AliasTableU64 {
    /// Build from weights. Outcomes are implicitly 0..N.
    pub fn from_weights(weights: &[f64]) -> Self {
        let n = weights.len();
        assert!(n > 0, "must have at least one outcome");

        let sum: f64 = weights.iter().sum();
        assert!(sum > 0.0, "total weight must be positive");

        let scale = n as f64 / sum;
        let mut scaled: Vec<f64> = weights.iter().map(|w| w * scale).collect();

        let mut small: VecDeque<usize> = VecDeque::new();
        let mut large: VecDeque<usize> = VecDeque::new();
        for (i, &w) in scaled.iter().enumerate() {
            if w < 1.0 {
                small.push_back(i);
            } else {
                large.push_back(i);
            }
        }

        let mut biases = vec![1.0f64; n];
        let primaries: Vec<u64> = (0..n as u64).collect();
        let mut aliases: Vec<u64> = (0..n as u64).collect();

        while let (Some(s), Some(l)) = (small.pop_front(), large.pop_front()) {
            biases[s] = scaled[s];
            aliases[s] = l as u64;

            scaled[l] -= 1.0 - scaled[s];
            if scaled[l] < 1.0 {
                small.push_back(l);
            } else {
                large.push_back(l);
            }
        }

        for &i in small.iter().chain(large.iter()) {
            biases[i] = 1.0;
        }

        Self {
            biases,
            primaries,
            aliases,
        }
    }

    /// Build a uniform table (all outcomes equally weighted).
    pub fn uniform(n: usize) -> Self {
        Self::from_weights(&vec![1.0; n])
    }

    /// Sample an outcome index from a uniform u64 input.
    ///
    /// Low bits select the slot, high bits test the bias.
    #[inline]
    pub fn sample(&self, input: u64) -> u64 {
        let n = self.biases.len();
        let slot_idx = (input as usize) % n;
        let frac = (input >> 32) as f64 / u32::MAX as f64;
        if frac < self.biases[slot_idx] {
            self.primaries[slot_idx]
        } else {
            self.aliases[slot_idx]
        }
    }

    /// Number of outcomes in the table.
    pub fn len(&self) -> usize {
        self.biases.len()
    }

    /// Whether the table has no outcomes.
    pub fn is_empty(&self) -> bool {
        self.biases.is_empty()
    }

    /// Access the bias array (for compiled kernel closure capture).
    pub fn biases(&self) -> &[f64] {
        &self.biases
    }

    /// Access the primary outcome array.
    pub fn primaries(&self) -> &[u64] {
        &self.primaries
    }

    /// Access the alias outcome array.
    pub fn aliases(&self) -> &[u64] {
        &self.aliases
    }
}

// -----------------------------------------------------------------
// Polydat node wrapping the alias table
// -----------------------------------------------------------------
//
// `AliasSample` flows through the `Const<Vec<f64>>` workload-list
// combinator (weights) and the `#[poly_const]` setup pattern
// (cached `AliasTableU64`). The node has no u64 kit
// (`Const<Vec<_>>`); it lowers through its slot kit, called from
// native code.

use polydat::derive_support::PolydatSetup;

impl PolydatSetup for AliasTableU64 {}

/// Build the alias table from raw weights. Single-call setup
/// invoked by the macro at construction time.
fn build_alias_table(weights: &[f64]) -> AliasTableU64 {
    AliasTableU64::from_weights(weights)
}

/// Sample an outcome index from a pre-built alias table over
/// `weights`. The input is a uniform u64 (hash upstream for
/// pseudo-random dispersion); the output is the chosen outcome
/// index.
#[polydat::polydat_node(category = Probability)]
fn alias_sample(
    input: u64,
    weights: Const<Vec<f64>>,
    #[poly_const(build_alias_table, from = weights)] table: &AliasTableU64,
) -> u64 {
    let _ = weights;
    table.sample(input)
}

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

    #[test]
    fn uniform_table_all_outcomes_reachable() {
        use xxhash_rust::xxh3::xxh3_64;

        let table = AliasTableU64::uniform(4);
        let mut seen = [false; 4];
        for i in 0..10_000u64 {
            let hashed = xxh3_64(&i.to_le_bytes());
            let outcome = table.sample(hashed) as usize;
            assert!(outcome < 4, "outcome {outcome} out of range");
            seen[outcome] = true;
        }
        for (i, &s) in seen.iter().enumerate() {
            assert!(s, "outcome {i} was never sampled");
        }
    }

    #[test]
    fn weighted_table_respects_distribution() {
        use xxhash_rust::xxh3::xxh3_64;

        // Heavily weighted: outcome 0 should dominate.
        // Inputs must be well-distributed (hashed), matching how the
        // Polydat uses alias tables — hash is always upstream.
        let table = AliasTableU64::from_weights(&[100.0, 1.0, 1.0]);
        let mut counts = [0u64; 3];
        let n = 100_000u64;
        for i in 0..n {
            let hashed = xxh3_64(&i.to_le_bytes());
            counts[table.sample(hashed) as usize] += 1;
        }
        // Outcome 0 has weight 100/102 ≈ 98%
        let ratio = counts[0] as f64 / n as f64;
        assert!(
            ratio > 0.90,
            "expected outcome 0 to dominate, got ratio {ratio} (counts: {counts:?})"
        );
    }

    #[test]
    fn deterministic() {
        let table = AliasTableU64::from_weights(&[1.0, 2.0, 3.0]);
        let a = table.sample(42);
        let b = table.sample(42);
        assert_eq!(a, b, "same input must produce same output");
    }

    #[test]
    fn generic_table_strings() {
        use xxhash_rust::xxh3::xxh3_64;

        let outcomes = vec!["alpha", "beta", "gamma"];
        let weights = vec![1.0, 1.0, 1.0];
        let table = AliasTable::from_weights(&outcomes, &weights);
        let mut seen = [false; 3];
        for i in 0..10_000u64 {
            let hashed = xxh3_64(&i.to_le_bytes());
            let result = *table.sample(hashed);
            match result {
                "alpha" => seen[0] = true,
                "beta" => seen[1] = true,
                "gamma" => seen[2] = true,
                other => panic!("unexpected outcome: {other}"),
            }
        }
        for (i, &s) in seen.iter().enumerate() {
            assert!(s, "outcome {i} never seen");
        }
    }

    #[test]
    fn polydat_node_eval() {
        use polydat::ast::PolydatNode;
        let node = AliasSample::new(vec![1.0, 1.0, 1.0, 1.0]);
        let mut out = [Value::None];
        node.eval(&[Value::U64(42)], &mut out);
        assert!(out[0].as_u64() < 4);
    }

    // `alias_sample` has no named JitOp; it runs through its slot
    // kit on every engine. The typed-eval path above covers
    // correctness.

    #[test]
    fn single_outcome() {
        let table = AliasTableU64::from_weights(&[1.0]);
        for i in 0..1000 {
            assert_eq!(table.sample(i), 0);
        }
    }

    #[test]
    fn two_outcomes_50_50() {
        use xxhash_rust::xxh3::xxh3_64;

        let table = AliasTableU64::from_weights(&[1.0, 1.0]);
        let mut counts = [0u64; 2];
        let n = 100_000u64;
        for i in 0..n {
            let hashed = xxh3_64(&i.to_le_bytes());
            counts[table.sample(hashed) as usize] += 1;
        }
        let ratio = counts[0] as f64 / n as f64;
        assert!(
            (0.40..0.60).contains(&ratio),
            "expected ~50/50, got ratio {ratio}"
        );
    }
}