mod node;
mod walker;
pub use node::{ParentRef, RecordId};
pub use walker::{BreadthFirstWalker, TreeWalker};
use std::collections::HashMap;
pub trait TreeRecord {
fn owner_index(&self) -> i32;
fn set_owner_index(&mut self, index: i32);
}
#[derive(Debug, Clone)]
pub struct RecordTree<R> {
records: Vec<R>,
parents: HashMap<RecordId, RecordId>,
children: HashMap<RecordId, Vec<RecordId>>,
roots: Vec<RecordId>,
}
impl<R: TreeRecord> RecordTree<R> {
pub fn from_records(records: Vec<R>) -> Self {
let mut tree = RecordTree {
records,
parents: HashMap::new(),
children: HashMap::new(),
roots: Vec::new(),
};
tree.rebuild_relationships();
tree
}
fn rebuild_relationships(&mut self) {
self.parents.clear();
self.children.clear();
self.roots.clear();
let len = self.records.len();
for (idx, record) in self.records.iter().enumerate() {
let child_id = RecordId::new(idx as u32);
let owner_index = record.owner_index();
if owner_index < 0 || (owner_index as usize) >= len {
self.roots.push(child_id);
} else {
let parent_id = RecordId::new(owner_index as u32);
self.parents.insert(child_id, parent_id);
self.children.entry(parent_id).or_default().push(child_id);
}
}
}
pub fn len(&self) -> usize {
self.records.len()
}
pub fn is_empty(&self) -> bool {
self.records.is_empty()
}
pub fn get(&self, id: RecordId) -> Option<&R> {
self.records.get(id.index() as usize)
}
pub fn get_mut(&mut self, id: RecordId) -> Option<&mut R> {
self.records.get_mut(id.index() as usize)
}
pub fn get_by_index(&self, index: usize) -> Option<&R> {
self.records.get(index)
}
pub fn iter(&self) -> impl Iterator<Item = (RecordId, &R)> {
self.records
.iter()
.enumerate()
.map(|(i, r)| (RecordId::new(i as u32), r))
}
pub fn iter_mut(&mut self) -> impl Iterator<Item = (RecordId, &mut R)> {
self.records
.iter_mut()
.enumerate()
.map(|(i, r)| (RecordId::new(i as u32), r))
}
pub fn roots(&self) -> impl Iterator<Item = (RecordId, &R)> {
self.roots
.iter()
.filter_map(move |&id| self.records.get(id.index() as usize).map(|r| (id, r)))
}
pub fn root_count(&self) -> usize {
self.roots.len()
}
pub fn is_root(&self, id: RecordId) -> bool {
self.roots.contains(&id)
}
pub fn parent(&self, id: RecordId) -> Option<(RecordId, &R)> {
self.parents.get(&id).and_then(|&parent_id| {
self.records
.get(parent_id.index() as usize)
.map(|r| (parent_id, r))
})
}
pub fn parent_id(&self, id: RecordId) -> Option<RecordId> {
self.parents.get(&id).copied()
}
pub fn children(&self, id: RecordId) -> impl Iterator<Item = (RecordId, &R)> {
self.children
.get(&id)
.map(|ids| ids.as_slice())
.unwrap_or(&[])
.iter()
.filter_map(move |&child_id| {
self.records
.get(child_id.index() as usize)
.map(|r| (child_id, r))
})
}
pub fn child_count(&self, id: RecordId) -> usize {
self.children.get(&id).map(|c| c.len()).unwrap_or(0)
}
pub fn has_children(&self, id: RecordId) -> bool {
self.children
.get(&id)
.map(|c| !c.is_empty())
.unwrap_or(false)
}
pub fn ancestors(&self, id: RecordId) -> impl Iterator<Item = (RecordId, &R)> {
AncestorIterator {
tree: self,
current: self.parent_id(id),
}
}
pub fn descendants(&self, id: RecordId) -> impl Iterator<Item = (RecordId, &R)> {
DescendantIterator {
tree: self,
stack: self.children.get(&id).cloned().unwrap_or_default(),
}
}
pub fn depth(&self, id: RecordId) -> usize {
let mut depth = 0;
let mut current = id;
while let Some(parent_id) = self.parent_id(current) {
depth += 1;
current = parent_id;
}
depth
}
pub fn walk_depth_first(&self) -> TreeWalker<'_, R> {
TreeWalker::new(self)
}
pub fn walk_from(&self, id: RecordId) -> TreeWalker<'_, R> {
TreeWalker::from_node(self, id)
}
pub fn find<F>(&self, predicate: F) -> impl Iterator<Item = (RecordId, &R)>
where
F: Fn(&R) -> bool,
{
self.records
.iter()
.enumerate()
.filter(move |(_, r)| predicate(r))
.map(|(i, r)| (RecordId::new(i as u32), r))
}
pub fn find_first<F>(&self, predicate: F) -> Option<(RecordId, &R)>
where
F: Fn(&R) -> bool,
{
self.records
.iter()
.enumerate()
.find(|(_, r)| predicate(r))
.map(|(i, r)| (RecordId::new(i as u32), r))
}
pub fn path_to_root(&self, id: RecordId) -> Vec<RecordId> {
let mut path = vec![id];
let mut current = id;
while let Some(parent_id) = self.parent_id(current) {
path.push(parent_id);
current = parent_id;
}
path
}
pub fn into_records(self) -> Vec<R> {
self.records
}
pub fn as_slice(&self) -> &[R] {
&self.records
}
pub fn add(&mut self, record: R) -> RecordId {
let id = RecordId::new(self.records.len() as u32);
let owner_index = record.owner_index();
self.records.push(record);
if owner_index < 0 || (owner_index as usize) >= self.records.len() - 1 {
self.roots.push(id);
} else {
let parent_id = RecordId::new(owner_index as u32);
self.parents.insert(id, parent_id);
self.children.entry(parent_id).or_default().push(id);
}
id
}
pub fn remove(&mut self, id: RecordId) -> Option<R> {
let index = id.index() as usize;
if index >= self.records.len() {
return None;
}
let record = self.records.remove(index);
self.rebuild_relationships();
Some(record)
}
pub fn reparent(&mut self, id: RecordId, new_parent: Option<RecordId>) {
if let Some(old_parent_id) = self.parents.remove(&id) {
if let Some(siblings) = self.children.get_mut(&old_parent_id) {
siblings.retain(|&child_id| child_id != id);
}
}
self.roots.retain(|&root_id| root_id != id);
if let Some(parent_id) = new_parent {
self.parents.insert(id, parent_id);
self.children.entry(parent_id).or_default().push(id);
if let Some(record) = self.records.get_mut(id.index() as usize) {
record.set_owner_index(parent_id.index() as i32);
}
} else {
self.roots.push(id);
if let Some(record) = self.records.get_mut(id.index() as usize) {
record.set_owner_index(-1);
}
}
}
}
impl<R> Default for RecordTree<R> {
fn default() -> Self {
RecordTree {
records: Vec::new(),
parents: HashMap::new(),
children: HashMap::new(),
roots: Vec::new(),
}
}
}
struct AncestorIterator<'a, R: TreeRecord> {
tree: &'a RecordTree<R>,
current: Option<RecordId>,
}
impl<'a, R: TreeRecord> Iterator for AncestorIterator<'a, R> {
type Item = (RecordId, &'a R);
fn next(&mut self) -> Option<Self::Item> {
let id = self.current?;
self.current = self.tree.parent_id(id);
self.tree.get(id).map(|r| (id, r))
}
}
struct DescendantIterator<'a, R: TreeRecord> {
tree: &'a RecordTree<R>,
stack: Vec<RecordId>,
}
impl<'a, R: TreeRecord> Iterator for DescendantIterator<'a, R> {
type Item = (RecordId, &'a R);
fn next(&mut self) -> Option<Self::Item> {
let id = self.stack.pop()?;
if let Some(children) = self.tree.children.get(&id) {
self.stack.extend(children.iter().rev());
}
self.tree.get(id).map(|r| (id, r))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Debug, Clone, Default)]
struct TestRecord {
owner_index: i32,
value: String,
}
impl TreeRecord for TestRecord {
fn owner_index(&self) -> i32 {
self.owner_index
}
fn set_owner_index(&mut self, index: i32) {
self.owner_index = index;
}
}
#[test]
fn test_build_tree() {
let records = vec![
TestRecord {
owner_index: -1,
value: "root".to_string(),
},
TestRecord {
owner_index: 0,
value: "child1".to_string(),
},
TestRecord {
owner_index: 0,
value: "child2".to_string(),
},
TestRecord {
owner_index: 1,
value: "grandchild".to_string(),
},
];
let tree = RecordTree::from_records(records);
assert_eq!(tree.len(), 4);
assert_eq!(tree.root_count(), 1);
}
#[test]
fn test_navigation() {
let records = vec![
TestRecord {
owner_index: -1,
value: "root".to_string(),
},
TestRecord {
owner_index: 0,
value: "child1".to_string(),
},
TestRecord {
owner_index: 0,
value: "child2".to_string(),
},
TestRecord {
owner_index: 1,
value: "grandchild".to_string(),
},
];
let tree = RecordTree::from_records(records);
let roots: Vec<_> = tree.roots().collect();
assert_eq!(roots.len(), 1);
assert_eq!(roots[0].1.value, "root");
let root_id = RecordId::new(0);
let children: Vec<_> = tree.children(root_id).collect();
assert_eq!(children.len(), 2);
let child_id = RecordId::new(1);
let parent = tree.parent(child_id);
assert!(parent.is_some());
assert_eq!(parent.unwrap().1.value, "root");
assert_eq!(tree.depth(RecordId::new(0)), 0);
assert_eq!(tree.depth(RecordId::new(1)), 1);
assert_eq!(tree.depth(RecordId::new(3)), 2);
}
#[test]
fn test_walk_depth_first() {
let records = vec![
TestRecord {
owner_index: -1,
value: "root".to_string(),
},
TestRecord {
owner_index: 0,
value: "child1".to_string(),
},
TestRecord {
owner_index: 0,
value: "child2".to_string(),
},
TestRecord {
owner_index: 1,
value: "grandchild".to_string(),
},
];
let tree = RecordTree::from_records(records);
let walked: Vec<_> = tree.walk_depth_first().collect();
assert_eq!(walked.len(), 4);
assert_eq!(walked[0].1.value, "root");
assert_eq!(walked[0].2, 0);
}
#[test]
fn test_ancestors() {
let records = vec![
TestRecord {
owner_index: -1,
value: "root".to_string(),
},
TestRecord {
owner_index: 0,
value: "child".to_string(),
},
TestRecord {
owner_index: 1,
value: "grandchild".to_string(),
},
];
let tree = RecordTree::from_records(records);
let ancestors: Vec<_> = tree.ancestors(RecordId::new(2)).collect();
assert_eq!(ancestors.len(), 2);
assert_eq!(ancestors[0].1.value, "child");
assert_eq!(ancestors[1].1.value, "root");
}
}