Use FST instead of Trie for consuming less memory

.gitignore

@@ -21,3 +21,5 @@ Cargo.lock
 *.temp.*
 
 heaptrack*
+
+**/*.fst

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

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!

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<String>,
+    pub entire_block_optional: Option<bool>,
+}
+
+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::<Vec<_>>();
+
+    words.sort();
+
+    for word in words {
+        let numbered_word = word.chars().map(|c| c as u8).collect::<Vec<u8>>();
+        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<String, Block> = 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::<Vec<_>>();
+
+    patterns.sort();
+
+    for pattern in patterns {
+        let numbered_word = pattern.chars().map(|c| c as u8).collect::<Vec<u8>>();
+        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();
+}

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);
 }

examples/example.rs

@@ -1,7 +1,6 @@
 use peak_alloc::PeakAlloc;
 use upodesh::suggest::Suggest;
 
-
 #[global_allocator]
 static PEAK_ALLOC: PeakAlloc = PeakAlloc;
 

src/fst.rs

@@ -0,0 +1,180 @@
+use fst::raw::{Fst, Node};
+
+#[derive(Clone)]
+pub struct FstTree<D: AsRef<[u8]>> {
+    fst: Fst<D>,
+}
+
+impl<D: AsRef<[u8]>> FstTree<D> {
+    pub fn from_fst(data: D) -> FstTree<D> {
+        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<FstNode<'a, D>> {
+        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<Vec<u8>> {
+    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::<Vec<u8>>();
+            builder.add(&numbered_word).unwrap();
+        }
+
+        Self {
+            fst: builder.into_fst(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct FstNode<'a, D: AsRef<[u8]>> {
+    fst: &'a Fst<D>,
+    node: Node<'a>,
+    word: String,
+}
+
+impl<'a, D: AsRef<[u8]>> FstNode<'a, D> {
+    pub fn get_matching_node(&self, suffix: &str) -> Option<FstNode<'a, D>> {
+        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<String> {
+        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);
+    }
+}

src/lib.rs

@@ -1,3 +1,3 @@
-mod trie;
-mod utils;
+mod fst;
 pub mod suggest;
+mod utils;