use crate::{path, Path};
use std::collections::BTreeMap;
#[derive(Debug, Clone)]
pub struct PathTrie<T> {
value: Option<T>,
children: BTreeMap<String, PathTrie<T>>,
}
impl<T> Default for PathTrie<T> {
fn default() -> Self {
Self {
value: None,
children: BTreeMap::new(),
}
}
}
impl<T> PathTrie<T> {
pub fn new() -> Self {
Self::default()
}
fn get_or_create_node(&mut self, path: &Path) -> &mut PathTrie<T> {
let mut current = self;
for component in path.iter() {
current = current.children.entry(component.to_string()).or_default();
}
current
}
fn get_node(&self, path: &Path) -> Option<&PathTrie<T>> {
let mut current = self;
for component in path.iter() {
current = current.children.get(component)?;
}
Some(current)
}
fn get_node_mut(&mut self, path: &Path) -> Option<&mut PathTrie<T>> {
let mut current = self;
for component in path.iter() {
current = current.children.get_mut(component)?;
}
Some(current)
}
pub fn insert(&mut self, path: &Path, value: T) -> Option<T> {
let node = self.get_or_create_node(path);
node.value.replace(value)
}
pub fn remove(&mut self, path: &Path) -> Option<T> {
self.get_node_mut(path)?.value.take()
}
pub fn remove_subtree(&mut self, path: &Path) -> Option<PathTrie<T>> {
if path.is_empty() {
let old = std::mem::take(self);
if old.value.is_some() || !old.children.is_empty() {
Some(old)
} else {
None
}
} else {
let parent_path = path.slice(0, path.len() - 1);
let child_name = &path[path.len() - 1];
let parent = self.get_node_mut(&parent_path)?;
parent.children.remove(child_name)
}
}
pub fn get(&self, path: &Path) -> Option<&T> {
self.get_node(path)?.value.as_ref()
}
pub fn get_mut(&mut self, path: &Path) -> Option<&mut T> {
self.get_node_mut(path)?.value.as_mut()
}
pub fn get_subtrie(&self, path: &Path) -> Option<&PathTrie<T>> {
self.get_node(path)
}
pub fn get_subtrie_mut(&mut self, path: &Path) -> Option<&mut PathTrie<T>> {
self.get_node_mut(path)
}
pub fn contains_value(&self, path: &Path) -> bool {
self.get(path).is_some()
}
pub fn len(&self) -> usize {
let self_count = if self.value.is_some() { 1 } else { 0 };
let children_count: usize = self.children.values().map(|child| child.len()).sum();
self_count + children_count
}
pub fn is_empty(&self) -> bool {
self.value.is_none() && self.children.values().all(|c| c.is_empty())
}
pub fn find_ancestor(&self, path: &Path) -> Option<(&T, Path)> {
let mut current = self;
let mut last_value: Option<&T> = self.value.as_ref();
let mut last_depth: usize = 0;
for (depth, component) in path.iter().enumerate() {
match current.children.get(component) {
Some(child) => {
current = child;
if child.value.is_some() {
last_value = child.value.as_ref();
last_depth = depth + 1;
}
}
None => break,
}
}
last_value.map(|v| {
let suffix = path.slice(last_depth, path.len());
(v, suffix)
})
}
pub fn find_ancestor_mut(&mut self, path: &Path) -> Option<(&mut T, Path)> {
let depth = {
let mut current = &*self;
let mut last_depth: usize = if self.value.is_some() { 0 } else { usize::MAX };
for (d, component) in path.iter().enumerate() {
match current.children.get(component) {
Some(child) => {
current = child;
if child.value.is_some() {
last_depth = d + 1;
}
}
None => break,
}
}
if last_depth == usize::MAX {
return None;
}
last_depth
};
let target_path = path.slice(0, depth);
let suffix = path.slice(depth, path.len());
self.get_mut(&target_path).map(|v| (v, suffix))
}
pub fn iter(&self) -> PathTrieIter<'_, T> {
PathTrieIter::new(self)
}
}
pub struct PathTrieIter<'a, T> {
stack: Vec<(Path, &'a PathTrie<T>)>,
}
impl<'a, T> PathTrieIter<'a, T> {
fn new(trie: &'a PathTrie<T>) -> Self {
Self {
stack: vec![(path!(""), trie)],
}
}
}
impl<'a, T> Iterator for PathTrieIter<'a, T> {
type Item = (Path, &'a T);
fn next(&mut self) -> Option<Self::Item> {
while let Some((path, node)) = self.stack.pop() {
for (name, child) in node.children.iter().rev() {
let child_path = if path.is_empty() {
Path::from_validated_components(vec![name.clone()])
} else {
let mut c: Vec<String> = path.iter().map(str::to_string).collect();
c.push(name.clone());
Path::from_validated_components(c)
};
self.stack.push((child_path, child));
}
if let Some(ref value) = node.value {
return Some((path, value));
}
}
None
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::path;
#[test]
fn new_trie_is_empty() {
let trie: PathTrie<i32> = PathTrie::new();
assert!(trie.is_empty());
assert_eq!(trie.len(), 0);
}
#[test]
fn insert_and_get() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a/b"), 42);
assert_eq!(trie.get(&path!("a/b")), Some(&42));
assert_eq!(trie.get(&path!("a")), None);
assert_eq!(trie.get(&path!("a/b/c")), None);
}
#[test]
fn insert_returns_previous() {
let mut trie: PathTrie<i32> = PathTrie::new();
assert_eq!(trie.insert(&path!("a"), 1), None);
assert_eq!(trie.insert(&path!("a"), 2), Some(1));
assert_eq!(trie.get(&path!("a")), Some(&2));
}
#[test]
fn remove_returns_value() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a/b"), 42);
assert_eq!(trie.remove(&path!("a/b")), Some(42));
assert_eq!(trie.get(&path!("a/b")), None);
}
#[test]
fn remove_keeps_children() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
trie.insert(&path!("a/b"), 2);
trie.remove(&path!("a"));
assert_eq!(trie.get(&path!("a")), None);
assert_eq!(trie.get(&path!("a/b")), Some(&2));
}
#[test]
fn remove_subtree() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
trie.insert(&path!("a/b"), 2);
trie.insert(&path!("c"), 3);
let subtree = trie.remove_subtree(&path!("a")).unwrap();
assert_eq!(subtree.get(&path!("")), Some(&1));
assert_eq!(subtree.get(&path!("b")), Some(&2));
assert_eq!(trie.get(&path!("a")), None);
assert_eq!(trie.get(&path!("a/b")), None);
assert_eq!(trie.get(&path!("c")), Some(&3));
}
#[test]
fn remove_subtree_at_root() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
trie.insert(&path!("b"), 2);
let subtree = trie.remove_subtree(&path!("")).unwrap();
assert!(trie.is_empty());
assert_eq!(subtree.get(&path!("a")), Some(&1));
assert_eq!(subtree.get(&path!("b")), Some(&2));
}
#[test]
fn remove_subtree_nonexistent() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
assert!(trie.remove_subtree(&path!("nonexistent")).is_none());
}
#[test]
fn get_mut() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
*trie.get_mut(&path!("a")).unwrap() = 42;
assert_eq!(trie.get(&path!("a")), Some(&42));
}
#[test]
fn contains_value() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a/b"), 1);
assert!(trie.contains_value(&path!("a/b")));
assert!(!trie.contains_value(&path!("a")));
assert!(!trie.contains_value(&path!("nonexistent")));
}
#[test]
fn len_and_is_empty() {
let mut trie: PathTrie<i32> = PathTrie::new();
assert!(trie.is_empty());
assert_eq!(trie.len(), 0);
trie.insert(&path!("a"), 1);
assert!(!trie.is_empty());
assert_eq!(trie.len(), 1);
trie.insert(&path!("a/b"), 2);
assert_eq!(trie.len(), 2);
trie.insert(&path!("c"), 3);
assert_eq!(trie.len(), 3);
}
#[test]
fn find_ancestor_basic() {
let mut trie: PathTrie<&str> = PathTrie::new();
trie.insert(&path!("data"), "data_store");
let (value, suffix) = trie.find_ancestor(&path!("data/users/1")).unwrap();
assert_eq!(*value, "data_store");
assert_eq!(suffix, path!("users/1"));
}
#[test]
fn find_ancestor_deeper_wins() {
let mut trie: PathTrie<&str> = PathTrie::new();
trie.insert(&path!("data"), "data_store");
trie.insert(&path!("data/cache"), "cache_store");
let (value, suffix) = trie.find_ancestor(&path!("data/cache/hot")).unwrap();
assert_eq!(*value, "cache_store");
assert_eq!(suffix, path!("hot"));
let (value, suffix) = trie.find_ancestor(&path!("data/users/1")).unwrap();
assert_eq!(*value, "data_store");
assert_eq!(suffix, path!("users/1"));
}
#[test]
fn find_ancestor_at_root() {
let mut trie: PathTrie<&str> = PathTrie::new();
trie.insert(&path!(""), "root_store");
let (value, suffix) = trie.find_ancestor(&path!("any/path")).unwrap();
assert_eq!(*value, "root_store");
assert_eq!(suffix, path!("any/path"));
}
#[test]
fn find_ancestor_no_match() {
let mut trie: PathTrie<&str> = PathTrie::new();
trie.insert(&path!("data"), "data_store");
assert!(trie.find_ancestor(&path!("other/path")).is_none());
}
#[test]
fn find_ancestor_exact_match() {
let mut trie: PathTrie<&str> = PathTrie::new();
trie.insert(&path!("data"), "data_store");
let (value, suffix) = trie.find_ancestor(&path!("data")).unwrap();
assert_eq!(*value, "data_store");
assert!(suffix.is_empty());
}
#[test]
fn find_ancestor_mut() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
trie.insert(&path!("a/b"), 2);
let (value, suffix) = trie.find_ancestor_mut(&path!("a/b/c")).unwrap();
assert_eq!(*value, 2);
assert_eq!(suffix, path!("c"));
*value = 42;
assert_eq!(trie.get(&path!("a/b")), Some(&42));
}
#[test]
fn find_ancestor_mut_no_match() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
assert!(trie.find_ancestor_mut(&path!("b/c")).is_none());
}
#[test]
fn iter_all_values() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
trie.insert(&path!("b"), 2);
trie.insert(&path!("a/c"), 3);
let mut items: Vec<_> = trie.iter().collect();
items.sort_by_key(|a| a.0.to_string());
assert_eq!(items.len(), 3);
assert_eq!(items[0], (path!("a"), &1));
assert_eq!(items[1], (path!("a/c"), &3));
assert_eq!(items[2], (path!("b"), &2));
}
#[test]
fn iter_empty() {
let trie: PathTrie<i32> = PathTrie::new();
assert_eq!(trie.iter().count(), 0);
}
#[test]
fn iter_root_value() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!(""), 42);
let items: Vec<_> = trie.iter().collect();
assert_eq!(items.len(), 1);
assert_eq!(items[0], (path!(""), &42));
}
#[test]
fn get_subtrie() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a/b"), 1);
trie.insert(&path!("a/c"), 2);
let subtrie = trie.get_subtrie(&path!("a")).unwrap();
assert_eq!(subtrie.get(&path!("b")), Some(&1));
assert_eq!(subtrie.get(&path!("c")), Some(&2));
}
#[test]
fn get_subtrie_nonexistent() {
let trie: PathTrie<i32> = PathTrie::new();
assert!(trie.get_subtrie(&path!("nonexistent")).is_none());
}
#[test]
fn clone_trie() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
let cloned = trie.clone();
assert_eq!(cloned.get(&path!("a")), Some(&1));
}
#[test]
fn default_trie() {
let trie: PathTrie<i32> = PathTrie::default();
assert!(trie.is_empty());
}
#[test]
fn debug_trie() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a"), 1);
let debug = format!("{:?}", trie);
assert!(debug.contains("PathTrie"));
}
#[test]
fn insert_at_root() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!(""), 42);
assert_eq!(trie.get(&path!("")), Some(&42));
assert_eq!(trie.len(), 1);
}
#[test]
fn remove_at_root() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!(""), 42);
assert_eq!(trie.remove(&path!("")), Some(42));
assert!(trie.is_empty());
}
#[test]
fn get_subtrie_mut() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a/b"), 1);
let subtrie = trie.get_subtrie_mut(&path!("a")).unwrap();
subtrie.insert(&path!("c"), 2);
assert_eq!(trie.get(&path!("a/c")), Some(&2));
}
#[test]
fn remove_subtree_empty_trie() {
let mut trie: PathTrie<i32> = PathTrie::new();
assert!(trie.remove_subtree(&path!("")).is_none());
}
#[test]
fn is_empty_with_only_structure() {
let mut trie: PathTrie<i32> = PathTrie::new();
trie.insert(&path!("a/b/c"), 1);
trie.remove(&path!("a/b/c"));
assert!(trie.is_empty());
}
}