wordle_suggest/
lib.rs

1use std::{
2    cmp::Ordering,
3    collections::{BinaryHeap, HashMap, HashSet},
4    iter,
5};
6
7use rand::{rngs::StdRng, Rng};
8
9mod weights {
10    include!(concat!(env!("OUT_DIR"), "/weights.rs"));
11}
12
13pub type Hint = [CharHint; 5];
14pub type Word = [char; 5];
15
16#[derive(Clone, Debug, PartialEq)]
17pub enum CharHint {
18    Here(char),
19    Elsewhere(char),
20    None(char),
21}
22
23#[derive(Debug, Ord, Eq, PartialEq)]
24struct WeightedWord {
25    word: Word,
26    weight: usize,
27    common: bool,
28}
29
30impl PartialOrd for WeightedWord {
31    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
32        if self.common == other.common {
33            self.weight.partial_cmp(&other.weight)
34        } else {
35            self.common.partial_cmp(&other.common)
36        }
37    }
38}
39
40impl Into<String> for WeightedWord {
41    fn into(self) -> String {
42        String::from_iter(self.word)
43    }
44}
45
46pub fn suggestions(
47    hints: &Vec<Hint>,
48    unique: bool,
49    mut random: Option<StdRng>,
50    limit: Option<usize>,
51) -> impl Iterator<Item = String> {
52    let mut heap: BinaryHeap<WeightedWord> = weights::WEIGHTS
53        .into_iter()
54        .filter(|(word, _, _)| satisfies_uniqueness(word, unique))
55        .filter(|(word, _, _)| satisfies_hints(word, hints))
56        .map(|(word, weight, common)| {
57            if let Some(rng) = random.as_mut() {
58                WeightedWord {
59                    word,
60                    weight: rng.gen::<usize>(),
61                    common: rng.gen::<bool>(),
62                }
63            } else {
64                WeightedWord {
65                    word,
66                    weight,
67                    common,
68                }
69            }
70        })
71        .collect();
72
73    let limit = limit.unwrap_or_else(|| heap.len());
74
75    iter::from_fn(move || heap.pop().map(Into::into)).take(limit)
76}
77
78fn satisfies_hints(word: &Word, hints: &Vec<Hint>) -> bool {
79    hints.iter().all(|hint| satisfies_hint(word, hint))
80}
81
82fn satisfies_hint(word: &Word, hint: &Hint) -> bool {
83    let mut known_char_counts = HashMap::new();
84
85    hint.iter().enumerate().all(|(i, ch)| match ch {
86        CharHint::Here(c) => {
87            *known_char_counts.entry(c).or_insert(0) += 1;
88            word[i] == *c
89        }
90        CharHint::Elsewhere(c) => {
91            *known_char_counts.entry(c).or_insert(0) += 1;
92            word[i] != *c && word.contains(c)
93        }
94        CharHint::None(_) => true
95    }) && hint.iter().all(|ch| match ch {
96        CharHint::None(c) => {
97            let expected_count = known_char_counts.get(c).cloned().unwrap_or(0);
98            let actual_count = word.iter().filter(|&wc| *wc == *c).count();
99
100            expected_count == actual_count
101        }
102        CharHint::Here(_) | CharHint::Elsewhere(_) => true,
103    })
104}
105
106fn satisfies_uniqueness(word: &Word, unique: bool) -> bool {
107    if unique {
108        HashSet::<char>::from_iter(*word).len() == word.len()
109    } else {
110        true
111    }
112}
113
114#[cfg(test)]
115mod test {
116    use crate::{satisfies_hint, CharHint, WeightedWord};
117
118    #[test]
119    fn test_weighted_word_ord() {
120        let mut lhs = WeightedWord {
121            word: ['a', 'b', 'c', 'd', 'e'],
122            weight: 0,
123            common: false,
124        };
125
126        let mut rhs = WeightedWord {
127            word: ['a', 'b', 'c', 'd', 'e'],
128            weight: 1,
129            common: false,
130        };
131
132        assert!(lhs < rhs);
133
134        lhs.common = true;
135        assert!(lhs > rhs);
136
137        rhs.common = true;
138        assert!(lhs < rhs);
139
140        rhs.weight = lhs.weight;
141        assert!(lhs == rhs);
142    }
143
144    #[test]
145    fn test_all_nones() {
146        assert!(
147            satisfies_hint(
148                &['m', 'o', 'n', 'e', 'y'],
149                &[
150                    CharHint::None('q'),
151                    CharHint::None('x'),
152                    CharHint::None('p'),
153                    CharHint::None('z'),
154                    CharHint::None('r'),
155                ]
156            ),
157            "All `None`s are satisfied by a word containing none of those letters"
158        );
159    }
160
161    #[test]
162    fn test_all_heres() {
163        assert!(
164            satisfies_hint(
165                &['m', 'o', 'n', 'e', 'y'],
166                &[
167                    CharHint::Here('m'),
168                    CharHint::Here('o'),
169                    CharHint::Here('n'),
170                    CharHint::Here('e'),
171                    CharHint::Here('y'),
172                ]
173            ),
174            "All `Here`s are satisified by the matching word"
175        );
176    }
177
178    #[test]
179    fn test_elsewhere() {
180        assert!(
181            satisfies_hint(
182                &['m', 'o', 'n', 'e', 'y'],
183                &[
184                    CharHint::Elsewhere('y'),
185                    CharHint::None('x'),
186                    CharHint::None('p'),
187                    CharHint::None('z'),
188                    CharHint::None('r'),
189                ]
190            ),
191            "An `Elsewhere` is satisfied with a letter in a different position"
192        );
193    }
194
195    #[test]
196    fn test_single_none() {
197        assert!(
198            !satisfies_hint(
199                &['a', 'p', 'n', 'i', 'c'],
200                &[
201                    CharHint::Elsewhere('p'),
202                    CharHint::None('a'), // <-
203                    CharHint::Here('n'),
204                    CharHint::Here('i'),
205                    CharHint::Here('c'),
206                ]
207            ),
208            "A single `None` rejects words containing that letter"
209        );
210    }
211
212    #[test]
213    fn test_repeated_hint_chars() {
214        assert!(
215            satisfies_hint(
216                &['b', 'o', 'a', 't', 's'],
217                &[
218                    CharHint::Here('b'),
219                    CharHint::Elsewhere('a'),
220                    CharHint::None('b'),
221                    CharHint::None('b'),
222                    CharHint::None('y'),
223                ]
224            ),
225            "Repeated hint characters can be marked `None`"
226        );
227    }
228
229    #[test]
230    fn test_here_after_none() {
231        assert!(satisfies_hint(
232            &['f', 'r', 'a', 'm', 'e'],
233            &[
234                CharHint::None('e'),
235                CharHint::Here('r'),
236                CharHint::Here('a'),
237                CharHint::None('s'),
238                CharHint::Here('e'),
239            ]
240        ),);
241    }
242
243    #[test]
244    fn test_belle() {
245        assert!(!satisfies_hint(
246            &['b', 'e', 'l', 'l', 'e'],
247            &[
248                CharHint::Here('b'),
249                CharHint::None('e'),
250                CharHint::None('l'),
251                CharHint::Here('l'),
252                CharHint::Here('e'),
253            ]
254        ),);
255    }
256
257    #[test]
258    fn test_here_elsewhere_none() {
259        assert!(satisfies_hint(
260            &['a', 'b', 'a', 'b', 'c'],
261            &[
262                CharHint::Here('a'),
263                CharHint::Elsewhere('a'),
264                CharHint::None('x'),
265                CharHint::None('a'),
266                CharHint::None('x'),
267            ]
268        ),);
269    }
270}