1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

//!A Rust implementation of Search Auto Completion
//! # Examples
//! ```
//! use search_autocompletion::AutoComplete;
//!
//! let mut com = AutoComplete::default();
//! com.insert(&("Hello", 9));
//! com.insert(&("Hell", 10));
//! com.insert(&("Ham", 1000));
//! com.insert(&("Hen", 54));
//!
//! let strings = com.get_strings_for_prefix("He").unwrap();
//! assert_eq!(strings, vec!["Hen", "Hell", "Hello"]);
//! ```
#![warn(clippy::all)]
#![warn(clippy::pedantic)]
#![warn(clippy::nursery)]
#![warn(clippy::cargo)]
#![deny(missing_docs)]
use std::collections::BTreeMap;

/// The [`AutoComplete`] struct, basically wrapper around [`Node`]
#[derive(Default)]
pub struct AutoComplete<T: PartialOrd + Ord + Clone + Default> {
    trie: Node<T>,
}

impl<T> AutoComplete<T>
where
    T: PartialOrd + Ord + Clone + Default,
{
    /// Create a new [`AutoComplete`] with pre-defined strings to autocomplete
    pub fn new(dict: &[(String, T)]) -> Self {
        let mut trie = Self::default();
        dict.iter().for_each(|v| trie.insert(&(&v.0, v.1.clone())));
        trie
    }

    /// Insert a string into the [`AutoComplete`] struct to autocomplete it
    pub fn insert(&mut self, value: &(&str, T)) {
        let mut curr = &mut self.trie;
        for (c, v) in value.0.char_indices() {
            curr = curr
                .children
                .entry(v)
                .and_modify(|f| {
                    if c == (value.0.len() - 1) {
                        f.weight = value.1.clone();
                        f.is_word = true;
                    }
                })
                .or_insert_with(|| Node::new(value.0[0..=c].to_string(), c == (value.0.len() - 1), value.1.clone()));
        }
    }

    /// Change weight of some string
    /// # Errors
    /// If Result is not OK then it failed to set the weight for this specific string for some reason
    pub fn change_weight(&mut self, value: &(&str, T)) -> Result<(), String> {
        let mut curr = &mut self.trie;

        for v in value.0.chars() {
            curr = match curr.children.get_mut(&v) {
                Some(a) => a,
                None => return Err(format!("Failed to set the weight for {}!", value.0)),
            };
        }
        curr.weight = value.1.clone();
        Ok(())
    }

    /// Get strings that fit the prefix
    pub fn get_strings_for_prefix(&self, prefix: &str) -> Option<Vec<String>> {
        let mut results: Vec<(String, T)> = Vec::new();
        let mut curr = &self.trie;

        for v in prefix.chars() {
            curr = match curr.children.get(&v) {
                Some(a) => a,
                None => return None,
            }
        }

        Self::find_all_child_words(curr, &mut results);
        results.sort_by(|(_, a), (_, b)| b.cmp(a));
        Some(results.into_iter().map(|(v, _)| v).collect())
    }

    /// Recursive function that searches for all children of a Vector of strings
    fn find_all_child_words(node: &Node<T>, results: &mut Vec<(String, T)>) {
        if node.is_word {
            results.push((node.prefix.clone(), node.weight.clone()));
        }

        node.children.keys().for_each(|v| Self::find_all_child_words(node.children.get(v).unwrap(), results));
    }
}

/// Node that is used in [`AutoComplete`] struct
#[derive(Default)]
struct Node<T>
where
    T: PartialOrd + Ord + Default + Clone,
{
    prefix: String,
    children: BTreeMap<char, Node<T>>,
    is_word: bool,
    weight: T,
}

impl<T> Node<T>
where
    T: PartialOrd + Ord + Default + Clone,
{
    /// Creates a new node with a prefix and indicator if it is a word
    fn new(prefix: String, is_word: bool, weight: T) -> Self {
        Self {
            prefix,
            children: BTreeMap::new(),
            is_word,
            weight: if is_word { weight } else { T::default() },
        }
    }
}