Repository: PyThaiNLP/nlpo3 (language: python)

Issue:
Clarify method in Tokenizer trait
From Tokenizer trait declaration:

```rust
pub trait Tokenizer {
    fn segment(&self, text: &str, safe: Option<bool>, parallel: Option<bool>) -> Vec<String>;

    fn segment_to_string(
        &self,
        text: &str,
        safe: Option<bool>,
        parallel: Option<bool>,
    ) -> Vec<String>;
}
```

It provides 2 methods that look do the same thing. And in `Newmm`, it do the same things for both methods (Ref [here](https://github.com/PyThaiNLP/nlpo3/blob/main/src/tokenizer/newmm_custom.rs#L366-L385)). I purpose to clarify which one we should use and clean up trait protocol. 


Discussion / hints:

Fix patch:
diff --git a/Cargo.toml b/Cargo.toml
index a19411f..85e1c8f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,9 +11,11 @@ homepage = "https://github.com/PyThaiNLP/nlpo3/"
 documentation = "https://github.com/PyThaiNLP/nlpo3/"
 repository = "https://github.com/PyThaiNLP/nlpo3/"
 readme = "README.md"
-exclude = [".gitignore", ".github/*", "build_tools/*", "tests/*"]
-
+exclude = [".gitignore", ".github/*", "build_tools/*", "tests/*", "nlpo3-cli/*", "nlpo3-nodejs/*", "nlpo3-python/*"]
 
+[profile.release]
+lto = true
+codegen-units = 1
 
 [lib]
 path = "src/lib.rs"
@@ -23,16 +25,16 @@ path = "src/lib.rs"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-ahash = "0.7.2"
+anyhow = "1.0.42"
 binary-heap-plus = "0.4.1"
 bytecount = "0.6.2"
 lazy_static = "1.4.0"
 rayon = "1.5"
 regex = "1.4.6"
-smol_str = "0.1.17"
+rustc-hash = "1.1.0"
 
 
 [[test]]
 name = "basic"
 path = "tests/test_tokenizer.rs"
-test = true
\ No newline at end of file
+test = true
diff --git a/src/fixed_bytes_str/four_bytes.rs b/src/fixed_bytes_str/four_bytes.rs
index 9537a6a..f492bce 100644
--- a/src/fixed_bytes_str/four_bytes.rs
+++ b/src/fixed_bytes_str/four_bytes.rs
@@ -262,7 +262,7 @@ impl CustomString {
         let mut output_content: Vec<u8> = Vec::with_capacity(input.len() / 100);
         for index in 0..input.chars_len() {
             let extracted_bytes =
-                trim_to_std_utf8(&input.slice_by_char_indice(index, index + 1)).unwrap();
+                trim_to_std_utf8(input.slice_by_char_indice(index, index + 1)).unwrap();
             match extracted_bytes {
                 (None, None, None, Some(first_byte)) => {
                     output_content.push(first_byte);
@@ -293,7 +293,7 @@ impl CustomString {
         let mut output_content: Vec<u8> = Vec::with_capacity(input.len() / 100);
         for index in 0..input.chars_len() {
             let extracted_bytes =
-                trim_to_std_utf8(&input.slice_by_char_indice(index, index + 1)).unwrap();
+                trim_to_std_utf8(input.slice_by_char_indice(index, index + 1)).unwrap();
             match extracted_bytes {
                 (None, None, None, Some(first_byte)) => {
                     output_content.push(first_byte);
@@ -329,7 +329,7 @@ pub trait FixedCharsLengthByteSlice {
 
 impl FixedCharsLengthByteSlice for &CustomStringBytesSlice {
     fn slice_by_char_indice(&self, start: usize, end: usize) -> Self {
-        unsafe { &self.get_unchecked((start * BYTES_PER_CHAR)..(end * BYTES_PER_CHAR)) }
+        unsafe { self.get_unchecked((start * BYTES_PER_CHAR)..(end * BYTES_PER_CHAR)) }
     }
     fn chars_len(&self) -> usize {
         self.len() / BYTES_PER_CHAR
@@ -368,7 +368,7 @@ fn check_slice() {
     let ex: &[u8] = &[255, 255, 255, 255, 0, 255, 111, 0];
     assert_eq!(ex.slice_by_char_indice(0, 1), &[255, 255, 255, 255]);
     assert_eq!(ex.slice_by_char_indice(1, 2), &[0, 255, 111, 0]);
-    assert_eq!("".is_empty(), true);
+    assert!("".is_empty());
 }
 
 #[test]
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
index 71f8b82..8b7edea 100644
--- a/src/tokenizer.rs
+++ b/src/tokenizer.rs
@@ -1,5 +1,5 @@
-pub mod dict_reader_custom;
+mod dict_reader_custom;
 pub mod newmm_custom;
-pub mod tcc_custom;
+mod tcc_custom;
 pub mod tokenizer_trait;
-pub mod trie_custom;
+mod trie_custom;
diff --git a/src/tokenizer/newmm_custom.rs b/src/tokenizer/newmm_custom.rs
index 01bfa82..11f0b20 100644
--- a/src/tokenizer/newmm_custom.rs
+++ b/src/tokenizer/newmm_custom.rs
@@ -1,3 +1,10 @@
+use super::super::fixed_bytes_str::four_bytes::{CustomStringBytesSlice, CustomStringBytesVec};
+use super::{
+    dict_reader_custom::{create_default_dict, create_dict_trie, DictSource},
+    tcc_custom,
+    tokenizer_trait::Tokenizer,
+    trie_custom::Trie,
+};
 /**
 Dictionary-based maximal matching word segmentation, constrained with
 Thai Character Cluster (TCC) boundaries.
@@ -11,23 +18,17 @@ with heuristic graph size limit added to avoid exponential wait time.
 
 Rust implementation: ["Thanathip Suntorntip"]
 */
-// TODO: use slice_by_chars_indice on &[u8]
 use crate::fixed_bytes_str::four_bytes::{
     rfind_space_char_index, CustomString, FixedCharsLengthByteSlice, BYTES_PER_CHAR,
 };
-
-use super::super::fixed_bytes_str::four_bytes::{CustomStringBytesSlice, CustomStringBytesVec};
-use super::{
-    dict_reader_custom::{create_default_dict, create_dict_trie, DictSource},
-    tcc_custom,
-    tokenizer_trait::Tokenizer,
-    trie_custom::{Trie},
-};
-use ahash::{AHashMap as HashMap, AHashSet as HashSet};
+use anyhow::Result as AnyResult;
 use binary_heap_plus::{BinaryHeap, MinComparator};
 use lazy_static::lazy_static;
 use rayon::prelude::*;
 use regex::bytes::Regex;
+use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
+use std::error::Error;
+use std::fmt::Display;
 use std::{collections::VecDeque, path::PathBuf};
 const MAX_GRAPH_SIZE: usize = 50;
 const USE_MULTITHREAD_THRESHOLD: usize = 10000;
@@ -66,6 +67,35 @@ lazy_static! {
 pub struct Newmm {
     dict: Box<Trie>,
 }
+#[derive(Clone, Debug)]
+struct BFSSearchError {
+    graph: HashMap<CharacterIndex, Vec<CharacterIndex>>,
+    start: CharacterIndex,
+    goal: CharacterIndex,
+}
+impl BFSSearchError {
+    pub fn new(
+        graph: &HashMap<CharacterIndex, Vec<CharacterIndex>>,
+        start: CharacterIndex,
+        goal: CharacterIndex,
+    ) -> Self {
+        Self {
+            graph: graph.clone(),
+            start,
+            goal,
+        }
+    }
+}
+impl Display for BFSSearchError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "Cannot find goal position {} with start position {} with graph {:?}",
+            self.goal, self.start, self.goal
+        )
+    }
+}
+impl Error for BFSSearchError {}
 
 impl Newmm {
     pub fn new(dict_path: Option<&str>) -> Self {
@@ -86,7 +116,7 @@ impl Newmm {
         start: CharacterIndex,
         goal: CharacterIndex,
         current_queue: &mut VecDeque<(usize, Vec<usize>)>,
-    ) -> Vec<CharacterIndex> {
+    ) -> AnyResult<Vec<CharacterIndex>> {
         current_queue.clear();
 
         // let mut current_queue: VecDeque<(usize, Vec<usize>)> = VecDeque::with_capacity(graph.len());
@@ -94,8 +124,8 @@ impl Newmm {
         let mut init_path: Vec<usize> = Vec::with_capacity(goal - start);
         init_path.push(start);
         current_queue.push_back((start, init_path));
-        while !current_queue.is_empty() {
-            let (vertex, path) = current_queue.pop_front().unwrap();
+
+        while let Some((vertex, path)) = current_queue.pop_front() {
             if let Some(idk) = graph.get(&vertex) {
                 for position in idk {
                     if *position != goal {
@@ -106,31 +136,33 @@ impl Newmm {
                         let mut appended_path = path;
                         appended_path.push(*position);
 
-                        return appended_path;
+                        return Ok(appended_path);
                     };
                 }
             };
         }
-        panic!("something wrong");
+
+        Err(BFSSearchError::new(graph, start, goal).into())
     }
 
     fn one_cut<'a, 'b>(
         input: &'a CustomStringBytesSlice,
         custom_dict: &'b Trie,
-    ) -> Vec<&'a CustomStringBytesSlice> {
+    ) -> AnyResult<Vec<&'a CustomStringBytesSlice>> {
         let text = input;
         let input_char_len = text.chars_len();
         let mut reused_queue: VecDeque<(usize, Vec<usize>)> = VecDeque::with_capacity(10);
         let mut graph_size: usize = 0;
-        let mut graph: HashMap<CharacterIndex, Vec<CharacterIndex>> =
-            HashMap::with_capacity(input_char_len / 100);
-        let mut result_str: Vec<&CustomStringBytesSlice> = Vec::with_capacity(input_char_len / 100);
+        let mut graph: HashMap<CharacterIndex, Vec<CharacterIndex>> = HashMap::default();
+        graph.reserve(input_char_len / 10);
+        let mut result_str: Vec<&CustomStringBytesSlice> = Vec::with_capacity(input_char_len / 10);
 
         // all position should be refered as character index
         let valid_position = tcc_custom::tcc_pos(input);
         let text_length = input_char_len;
         let mut position_list: BinaryHeap<CharacterIndex, MinComparator> = BinaryHeap::new_min();
-        let mut existing_candidate: HashSet<CharacterIndex> = HashSet::with_capacity(50);
+        let mut existing_candidate: HashSet<CharacterIndex> = HashSet::default();
+        existing_candidate.reserve(input_char_len / 10);
         position_list.push(0);
         existing_candidate.insert(0);
         let mut end_position: CharacterIndex = 0;
@@ -143,7 +175,7 @@ impl Newmm {
         } {
             if let Some(begin_position) = position_list.pop() {
                 let sub_text_prefix = text.slice_by_char_indice(begin_position, text.chars_len());
-                let prefixes = Trie::prefix_ref(sub_text_prefix, &custom_dict);
+                let prefixes = Trie::prefix_ref(sub_text_prefix, custom_dict);
                 for word in prefixes {
                     let word_length = word.chars_len();
                     let end_position_candidate = begin_position + word_length;
@@ -177,7 +209,7 @@ impl Newmm {
                             end_position,
                             *first_position_list,
                             &mut reused_queue,
-                        );
+                        )?;
                         graph_size = 0; // reset our graph
 
                         for position in group_of_end_position_candidate.iter().skip(1) {
@@ -211,12 +243,12 @@ impl Newmm {
                                 if valid_position.contains(&position) {
                                     let prefix = &text.slice_by_char_indice(position, text_length);
 
-                                    let list_of_prefixes = Trie::prefix_ref(prefix, &custom_dict);
+                                    let list_of_prefixes = Trie::prefix_ref(prefix, custom_dict);
                                     let valid_word_filter = |word: &&[u8]| {
                                         let new_position = position + word.chars_len();
                                         let is_valid = valid_position.contains(&new_position);
                                         let is_two_thai_chars =
-                                            THAI_TWOCHARS_PATTERN.is_match(&word);
+                                            THAI_TWOCHARS_PATTERN.is_match(word);
                                         is_valid && !is_two_thai_chars
                                     };
                                     let valid_words: Vec<&[u8]> =
@@ -237,7 +269,7 @@ impl Newmm {
                                         finish_without_break = false;
                                         break;
                                     };
-                                    if NON_THAI_PATTERN.is_match(&prefix) {
+                                    if NON_THAI_PATTERN.is_match(prefix) {
                                         end_position = position;
                                         finish_without_break = false;
                                         break;
@@ -249,32 +281,28 @@ impl Newmm {
                             }
                         }
                     }
-                    let current_graph_opt = graph.get_mut(&begin_position);
-                    match current_graph_opt {
-                        Some(existing_path) => {
-                            existing_path.push(end_position);
-                            graph_size += 1;
-                            let token = text.slice_by_char_indice(begin_position, end_position);
+                    if let Some(existing_path) = graph.get_mut(&begin_position) {
+                        existing_path.push(end_position);
+                        graph_size += 1;
+                        let token = text.slice_by_char_indice(begin_position, end_position);
 
-                            result_str.push(token);
-                            position_list.push(end_position);
-                            existing_candidate.insert(end_position);
-                        }
-                        None => {
-                            let mut graph_elem: Vec<usize> = Vec::with_capacity(10);
-                            graph_elem.push(end_position);
-                            graph.insert(begin_position, graph_elem);
-                            graph_size += 1;
-                            let token = text.slice_by_char_indice(begin_position, end_position);
-                            result_str.push(token);
-                            position_list.push(end_position);
-                            existing_candidate.insert(end_position);
-                        }
+                        result_str.push(token);
+                        position_list.push(end_position);
+                        existing_candidate.insert(end_position);
+                    } else {
+                        let mut graph_elem: Vec<usize> = Vec::with_capacity(10);
+                        graph_elem.push(end_position);
+                        graph.insert(begin_position, graph_elem);
+                        graph_size += 1;
+                        let token = text.slice_by_char_indice(begin_position, end_position);
+                        result_str.push(token);
+                        position_list.push(end_position);
+                        existing_candidate.insert(end_position);
                     }
                 }
             }
         }
-        result_str
+        Ok(result_str)
     }
 
     pub fn internal_segment(
@@ -282,27 +310,27 @@ impl Newmm {
         custom_dict: &Trie,
         safe: bool,
         parallel: bool,
-    ) -> Vec<String> {
-        if input.len() == 0 {
-            return vec![];
+    ) -> AnyResult<Vec<String>> {
+        if input.is_empty() {
+            return Ok(vec![]);
         }
         if !safe || input.chars_len() < TEXT_SCAN_END {
-            let result = Self::one_cut(input.raw_content(), custom_dict);
-            if parallel {
+            let result = Self::one_cut(input.raw_content(), custom_dict)?;
+            Ok(if parallel {
                 result
                     .into_par_iter()
                     .map(|custom_string_bytes| {
-                        CustomString::convert_raw_bytes_to_std_string(&custom_string_bytes)
+                        CustomString::convert_raw_bytes_to_std_string(custom_string_bytes)
                     })
                     .collect()
             } else {
                 result
                     .into_iter()
                     .map(|custom_string_bytes| {
-                        CustomString::convert_raw_bytes_to_std_string(&custom_string_bytes)
+                        CustomString::convert_raw_bytes_to_std_string(custom_string_bytes)
                     })
                     .collect()
-            }
+            })
         } else {
             let mut txt = input.raw_content();
             let mut txt_parts: Vec<CustomStringBytesVec> = Vec::with_capacity(txt.len() / 10);
@@ -316,7 +344,7 @@ impl Newmm {
                 if let Some(space_char_index) = space_char_index {
                     cut_pos = space_char_index + 1;
                 } else {
-                    let word_tokens = Self::one_cut(sample, &custom_dict);
+                    let word_tokens = Self::one_cut(sample, custom_dict)?;
                     let mut token_max_index = 0;
                     let mut token_max_length = 0;
                     for (idx, token) in word_tokens.iter().enumerate() {
@@ -338,33 +366,41 @@ impl Newmm {
                 txt_parts.push(txt.to_owned());
             }
 
-            if parallel {
+            Ok(if parallel {
                 txt_parts
-                    .into_par_iter()
-                    .flat_map(|part| {
-                        Self::one_cut(&part, &custom_dict)
+                    .par_iter()
+                    .flat_map(|part| -> AnyResult<_> {
+                        let words = Self::one_cut(part, custom_dict)?;
+                        Ok(words
                             .into_par_iter()
-                            .map(|word| CustomString::convert_raw_bytes_to_std_string(&word))
-                            .collect::<Vec<String>>()
+                            .map(|word| CustomString::convert_raw_bytes_to_std_string(word))
+                            .collect::<Vec<String>>())
                     })
-                    .collect::<Vec<String>>()
+                    .flatten()
+                    .collect()
             } else {
                 txt_parts
                     .iter()
-                    .flat_map(|part| {
-                        Self::one_cut(&part, &custom_dict)
+                    .flat_map(|part| -> AnyResult<_> {
+                        Ok(Self::one_cut(part, custom_dict)?
                             .iter()
-                            .map(|word| CustomString::convert_raw_bytes_to_std_string(&word))
-                            .collect::<Vec<String>>()
+                            .map(|word| CustomString::convert_raw_bytes_to_std_string(word))
+                            .collect::<Vec<String>>())
                     })
-                    .collect::<Vec<String>>()
-            }
+                    .flatten()
+                    .collect()
+            })
         }
     }
 }
 
 impl Tokenizer for Newmm {
-    fn segment(&self, text: &str, safe: Option<bool>, parallel: Option<bool>) -> Vec<String> {
+    fn segment(
+        &self,
+        text: &str,
+        safe: Option<bool>,
+        parallel: Option<bool>,
+    ) -> AnyResult<Vec<String>> {
         let safe_flag = safe.unwrap_or(false);
         let parallel_flag = match parallel {
             Some(val) => val,
@@ -380,6 +416,6 @@ impl Tokenizer for Newmm {
         safe: Option<bool>,
         parallel: Option<bool>,
     ) -> Vec<String> {
-        self.segment(text, safe, parallel)
+        self.segment(text, safe, parallel).unwrap()
     }
 }
diff --git a/src/tokenizer/tcc_custom.rs b/src/tokenizer/tcc_custom.rs
index fd3a6f7..ad51ba3 100644
--- a/src/tokenizer/tcc_custom.rs
+++ b/src/tokenizer/tcc_custom.rs
@@ -10,7 +10,7 @@ Credits:
     * Python code: Korakot Chaovavanich
     * Rust Code Translation: Thanathip Suntorntip
 */
-use ahash::AHashSet as HashSet;
+use rustc_hash::FxHashSet as HashSet;
 use lazy_static::lazy_static;
 use regex::bytes::Regex;
 
@@ -65,28 +65,48 @@ lazy_static! {
     )
     .unwrap();
 }
-
-pub fn tcc_pos(custom_text_type: &CustomStringBytesSlice) -> HashSet<usize> {
-    let mut set: HashSet<usize> = HashSet::with_capacity(custom_text_type.chars_len() / 10);
-    if custom_text_type.is_empty() {
-        set
-    } else {
-        let mut position: usize = 0;
-        let four_bytes_chars_segment = segment(custom_text_type);
-        for segment in four_bytes_chars_segment.into_iter() {
-            let segment_size = segment.chars_len();
+pub fn tcc_pos(custom_text_type: &CustomStringBytesSlice) -> HashSet<usize>{
+    let mut set: HashSet<usize> = HashSet::default();
+    set.reserve(custom_text_type.chars_len() / 10);
+    let mut txt = custom_text_type;
+    let mut position: usize = 0;
+    while !txt.is_empty() {
+        if let Some(result) = NON_LOOKAHEAD_TCC.find(txt) {
+            let mut matched = &txt[result.start()..result.end()];
+            let match_length = matched.len();
+            if LOOKAHEAD_TCC.is_match(matched) {
+                // trim one more char to the right.
+                let end_bytes_index = match_length - (1 * BYTES_PER_CHAR);
+                matched = &matched[0..end_bytes_index];
+                let segment_size = matched.chars_len();
+                position += segment_size;
+                set.insert(position);
+                txt = &txt[end_bytes_index..];
+            } else {
+                let segment_size = matched.chars_len();
+                position += segment_size;
+                set.insert(position);
+                let end_bytes_index = match_length;
+                txt = &txt[end_bytes_index..];
+            }
+        } else {
+            // not thai
+            let first_char = txt.slice_by_char_indice(0, 1);
+            let segment_size = first_char.chars_len();
             position += segment_size;
             set.insert(position);
+            txt = txt.slice_by_char_indice(1, txt.chars_len());
         }
-        set
     }
+    set
 }
 
-pub fn segment(custom_text_type: &CustomStringBytesSlice) -> Vec<&CustomStringBytesSlice> {
+#[allow(dead_code)]
+pub fn tcc_segment(custom_text_type: &CustomStringBytesSlice) -> Vec<&CustomStringBytesSlice> {
     let mut txt = custom_text_type;
     let mut tcc_result: Vec<&[u8]> = Vec::with_capacity(txt.len() / 10);
     while !txt.is_empty() {
-        if let Some(result) = NON_LOOKAHEAD_TCC.find(&txt) {
+        if let Some(result) = NON_LOOKAHEAD_TCC.find(txt) {
             let mut matched = &txt[result.start()..result.end()];
             let match_length = matched.len();
             if LOOKAHEAD_TCC.is_match(matched) {
diff --git a/src/tokenizer/tokenizer_trait.rs b/src/tokenizer/tokenizer_trait.rs
index 17dfa6f..a652590 100644
--- a/src/tokenizer/tokenizer_trait.rs
+++ b/src/tokenizer/tokenizer_trait.rs
@@ -1,8 +1,11 @@
-/**
-   This should be the only part exposed to lib.rs
-*/
+use anyhow::Result as AnyResult;
 pub trait Tokenizer {
-    fn segment(&self, text: &str, safe: Option<bool>, parallel: Option<bool>) -> Vec<String>;
+    fn segment(
+        &self,
+        text: &str,
+        safe: Option<bool>,
+        parallel: Option<bool>,
+    ) -> AnyResult<Vec<String>>;
 
     fn segment_to_string(
         &self,
diff --git a/src/tokenizer/trie_custom.rs b/src/tokenizer/trie_custom.rs
index f359846..49f5b2f 100644
--- a/src/tokenizer/trie_custom.rs
+++ b/src/tokenizer/trie_custom.rs
@@ -2,9 +2,8 @@ use crate::fixed_bytes_str::four_bytes::{
     CustomString, CustomStringBytesSlice, CustomStringBytesVec, FixedCharsLengthByteSlice,
     BYTES_PER_CHAR,
 };
-use ahash::{AHashMap as HashMap, AHashSet as HashSet};
+use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
 use std::borrow::BorrowMut;
-use std::iter::Iterator;
 /**
 This module is meant to be a direct implementation of Dict Trie in PythaiNLP.
 
@@ -31,7 +30,7 @@ impl TrieNode {
     pub fn new() -> Self {
         Self {
             // text: None,
-            children: HashMap::with_capacity(100),
+            children: HashMap::default(),
             end: false,
         }
     }
@@ -76,7 +75,7 @@ impl TrieNode {
                 if char_count == 1 {
                     child.set_not_end();
                 }
-                child.remove_word_from_node(&word);
+                child.remove_word_from_node(word);
 
                 if !child.end && child.children.is_empty() {
                     self.remove_child(character);
@@ -85,8 +84,8 @@ impl TrieNode {
         }
     }
 
-    pub fn list_prefix(&self, prefix: &CustomStringBytesSlice) -> Vec<CustomStringBytesVec> {
-        let mut result: Vec<CustomStringBytesVec> = Vec::with_capacity(100);
+    pub fn list_prefix<'d,'p>(&'d self, prefix: &'p CustomStringBytesSlice) -> Vec<&'p CustomStringBytesSlice> {
+        let mut result: Vec<&CustomStringBytesSlice> = Vec::with_capacity(100);
         let prefix_cpy = prefix;
         let mut current_index = 0;
         let mut current_node_wrap = Some(self);
@@ -97,7 +96,7 @@ impl TrieNode {
                     if child.end {
                         let substring_of_prefix =
                             prefix_cpy.slice_by_char_indice(0, current_index + 1);
-                        result.push(substring_of_prefix.to_vec());
+                        result.push(substring_of_prefix);
                     }
                     current_node_wrap = Some(child);
                 } else {
@@ -118,12 +117,12 @@ pub struct Trie {
 impl Trie {
     pub fn new(words: &[CustomString]) -> Self {
         let mut instance = Self {
-            words: HashSet::with_capacity(words.len() / 10),
+            words: HashSet::default(),
             root: TrieNode::new(),
         };
         for word in words.iter() {
             if !word.is_empty() {
-                instance.add(&word);
+                instance.add(word);
             }
         }
         instance
@@ -137,7 +136,7 @@ impl Trie {
         let stripped_word = word.trim();
         self.words.insert(stripped_word.raw_content().into());
         let current_cursor = self.root.borrow_mut();
-        current_cursor.add_word(&stripped_word.raw_content());
+        current_cursor.add_word(stripped_word.raw_content());
     }
 
     pub fn remove(&mut self, word: &CustomString) {
@@ -149,7 +148,7 @@ impl Trie {
         }
     }
 
-    pub fn prefix(&self, prefix: &CustomStringBytesSlice) -> Vec<Vec<u8>> {
+    pub fn prefix<'d,'p>(&'d self, prefix: &'p CustomStringBytesSlice) -> Vec<&'p CustomStringBytesSlice> {
         self.root.list_prefix(prefix)
     }
 
@@ -160,6 +159,7 @@ impl Trie {
 
        The only downside I can think about this function is its non-descriptive name..
     */
+
     pub fn prefix_ref<'p, 't>(
         prefix: &'p CustomStringBytesSlice,
         dict_trie: &'t Self,
@@ -196,6 +196,6 @@ impl Trie {
     }
 
     pub fn amount_of_words(&self) -> usize {
-        self.words.iter().count()
+        self.words.len()
     }
 }


Test patch:
diff --git a/tests/test_tokenizer.rs b/tests/test_tokenizer.rs
index 4a419d4..006d606 100644
--- a/tests/test_tokenizer.rs
+++ b/tests/test_tokenizer.rs
@@ -150,9 +150,9 @@ fn test_long_text_byte_tokenizer() {
     .join("");
 
     let newmm_default_dict = NewmmCustom::new(None);
-    let result = newmm_default_dict.segment(&long_text, None, Some(true));
+    let result = newmm_default_dict.segment(&long_text, None, Some(true)).unwrap();
 
-    let safe_result = newmm_default_dict.segment(&long_text, Some(true), Some(true));
+    let safe_result = newmm_default_dict.segment(&long_text, Some(true), Some(true)).unwrap();
     assert_eq!(result.len(), 1889);
     assert_eq!(safe_result.len(), 1991);
 }