optimle/wordlelike/src/constraints.rs

use std::collections::HashMap;
use std::convert::TryInto;
use crate::guess::*;

pub trait Constraints<const N: usize> {
    fn push(&mut self, result: GuessResult<N>);
    fn check(&self, word: Word<N>) -> bool;
}

#[derive(Clone)]
pub struct LetterCount<const N: usize> {
    letter_minimums: HashMap<char, usize>,
    letter_maximums: HashMap<char, usize>,
    empty: bool,
}

impl<const N: usize> LetterCount<N> {
    pub fn new() -> Self {
        LetterCount {
            letter_minimums: HashMap::new(),
            letter_maximums: HashMap::new(),
            empty: true,
        }
    }
}

impl<const N: usize> Constraints<N> for LetterCount<N> {
    fn push(&mut self, result: GuessResult<N>) {
        self.empty = false;

        let mut letter_matches = HashMap::new();
        for a in 0..N {
            letter_matches
                .entry(result.word[a])
                .or_insert(Vec::new())
                .push(result.matches[a]);
        }

        for (letter, matches) in letter_matches {
            let nowheres = matches.iter().filter(|&m| *m == LetterMatch::NOWHERE).count();
            let somewheres = matches.len() - nowheres;
            self.letter_minimums.insert(
                letter,
                std::cmp::max(
                    *self.letter_minimums.get(&letter).unwrap_or(&0),
                    somewheres,
                ),
            );
            if nowheres > 0 {
                self.letter_maximums.insert(
                    letter,
                    std::cmp::min(
                        *self.letter_maximums.get(&letter).unwrap_or(&N),
                        somewheres,
                    ),
                );
            }
        }
    }

    fn check(&self, word: Word<N>) -> bool {
        if self.empty { return true; }

        let mut letter_counts = HashMap::new();
        for a in 0..N {
            *letter_counts.entry(word[a]).or_insert(0) += 1;
        }

        for (letter, minimum) in &self.letter_minimums {
            let letter_count = letter_counts.get(&letter).unwrap_or(&0);
            if letter_count < minimum {
                return false;
            }
        }
        for (letter, maximum) in &self.letter_maximums {
            let letter_count = letter_counts.get(&letter).unwrap_or(&0);
            if letter_count > maximum {
                return false;
            }
        }

        true
    }
}

#[derive(Clone)]
pub struct LetterPosition<const N: usize> {
    heres: [Option<char>; N],
    elsewheres: [Vec<char>; N],
    empty: bool,
}

impl<const N: usize> LetterPosition<N> {
    pub fn new() -> Self {
        Self {
            heres: [None; N],
            elsewheres: vec![Vec::new(); N].try_into().expect("failed to initialize [Vec<char>; N]"),
            empty: true,
        }
    }
}

impl<const N: usize> Constraints<N> for LetterPosition<N> {
    fn push(&mut self, result: GuessResult<N>) {
        self.empty = false;

        for a in 0..N {
            match result.matches[a] {
                LetterMatch::HERE => self.heres[a] = Some(result.word[a]),
                LetterMatch::ELSEWHERE => {
                    self.elsewheres[a].push(result.word[a]);
                    for b in 0..N {
                        if a != b && result.word[b] == result.word[a] && result.matches[b] == LetterMatch::NOWHERE {
                            self.elsewheres[b].push(result.word[a]);
                        }
                    }
                },
                LetterMatch::NOWHERE => (),
            }
        }
    }

    fn check(&self, word: Word<N>) -> bool {
        if self.empty { return true; }

        for a in 0..N {
            if let Some(here) = self.heres[a] {
                if word[a] != here {
                    return false;
                }
            }
            if self.elsewheres[a].contains(&word[a]) {
                return false;
            }
        }

        true
    }
}

#[derive(Clone)]
pub struct Total<const N: usize> {
    count: LetterCount<N>,
    position: LetterPosition<N>,
}

impl<const N: usize> Total<N> {
    pub fn new() -> Self {
        Self {
            count: LetterCount::new(),
            position: LetterPosition::new(),
        }
    }
}

impl<const N: usize> Constraints<N> for Total<N> {
    fn push(&mut self, result: GuessResult<N>) {
        self.count.push(result);
        self.position.push(result);
    }

    fn check(&self, word: Word<N>) -> bool {
        self.count.check(word) && self.position.check(word)
    }
}


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

    #[test]
    fn it_constrains_nowheres() {
        let mut constraints = Total::new();
        constraints.push(GuessResult {
            word: word!("loose"),
            matches: [LetterMatch::NOWHERE; 5],
        });

        assert_eq!(false, constraints.check(word!("blink"))); // no L
        assert_eq!(false, constraints.check(word!("ghost"))); // no O
        assert_eq!(true, constraints.check(word!("chain")));
    }
    
    #[test]
    fn it_constraints_elsewheres() {
        let mut constraints = Total::new();
        constraints.push(GuessResult {
            word: word!("loose"),
            matches: [LetterMatch::NOWHERE, LetterMatch::ELSEWHERE, LetterMatch::NOWHERE, LetterMatch::NOWHERE, LetterMatch::NOWHERE],
        });

        assert_eq!(false, constraints.check(word!("brood"))); // too many O's
        assert_eq!(false, constraints.check(word!("going"))); // O can't go here
        assert_eq!(false, constraints.check(word!("abort"))); // O can't go here
        assert_eq!(true, constraints.check(word!("favor")));
    }

    #[test]
    fn it_constrains_heres() {
        let mut constraints = Total::new();
        constraints.push(GuessResult {
            word: word!("loose"),
            matches: [LetterMatch::NOWHERE, LetterMatch::HERE, LetterMatch::HERE, LetterMatch::HERE, LetterMatch::HERE],
        });

        assert_eq!(false, constraints.check(word!("house"))); // no U here
        assert_eq!(false, constraints.check(word!("loose"))); // no L here
        assert_eq!(true, constraints.check(word!("goose")));
    }
}