yahtzee-engine 0.1.0

//! Dice multisets: a full roll of five dice and a kept subset.

use core::fmt;
use core::str::FromStr;

/// A roll of exactly five six-sided dice, stored as face counts.
///
/// Order never matters in Yahtzee, so the multiset is the canonical
/// representation: `counts[f - 1]` dice show face `f`.  The five-dice
/// invariant is enforced by every constructor, which is why scoring
/// functions never have to handle short or long rolls.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Dice {
    counts: [u8; 6],
}

/// A sub-multiset of zero to five dice held between rolls.
///
/// A separate type from [`Dice`] keeps "a full roll" and "what I am
/// keeping" from being confused: every function that needs five dice
/// takes [`Dice`], and the type system rules out partial rolls there.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct Keep {
    counts: [u8; 6],
}

/// Error parsing dice from text.
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[non_exhaustive]
pub enum ParseDiceError {
    /// A character other than `1`–`6` or whitespace was found.
    #[error("invalid die face (expected digits 1-6)")]
    InvalidFace,
    /// The number of dice does not fit the target type.
    #[error("wrong number of dice")]
    WrongCount,
}

const fn parse_counts(s: &str) -> Result<([u8; 6], u8), ParseDiceError> {
    let mut counts = [0u8; 6];
    let mut total = 0u8;
    let bytes = s.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        match bytes[i] {
            b'1'..=b'6' => {
                counts[(bytes[i] - b'1') as usize] += 1;
                total += 1;
                if total > 5 {
                    return Err(ParseDiceError::WrongCount);
                }
            }
            b' ' | b'\t' | b'\n' | b'\r' => {}
            _ => return Err(ParseDiceError::InvalidFace),
        }
        i += 1;
    }
    Ok((counts, total))
}

fn fmt_counts(counts: &[u8; 6], f: &mut fmt::Formatter<'_>) -> fmt::Result {
    for (i, &n) in counts.iter().enumerate() {
        for _ in 0..n {
            write!(f, "{}", i + 1)?;
        }
    }
    Ok(())
}

impl Dice {
    /// Builds a roll from five faces in any order.
    ///
    /// Returns [`None`] if any face is outside `1..=6`.
    #[must_use]
    pub const fn from_faces(faces: [u8; 5]) -> Option<Self> {
        let mut counts = [0u8; 6];
        let mut i = 0;
        while i < 5 {
            if faces[i] < 1 || faces[i] > 6 {
                return None;
            }
            counts[(faces[i] - 1) as usize] += 1;
            i += 1;
        }
        Some(Self { counts })
    }

    /// Builds a roll from face counts (`counts[f - 1]` dice show `f`).
    ///
    /// Returns [`None`] unless the counts sum to exactly five.
    #[must_use]
    pub const fn from_counts(counts: [u8; 6]) -> Option<Self> {
        // A u16 sum: u8 would wrap on adversarial counts in release
        // builds and let a 261-die "roll" through as five.
        let mut sum = 0u16;
        let mut i = 0;
        while i < 6 {
            sum += counts[i] as u16;
            i += 1;
        }
        if sum == 5 {
            Some(Self { counts })
        } else {
            None
        }
    }

    /// How many dice show `face` (`1..=6`).
    ///
    /// # Panics
    ///
    /// Panics if `face` is outside `1..=6`.
    #[must_use]
    pub const fn count(self, face: u8) -> u8 {
        self.counts[(face - 1) as usize]
    }

    /// The face counts, indexed by `face - 1`.
    #[must_use]
    pub const fn counts(self) -> [u8; 6] {
        self.counts
    }

    /// The pip total of all five dice.
    #[must_use]
    pub const fn sum(self) -> u8 {
        let mut sum = 0u8;
        let mut f = 0;
        while f < 6 {
            sum += self.counts[f] * (f as u8 + 1);
            f += 1;
        }
        sum
    }

