150 lines
4.7 KiB
Rust
150 lines
4.7 KiB
Rust
use crate::*;
|
|
use pyo3::prelude::*;
|
|
use rayon::prelude::*;
|
|
|
|
// Node in our character trie
|
|
#[derive(Default)]
|
|
struct CharTrieNode {
|
|
// Store character and node pairs, keep sorted for binary search
|
|
children: Vec<(char, CharTrieNode)>,
|
|
// Store token IDs that are valid at this point
|
|
token_id: Option<Py<PyAny>>,
|
|
}
|
|
|
|
impl CharTrieNode {
|
|
// 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::default()));
|
|
&mut self.children[index].1
|
|
}
|
|
}
|
|
}
|
|
|
|
// Add a token id to this node
|
|
fn add_token(&mut self, token_id: Py<PyAny>) {
|
|
self.token_id = Some(token_id);
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
// 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 {
|
|
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
|
|
.iter()
|
|
.map(|(ch, node)| (*ch, node.clone_ref(py)))
|
|
.collect(),
|
|
token_id: self
|
|
.token_id
|
|
.iter()
|
|
.map(|token_id| token_id.clone_ref(py))
|
|
.nth(0),
|
|
}
|
|
}
|
|
}
|
|
|
|
// The main trie structure
|
|
#[derive(Default)]
|
|
pub struct CharTrie {
|
|
root: CharTrieNode,
|
|
}
|
|
|
|
impl CharTrie {
|
|
// 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;
|
|
|
|
// Add each character in the token text
|
|
for c in token_text.chars() {
|
|
current = current.add_child(c);
|
|
}
|
|
|
|
// Mark this node as representing token_id
|
|
current.add_token(token_id);
|
|
}
|
|
|
|
// Find all tokens that could be invalid given the current XML state
|
|
pub fn find_invalid_tokens(
|
|
&self,
|
|
py: Python<'_>,
|
|
xml_validator: &XmlSchemaValidator,
|
|
) -> Vec<Py<PyAny>> {
|
|
// 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 {
|
|
Self {
|
|
root: self.root.clone_ref(py),
|
|
}
|
|
}
|
|
} |