werewolves/werewolves-proto/src/bag.rs

// Copyright (C) 2025 Emilis Bliūdžius
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use rand::{SeedableRng, rngs::SmallRng, seq::SliceRandom};
// #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
// pub enum BagItem<T, V> {
//     Left(T),
//     Right(V),
// }
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Bag<T>(Vec<T>);

impl<T> Bag<T> {
    pub fn new(items: impl IntoIterator<Item = T>) -> Self {
        Self(items.into_iter().collect())
    }

    pub fn pull(&mut self) -> Option<T> {
        self.0.pop()
    }

    pub fn peek(&self) -> Option<&T> {
        self.0.last()
    }

    pub const fn len(&self) -> usize {
        self.0.len()
    }
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Default)]
pub enum DrunkRoll {
    Drunk,
    #[default]
    Sober,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct DrunkBag {
    #[serde(skip)]
    rng: SmallRng,
    seed: u64,
    bag_number: usize,
    bag: Bag<DrunkRoll>,
}

impl<'de> Deserialize<'de> for DrunkBag {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct DrunkBagNoRng {
            seed: u64,
            bag_number: usize,
            bag: Bag<DrunkRoll>,
        }
        let DrunkBagNoRng {
            seed,
            bag_number,
            bag,
        } = DrunkBagNoRng::deserialize(deserializer)?;
        let mut rng = SmallRng::seed_from_u64(seed);
        // Shuffle the default bag bag_number of times to get the smallrng to the same state
        for _ in 0..bag_number {
            Self::DEFAULT_BAG
                .iter()
                .copied()
                .collect::<Box<[_]>>()
                .shuffle(&mut rng);
        }
        Ok(Self {
            rng,
            seed,
            bag_number,
            bag,
        })
    }
}

impl Default for DrunkBag {
    fn default() -> Self {
        Self::new()
    }
}

impl DrunkBag {
    const DEFAULT_BAG: &[DrunkRoll] = &[
        DrunkRoll::Drunk,
        DrunkRoll::Drunk,
        DrunkRoll::Sober,
        DrunkRoll::Sober,
        DrunkRoll::Sober,
    ];

    #[cfg(test)]
    #[doc(hidden)]
    pub fn all_drunk() -> Self {
        Self {
            rng: SmallRng::seed_from_u64(0),
            seed: 0,
            bag_number: 1,
            bag: Bag::new([
                DrunkRoll::Drunk,
                DrunkRoll::Drunk,
                DrunkRoll::Drunk,
                DrunkRoll::Drunk,
                DrunkRoll::Drunk,
            ]),
        }
    }

    pub fn new() -> Self {
        let seed = rand::random();
        let mut rng = SmallRng::seed_from_u64(seed);
        let mut starting_bag = Self::DEFAULT_BAG.iter().copied().collect::<Box<[_]>>();
        starting_bag.shuffle(&mut rng);
        let bag = Bag::new(starting_bag);

        Self {
            rng,
            seed,
            bag,
            bag_number: 1,
        }
    }

    pub fn peek(&self) -> DrunkRoll {
        self.bag.peek().copied().unwrap_or_default()
    }

    fn next_bag(&mut self) {
        let mut starting_bag = Self::DEFAULT_BAG.iter().copied().collect::<Box<[_]>>();
        starting_bag.shuffle(&mut self.rng);
        self.bag = Bag::new(starting_bag);
        self.bag_number += 1;
    }

    pub fn pull(&mut self) -> DrunkRoll {
        if self.bag.len() < 2 {
            *self = Self::new();
        } else if self.bag.len() == 2 {
            let pulled = self.bag.pull().unwrap_or_default();
            self.next_bag();
            return pulled;
        }
        self.bag.pull().unwrap_or_default()
    }
}