use pyo3::prelude::*;
use std::collections::HashMap;
#[pyclass]
pub struct RWayTrie {
root: Option<Box<Node>>,
}
struct Node {
value: Option<String>,
children: HashMap<u8, Box<Node>>,
}
impl Node {
fn new() -> Self {
Node {
value: None,
children: HashMap::new(),
}
}
}
#[pymethods]
impl RWayTrie {
#[new]
fn new() -> Self {
RWayTrie { root: None }
}
fn insert(&mut self, key: &str, value: &str) {
if self.root.is_none() {
self.root = Some(Box::new(Node::new()));
}
let mut current = self.root.as_mut().unwrap();
for &byte in key.as_bytes() {
current = current.children.entry(byte).or_insert_with(|| Box::new(Node::new()));
}
current.value = Some(value.to_string());
}
fn search(&self, key: &str) -> Option<String> {
let mut current = self.root.as_ref()?;
for &byte in key.as_bytes() {
current = current.children.get(&byte)?;
}
current.value.clone()
}
fn keys_with_prefix(&self, prefix: &str) -> PyResult<HashMap<String, String>> {
let mut results = HashMap::new();
let mut current = self.root.as_deref();
for &byte in prefix.as_bytes() {
match current {
Some(node) => current = node.children.get(&byte).map(|boxed_node| boxed_node.as_ref()),
None => return Ok(results),
}
}
if let Some(node) = current {
self.collect(node, &mut String::from(prefix), &mut results);
}
Ok(results)
}
}
impl RWayTrie {
fn collect(&self, node: &Node, prefix: &mut String, results: &mut HashMap<String, String>) {
if let Some(value) = &node.value {
results.insert(prefix.clone(), value.clone());
}
for (&byte, child) in &node.children {
prefix.push(byte as char);
self.collect(child, prefix, results);
prefix.pop();
}
}
}