use std::{
fmt::{self, Display},
mem,
};
use ego_tree::{NodeId, NodeRef, Tree};
pub struct TreeWrapper<'a, T, F> {
tree: &'a Tree<T>,
root: NodeId,
f: F,
}
impl<'a, T, F> TreeWrapper<'a, T, F>
where
F: Fn(&NodeRef<T>) -> String,
{
pub fn new(tree: &'a Tree<T>, root: NodeId, f: F) -> Self {
Self { tree, f, root }
}
}
impl<'a, T, F: Fn(&NodeRef<T>) -> String> Display for TreeWrapper<'a, T, F> {
fn fmt(&self, w: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(w, "digraph {{")?;
for node in self.tree.nodes() {
if node.ancestors().any(|x| x.id() == self.root) {
let id: usize = unsafe { mem::transmute(node.id()) };
writeln!(
w,
"N{id} [ label = \"{}\"]",
(self.f)(&node).escape_default()
)?;
for child in node.children() {
let c_id: usize = unsafe { mem::transmute(child.id()) };
writeln!(w, "N{id} -> N{c_id}")?;
}
}
}
writeln!(w, "}}")?;
Ok(())
}
}