1use std::collections::HashMap;
2
3#[derive(Default, Debug)]
4pub struct TrieNode {
5 pub(crate) is_end_of_word: bool,
6 pub(crate) children: HashMap<char, TrieNode>,
7}
8
9impl TrieNode {
10 pub(crate) fn len(&self) -> usize {
16 let mut sum = self.children.len();
17 for a in self.children.values() {
18 sum += a.len();
19 }
20 sum
21 }
22
23 pub(crate) fn get_suggestions(&self, current: &str) -> Vec<String> {
27 let mut words: Vec<String> = Vec::new();
28 if self.is_end_of_word {
29 words.push(current.to_string());
30 eprintln!("Inserting {current}")
31 }
32 for (c, node) in self.children.iter() {
33 words.append(
34 node.get_suggestions(format!("{current}{c}").as_str())
35 .as_mut(),
36 );
37 }
38 eprintln!("returning {words:#?}");
39 words
40 }
41}
42
43#[cfg(test)]
44mod tests {
45 use super::*;
46
47 #[test]
48 fn it_should_create_a_new_node() {
49 let result = TrieNode::default();
50 assert!(!result.is_end_of_word);
51 assert_eq!(result.children.len(), 0);
52 }
53
54 #[test]
55 fn it_should_have_correct_length() {
56 let mut node = TrieNode::default();
57 assert_eq!(node.len(), 0);
58 let mut child = TrieNode::default();
59 child.is_end_of_word = true;
60 node.children.insert('a', child);
61 assert_eq!(node.len(), 1);
62 }
63
64 #[test]
65 fn it_should_get_suggestions() {
66 let mut node = TrieNode::default();
67 let mut child = TrieNode::default();
68 child.is_end_of_word = true;
69 node.children.insert('a', child);
70 let words = node.get_suggestions("");
71 assert_eq!(words, vec!["a"]);
72 }
73
74 #[test]
75 fn it_should_get_suggestions_independantly_from_prefix() {
76 let mut node = TrieNode::default();
77 let mut child = TrieNode::default();
78 child.is_end_of_word = true;
79 node.children.insert('a', child);
80 let prefixes = ["some", "prefix"];
81 for prefix in prefixes.iter() {
82 let words = node.get_suggestions(prefix);
83 assert_eq!(words, vec![format!("{prefix}a")]);
84 }
85 }
86
87 #[test]
88 fn it_should_get_suggestions_with_multiple_children() {
89 let mut node = TrieNode::default();
90 let mut child = TrieNode::default();
91 child.is_end_of_word = true;
92 node.children.insert('a', child);
93 let mut child = TrieNode::default();
94 child.is_end_of_word = true;
95 node.children.insert('b', child);
96 let words = node.get_suggestions("");
97 assert!(words.contains(&"a".to_string()));
98 assert!(words.contains(&"b".to_string()));
99 }
100}