arbor_cli/common/
autocomplete.rs

1use std::error::Error;
2
3use crate::util::{app_data::AppData, backup::Backup};
4
5use super::trie::Trie;
6
7pub struct Autocomplete {
8    app_data: AppData,
9    backup: Option<Backup>,
10    trie: Trie,
11}
12
13impl Autocomplete {
14    pub async fn build(
15        language: Option<String>,
16        thread_count: Option<u8>,
17        max_suggestion: Option<u8>,
18        has_backup: bool,
19        backup_path: Option<&str>,
20    ) -> Result<Self, Box<dyn Error>> {
21        let app_data = AppData::build(language, thread_count, max_suggestion)?;
22        let backup = if has_backup {
23            Some(Backup::build(backup_path).await?)
24        } else {
25            None
26        };
27        let trie = Trie::new();
28
29        Ok(Self {
30            app_data,
31            backup,
32            trie,
33        })
34    }
35
36    pub async fn load_backup(&mut self) -> Result<(), Box<dyn Error>> {
37        if self.backup.is_some() {
38            let backup = self.backup.as_ref().unwrap();
39
40            let backup_data = backup.load_data().await?;
41
42            for word in backup_data {
43                Trie::insert(word, &mut self.trie.root, 0)?;
44            }
45        }
46
47        Ok(())
48    }
49
50    pub async fn insert_word(&mut self, word: String) -> Result<(), Box<dyn Error>> {
51        Trie::insert(word.clone(), &mut self.trie.root, 0)?;
52
53        if self.backup.is_some() {
54            self.backup
55                .as_mut()
56                .unwrap()
57                .save_data(Vec::from([word]))
58                .await?;
59        }
60
61        Ok(())
62    }
63
64    pub async fn suggest_word(&self, prefix: &str) -> Result<Vec<String>, Box<dyn Error>> {
65        let suggestions = self.trie.suggest(prefix)?;
66
67        let limit = self.app_data.get_max_suggestion() as usize;
68
69        Ok(suggestions.iter().take(limit).cloned().collect())
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use std::error::Error;
76
77    use super::*;
78
79    #[tokio::test]
80    async fn it_creates_autocomplete_instance() -> Result<(), Box<dyn Error>> {
81        let ac = Autocomplete::build(None, None, None, false, None).await?;
82
83        assert_eq!(ac.app_data.get_language(), "en-US");
84        assert!(ac.backup.is_none());
85
86        Ok(())
87    }
88
89    #[tokio::test]
90    async fn it_loads_backup() -> Result<(), Box<dyn Error>> {
91        let backup = Backup::build(None).await?;
92
93        let words = Vec::from(["hello".to_string(), "hi".to_string(), "hey".to_string()]);
94
95        backup.save_data(words.clone()).await?;
96
97        let mut ac = Autocomplete::build(None, None, None, true, None).await?;
98
99        ac.load_backup().await?;
100
101        let word_1 = ac.suggest_word("hell").await?;
102        let word_2 = ac.suggest_word("hi").await?;
103        let word_3 = ac.suggest_word("hey").await?;
104
105        assert_eq!(
106            Vec::from([
107                word_1.iter().nth(0).unwrap().to_owned(),
108                word_2.iter().nth(0).unwrap().to_owned(),
109                word_3.iter().nth(0).unwrap().to_owned(),
110            ]),
111            words
112        );
113
114        std::fs::remove_file(backup.file_path).unwrap();
115
116        Ok(())
117    }
118
119    #[tokio::test]
120    async fn it_inserts_word_and_suggests() -> Result<(), Box<dyn Error>> {
121        let mut ac = Autocomplete::build(None, None, None, false, None).await?;
122
123        let word = "test".to_string();
124
125        ac.insert_word(word.clone()).await?;
126
127        let suggestion = ac.suggest_word(word.as_str()).await?;
128
129        let suggestion = suggestion.iter().nth(0).unwrap().to_owned();
130
131        assert_eq!(word, suggestion);
132
133        Ok(())
134    }
135}