1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

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,
{
    /// Creates a new TreeWrapper
    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(())
    }
}