yahtzee-engine 0.1.0

//! Canonical indices and probability tables for the 252 rolls and 462
//! keeps, the combinatorial skeleton of the solver.
//!
//! Multisets are enumerated by size and then lexicographically, so the
//! size-5 block sits at the end: `keep_id = roll_id + 210`.  Transitions
//! are stored compressed-sparse: every keep has exactly the rolls that
//! contain it as successors — 4368 pairs in total — and a full keep has
//! itself with probability one, which makes "stop rolling" an ordinary
//! lattice edge instead of a special case.

use crate::{Dice, Keep};

/// The number of distinct five-die rolls.
pub(crate) const ROLLS: usize = 252;

/// The number of distinct keeps of zero to five dice.
pub(crate) const KEEPS: usize = 462;

/// Keeps smaller than five dice; also the id offset of the size-5 block.
pub(crate) const SIZE5_OFFSET: usize = 210;

/// `n!` for `n <= 5`.
const FACTORIAL: [f64; 6] = [1.0, 1.0, 2.0, 6.0, 24.0, 120.0];

/// Ways to roll this multiset of faces, over `6^n` equally likely
/// ordered outcomes.
fn multinomial(counts: [u8; 6]) -> f64 {
    let n: u8 = counts.iter().sum();
    counts.iter().fold(FACTORIAL[usize::from(n)], |acc, &c| {
        acc / FACTORIAL[usize::from(c)]
    })
}

/// Packs face counts into a base-6 key for reverse lookup.
fn key(counts: [u8; 6]) -> usize {
    counts
        .iter()
        .rev()
        .fold(0, |acc, &c| acc * 6 + usize::from(c))
}

/// Precomputed dice combinatorics shared by all solver queries.
#[derive(Debug, Clone)]
pub(crate) struct DiceTables {
    /// The 252 rolls in canonical order.
    rolls: Vec<Dice>,
    /// The 462 keeps in canonical order; the last 252 mirror `rolls`.
    keeps: Vec<Keep>,
    /// Base-6 count key to keep id; `u16::MAX` marks impossible keys.
    index_of: Vec<u16>,
    /// Probability of each roll from five fresh dice.
    roll_prob: [f64; ROLLS],
    /// CSR offsets into `keep_succ`, one slice per keep.
    keep_succ_off: Vec<u32>,
    /// (roll id, probability) successors of each keep.
    keep_succ: Vec<(u16, f64)>,
    /// CSR offsets into `roll_keeps`, one slice per roll.
    roll_keeps_off: Vec<u32>,
    /// Keep ids of every sub-multiset of each roll.
    roll_keeps: Vec<u16>,
}

/// All face-count multisets of `size` dice, lexicographic by ascending
/// face string.
fn multisets(size: u8) -> Vec<[u8; 6]> {
    fn rec(counts: &mut [u8; 6], min_face: usize, left: u8, out: &mut Vec<[u8; 6]>) {
        if left == 0 {
            out.push(*counts);
            return;
        }
        for face in min_face..6 {
            counts[face] += 1;
            rec(counts, face, left - 1, out);
            counts[face] -= 1;
        }
    }
    let mut out = Vec::new();
    rec(&mut [0; 6], 0, size, &mut out);
    out
}

impl DiceTables {
    pub(crate) fn new() -> Self {
        let counts_by_id: Vec<[u8; 6]> = (0..=5).flat_map(multisets).collect();
        debug_assert_eq!(counts_by_id.len(), KEEPS);

        let keeps: Vec<Keep> = counts_by_id
            .iter()
            .map(|&c| Keep::from_counts(c).expect("at most five dice"))
            .collect();
        let rolls: Vec<Dice> = counts_by_id[SIZE5_OFFSET..]
            .iter()
            .map(|&c| Dice::from_counts(c).expect("exactly five dice"))
            .collect();

        let mut index_of = vec![u16::MAX; 6usize.pow(6)];
        for (id, &counts) in counts_by_id.iter().enumerate() {
            index_of[key(counts)] = id as u16;
        }

        let mut roll_prob = [0.0; ROLLS];
        for (r, dice) in rolls.iter().enumerate() {
            roll_prob[r] = multinomial(dice.counts()) / 7776.0;
        }

        let mut keep_succ_off = vec![0u32; KEEPS + 1];
        let mut keep_succ = Vec::new();
        for (k, keep) in keeps.iter().enumerate() {
            let rerolled = 5 - keep.len();
            let outcomes = 6f64.powi(i32::from(rerolled));
            for (r, roll) in rolls.iter().enumerate() {
                if !roll.contains(*keep) {
                    continue;
                }
                let mut diff = roll.counts();
                for (d, kept) in diff.iter_mut().zip(keep.counts()) {
                    *d -= kept;
                }
                keep_succ.push((r as u16, multinomial(diff) / outcomes));
            }
            keep_succ_off[k + 1] = keep_succ.len() as u32;
        }

        let mut roll_keeps_off = vec![0u32; ROLLS + 1];
        let mut roll_keeps = Vec::new();
        for (r, roll) in rolls.iter().enumerate() {
            let mut ids: Vec<u16> = roll.keeps().map(|k| index_of[key(k.counts())]).collect();
            ids.sort_unstable();
            roll_keeps.extend_from_slice(&ids);
            roll_keeps_off[r + 1] = roll_keeps.len() as u32;
        }

        Self {
            rolls,
            keeps,
            index_of,
            roll_prob,
            keep_succ_off,
            keep_succ,
            roll_keeps_off,
            roll_keeps,
        }
    }

