augmented-rbtree 0.3.0

An augmented red-black tree with generic, user-defined per-node statistics — enables interval trees, order-statistics trees, range-sum trees, and more.
Documentation 
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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use std::{fmt::Write, fs::File, path::Path};

use augmented_rbtree::{Allocator, AugmentedRBTree, Color};

pub(crate) fn dump_tree_to_svg<K, V, G, S, A: Allocator, P: AsRef<Path>>(
    tree: &AugmentedRBTree<K, V, G, A>,
    path: P,
    show_nulls: bool,
) -> std::io::Result<()>
where
    K: std::fmt::Display + Ord,
    V: std::fmt::Display,
    G: augmented_rbtree::Augment<K, V, Stats = S>,
    S: std::fmt::Display,
{
    let mut dot = String::new();
    dot.push_str("digraph RBT {\n");
    dot.push_str("    node [fontname=\"Arial\", shape=Mrecord, style=filled, color=\"white\", fontcolor=white];\n");

    let mut nil_counter = 0;

    tree.visit_topology(|key, color, left, right| {
        let color_str = match color {
            Color::Red => "red",
            Color::Black => "black",
        };

        // Fetch the value for the current key
        let (value_str, stats_str) = if let Some((val, stats)) = tree.get_value_stats(key) {
            (format!("{val}"), format!("{stats}"))
        } else {
            ("None".to_string(), "None".to_string())
        };

        // Escape special Graphviz characters including double quotes inside labels
        let escaped_key = format!("{key}")
            .replace('"', "\\\"")
            .replace('|', "\\|")
            .replace('{', "\\{")
            .replace('}', "\\}");
        let escaped_val = value_str
            .replace('"', "\\\"")
            .replace('|', "\\|")
            .replace('{', "\\{")
            .replace('}', "\\}");
        let escaped_stat = stats_str
            .replace('"', "\\\"")
            .replace('|', "\\|")
            .replace('{', "\\{")
            .replace('}', "\\}");

        let label = format!("{{ {escaped_key} | {{ {escaped_val} | {escaped_stat} }} }}");

        // Safely escape the key string to use it as a Node Identifier
        let id_key = format!("{key}").replace('"', "\\\"");

        // Wrap the node ID in double quotes to handle spaces/brackets securely
        let _ = writeln!(
            dot,
            "    \"{id_key}\" [label=\"{label}\", fillcolor={color_str}, color=\"white\"];"
        );

        if let Some(l_key) = left {
            let id_l_key = format!("{l_key}").replace('"', "\\\"");
            let _ = writeln!(dot, "    \"{id_key}\" -> \"{id_l_key}\";");
        } else if show_nulls {
            nil_counter += 1;
            let nil_name = format!("nil_{nil_counter}");
            let _ = writeln!(
                dot,
                "    {nil_name} [label=\"NIL\", shape=box, fillcolor=black, color=white, fontcolor=white];"
            );
            let _ = writeln!(dot, "    \"{id_key}\" -> {nil_name};");
        }

        if let Some(r_key) = right {
            let id_r_key = format!("{r_key}").replace('"', "\\\"");
            let _ = writeln!(dot, "    \"{id_key}\" -> \"{id_r_key}\";");
        } else if show_nulls {
            nil_counter += 1;
            let nil_name = format!("nil_{nil_counter}");
            let _ = writeln!(
                dot,
                "    {nil_name} [label=\"NIL\", shape=box, fillcolor=black, color=white, fontcolor=white];"
            );
            let _ = writeln!(dot, "    \"{id_key}\" -> {nil_name};");
        }
    });

    dot.push_str("}\n");

    let parsed_graph = graphviz_rust::parse(&dot)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

    let svg_bytes = graphviz_rust::exec(
        parsed_graph,
        &mut graphviz_rust::printer::PrinterContext::default(),
        vec![graphviz_rust::cmd::CommandArg::Format(
            graphviz_rust::cmd::Format::Svg,
        )],
    )
    .map_err(std::io::Error::other)?;

    let mut file = File::create(path)?;
    std::io::Write::write_all(&mut file, svg_bytes.as_ref())?;

    Ok(())
}