    /// The face of a five-of-a-kind, or [`None`] if this roll is not one.
    #[must_use]
    pub const fn yahtzee_face(self) -> Option<u8> {
        let mut f = 0;
        while f < 6 {
            if self.counts[f] == 5 {
                return Some(f as u8 + 1);
            }
            f += 1;
        }
        None
    }

    /// The five faces in ascending order.
    pub fn faces(self) -> impl Iterator<Item = u8> {
        (1..=6u8).flat_map(move |f| core::iter::repeat_n(f, usize::from(self.count(f))))
    }

    /// All distinct sub-multisets of this roll, from [`Keep::EMPTY`] to
    /// keeping everything.  A roll has at most 32 keeps (all faces
    /// distinct) and as few as 6 (five of a kind).
    pub fn keeps(self) -> impl Iterator<Item = Keep> {
        Keeps {
            limits: self.counts,
            next: Some([0; 6]),
        }
    }

    /// Whether `keep` is a sub-multiset of this roll.
    #[must_use]
    pub const fn contains(self, keep: Keep) -> bool {
        let mut f = 0;
        while f < 6 {
            if keep.counts[f] > self.counts[f] {
                return false;
            }
            f += 1;
        }
        true
    }
}

impl Keep {
    /// Keeping no dice: reroll everything.
    pub const EMPTY: Self = Self { counts: [0; 6] };

    /// Builds a keep from face counts (`counts[f - 1]` dice show `f`).
    ///
    /// Returns [`None`] if the counts sum to more than five.
    #[must_use]
    pub const fn from_counts(counts: [u8; 6]) -> Option<Self> {
        // A u16 sum, like `Dice::from_counts`: u8 would wrap in release.
        let mut sum = 0u16;
        let mut i = 0;
        while i < 6 {
            sum += counts[i] as u16;
            i += 1;
        }
        if sum <= 5 {
            Some(Self { counts })
        } else {
            None
        }
    }

    /// How many dice show `face` (`1..=6`).
    ///
    /// # Panics
    ///
    /// Panics if `face` is outside `1..=6`.
    #[must_use]
    pub const fn count(self, face: u8) -> u8 {
        self.counts[(face - 1) as usize]
    }

    /// The face counts, indexed by `face - 1`.
    #[must_use]
    pub const fn counts(self) -> [u8; 6] {
        self.counts
    }

    /// The number of dice kept.
    #[must_use]
    pub const fn len(self) -> u8 {
        let mut sum = 0u8;
        let mut i = 0;
        while i < 6 {
            sum += self.counts[i];
            i += 1;
        }
        sum
    }

    /// Whether no dice are kept.
    #[must_use]
    pub const fn is_empty(self) -> bool {
        self.len() == 0
    }
}

impl From<Dice> for Keep {
    /// Keeping all five dice.
    fn from(dice: Dice) -> Self {
        Self {
            counts: dice.counts,
        }
    }
}

impl TryFrom<&[u8]> for Keep {
    type Error = ParseDiceError;

    /// Builds a keep from at most five faces in any order.
    fn try_from(faces: &[u8]) -> Result<Self, ParseDiceError> {
        if faces.len() > 5 {
            return Err(ParseDiceError::WrongCount);
        }
        let mut counts = [0u8; 6];
        for &face in faces {
            if !(1..=6).contains(&face) {
                return Err(ParseDiceError::InvalidFace);
            }
            counts[usize::from(face - 1)] += 1;
        }
        Ok(Self { counts })
    }
}

impl FromStr for Dice {
    type Err = ParseDiceError;

    /// Parses five digits `1`–`6`, ignoring whitespace: `"13446"`.
    fn from_str(s: &str) -> Result<Self, ParseDiceError> {
        let (counts, total) = parse_counts(s)?;
        if total == 5 {
            Ok(Self { counts })
        } else {
            Err(ParseDiceError::WrongCount)
        }
    }
}

impl FromStr for Keep {
    type Err = ParseDiceError;

    /// Parses up to five digits `1`–`6`, ignoring whitespace.  The empty
    /// string is [`Keep::EMPTY`].
    fn from_str(s: &str) -> Result<Self, ParseDiceError> {
        let (counts, _) = parse_counts(s)?;
        Ok(Self { counts })
    }
}

