Multi-thread chartie

lib/xml_schema_validator/Cargo.toml

@@ -12,6 +12,7 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 pyo3 = { version = "0.24", features = ["auto-initialize"] }
 quick-xml = { version = "0.28.1", features = ["serde", "serialize"] }
+rayon = "1.8"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1.0"
 thiserror = "1.0.38"

lib/xml_schema_validator/src/char_trie.rs

@@ -1,5 +1,6 @@
 use crate::*;
 use pyo3::prelude::*;
+use rayon::prelude::*;
 
 // Node in our character trie
 #[derive(Default)]
@@ -51,6 +52,39 @@ impl CharTrieNode {
         res
     }
 
+    // Parallel version of find_invalid_tokens for the first level of recursion
+    fn find_invalid_tokens_parallel(
+        &self,
+        _py: Python<'_>,
+        xml_validator: &XmlSchemaValidator,
+    ) -> Vec<Py<PyAny>> {
+        // We need to use the Python GIL to share PyAny objects across threads
+        Python::with_gil(|_py| {
+            // Convert children to a parallel iterator
+            self.children
+                .par_iter() // Use parallel iterator from rayon
+                .map(|(c, n)| {
+                    // Clone validator for this thread
+                    let mut node_validator = xml_validator.clone();
+                    
+                    // Process this character
+                    if node_validator.append_char(*c).is_ok() {
+                        // For valid nodes, use the regular (non-parallel) method for deeper levels
+                        Python::with_gil(|py| {
+                            n.find_invalid_tokens(py, &node_validator)
+                        })
+                    } else {
+                        // For invalid nodes, get all tokens
+                        Python::with_gil(|py| {
+                            n.all_tokens(py)
+                        })
+                    }
+                })
+                .flatten()
+                .collect()
+        })
+    }
+
     fn all_tokens(&self, py: Python<'_>) -> Vec<Py<PyAny>> {
         let mut tokens = Vec::new();
         if let Some(token_id) = &self.token_id {
@@ -104,7 +138,8 @@ impl CharTrie {
         py: Python<'_>,
         xml_validator: &XmlSchemaValidator,
     ) -> Vec<Py<PyAny>> {
-        self.root.find_invalid_tokens(py, xml_validator)
+        // Use the parallel version starting at the root
+        self.root.find_invalid_tokens_parallel(py, xml_validator)
     }
 
     pub fn clone_ref(&self, py: Python<'_>) -> Self {
@@ -112,4 +147,4 @@ impl CharTrie {
             root: self.root.clone_ref(py),
         }
     }
-}
+}