merkle_search_tree/visitor/
dot.rs

1//! Output the structure of a [`MerkleSearchTree`] formatted using [Graphviz]
2//! DOT grammar.
3//!
4//! [`MerkleSearchTree`]: crate::MerkleSearchTree
5//! [Graphviz]: https://graphviz.org/
6
7use std::fmt::{Display, Write};
8
9use super::Visitor;
10use crate::{node::Node, page::Page};
11
12#[derive(Debug, Clone)]
13enum Parent {
14    Node(String),
15    Page(String, usize),
16}
17
18/// Serialise a tree into [Graphviz DOT language][dot] output to be rendered as
19/// a diagram.
20///
21/// [dot]: https://graphviz.org/doc/info/lang.html
22#[derive(Debug)]
23pub struct DotVisitor {
24    buf: String,
25
26    /// Total number of pages visited so far (1-based)
27    page_count: usize,
28
29    /// The stack of parent node keys / page records, most recently visited
30    /// last.
31    link_stack: Vec<Parent>,
32
33    /// A set of per-page buffers, populated incrementally and merged into `buf`
34    /// once complete.
35    page_bufs: Vec<String>,
36}
37
38impl Default for DotVisitor {
39    fn default() -> Self {
40        Self {
41            buf: "digraph g {\n".to_string(),
42            page_count: Default::default(),
43            link_stack: vec![],
44            page_bufs: vec![],
45        }
46    }
47}
48
49impl<'a, const N: usize, K> Visitor<'a, N, K> for DotVisitor
50where
51    K: Display,
52{
53    fn visit_page(&mut self, page: &'a Page<N, K>, high_page: bool) -> bool {
54        let mut buf = String::new();
55
56        self.page_count += 1;
57
58        // Write the link directly to the buffer.
59        match self.link_stack.last() {
60            // If this is the first page, render a pseudo root node
61            None if self.page_count == 1 => {
62                writeln!(buf, "\troot [shape=diamond style=dotted];").unwrap();
63                writeln!(buf, "\troot -> page_{}:head", self.page_count)
64            }
65            Some(Parent::Page(p, _)) => writeln!(
66                buf,
67                "\t{} -> page_{}:high_page [fontcolor=red color=red label=\"high page\"];",
68                p, self.page_count
69            ),
70            Some(Parent::Node(n)) if !high_page => {
71                writeln!(buf, "\t{} -> page_{}:head;", n, self.page_count)
72            }
73            _ => Ok(()),
74        }
75        .unwrap();
76
77        write!(
78            buf,
79            "\tpage_{} [shape=record, label=\"<head>Level {}|",
80            self.page_count,
81            page.level(),
82        )
83        .unwrap();
84
85        self.link_stack.push(Parent::Page(
86            format!("page_{}:head", self.page_count),
87            self.page_count,
88        ));
89
90        self.page_bufs.push(buf);
91
92        true
93    }
94
95    fn post_visit_page(&mut self, page: &'a Page<N, K>) -> bool {
96        let mut buf = self.page_bufs.pop().unwrap();
97
98        // Remove the trailing | from the node field
99        let _ = buf.pop();
100
101        let me = match self.link_stack.pop().unwrap() {
102            Parent::Node(_) => panic!("pop should yield visited page"),
103            Parent::Page(_, id) => id,
104        };
105
106        // If this page has a high page, it'll be visited next.
107        if page.high_page().is_some() {
108            // Add a high page to this record
109            writeln!(buf, r#"|<high_page>·"]"#).unwrap();
110
111            // Link the high page to the referenced page
112            writeln!(
113                buf,
114                "\tpage_{}:high_page -> page_{}:head [fontcolor=red color=red label=\"high \
115                 page\"];",
116                me,
117                self.page_count + 1,
118            )
119            .unwrap();
120        } else {
121            // No high page, terminate record without it.
122            writeln!(buf, r#""]"#).unwrap();
123        }
124
125        writeln!(self.buf, "{}", buf).unwrap();
126
127        true
128    }
129
130    fn pre_visit_node(&mut self, node: &'a Node<N, K>) -> bool {
131        // Find the ID of the last visited page, which will be the parent of
132        // this node.
133        let page_id = self
134            .link_stack
135            .iter()
136            .rev()
137            .filter_map(|v| match v {
138                Parent::Node(_) => None,
139                Parent::Page(_, id) => Some(id),
140            })
141            .next()
142            .unwrap();
143
144        let name = clean_name(node.key());
145        self.link_stack
146            .push(Parent::Node(format!("page_{}:{}", page_id, &name)));
147
148        true
149    }
150
151    fn visit_node(&mut self, node: &'a Node<N, K>) -> bool {
152        let buf = self.page_bufs.last_mut().unwrap();
153
154        // Add this node to the page record
155        let name = clean_name(node.key());
156        write!(buf, "<{}>·|{}|", &name, name).unwrap();
157
158        true
159    }
160
161    fn post_visit_node(&mut self, _node: &'a Node<N, K>) -> bool {
162        self.link_stack.pop();
163        true
164    }
165}
166
167impl DotVisitor {
168    /// Consume this visitor, yielding the generated DOT representation.
169    pub fn finalise(self) -> String {
170        assert!(self.page_bufs.is_empty());
171        assert!(self.link_stack.is_empty());
172
173        format!("{}}}\n", self.buf)
174    }
175}
176
177fn clean_name<'a, T>(name: T) -> String
178where
179    T: std::fmt::Display + 'a,
180{
181    name.to_string()
182        .chars()
183        .map(|v| match v {
184            'a'..='z' | 'A'..='Z' | '0'..='9' | '.' | '_' => v,
185            _ => '_',
186        })
187        .collect()
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use crate::{
194        assert_tree,
195        builder::Builder,
196        digest::{
197            mock::{LevelKey, MockHasher},
198            Digest, ValueDigest,
199        },
200    };
201
202    const MOCK_VALUE: ValueDigest<32> = ValueDigest::new(Digest::new([0; 32]));
203
204    #[test]
205    fn test_dot_flat() {
206        let p = Page::new(
207            42,
208            vec![
209                Node::new(&"k1", MOCK_VALUE, None),
210                Node::new(&"k2", MOCK_VALUE, None),
211            ],
212        );
213
214        assert_tree!(page = p);
215    }
216
217    #[test]
218    fn test_dot_high_page() {
219        let h = Page::new(0, vec![Node::new(&"z_high1", MOCK_VALUE, None)]);
220
221        let mut p = Page::new(
222            42,
223            vec![
224                Node::new(&"k1", MOCK_VALUE, None),
225                Node::new(&"k2", MOCK_VALUE, None),
226            ],
227        );
228        p.insert_high_page(Box::new(h));
229
230        assert_tree!(page = p);
231    }
232
233    #[test]
234    fn test_dot_lt_pointer() {
235        let lt_page_1 = Page::new(1, vec![Node::new(&"lt1", MOCK_VALUE, None)]);
236        let lt_page_2 = Page::new(
237            2,
238            vec![Node::new(&"lt2", MOCK_VALUE, Some(Box::new(lt_page_1)))],
239        );
240
241        let p = Page::new(
242            42,
243            vec![
244                Node::new(&"z_k1", MOCK_VALUE, Some(Box::new(lt_page_2))),
245                Node::new(&"z_k2", MOCK_VALUE, None),
246            ],
247        );
248
249        assert_tree!(page = p);
250    }
251
252    #[test]
253    fn test_dot_high_page_lt_pointer() {
254        let lt_page_1 = Page::new(10, vec![Node::new(&"lt1", MOCK_VALUE, None)]);
255        let lt_page_2 = Page::new(
256            11,
257            vec![Node::new(&"lt2", MOCK_VALUE, Some(Box::new(lt_page_1)))],
258        );
259
260        let h1 = Page::new(0, vec![Node::new(&"zz_h1", MOCK_VALUE, None)]);
261        let h2 = Page::new(1, vec![Node::new(&"zz_h2", MOCK_VALUE, Some(Box::new(h1)))]);
262
263        let mut p = Page::new(
264            42,
265            vec![
266                Node::new(&"z_k1", MOCK_VALUE, Some(Box::new(lt_page_2))),
267                Node::new(&"z_k2", MOCK_VALUE, None),
268            ],
269        );
270        p.insert_high_page(Box::new(h2));
271
272        assert_tree!(page = p);
273    }
274
275    #[test]
276    fn test_parent_lookup() {
277        const MOCK_VALUE_1: ValueDigest<1> = ValueDigest::new(Digest::new([0; 1]));
278
279        let mut p = Page::new(1, vec![Node::new(4, MOCK_VALUE_1, None)]);
280
281        p.upsert(3, 0, MOCK_VALUE_1);
282        p.upsert(1, 0, MOCK_VALUE_1);
283        p.upsert(2, 1, MOCK_VALUE_1);
284
285        assert_tree!(page = p);
286    }
287
288    #[test]
289    fn test_linear_children() {
290        let mut t = Builder::default()
291            .with_hasher(MockHasher::default())
292            .build();
293
294        t.upsert(LevelKey::new("I", 2), &"bananas");
295        t.upsert(LevelKey::new("E", 1), &"bananas");
296        t.upsert(LevelKey::new("F", 0), &"bananas");
297
298        assert_tree!(t);
299    }
300}
merkle_search_tree/visitor/dot.rs

merkle_search_tree/visitor/
dot.rs
