Use CharTrie and fix multi-byte char issue

lib/xml_schema_validator/src/char_trie.rs

@@ -2,6 +2,7 @@ use crate::*;
 use pyo3::prelude::*;
 
 // Node in our character trie
+#[derive(Default)]
 struct CharTrieNode {
     // Store character and node pairs, keep sorted for binary search
     children: Vec<(char, CharTrieNode)>,
@@ -10,20 +11,13 @@ struct CharTrieNode {
 }
 
 impl CharTrieNode {
-    fn new() -> Self {
-        Self {
-            children: Vec::new(),
-            token_id: None,
-        }
-    }
-
     // Add a child for character c, maintaining sorted order
     fn add_child(&mut self, c: char) -> &mut CharTrieNode {
         match self.children.binary_search_by_key(&c, |(ch, _)| *ch) {
             Ok(index) => &mut self.children[index].1,
             Err(index) => {
                 // Insert at the right position to maintain sort order
-                self.children.insert(index, (c, CharTrieNode::new()));
+                self.children.insert(index, (c, CharTrieNode::default()));
                 &mut self.children[index].1
             }
         }
@@ -34,7 +28,38 @@ impl CharTrieNode {
         self.token_id = Some(token_id);
     }
 
-    pub fn clone_ref(&self, py: Python<'_>) -> Self {
+    fn find_invalid_tokens(
+        &self,
+        py: Python<'_>,
+        xml_validator: &XmlSchemaValidator,
+    ) -> Vec<Py<PyAny>> {
+        let res = self.children.iter().map(|(c, n)|{
+            let mut node_validator = xml_validator.clone();
+            if node_validator.append_char(*c).is_ok() {
+                // If the node is valid, recursively check its children
+                n.find_invalid_tokens(py, &node_validator)
+            } else {
+                // If the node is invalid, add it and its children to the list of invalid tokens
+                n.all_tokens(py)
+            }
+        })
+            .flatten()
+            .collect();
+        res
+    }
+
+    fn all_tokens(&self, py: Python<'_>) -> Vec<Py<PyAny>> {
+        let mut tokens = Vec::new();
+        if let Some(token_id) = &self.token_id {
+            tokens.push(token_id.clone_ref(py));
+        }
+        for (_, child) in &self.children {
+            tokens.extend(child.all_tokens(py));
+        }
+        tokens
+    }
+
+    fn clone_ref(&self, py: Python<'_>) -> Self {
         Self {
             children: self
                 .children
@@ -51,17 +76,12 @@ impl CharTrieNode {
 }
 
 // The main trie structure
+#[derive(Default)]
 pub struct CharTrie {
     root: CharTrieNode,
 }
 
 impl CharTrie {
-    pub fn new() -> Self {
-        Self {
-            root: CharTrieNode::new(),
-        }
-    }
-
     // Insert a token's character sequence into the trie
     pub fn insert(&mut self, token_text: &str, token_id: Py<PyAny>) {
         let mut current = &mut self.root;
@@ -80,32 +100,8 @@ impl CharTrie {
         &self,
         py: Python<'_>,
         xml_validator: &XmlSchemaValidator,
-        token_map: &Vec<(Py<PyAny>, String)>,
     ) -> Vec<Py<PyAny>> {
-        let mut invalid_tokens = Vec::new();
-
-        for (token_id, token_text) in token_map {
-            // Check if this token would make the XML invalid
-            if !self.is_valid_continuation(xml_validator, token_text) {
-                invalid_tokens.push(token_id.clone_ref(py));
-            }
-        }
-
-        invalid_tokens
-    }
-
-    // Check if a token would be a valid continuation
-    fn is_valid_continuation(&self, xml_validator: &XmlSchemaValidator, token_text: &str) -> bool {
-        // First, quick check with just the first character
-        if let Some(first_char) = token_text.chars().next() {
-            if !xml_validator.can_continue_with(first_char) {
-                return false;
-            }
-        }
-
-        // If first character is valid, try the full token
-        let mut validator_clone = xml_validator.clone();
-        validator_clone.append(token_text).is_ok()
+        self.root.find_invalid_tokens(py, xml_validator)
     }
 
     pub fn clone_ref(&self, py: Python<'_>) -> Self {
@@ -113,4 +109,4 @@ impl CharTrie {
             root: self.root.clone_ref(py),
         }
     }
-}
+}