use std::cell::RefCell;
use std::rc::Rc;
use slotmap::SlotMap;
use teksilo_core::ObserverHandle;
use crate::tree_change::{NodeId, TreeChange};
struct TreeNode<T> {
data: T,
parent: Option<NodeId>,
children: Vec<NodeId>,
}
struct TreeObserverEntry {
id: u64,
callback: Rc<dyn Fn(&TreeChange)>,
}
struct TreeModelInner<T> {
arena: SlotMap<slotmap::DefaultKey, TreeNode<T>>,
roots: Vec<NodeId>,
observers: Vec<TreeObserverEntry>,
next_observer_id: u64,
#[cfg(debug_assertions)]
debug_adapter: Option<Rc<dyn crate::debug_registry::ModelDebug>>,
}
pub struct TreeModel<T: 'static> {
inner: Rc<RefCell<TreeModelInner<T>>>,
}
impl<T: 'static> TreeModel<T> {
pub fn new() -> Self {
Self {
inner: Rc::new(RefCell::new(TreeModelInner {
arena: SlotMap::new(),
roots: Vec::new(),
observers: Vec::new(),
next_observer_id: 1,
#[cfg(debug_assertions)]
debug_adapter: None,
})),
}
}
pub fn root_count(&self) -> usize {
self.inner.borrow().roots.len()
}
pub fn root(&self, index: usize) -> NodeId {
self.inner.borrow().roots[index]
}
pub fn child_count(&self, parent: NodeId) -> usize {
let guard = self.inner.borrow();
guard
.arena
.get(parent.key())
.map(|n| n.children.len())
.unwrap_or(0)
}
pub fn child(&self, parent: NodeId, index: usize) -> NodeId {
let guard = self.inner.borrow();
guard.arena[parent.key()].children[index]
}
pub fn parent(&self, node: NodeId) -> Option<NodeId> {
let guard = self.inner.borrow();
guard.arena.get(node.key()).and_then(|n| n.parent)
}
pub fn depth(&self, node: NodeId) -> usize {
let guard = self.inner.borrow();
let mut depth = 0;
let mut current = guard.arena.get(node.key()).and_then(|n| n.parent);
while let Some(pid) = current {
depth += 1;
current = guard.arena.get(pid.key()).and_then(|n| n.parent);
}
depth
}
pub fn has_children(&self, node: NodeId) -> bool {
self.child_count(node) > 0
}
pub fn children(&self, node: NodeId) -> Vec<NodeId> {
let guard = self.inner.borrow();
guard
.arena
.get(node.key())
.map(|n| n.children.clone())
.unwrap_or_default()
}
pub fn with_item<R>(&self, node: NodeId, f: impl FnOnce(&T) -> R) -> Option<R> {
let guard = self.inner.borrow();
guard.arena.get(node.key()).map(|n| f(&n.data))
}
pub fn find_by(&self, predicate: impl Fn(&T) -> bool) -> Option<NodeId> {
let guard = self.inner.borrow();
let mut stack: Vec<NodeId> = guard.roots.iter().rev().copied().collect();
while let Some(nid) = stack.pop() {
if let Some(node) = guard.arena.get(nid.key()) {
if predicate(&node.data) {
return Some(nid);
}
for &child_id in node.children.iter().rev() {
stack.push(child_id);
}
}
}
None
}
pub fn insert_root(&self, index: usize, item: T) -> NodeId {
let node_id = {
let mut guard = self.inner.borrow_mut();
let key = guard.arena.insert(TreeNode {
data: item,
parent: None,
children: Vec::new(),
});
let node_id = NodeId::from_key(key);
guard.roots.insert(index, node_id);
node_id
};
self.notify(TreeChange::NodeInserted {
parent: None,
index,
node: node_id,
});
node_id
}
pub fn insert_child(&self, parent: NodeId, index: usize, item: T) -> NodeId {
let node_id = {
let mut guard = self.inner.borrow_mut();
let key = guard.arena.insert(TreeNode {
data: item,
parent: Some(parent),
children: Vec::new(),
});
let node_id = NodeId::from_key(key);
guard.arena[parent.key()].children.insert(index, node_id);
node_id
};
self.notify(TreeChange::NodeInserted {
parent: Some(parent),
index,
node: node_id,
});
node_id
}
pub fn remove(&self, node: NodeId) {
let parent = {
let mut guard = self.inner.borrow_mut();
let parent = guard.arena[node.key()].parent;
if let Some(pid) = parent {
guard.arena[pid.key()].children.retain(|&c| c != node);
} else {
guard.roots.retain(|&r| r != node);
}
Self::remove_subtree(&mut guard.arena, node);
parent
};
self.notify(TreeChange::NodeRemoved { parent, node });
}
pub fn move_node(&self, node: NodeId, new_parent: NodeId, new_index: usize) {
let old_parent = {
let mut guard = self.inner.borrow_mut();
assert!(
!Self::is_descendant_of(&guard.arena, new_parent, node),
"cannot move a node into its own subtree"
);
let old_parent = guard.arena[node.key()].parent;
if let Some(pid) = old_parent {
guard.arena[pid.key()].children.retain(|&c| c != node);
} else {
guard.roots.retain(|&r| r != node);
}
guard.arena[node.key()].parent = Some(new_parent);
guard.arena[new_parent.key()]
.children
.insert(new_index, node);
old_parent
};
self.notify(TreeChange::NodeMoved {
node,
old_parent,
new_parent: Some(new_parent),
new_index,
});
}
pub fn move_to_root(&self, node: NodeId, new_index: usize) {
let old_parent = {
let mut guard = self.inner.borrow_mut();
let old_parent = guard.arena[node.key()].parent;
if let Some(pid) = old_parent {
guard.arena[pid.key()].children.retain(|&c| c != node);
} else {
guard.roots.retain(|&r| r != node);
}
guard.arena[node.key()].parent = None;
guard.roots.insert(new_index, node);
old_parent
};
self.notify(TreeChange::NodeMoved {
node,
old_parent,
new_parent: None,
new_index,
});
}
pub fn update(&self, node: NodeId, item: T) {
{
let mut guard = self.inner.borrow_mut();
guard.arena[node.key()].data = item;
}
self.notify(TreeChange::NodeUpdated { node });
}
pub fn observe_changes(&self, f: impl Fn(&TreeChange) + 'static) -> ObserverHandle {
let mut guard = self.inner.borrow_mut();
let id = guard.next_observer_id;
guard.next_observer_id += 1;
guard.observers.push(TreeObserverEntry {
id,
callback: Rc::new(f),
});
let inner = self.inner.clone();
ObserverHandle::new(
self.inner.clone(),
id,
Rc::new(move |observer_id| {
inner.borrow_mut().observers.retain(|e| e.id != observer_id);
}),
)
}
fn notify(&self, change: TreeChange) {
let callbacks: Vec<Rc<dyn Fn(&TreeChange)>> = self
.inner
.borrow()
.observers
.iter()
.map(|e| e.callback.clone())
.collect();
for cb in &callbacks {
cb(&change);
}
}
fn remove_subtree(arena: &mut SlotMap<slotmap::DefaultKey, TreeNode<T>>, node: NodeId) {
let mut stack = vec![node];
let mut to_remove = Vec::new();
while let Some(current) = stack.pop() {
if let Some(n) = arena.get(current.key()) {
stack.extend(n.children.iter().copied());
}
to_remove.push(current);
}
for id in to_remove {
arena.remove(id.key());
}
}
fn is_descendant_of(
arena: &SlotMap<slotmap::DefaultKey, TreeNode<T>>,
candidate: NodeId,
ancestor: NodeId,
) -> bool {
let mut current = Some(candidate);
while let Some(nid) = current {
if nid == ancestor {
return true;
}
current = arena.get(nid.key()).and_then(|n| n.parent);
}
false
}
}
impl<T: std::fmt::Debug + 'static> TreeModel<T> {
pub fn debug_named(self, _name: impl Into<String>) -> Self {
#[cfg(debug_assertions)]
{
let weak = Rc::downgrade(&self.inner);
let adapter: Rc<dyn crate::debug_registry::ModelDebug> =
Rc::new(TreeModelDebug::<T> { weak });
let name = _name.into();
crate::debug_registry::register(name, Rc::downgrade(&adapter));
self.inner.borrow_mut().debug_adapter = Some(adapter);
}
self
}
}
#[cfg(debug_assertions)]
struct TreeModelDebug<T> {
weak: std::rc::Weak<RefCell<TreeModelInner<T>>>,
}
#[cfg(debug_assertions)]
impl<T: std::fmt::Debug + 'static> crate::debug_registry::ModelDebug for TreeModelDebug<T> {
fn kind(&self) -> &'static str {
"TreeModel"
}
fn len(&self) -> usize {
self.weak
.upgrade()
.map(|inner| inner.borrow().arena.len())
.unwrap_or(0)
}
fn debug_dump(&self, out: &mut dyn std::fmt::Write) {
let Some(inner) = self.weak.upgrade() else {
return;
};
let guard = inner.borrow();
let roots: Vec<NodeId> = guard.roots.clone();
for root in roots {
dump_subtree(&guard, root, 0, out);
}
}
}
#[cfg(debug_assertions)]
fn dump_subtree<T: std::fmt::Debug>(
guard: &TreeModelInner<T>,
node: NodeId,
depth: usize,
out: &mut dyn std::fmt::Write,
) {
let Some(n) = guard.arena.get(node.key()) else {
return;
};
let _ = writeln!(out, "{:indent$}{:?}", "", n.data, indent = depth * 2);
let children = n.children.clone();
for child in children {
dump_subtree(guard, child, depth + 1, out);
}
}
impl<T: 'static> Default for TreeModel<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: 'static> Clone for TreeModel<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
impl<T: std::fmt::Debug + 'static> std::fmt::Debug for TreeModel<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let guard = self.inner.borrow();
f.debug_struct("TreeModel")
.field("root_count", &guard.roots.len())
.field("total_nodes", &guard.arena.len())
.finish()
}
}
#[cfg(test)]
mod tests {
use std::cell::Cell;
use super::*;
fn sample_tree() -> (TreeModel<&'static str>, NodeId, NodeId, NodeId, NodeId) {
let tree = TreeModel::new();
let a = tree.insert_root(0, "A");
let b = tree.insert_root(1, "B");
let a1 = tree.insert_child(a, 0, "A1");
let a2 = tree.insert_child(a, 1, "A2");
(tree, a, b, a1, a2)
}
#[test]
fn empty_tree() {
let tree: TreeModel<i32> = TreeModel::new();
assert_eq!(tree.root_count(), 0);
}
#[test]
fn insert_roots() {
let tree = TreeModel::new();
let a = tree.insert_root(0, "A");
let b = tree.insert_root(1, "B");
assert_eq!(tree.root_count(), 2);
assert_eq!(tree.root(0), a);
assert_eq!(tree.root(1), b);
}
#[test]
fn insert_children() {
let (tree, a, _, a1, a2) = sample_tree();
assert_eq!(tree.child_count(a), 2);
assert_eq!(tree.child(a, 0), a1);
assert_eq!(tree.child(a, 1), a2);
}
#[test]
fn parent_and_depth() {
let (tree, a, _, a1, _) = sample_tree();
assert_eq!(tree.parent(a), None);
assert_eq!(tree.parent(a1), Some(a));
assert_eq!(tree.depth(a), 0);
assert_eq!(tree.depth(a1), 1);
}
#[test]
fn has_children_query() {
let (tree, a, b, _, _) = sample_tree();
assert!(tree.has_children(a));
assert!(!tree.has_children(b));
}
#[test]
fn with_item() {
let (tree, a, _, _, _) = sample_tree();
assert_eq!(tree.with_item(a, |v| *v), Some("A"));
}
#[test]
fn find_by() {
let (tree, _, _, a1, _) = sample_tree();
let found = tree.find_by(|v| *v == "A1");
assert_eq!(found, Some(a1));
let not_found = tree.find_by(|v| *v == "Z");
assert_eq!(not_found, None);
}
#[test]
fn insert_root_emits_change() {
let tree = TreeModel::new();
let changes: Rc<RefCell<Vec<TreeChange>>> = Rc::new(RefCell::new(Vec::new()));
let c = changes.clone();
let _handle = tree.observe_changes(move |change| {
c.borrow_mut().push(change.clone());
});
let a = tree.insert_root(0, "A");
let log = changes.borrow();
assert_eq!(
log[0],
TreeChange::NodeInserted {
parent: None,
index: 0,
node: a
}
);
}
#[test]
fn insert_child_emits_change() {
let tree = TreeModel::new();
let a = tree.insert_root(0, "A");
let changes: Rc<RefCell<Vec<TreeChange>>> = Rc::new(RefCell::new(Vec::new()));
let c = changes.clone();
let _handle = tree.observe_changes(move |change| {
c.borrow_mut().push(change.clone());
});
let a1 = tree.insert_child(a, 0, "A1");
let log = changes.borrow();
assert_eq!(log.len(), 1, "insert_child should emit exactly one change");
assert_eq!(
log[0],
TreeChange::NodeInserted {
parent: Some(a),
index: 0,
node: a1
}
);
}
#[test]
fn remove_emits_change() {
let (tree, a, _, a1, _) = sample_tree();
let changes: Rc<RefCell<Vec<TreeChange>>> = Rc::new(RefCell::new(Vec::new()));
let c = changes.clone();
let _handle = tree.observe_changes(move |change| {
c.borrow_mut().push(change.clone());
});
tree.remove(a1);
assert_eq!(tree.child_count(a), 1);
let log = changes.borrow();
assert_eq!(log.len(), 1, "remove should emit exactly one change");
assert_eq!(
log[0],
TreeChange::NodeRemoved {
parent: Some(a),
node: a1
}
);
}
#[test]
fn remove_subtree() {
let tree = TreeModel::new();
let a = tree.insert_root(0, "A");
let a1 = tree.insert_child(a, 0, "A1");
let a1a = tree.insert_child(a1, 0, "A1a");
tree.remove(a1);
assert_eq!(tree.child_count(a), 0);
assert_eq!(tree.with_item(a1, |_| ()), None, "a1 should be removed");
assert_eq!(
tree.with_item(a1a, |_| ()),
None,
"a1a (grandchild) should also be removed"
);
}
#[test]
fn remove_root() {
let (tree, a, b, _, _) = sample_tree();
tree.remove(a);
assert_eq!(tree.root_count(), 1);
assert_eq!(tree.root(0), b);
}
#[test]
fn move_node() {
let (tree, a, b, a1, _) = sample_tree();
let changes: Rc<RefCell<Vec<TreeChange>>> = Rc::new(RefCell::new(Vec::new()));
let c = changes.clone();
let _handle = tree.observe_changes(move |change| {
c.borrow_mut().push(change.clone());
});
tree.move_node(a1, b, 0);
assert_eq!(tree.child_count(a), 1); assert_eq!(tree.child_count(b), 1); assert_eq!(tree.child(b, 0), a1);
assert_eq!(tree.parent(a1), Some(b));
let log = changes.borrow();
assert_eq!(
log[0],
TreeChange::NodeMoved {
node: a1,
old_parent: Some(a),
new_parent: Some(b),
new_index: 0,
}
);
}
#[test]
fn move_to_root() {
let (tree, a, _, a1, _) = sample_tree();
tree.move_to_root(a1, 0);
assert_eq!(tree.root_count(), 3); assert_eq!(tree.root(0), a1);
assert_eq!(tree.parent(a1), None);
assert_eq!(tree.child_count(a), 1); }
#[test]
#[should_panic(expected = "cannot move a node into its own subtree")]
fn move_into_own_subtree_panics() {
let (tree, a, _, a1, _) = sample_tree();
tree.move_node(a, a1, 0); }
#[test]
fn update_emits_change() {
let (tree, a, _, _, _) = sample_tree();
let changes: Rc<RefCell<Vec<TreeChange>>> = Rc::new(RefCell::new(Vec::new()));
let c = changes.clone();
let _handle = tree.observe_changes(move |change| {
c.borrow_mut().push(change.clone());
});
tree.update(a, "A-updated");
assert_eq!(tree.with_item(a, |v| *v), Some("A-updated"));
let log = changes.borrow();
assert_eq!(log[0], TreeChange::NodeUpdated { node: a });
}
#[test]
fn observer_removed_on_handle_drop() {
let tree = TreeModel::new();
let count = Rc::new(Cell::new(0));
let c = count.clone();
let handle = tree.observe_changes(move |_| c.set(c.get() + 1));
tree.insert_root(0, "A");
assert_eq!(count.get(), 1);
drop(handle);
tree.insert_root(1, "B");
assert_eq!(count.get(), 1); }
#[test]
fn clone_shares_data() {
let tree = TreeModel::new();
let a = tree.insert_root(0, "A");
let clone = tree.clone();
assert_eq!(clone.root_count(), 1);
assert_eq!(clone.with_item(a, |v| *v), Some("A"));
clone.insert_root(1, "B");
assert_eq!(tree.root_count(), 2);
}
#[test]
fn deep_tree_depth() {
let tree = TreeModel::new();
let r = tree.insert_root(0, "r");
let c1 = tree.insert_child(r, 0, "c1");
let c2 = tree.insert_child(c1, 0, "c2");
let c3 = tree.insert_child(c2, 0, "c3");
assert_eq!(tree.depth(r), 0);
assert_eq!(tree.depth(c1), 1);
assert_eq!(tree.depth(c2), 2);
assert_eq!(tree.depth(c3), 3);
}
#[test]
fn children_returns_correct_ids() {
let (tree, a, _, a1, a2) = sample_tree();
let children = tree.children(a);
assert_eq!(children, vec![a1, a2]);
}
#[test]
fn remove_deep_chain_does_not_overflow() {
const DEPTH: usize = 50_000;
let tree = TreeModel::new();
let root = tree.insert_root(0, 0usize);
let mut leaf = root;
for i in 1..DEPTH {
leaf = tree.insert_child(leaf, 0, i);
}
tree.remove(root);
assert_eq!(tree.root_count(), 0);
assert_eq!(tree.with_item(leaf, |_| ()), None);
}
}