impl fmt::Display for Dice {
    /// The five faces in ascending order: `"13446"`.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_counts(&self.counts, f)
    }
}

impl fmt::Display for Keep {
    /// The kept faces in ascending order, empty when nothing is kept.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_counts(&self.counts, f)
    }
}

/// Odometer over all sub-multisets of a roll.
struct Keeps {
    limits: [u8; 6],
    next: Option<[u8; 6]>,
}

impl Iterator for Keeps {
    type Item = Keep;

    fn next(&mut self) -> Option<Keep> {
        let current = self.next?;
        let mut counts = current;
        self.next = 'carry: {
            for f in 0..6 {
                if counts[f] < self.limits[f] {
                    counts[f] += 1;
                    break 'carry Some(counts);
                }
                counts[f] = 0;
            }
            None
        };
        Some(Keep { counts: current })
    }
}

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

    #[test]
    fn constructors_enforce_invariants() {
        assert_eq!(Dice::from_faces([1, 3, 4, 4, 6]), "13446".parse().ok());
        assert_eq!(Dice::from_faces([0, 3, 4, 4, 6]), None);
        assert_eq!(Dice::from_faces([7, 3, 4, 4, 6]), None);
        assert_eq!(Dice::from_counts([5, 0, 0, 0, 0, 0]), "11111".parse().ok());
        assert_eq!(Dice::from_counts([4, 0, 0, 0, 0, 0]), None);
        assert_eq!(Keep::from_counts([6, 0, 0, 0, 0, 0]), None);
        // Counts summing to 5 mod 256 must not sneak past validation.
        assert_eq!(Dice::from_counts([200, 61, 0, 0, 0, 0]), None);
        assert_eq!(Keep::from_counts([255, 6, 0, 0, 0, 0]), None);
        assert_eq!("1344".parse::<Dice>(), Err(ParseDiceError::WrongCount));
        assert_eq!("134467".parse::<Dice>(), Err(ParseDiceError::InvalidFace));
        assert_eq!("".parse::<Keep>(), Ok(Keep::EMPTY));
        assert_eq!("1 3 4\t46".parse::<Dice>(), "13446".parse());
    }

    #[test]
    fn accessors() {
        let dice: Dice = "13446".parse().expect("a valid roll");
        assert_eq!(dice.count(4), 2);
        assert_eq!(dice.sum(), 18);
        assert_eq!(dice.yahtzee_face(), None);
        assert_eq!(dice.faces().collect::<Vec<_>>(), [1, 3, 4, 4, 6]);
        assert_eq!(dice.to_string(), "13446");

        let yahtzee: Dice = "22222".parse().expect("a valid roll");
        assert_eq!(yahtzee.yahtzee_face(), Some(2));
    }

    #[test]
    fn keep_containment() {
        let dice: Dice = "13446".parse().expect("a valid roll");
        assert!(dice.contains("44".parse().expect("a valid keep")));
        assert!(dice.contains(Keep::EMPTY));
        assert!(dice.contains(Keep::from(dice)));
        assert!(!dice.contains("444".parse().expect("a valid keep")));
        assert!(!dice.contains("2".parse().expect("a valid keep")));
    }

    #[test]
    fn keeps_enumerates_every_sub_multiset() {
        let distinct: Dice = "12345".parse().expect("a valid roll");
        assert_eq!(distinct.keeps().count(), 32);

        let yahtzee: Dice = "66666".parse().expect("a valid roll");
        assert_eq!(yahtzee.keeps().count(), 6);

        let dice: Dice = "13446".parse().expect("a valid roll");
        let keeps: Vec<Keep> = dice.keeps().collect();
        assert_eq!(keeps.len(), 24); // 2 * 2 * 3 * 2
        assert!(keeps.iter().all(|&k| dice.contains(k)));
        // No duplicates: every keep is a distinct multiset.
        for (i, a) in keeps.iter().enumerate() {
            assert!(keeps[i + 1..].iter().all(|b| a != b));
        }
    }
}