search_autocompletion/
lib.rs

1//!A Rust implementation of Search Auto Completion
2//! # Examples
3//! ```
4//! use search_autocompletion::AutoComplete;
5//!
6//! let mut com = AutoComplete::default();
7//! com.insert(&("Hello", 9));
8//! com.insert(&("Hell", 10));
9//! com.insert(&("Ham", 1000));
10//! com.insert(&("Hen", 54));
11//!
12//! let strings = com.get_strings_for_prefix("He").unwrap();
13//! assert_eq!(strings, vec!["Hen", "Hell", "Hello"]);
14//! ```
15#![warn(clippy::all)]
16#![warn(clippy::pedantic)]
17#![warn(clippy::nursery)]
18#![warn(clippy::cargo)]
19#![deny(missing_docs)]
20use std::collections::BTreeMap;
21
22/// The [`AutoComplete`] struct, basically wrapper around [`Node`]
23#[derive(Default)]
24pub struct AutoComplete<T: PartialOrd + Ord + Clone + Default> {
25    trie: Node<T>,
26}
27
28impl<T> AutoComplete<T>
29where
30    T: PartialOrd + Ord + Clone + Default,
31{
32    /// Create a new [`AutoComplete`] with pre-defined strings to autocomplete
33    pub fn new(dict: &[(String, T)]) -> Self {
34        let mut trie = Self::default();
35        dict.iter().for_each(|v| trie.insert(&(&v.0, v.1.clone())));
36        trie
37    }
38
39    /// Insert a string into the [`AutoComplete`] struct to autocomplete it
40    pub fn insert(&mut self, value: &(&str, T)) {
41        let mut curr = &mut self.trie;
42        for (c, v) in value.0.char_indices() {
43            curr = curr
44                .children
45                .entry(v)
46                .and_modify(|f| {
47                    if c == (value.0.len() - 1) {
48                        f.weight = value.1.clone();
49                        f.is_word = true;
50                    }
51                })
52                .or_insert_with(|| Node::new(value.0[0..=c].to_string(), c == (value.0.len() - 1), value.1.clone()));
53        }
54    }
55
56    /// Change weight of some string
57    /// # Errors
58    /// If Result is not OK then it failed to set the weight for this specific string for some reason
59    pub fn change_weight(&mut self, value: &(&str, T)) -> Result<(), String> {
60        let mut curr = &mut self.trie;
61
62        for v in value.0.chars() {
63            curr = match curr.children.get_mut(&v) {
64                Some(a) => a,
65                None => return Err(format!("Failed to set the weight for {}!", value.0)),
66            };
67        }
68        curr.weight = value.1.clone();
69        Ok(())
70    }
71
72    /// Get strings that fit the prefix
73    pub fn get_strings_for_prefix(&self, prefix: &str) -> Option<Vec<String>> {
74        let mut results: Vec<(String, T)> = Vec::new();
75        let mut curr = &self.trie;
76
77        for v in prefix.chars() {
78            curr = match curr.children.get(&v) {
79                Some(a) => a,
80                None => return None,
81            }
82        }
83
84        Self::find_all_child_words(curr, &mut results);
85        results.sort_by(|(_, a), (_, b)| b.cmp(a));
86        Some(results.into_iter().map(|(v, _)| v).collect())
87    }
88
89    /// Recursive function that searches for all children of a Vector of strings
90    fn find_all_child_words(node: &Node<T>, results: &mut Vec<(String, T)>) {
91        if node.is_word {
92            results.push((node.prefix.clone(), node.weight.clone()));
93        }
94
95        node.children.keys().for_each(|v| Self::find_all_child_words(node.children.get(v).unwrap(), results));
96    }
97}
98
99/// Node that is used in [`AutoComplete`] struct
100#[derive(Default)]
101struct Node<T>
102where
103    T: PartialOrd + Ord + Default + Clone,
104{
105    prefix: String,
106    children: BTreeMap<char, Node<T>>,
107    is_word: bool,
108    weight: T,
109}
110
111impl<T> Node<T>
112where
113    T: PartialOrd + Ord + Default + Clone,
114{
115    /// Creates a new node with a prefix and indicator if it is a word
116    fn new(prefix: String, is_word: bool, weight: T) -> Self {
117        Self {
118            prefix,
119            children: BTreeMap::new(),
120            is_word,
121            weight: if is_word { weight } else { T::default() },
122        }
123    }
124}