    /// The canonical roll for an id below [`ROLLS`].
    pub(crate) fn roll(&self, roll_id: usize) -> Dice {
        self.rolls[roll_id]
    }

    /// The canonical keep for an id below [`KEEPS`].
    pub(crate) fn keep(&self, keep_id: usize) -> Keep {
        self.keeps[keep_id]
    }

    /// The canonical id of a keep.
    pub(crate) fn keep_id(&self, keep: Keep) -> usize {
        usize::from(self.index_of[key(keep.counts())])
    }

    /// The canonical id of a roll.
    pub(crate) fn roll_id(&self, dice: Dice) -> usize {
        usize::from(self.index_of[key(dice.counts())]) - SIZE5_OFFSET
    }

    /// Probability of each roll from five fresh dice.
    pub(crate) fn roll_prob(&self) -> &[f64; ROLLS] {
        &self.roll_prob
    }

    /// The rolls reachable from a keep, with probabilities summing to 1.
    pub(crate) fn keep_successors(&self, keep_id: usize) -> &[(u16, f64)] {
        let (lo, hi) = (self.keep_succ_off[keep_id], self.keep_succ_off[keep_id + 1]);
        &self.keep_succ[lo as usize..hi as usize]
    }

    /// The keep ids of every sub-multiset of a roll, ascending.
    pub(crate) fn roll_keeps(&self, roll_id: usize) -> &[u16] {
        let (lo, hi) = (
            self.roll_keeps_off[roll_id],
            self.roll_keeps_off[roll_id + 1],
        );
        &self.roll_keeps[lo as usize..hi as usize]
    }
}

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

    #[test]
    fn canonical_ids_line_up() {
        let tables = DiceTables::new();
        for keep_id in 0..KEEPS {
            assert_eq!(tables.keep_id(tables.keep(keep_id)), keep_id);
        }
        for roll_id in 0..ROLLS {
            let dice = tables.roll(roll_id);
            assert_eq!(tables.roll_id(dice), roll_id);
            // The size-5 block of keeps mirrors the rolls exactly.
            assert_eq!(tables.keep_id(Keep::from(dice)), roll_id + SIZE5_OFFSET);
        }
    }

    #[test]
    fn probabilities_are_distributions() {
        let tables = DiceTables::new();
        let total: f64 = tables.roll_prob().iter().sum();
        assert!((total - 1.0).abs() < 1e-12);
        for keep_id in 0..KEEPS {
            let successors = tables.keep_successors(keep_id);
            let total: f64 = successors.iter().map(|&(_, p)| p).sum();
            assert!(
                (total - 1.0).abs() < 1e-12,
                "keep {keep_id} sums to {total}"
            );
        }
        // A full keep has itself as its only successor.
        for roll_id in 0..ROLLS {
            let successors = tables.keep_successors(roll_id + SIZE5_OFFSET);
            assert_eq!(successors.len(), 1);
            assert_eq!(successors[0], (roll_id as u16, 1.0));
        }
    }

    #[test]
    fn transition_pairs_total_4368() {
        let tables = DiceTables::new();
        let succ: usize = (0..KEEPS).map(|k| tables.keep_successors(k).len()).sum();
        let subs: usize = (0..ROLLS).map(|r| tables.roll_keeps(r).len()).sum();
        assert_eq!(succ, 4368);
        assert_eq!(subs, 4368);
    }
}