diff --git a/.gitignore b/.gitignore index 60e01bb..edacbd0 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,5 @@ Cargo.lock *.temp.* heaptrack* + +**/*.fst diff --git a/Cargo.toml b/Cargo.toml index 0db8dc1..85fd9d8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,8 +2,10 @@ name = "upodesh" version = "0.1.0" edition = "2024" +build = "build.rs" [dependencies] +fst = "0.4" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" @@ -14,6 +16,11 @@ regex = "1" okkhor = { version = "0.7", features = ["regex"] } peak_alloc = "0.2" +[build-dependencies] +fst = "0.4" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" + [[bench]] name = "suggestions" harness = false diff --git a/README.md b/README.md index 1e44ba1..ac9fcc6 100644 --- a/README.md +++ b/README.md @@ -1,21 +1,21 @@ # upodesh `upodesh` (In Bangla, `উপদেশ` is a synonym of the word `পরামর্শ` which means `suggestion`) is a Bangla word suggestion library. -This implementation uses an approach based on the Trie data structure which is substantially faster than the Regular Expression based approach. This is a Rust port of the Go project [`libavrophonetic`](https://github.com/mugli/libavrophonetic/) of Mehdi Hasan Khan. +This implementation uses an approach based on the Finite State Transducer (FST) data structure which is substantially faster than the Regular Expression based approach. This approach is inspired by the Go project [`libavrophonetic`](https://github.com/mugli/libavrophonetic/) of Mehdi Hasan Khan which used Trie data structure. ## Benchmarks -`upodesh` is around **~5× faster** than a heavily optimized regex based search approach previously used in OpenBangla Keyboard. And in cases where the old implementation struggled with large regex patterns, upodesh is a staggering **~80×** faster! - +`upodesh` is significantly faster than the previously used heavily optimized regex-based search approach in OpenBangla Keyboard. Based on recent benchmarks, it is approximately ~21× to ~58× faster, depending on the input. This demonstrates a substantial performance gain over regex, especially in cases where large patterns previously caused bottlenecks. ### 📊 Summary of the Benchmark This benchmark was performed on a Apple MacBook Air M1: | Word | `upodesh` Time | `regex` Time | Speedup | | --------- | -------------- | ------------ | --------------- | -| `a` | 39.190 µs | 193.50 µs | **\~4.94× faster** | -| `arO` | 45.942 µs | 247.68 µs | **\~5.39× faster** | -| `bistari` | 4.4495 µs | 355.04 µs | **\~79.79× faster** | +| `a` | ~3.341 µs | ~194.34 µs | **\~58× faster** | +| `arO` | ~11.840 µs | ~246.53 µs | **\~20.8× faster** | +| `bistari` | ~9.734 µs | ~353.74 µs | **\~36.3× faster** | ## Acknowledgement * [Mehdi Hasan Khan](https://github.com/mugli) and [Tahmid Sadik](https://github.com/tahmidsadik/) for their [`libavrophonetic`](https://github.com/mugli/libavrophonetic/) project. +* [Andrew Gallant](https://github.com/BurntSushi) for his amazing [`fst`](https://github.com/BurntSushi/fst) crate and [Index 1,600,000,000 Keys with Automata and Rust](https://burntsushi.net/transducers/) blog post! diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..2146136 --- /dev/null +++ b/build.rs @@ -0,0 +1,63 @@ +use std::{ + collections::HashMap, + fs::{File, read, read_to_string}, + io::BufWriter, +}; + +use fst::raw::Builder; +use serde::Deserialize; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Block { + pub transliterate: Vec, + pub entire_block_optional: Option, +} + +fn generate_words_fst() { + let file = File::create("src/words.fst").expect("Failed to create words.fst"); + let writer = BufWriter::new(file); + + let mut fst = Builder::new(writer).unwrap(); + let words = read_to_string("data/source-words.txt").expect("Failed to read source words file"); + + let mut words = words.lines().map(str::trim).collect::>(); + + words.sort(); + + for word in words { + let numbered_word = word.chars().map(|c| c as u8).collect::>(); + fst.add(&numbered_word).expect("Failed to add word to FST"); + } + + fst.finish().expect("Failed to finish words FST generation"); +} + +fn generate_patterns_fst() { + let file = File::create("src/patterns.fst").expect("Failed to create patterns.fst"); + let writer = BufWriter::new(file); + + let mut fst = Builder::new(writer).unwrap(); + let patterns: HashMap = serde_json::from_slice( + &read("data/preprocessed-patterns.json").expect("Failed to read source patterns file"), + ) + .unwrap(); + + let mut patterns = patterns.keys().map(|s| s.as_str()).collect::>(); + + patterns.sort(); + + for pattern in patterns { + let numbered_word = pattern.chars().map(|c| c as u8).collect::>(); + fst.add(&numbered_word) + .expect("Failed to add pattern to FST"); + } + + fst.finish() + .expect("Failed to finish patterns FST generation"); +} + +fn main() { + generate_words_fst(); + generate_patterns_fst(); +} diff --git a/examples/alloc.rs b/examples/alloc.rs index 0d1fc4b..fe1e1b5 100644 --- a/examples/alloc.rs +++ b/examples/alloc.rs @@ -1,5 +1,5 @@ -use upodesh::suggest::Suggest; use peak_alloc::PeakAlloc; +use upodesh::suggest::Suggest; #[global_allocator] static PEAK_ALLOC: PeakAlloc = PeakAlloc; @@ -8,7 +8,7 @@ fn main() { let _suggest = Suggest::new(); let current_mem = PEAK_ALLOC.current_usage_as_mb(); - println!("This program currently uses {} MB of RAM.", current_mem); - let peak_mem = PEAK_ALLOC.peak_usage_as_mb(); - println!("The max amount that was used {} MB.", peak_mem); + println!("This program currently uses {} MB of RAM.", current_mem); + let peak_mem = PEAK_ALLOC.peak_usage_as_mb(); + println!("The max amount that was used {} MB.", peak_mem); } diff --git a/examples/example.rs b/examples/example.rs index 470732e..6bee2ae 100644 --- a/examples/example.rs +++ b/examples/example.rs @@ -1,7 +1,6 @@ use peak_alloc::PeakAlloc; use upodesh::suggest::Suggest; - #[global_allocator] static PEAK_ALLOC: PeakAlloc = PeakAlloc; diff --git a/src/fst.rs b/src/fst.rs new file mode 100644 index 0000000..b6a99d1 --- /dev/null +++ b/src/fst.rs @@ -0,0 +1,180 @@ +use fst::raw::{Fst, Node}; + +#[derive(Clone)] +pub struct FstTree> { + fst: Fst, +} + +impl> FstTree { + pub fn from_fst(data: D) -> FstTree { + let fst = Fst::new(data).expect("Failed to create FST from bytes"); + Self { fst } + } + + pub fn match_longest_common_prefix<'a>(&self, prefix: &'a str) -> (&'a str, &'a str, bool) { + let mut iter = prefix.chars(); + let mut index = 0; + let mut node = self.fst.root(); + + while let Some(c) = iter.next() { + match node.find_input(c as u8) { + Some(addr) => { + node = self.fst.node(node.transition_addr(addr)); + index += c.len_utf8(); + } + None => break, + } + } + + let (matched_prefix, remaining) = prefix.split_at(index); + + (matched_prefix, remaining, node.is_final()) + } + + pub fn matching_node<'a>(&'a self, word: &str) -> Option> { + let mut iter = word.chars(); + let mut node = self.fst.root(); + + while let Some(c) = iter.next() { + match node.find_input(c as u8) { + Some(addr) => { + node = self.fst.node(node.transition_addr(addr)); + } + None => return None, + } + } + + Some(FstNode { + fst: &self.fst, + node, + word: word.to_string(), + }) + } +} + +#[cfg(test)] +impl FstTree> { + fn from_strings(mut set: Vec<&str>) -> Self { + set.sort(); + + let mut builder = fst::raw::Builder::memory(); + + for word in set { + // Convert each Bengali character to its single byte representation + let numbered_word = word.chars().map(|c| c as u8).collect::>(); + builder.add(&numbered_word).unwrap(); + } + + Self { + fst: builder.into_fst(), + } + } +} + +#[derive(Clone)] +pub struct FstNode<'a, D: AsRef<[u8]>> { + fst: &'a Fst, + node: Node<'a>, + word: String, +} + +impl<'a, D: AsRef<[u8]>> FstNode<'a, D> { + pub fn get_matching_node(&self, suffix: &str) -> Option> { + let mut iter = suffix.chars(); + let mut node = self.node; + + while let Some(c) = iter.next() { + match node.find_input(c as u8) { + Some(addr) => { + node = self.fst.node(node.transition_addr(addr)); + } + None => return None, + } + } + + let word = self.word.clone() + suffix; + + Some(FstNode { + fst: self.fst, + node, + word, + }) + } + + pub fn get_word(self) -> Option { + if self.node.is_final() { + Some(self.word) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_match_longest_common_prefix() { + let fst = FstTree::from_strings(vec![ + "ক", + "কখগ", + "কখগঘঙ", + "চ", + "চছজ", + "চছজঝঞ", + "১", + "a", + "bc", + "abcd", + ]); + + assert_eq!(fst.match_longest_common_prefix("ক"), ("ক", "", true)); + assert_eq!( + fst.match_longest_common_prefix("ক1234"), + ("ক", "1234", true) + ); + assert_eq!(fst.match_longest_common_prefix("1234"), ("", "1234", false)); + assert_eq!( + fst.match_longest_common_prefix("কখগঘঙচছজঝঞ"), + ("কখগঘঙ", "চছজঝঞ", true) + ); + + assert_eq!(fst.match_longest_common_prefix("a"), ("a", "", true)); + assert_eq!(fst.match_longest_common_prefix("a123"), ("a", "123", true)); + assert_eq!( + fst.match_longest_common_prefix("abcdefg"), + ("abcd", "efg", true) + ); + } + + #[test] + fn test_find_matching_node() { + let fst = FstTree::from_strings(vec!["ক", "কখ", "কখগঘঙচছ"]); + + let n1 = fst.matching_node("ক").unwrap(); + + let n2 = n1.get_matching_node("খ").unwrap(); + + _ = n2.get_matching_node("গঘ").unwrap(); + + _ = fst.matching_node("কখগঘ").unwrap(); + } + + #[test] + fn test_get_word() { + let trie = FstTree::from_strings(vec!["ক", "কখ", "কখগঘঙচছ"]); + + let n1 = trie.matching_node("ক").unwrap(); + assert_eq!(n1.clone().get_word(), Some("ক".to_string())); + + let n2 = n1.get_matching_node("খ").unwrap(); + assert_eq!(n2.clone().get_word(), Some("কখ".to_string())); + + let n3 = n2.get_matching_node("গঘ").unwrap(); + assert_eq!(n3.get_word(), None); + + let n4 = trie.matching_node("কখগঘ").unwrap(); + assert_eq!(n4.get_word(), None); + } +} diff --git a/src/lib.rs b/src/lib.rs index 9abe03e..4d4f8a3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,3 @@ -mod trie; -mod utils; +mod fst; pub mod suggest; +mod utils; diff --git a/src/suggest.rs b/src/suggest.rs index b8612a0..923ccbb 100644 --- a/src/suggest.rs +++ b/src/suggest.rs @@ -1,8 +1,17 @@ -use std::collections::{HashMap, HashSet}; +use std::{ + collections::{HashMap, HashSet}, + sync::LazyLock, +}; use serde::Deserialize; -use crate::{trie::Trie, utils::fix_string}; +use crate::{fst::FstTree, utils::fix_string}; + +static WORDS: LazyLock> = + LazyLock::new(|| FstTree::from_fst(include_bytes!("words.fst"))); + +static PATTERNS: LazyLock> = + LazyLock::new(|| FstTree::from_fst(include_bytes!("patterns.fst"))); #[derive(Deserialize)] #[serde(rename_all = "camelCase")] @@ -13,39 +22,34 @@ pub struct Block { pub struct Suggest { patterns: HashMap, - patterns_trie: Trie, - words: Trie, common_suffixes: Vec, } impl Suggest { pub fn new() -> Self { let patterns_data = include_bytes!("../data/preprocessed-patterns.json"); - let words_data = include_str!("../data/source-words.txt"); let common_data = include_bytes!("../data/source-common-patterns.json"); let patterns: HashMap = serde_json::from_slice(patterns_data).unwrap(); - let patterns_trie = Trie::from_strings(patterns.keys().map(|s| s.as_str())); - let words = Trie::from_strings(words_data.lines().map(|s| s.trim())); let common_suffixes = serde_json::from_slice(common_data).unwrap(); Suggest { patterns, - patterns_trie, - words, common_suffixes, } } pub fn suggest(&self, input: &str) -> Vec { + let words = LazyLock::force(&WORDS); + let patterns = LazyLock::force(&PATTERNS); let input = fix_string(input); - let (matched, mut remaining, _) = self.patterns_trie.match_longest_common_prefix(&input); + let (matched, mut remaining, _) = patterns.match_longest_common_prefix(&input); let matched_patterns = &self.patterns.get(matched).unwrap().transliterate; let mut matched_nodes = matched_patterns .iter() - .filter_map(|p| self.words.matching_node(p)) + .filter_map(|p| words.matching_node(p)) .collect::>(); let additional_nodes = matched_nodes @@ -56,17 +60,17 @@ impl Suggest { .filter_map(|suffix| node.get_matching_node(suffix)) }) .collect::>(); + matched_nodes.extend(additional_nodes); while !remaining.is_empty() { - let (mut new_matched, mut new_remaining, mut complete) = - self.patterns_trie.match_longest_common_prefix(&remaining); + let (mut new_matched, new_remaining, mut complete) = + patterns.match_longest_common_prefix(&remaining); if !complete { for i in (0..remaining.len()).rev() { - (new_matched, new_remaining, complete) = self - .patterns_trie - .match_longest_common_prefix(&remaining[..i]); + (new_matched, _, complete) = + patterns.match_longest_common_prefix(&remaining[..i]); if complete { remaining = &remaining[i..]; @@ -111,8 +115,11 @@ impl Suggest { matched_nodes.extend(additional_matched_nodes); } - let suggestions: HashSet<_> = matched_nodes.iter().filter_map(|n| n.get_word()).collect(); - suggestions.into_iter().map(|s| s.to_string()).collect() + let suggestions: HashSet<_> = matched_nodes + .into_iter() + .filter_map(|n| n.get_word()) + .collect(); + suggestions.into_iter().collect() } } diff --git a/src/trie.rs b/src/trie.rs deleted file mode 100644 index 2081584..0000000 --- a/src/trie.rs +++ /dev/null @@ -1,203 +0,0 @@ -use std::collections::BTreeMap; - -#[derive(Debug, Clone, PartialEq)] -pub struct TrieNode { - children: BTreeMap, - word: Option, // Store the complete word at the end of the node -} - -impl<'a> TrieNode { - fn new() -> Self { - TrieNode { - children: BTreeMap::new(), - word: None, - } - } - - #[cfg(test)] - fn is_complete_word(&self) -> bool { - self.word.is_some() - } - - pub fn get_word(&'a self) -> Option<&'a str> { - self.word.as_deref() - } - - #[cfg(test)] - fn find_complete_words(&'a self) -> Vec<&'a str> { - let mut words = Vec::new(); - for node in self.children.values() { - if let Some(word) = &node.word { - words.push(word.as_str()); - } - - let child_words = node.find_complete_words(); - words.extend(child_words); - } - - words - } - - /// FindMatchingNode - pub fn get_matching_node(&self, word: &str) -> Option<&TrieNode> { - let mut current_node = self; - for ch in word.chars() { - match current_node.children.get(&ch) { - Some(node) => current_node = node, - None => return None, - } - } - Some(current_node) - } -} - -pub struct Trie { - root: TrieNode, -} - -impl<'a> Trie { - fn new() -> Self { - Trie { - root: TrieNode::new(), - } - } - - fn insert(&mut self, word: &str) { - let mut current_node = &mut self.root; - for ch in word.chars() { - current_node = current_node.children.entry(ch).or_insert(TrieNode::new()); - } - - current_node.word = Some(word.to_string()); // Store the word at the end of the node - } - - pub fn from_strings(words: impl Iterator) -> Self { - let mut trie = Trie::new(); - for word in words { - trie.insert(word); - } - trie - } - - /// FindMatchingNode - pub fn matching_node(&self, word: &str) -> Option<&TrieNode> { - self.root.get_matching_node(word) - } - - /// findLongestPrefixNode - fn longest_prefix(&self, word: &str) -> (&TrieNode, usize) { - let mut current_node = &self.root; - let mut prefix_len: usize = 0; - - for ch in word.chars() { - match current_node.children.get(&ch) { - Some(node) => { - prefix_len += ch.len_utf8(); - current_node = node; - } - None => break, - } - } - - (current_node, prefix_len) - } - - /// MatchPrefix - #[cfg(test)] - pub fn match_prefix(&'a self, prefix: &'a str) -> Vec<&'a str> { - let mut result = Vec::new(); - - if prefix.is_empty() { - return result; - } - - let (node, prefix_len) = self.longest_prefix(prefix); - if prefix_len == 0 { - return result; - }; - - result.extend(node.find_complete_words()); - - if let Some(word) = &node.word { - result.push(word); - } - - result - } - - /// MatchLongestCommonPrefix - pub fn match_longest_common_prefix(&self, prefix: &'a str) -> (&'a str, &'a str, bool) { - if prefix.is_empty() { - return ("", "", false); - } - - let (node, matched_prefix_len) = self.longest_prefix(prefix); - let (matched_prefix, remaining) = prefix.split_at(matched_prefix_len); - (matched_prefix, remaining, node.word.is_some()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_find_matching_node() { - let trie = Trie::from_strings(["ক", "কখ", "কখগঘঙচছ"].into_iter()); - - let n1 = trie.matching_node("ক").unwrap(); - - let n2 = n1.get_matching_node("খ").unwrap(); - - _ = n2.get_matching_node("গঘ").unwrap(); - - _ = trie.matching_node("কখগঘ").unwrap(); - } - - #[test] - fn test_is_complete_word() { - let trie = Trie::from_strings(["ক", "কখ", "কখগঘঙচছ"].into_iter()); - - let n1 = trie.matching_node("ক").unwrap(); - assert!(n1.is_complete_word()); - - let n2 = n1.get_matching_node("খ").unwrap(); - assert!(n2.is_complete_word()); - - let n3 = n2.get_matching_node("গঘ").unwrap(); - assert!(!n3.is_complete_word()); - - let n4 = trie.matching_node("কখগঘ").unwrap(); - assert!(!n4.is_complete_word()); - } - - #[test] - fn test_match_prefix() { - let trie = Trie::from_strings(["ক", "কখগ", "কখগঘঙ", "চ", "চছজ", "চছজঝঞ", "১"].into_iter()); - - assert_eq!(trie.match_prefix("ক"), vec!["কখগ", "কখগঘঙ", "ক"]); - assert_eq!(trie.match_prefix("কখ"), vec!["কখগ", "কখগঘঙ"]); - assert_eq!(trie.match_prefix("চছজঝঞ"), vec!["চছজঝঞ"]); - assert_eq!(trie.match_prefix("২"), Vec::::new()); - assert_eq!(trie.match_prefix(""), Vec::::new()); - } - - #[test] - fn test_match_longest_common_prefix() { - let trie = Trie::from_strings(["ক", "কখগ", "কখগঘঙ", "চ", "চছজ", "চছজঝঞ", "১"].into_iter()); - - assert_eq!(trie.match_longest_common_prefix("ক"), ("ক", "", true)); - assert_eq!( - trie.match_longest_common_prefix("ক1234"), - ("ক", "1234", true) - ); - assert_eq!( - trie.match_longest_common_prefix("1234"), - ("", "1234", false) - ); - assert_eq!( - trie.match_longest_common_prefix("কখগঘঙচছজঝঞ"), - ("কখগঘঙ", "চছজঝঞ", true) - ); - } -}