fix(newmm): remove quadratic time in graph reset and BFS path copy

CHANGELOG.md

@@ -41,6 +41,10 @@ Major release with breaking changes.
 ### Fixed
 
 - Fixed major slowdown risks on highly ambiguous tokenization paths.
+- Fixed superlinear (quadratic) tokenization time in `newmm` on long text:
+  the ambiguity graph is no longer cleared with `HashMap::clear()` (which costs
+  O(capacity) per resolution), and BFS rebuilds its path from a predecessor map
+  instead of copying the path at every step.
 
 ### Migration
 

src/tokenizer/newmm.rs

@@ -306,32 +306,35 @@ impl<D: DictBackend> NewmmTokenizer<D> {
         graph: &HashMap<CharacterIndex, Vec<CharacterIndex>>,
         start: CharacterIndex,
         goal: CharacterIndex,
-        current_queue: &mut VecDeque<(usize, Vec<usize>)>,
+        current_queue: &mut VecDeque<CharacterIndex>,
     ) -> AnyResult<Vec<CharacterIndex>> {
         current_queue.clear();
 
-        // The visited set ensures each vertex is enqueued at most once.
-        // This avoids the exponential blow-up from revisiting the same nodes.
-        let mut visited: HashSet<CharacterIndex> =
-            HashSet::with_capacity_and_hasher(goal - start, Default::default());
-        visited.insert(start);
+        // Each vertex is enqueued at most once (this also avoids the
+        // exponential blow-up from revisiting nodes). Instead of copying the
+        // whole path at every step, record one predecessor per vertex and
+        // rebuild the path once, when the goal is reached.
+        let mut predecessor: HashMap<CharacterIndex, CharacterIndex> = HashMap::default();
+        predecessor.insert(start, start);
+        current_queue.push_back(start);
 
-        let mut init_path: Vec<usize> = Vec::with_capacity(goal - start);
-        init_path.push(start);
-        current_queue.push_back((start, init_path));
-
-        while let Some((vertex, path)) = current_queue.pop_front() {
+        while let Some(vertex) = current_queue.pop_front() {
             if let Some(idk) = graph.get(&vertex) {
                 for &position in idk {
                     if position == goal {
-                        let mut appended_path = path;
-                        appended_path.push(position);
-                        return Ok(appended_path);
-                    } else if !visited.contains(&position) {
-                        visited.insert(position);
-                        let mut appended_path = path.clone();
-                        appended_path.push(position);
-                        current_queue.push_back((position, appended_path));
+                        let mut path = vec![goal, vertex];
+                        let mut current = vertex;
+                        while current != start {
+                            current = predecessor[&current];
+                            path.push(current);
+                        }
+                        path.reverse();
+                        return Ok(path);
+                    } else if let std::collections::hash_map::Entry::Vacant(entry) =
+                        predecessor.entry(position)
+                    {
+                        entry.insert(vertex);
+                        current_queue.push_back(position);
                     }
                 }
             };
@@ -343,10 +346,9 @@ impl<D: DictBackend> NewmmTokenizer<D> {
     fn one_cut<'a>(input: &'a CharString, custom_dict: &D) -> AnyResult<Vec<&'a str>> {
         let text = input;
         let input_char_len = text.chars_len();
-        let mut reused_queue: VecDeque<(usize, Vec<usize>)> = VecDeque::with_capacity(10);
+        let mut reused_queue: VecDeque<CharacterIndex> = VecDeque::with_capacity(10);
         let mut graph_size: usize = 0;
         let mut graph: HashMap<CharacterIndex, Vec<CharacterIndex>> = HashMap::default();
-        graph.reserve(input_char_len / 10);
         let mut result_str: Vec<&str> = Vec::with_capacity(input_char_len / 10);
 
         // All positions are character indices (not byte offsets).
@@ -395,7 +397,9 @@ impl<D: DictBackend> NewmmTokenizer<D> {
                         &mut reused_queue,
                     )?;
                     graph_size = 0;
-                    graph.clear();
+                    // Not `clear()`: it costs O(capacity), which grows with the
+                    // largest graph seen, and this runs once per resolution.
+                    graph = HashMap::default();
 
                     for &position in group_of_end_position_candidate.iter().skip(1) {
                         let token = text.substring_as_str(end_position, position);
@@ -465,7 +469,7 @@ impl<D: DictBackend> NewmmTokenizer<D> {
                 }
 
                 graph_size = 0;
-                graph.clear();
+                graph = HashMap::default();
                 let token = text.substring_as_str(begin_position, end_position);
                 result_str.push(token);
                 position_list.push(end_position);

tests/test_tokenizer.rs

@@ -391,3 +391,53 @@ fn test_newmm_ambiguous_performance() {
         "tokenization took {elapsed:?}; BFS path explosion may still be present"
     );
 }
+
+/// Best of three wall-clock timings of `segment` on `text`.
+fn best_segment_time(tokenizer: &NewmmTokenizer, text: &str) -> std::time::Duration {
+    (0..3)
+        .map(|_| {
+            let start = std::time::Instant::now();
+            let result = tokenizer.segment_with_options(text, false, None).unwrap();
+            let elapsed = start.elapsed();
+            assert_eq!(result.concat(), text);
+            elapsed
+        })
+        .min()
+        .unwrap()
+}
+
+/// Time must grow roughly linearly with text length, so a 4x longer text may
+/// take about 4x as long; the 10x limit catches quadratic behavior (16x).
+#[test]
+fn test_newmm_persistent_ambiguity_scaling() {
+    let mut path = env!("CARGO_MANIFEST_DIR").to_string();
+    path.push_str(DEFAULT_DICT_PATH);
+    let tokenizer = NewmmTokenizer::new(&path).unwrap();
+
+    // "กรรมกร" contains several overlapping dictionary words, so the
+    // ambiguity graph keeps growing across the whole text.
+    let small = best_segment_time(&tokenizer, &"กรรมกร".repeat(3_000));
+    let large = best_segment_time(&tokenizer, &"กรรมกร".repeat(12_000));
+    let limit = small.max(std::time::Duration::from_millis(10)) * 10;
+    assert!(
+        large < limit,
+        "4x longer text took {large:?} vs {small:?}; expected less than {limit:?}"
+    );
+}
+
+/// Ordinary long text must also scale linearly (10x longer text should take
+/// about 10x as long; the 25x limit catches quadratic behavior).
+#[test]
+fn test_newmm_long_text_scaling() {
+    let mut path = env!("CARGO_MANIFEST_DIR").to_string();
+    path.push_str(DEFAULT_DICT_PATH);
+    let tokenizer = NewmmTokenizer::new(&path).unwrap();
+
+    let small = best_segment_time(&tokenizer, &SECOND_TEXT.repeat(1_000));
+    let large = best_segment_time(&tokenizer, &SECOND_TEXT.repeat(10_000));
+    let limit = small.max(std::time::Duration::from_millis(10)) * 25;
+    assert!(
+        large < limit,
+        "10x longer text took {large:?} vs {small:?}; expected less than {limit:?}"
+    );
+}