redact_composer_core/render/
tree.rs

1use std::fmt::Debug;
2use std::ops::Index;
3
4#[cfg(feature = "serde")]
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug)]
8/// An index based n-ary tree node.
9pub struct Node<T> {
10    /// Index of this node.
11    pub idx: usize,
12    /// Contained value of this node.
13    pub value: T,
14    /// This node's parent index as an [`Option`], or [`None`] if it is a root.
15    pub parent: Option<usize>,
16    /// The indices of this node's children.
17    pub children: Vec<usize>,
18}
19
20#[derive(Debug)]
21/// An n-ary index-based tree.
22pub struct Tree<T> {
23    nodes: Vec<Node<T>>,
24}
25
26impl<T> Tree<T> {
27    /// Creates an empty tree.
28    pub fn new() -> Tree<T> {
29        Tree { nodes: vec![] }
30    }
31
32    /// The number of nodes in the tree.
33    pub fn len(&self) -> usize {
34        self.nodes.len()
35    }
36
37    /// Returns `true` iff the tree is empty.
38    pub fn is_empty(&self) -> bool {
39        self.nodes.is_empty()
40    }
41
42    /// Returns this tree's root, or [`None`] if it is empty.
43    pub fn root(&self) -> Option<&Node<T>> {
44        self.get(0)
45    }
46
47    /// Gets a node by it's index.
48    pub fn get(&self, idx: usize) -> Option<&Node<T>> {
49        self.nodes.get(idx)
50    }
51
52    /// Iterates a subtree starting from a given node. Order is not guaranteed.
53    pub fn node_iter<'a>(&'a self, start: &'a Node<T>) -> NodeIter<T> {
54        NodeIter {
55            tree: self,
56            idx_idx: 0,
57            idxs: vec![&start.idx],
58            skip: None,
59        }
60    }
61
62    /// Iterates a subtree, skipping indices contained in `skip`.
63    pub fn node_iter_with_skip<'a>(&'a self, start: &'a Node<T>, skip: Vec<usize>) -> NodeIter<T> {
64        NodeIter {
65            tree: self,
66            idx_idx: 0,
67            idxs: vec![&start.idx],
68            skip: Some(skip),
69        }
70    }
71
72    /// Returns a node iterator starting from the tree's root. Order is not guaranteed.
73    pub fn iter(&self) -> NodeIter<T> {
74        match self.root() {
75            Some(root) => self.node_iter(root),
76            None => NodeIter {
77                tree: self,
78                idx_idx: 0,
79                idxs: vec![],
80                skip: None,
81            },
82        }
83    }
84
85    /// Inserts a new value in this tree as a child of the `parent_idx` node.
86    pub fn insert(&mut self, item: T, parent_idx: Option<usize>) -> usize {
87        let new_idx = self.nodes.len();
88        self.nodes.push(Node {
89            idx: new_idx,
90            parent: parent_idx,
91            value: item,
92            children: vec![],
93        });
94
95        if let Some(parent_idx) = parent_idx {
96            self.nodes[parent_idx].children.push(new_idx);
97        }
98
99        new_idx
100    }
101}
102
103impl<T> Default for Tree<T> {
104    fn default() -> Self {
105        Self::new()
106    }
107}
108
109impl<'a, T> IntoIterator for &'a Tree<T> {
110    type Item = &'a Node<T>;
111
112    type IntoIter = NodeIter<'a, T>;
113
114    fn into_iter(self) -> Self::IntoIter {
115        self.iter()
116    }
117}
118
119/// A node iterator with optional `skip` indices.
120#[derive(Debug)]
121pub struct NodeIter<'a, T> {
122    tree: &'a Tree<T>,
123    idx_idx: usize,
124    idxs: Vec<&'a usize>,
125    skip: Option<Vec<usize>>,
126}
127
128impl<'a, T> Iterator for NodeIter<'a, T> {
129    type Item = &'a Node<T>;
130
131    fn next(&mut self) -> Option<Self::Item> {
132        match self.idxs.get(self.idx_idx) {
133            Some(idx) => {
134                let ret = &self.tree.nodes[**idx];
135                self.idxs.append(
136                    &mut ret
137                        .children
138                        .iter()
139                        .filter(|n| {
140                            if let Some(skip) = &self.skip {
141                                !skip.contains(n)
142                            } else {
143                                true
144                            }
145                        })
146                        .collect(),
147                );
148                self.idx_idx += 1;
149
150                Some(ret)
151            }
152            None => None,
153        }
154    }
155}
156
157impl<Idx: std::slice::SliceIndex<[Node<T>]>, T> Index<Idx> for Tree<T> {
158    type Output = Idx::Output;
159
160    fn index(&self, index: Idx) -> &Self::Output {
161        &self.nodes[index]
162    }
163}
164
165impl<Idx: std::slice::SliceIndex<[Node<T>]>, T> std::ops::IndexMut<Idx> for Tree<T> {
166    fn index_mut(&mut self, index: Idx) -> &mut Self::Output {
167        &mut self.nodes[index]
168    }
169}
170
171#[cfg(feature = "serde")]
172impl<T> Serialize for Tree<T>
173where
174    T: Serialize,
175{
176    fn serialize<S>(&self, serializer: S) -> crate::render::Result<S::Ok, S::Error>
177    where
178        S: serde::Serializer,
179    {
180        SerializeHelperNode::from((&self[0], self)).serialize(serializer)
181    }
182}
183
184#[cfg(feature = "serde")]
185impl<'de, T> Deserialize<'de> for Tree<T>
186where
187    T: Deserialize<'de> + Debug,
188{
189    fn deserialize<D>(deserializer: D) -> crate::render::Result<Self, D::Error>
190    where
191        D: serde::Deserializer<'de>,
192    {
193        Ok(DeserializeHelperNode::deserialize(deserializer)?.into())
194    }
195}
196
197// Private serialization helper struct
198#[cfg(feature = "serde")]
199#[cfg_attr(feature = "serde", derive(Serialize))]
200struct SerializeHelperNode<'a, T> {
201    #[cfg_attr(feature = "serde", serde(flatten))]
202    pub val: &'a T,
203    #[cfg_attr(feature = "serde", serde(skip_serializing_if = "Vec::is_empty"))]
204    pub children: Vec<SerializeHelperNode<'a, T>>,
205}
206
207#[cfg(feature = "serde")]
208#[cfg_attr(feature = "serde", derive(Deserialize))]
209struct DeserializeHelperNode<T> {
210    #[cfg_attr(feature = "serde", serde(flatten))]
211    pub val: T,
212    #[cfg_attr(feature = "serde", serde(default = "Vec::new"))]
213    pub children: Vec<DeserializeHelperNode<T>>,
214}
215
216#[cfg(feature = "serde")]
217impl<'a, T> From<(&'a Node<T>, &'a Tree<T>)> for SerializeHelperNode<'a, T> {
218    fn from(value: (&'a Node<T>, &'a Tree<T>)) -> SerializeHelperNode<'a, T> {
219        let (node, tree) = value;
220        SerializeHelperNode {
221            val: &node.value,
222            children: node
223                .children
224                .iter()
225                .map(|n| SerializeHelperNode::from((&tree[*n], tree)))
226                .collect::<Vec<_>>(),
227        }
228    }
229}
230
231#[cfg(feature = "serde")]
232impl<T> From<DeserializeHelperNode<T>> for Tree<T> {
233    fn from(value: DeserializeHelperNode<T>) -> Self {
234        let mut nodes_to_add = vec![(0_usize, value, None)];
235        let mut nodes = vec![];
236        let mut id_counter = 1;
237
238        while !nodes_to_add.is_empty() {
239            let mut next_nodes = nodes_to_add
240                .drain(..)
241                .flat_map(|(idx, n, parent)| {
242                    let (value, children) = (n.val, n.children);
243
244                    let child_idx_range = id_counter..(id_counter + children.len());
245                    id_counter += children.len();
246
247                    nodes.push(Node {
248                        idx,
249                        value,
250                        parent,
251                        children: child_idx_range.clone().collect(),
252                    });
253
254                    child_idx_range
255                        .zip(children.into_iter())
256                        .map(move |(child_idx, child_node)| (child_idx, child_node, Some(idx)))
257                })
258                .collect::<Vec<_>>();
259
260            nodes_to_add.append(&mut next_nodes);
261        }
262
263        Tree { nodes }
264    }
265}