columnar/adts/
tree.rs

1//! Columnar representations of tree-structured data.
2//!
3//! A `Tree<T>` is a node with a value of type `T` and a list of children.
4//! `Trees<TC>` stores a collection of trees with columnar storage for node values.
5use alloc::{vec::Vec, string::String};
6
7use crate::{Borrow, Index, IndexAs, Len, Clear, Push};
8
9/// A tree node with a value and children.
10#[derive(Clone)]
11pub struct Tree<T> {
12    pub data: T,
13    pub kids: Vec<Tree<T>>,
14}
15
16impl Tree<usize> {
17    pub fn sum(&self) -> usize {
18        self.data + self.kids.iter().map(|x| x.sum()).sum::<usize>()
19    }
20}
21
22/// A stand-in for `Vec<Tree<T>>`, with columnar storage for node values.
23///
24/// Nodes are stored in BFS order. `groups` tracks tree boundaries (cumulative
25/// node counts with a leading 0). `bounds` tracks child ranges per node
26/// (also with a leading 0 sentinel).
27#[derive(Copy, Clone)]
28pub struct Trees<TC, BC = Vec<u64>> {
29    /// Cumulative node counts: tree `i` starts at node `groups[i]`.
30    pub groups: BC,
31    /// Child delimiters: node `j`'s children are at `bounds[j]..bounds[j+1]`
32    /// (root nodes use `bounds[j]+1..bounds[j+1]` to skip themselves).
33    pub bounds: BC,
34    /// Columnar container for node values in BFS order.
35    pub values: TC,
36}
37
38impl<TC: Default> Default for Trees<TC> {
39    fn default() -> Self {
40        Self {
41            groups: vec![0u64],
42            bounds: vec![0u64],
43            values: TC::default(),
44        }
45    }
46}
47
48/// A reference to a single node within a `Trees` container.
49///
50/// Holds a copy of the borrowed values and bounds containers,
51/// plus the node's index and child range. Navigation to children
52/// constructs new `TreesRef` values.
53pub struct TreesRef<V, B> {
54    index: usize,
55    lower: usize,
56    upper: usize,
57    values: V,
58    bounds: B,
59}
60
61impl<V: Copy, B: Copy> Clone for TreesRef<V, B> {
62    fn clone(&self) -> Self { *self }
63}
64impl<V: Copy, B: Copy> Copy for TreesRef<V, B> {}
65
66impl<V: Index, B: IndexAs<u64>> TreesRef<V, B> {
67    /// The value at this node.
68    #[inline(always)]
69    pub fn value(&self) -> V::Ref {
70        self.values.get(self.index)
71    }
72    /// The number of children of this node.
73    #[inline(always)]
74    pub fn kids(&self) -> usize {
75        self.upper - self.lower
76    }
77}
78impl<V: Index + Copy, B: IndexAs<u64> + Copy> TreesRef<V, B> {
79    /// A reference to the `index`-th child of this node.
80    #[inline(always)]
81    pub fn child(&self, index: usize) -> Self {
82        assert!(index < self.upper - self.lower);
83        let child = self.lower + index;
84        TreesRef {
85            index: child,
86            lower: self.bounds.index_as(child) as usize,
87            upper: self.bounds.index_as(child + 1) as usize,
88            values: self.values,
89            bounds: self.bounds,
90        }
91    }
92}
93
94impl<TC, BC: Len> Len for Trees<TC, BC> {
95    #[inline(always)]
96    fn len(&self) -> usize { self.groups.len() - 1 }
97}
98
99impl<TC: Index + Copy, BC: IndexAs<u64> + Len + Copy> Index for Trees<TC, BC> {
100    type Ref = TreesRef<TC, BC>;
101    #[inline(always)]
102    fn get(&self, index: usize) -> Self::Ref {
103        let root = self.groups.index_as(index) as usize;
104        TreesRef {
105            index: root,
106            lower: self.bounds.index_as(root) as usize + 1,
107            upper: self.bounds.index_as(root + 1) as usize,
108            values: self.values,
109            bounds: self.bounds,
110        }
111    }
112}
113
114impl<'a, TC, BC: IndexAs<u64> + Len> Index for &'a Trees<TC, BC>
115where
116    &'a TC: Index,
117    &'a BC: IndexAs<u64>,
118{
119    type Ref = TreesRef<&'a TC, &'a BC>;
120    #[inline(always)]
121    fn get(&self, index: usize) -> Self::Ref {
122        let root = self.groups.index_as(index) as usize;
123        TreesRef {
124            index: root,
125            lower: self.bounds.index_as(root) as usize + 1,
126            upper: self.bounds.index_as(root + 1) as usize,
127            values: &self.values,
128            bounds: &self.bounds,
129        }
130    }
131}
132
133impl<TC: Borrow> Borrow for Trees<TC> {
134    type Ref<'a> = TreesRef<TC::Borrowed<'a>, &'a [u64]> where TC: 'a;
135    type Borrowed<'a> = Trees<TC::Borrowed<'a>, &'a [u64]> where TC: 'a;
136    #[inline(always)]
137    fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
138        Trees {
139            groups: &self.groups[..],
140            bounds: &self.bounds[..],
141            values: self.values.borrow(),
142        }
143    }
144    #[inline(always)]
145    fn reborrow<'b, 'a: 'b>(thing: Self::Borrowed<'a>) -> Self::Borrowed<'b> where TC: 'a {
146        Trees {
147            groups: thing.groups,
148            bounds: thing.bounds,
149            values: TC::reborrow(thing.values),
150        }
151    }
152    #[inline(always)]
153    fn reborrow_ref<'b, 'a: 'b>(thing: Self::Ref<'a>) -> Self::Ref<'b> where Self: 'a {
154        TreesRef {
155            index: thing.index,
156            lower: thing.lower,
157            upper: thing.upper,
158            values: TC::reborrow(thing.values),
159            bounds: thing.bounds,
160        }
161    }
162}
163
164impl<TC: Clear> Clear for Trees<TC> {
165    fn clear(&mut self) {
166        self.groups.clear();
167        self.groups.push(0u64);
168        self.bounds.clear();
169        self.bounds.push(0u64);
170        self.values.clear();
171    }
172}
173
174impl<TC: Len> Trees<TC> {
175    /// Pushes a tree into the container, storing nodes in BFS order.
176    pub fn push_tree<T>(&mut self, tree: Tree<T>) where TC: for<'a> Push<&'a T> {
177        let mut todo = alloc::collections::VecDeque::default();
178        todo.push_back(tree);
179        while let Some(node) = todo.pop_front() {
180            let cursor = self.values.len() + todo.len() + 1;
181            self.values.push(&node.data);
182            self.bounds.push((cursor + node.kids.len()) as u64);
183            for child in node.kids.into_iter() {
184                todo.push_back(child);
185            }
186        }
187        self.groups.push(self.values.len() as u64);
188    }
189}
190
191impl<'a, TC: crate::AsBytes<'a>, BC: crate::AsBytes<'a>> crate::AsBytes<'a> for Trees<TC, BC> {
192    const SLICE_COUNT: usize = BC::SLICE_COUNT + BC::SLICE_COUNT + TC::SLICE_COUNT;
193    #[inline]
194    fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
195        debug_assert!(index < Self::SLICE_COUNT);
196        if index < BC::SLICE_COUNT {
197            self.groups.get_byte_slice(index)
198        } else if index < BC::SLICE_COUNT + BC::SLICE_COUNT {
199            self.bounds.get_byte_slice(index - BC::SLICE_COUNT)
200        } else {
201            self.values.get_byte_slice(index - BC::SLICE_COUNT - BC::SLICE_COUNT)
202        }
203    }
204}
205
206impl<'a, TC: crate::FromBytes<'a>, BC: crate::FromBytes<'a>> crate::FromBytes<'a> for Trees<TC, BC> {
207    const SLICE_COUNT: usize = BC::SLICE_COUNT + BC::SLICE_COUNT + TC::SLICE_COUNT;
208    #[inline(always)]
209    fn from_bytes(bytes: &mut impl Iterator<Item=&'a [u8]>) -> Self {
210        Self {
211            groups: crate::FromBytes::from_bytes(bytes),
212            bounds: crate::FromBytes::from_bytes(bytes),
213            values: crate::FromBytes::from_bytes(bytes),
214        }
215    }
216    #[inline(always)]
217    fn from_store(store: &crate::bytes::indexed::DecodedStore<'a>, offset: &mut usize) -> Self {
218        Self {
219            groups: BC::from_store(store, offset),
220            bounds: BC::from_store(store, offset),
221            values: TC::from_store(store, offset),
222        }
223    }
224    #[inline(always)]
225    fn element_sizes(sizes: &mut Vec<usize>) -> Result<(), String> {
226        BC::element_sizes(sizes)?;
227        BC::element_sizes(sizes)?;
228        TC::element_sizes(sizes)?;
229        Ok(())
230    }
231}
232
233/// LOUDS (level ordered unary degree sequence) is a succinct tree representation.
234mod louds {
235
236    // The tree is encoded by traversing it in a BFS order, with each node producing
237    // as many `1`s as it has children, followed by a `0`. There is also a `1` at the
238    // beginning of the sequence, for the root.
239    //
240    // The logic to determine the child of a node is as follows:
241    //
242    //      child(x, i) = select0(rank1(x) + i − 1) + 1
243    //
244    // It is possible that `i` here starts at 1, which we should fix to be `0`.
245
246}
247
248#[cfg(test)]
249mod test {
250
251    use alloc::{vec, vec::Vec, string::ToString};
252    use crate::common::{Index, Len, Clear};
253    use crate::{Borrow, AsBytes, FromBytes};
254    use super::{Tree, Trees};
255
256    fn leaf<T>(data: T) -> Tree<T> {
257        Tree { data, kids: vec![] }
258    }
259    fn branch<T>(data: T, kids: Vec<Tree<T>>) -> Tree<T> {
260        Tree { data, kids }
261    }
262
263    #[test]
264    fn push_and_index() {
265        let mut trees: Trees<Vec<u64>> = Default::default();
266        // Tree: 10 -> [20, 30 -> [40]]
267        let tree = branch(10u64, vec![leaf(20), branch(30, vec![leaf(40)])]);
268        trees.push_tree(tree);
269
270        assert_eq!(trees.len(), 1);
271
272        let borrowed = trees.borrow();
273        let root = borrowed.get(0);
274        assert_eq!(*root.value(), 10);
275        assert_eq!(root.kids(), 2);
276
277        let c0 = root.child(0);
278        assert_eq!(*c0.value(), 20);
279        assert_eq!(c0.kids(), 0);
280
281        let c1 = root.child(1);
282        assert_eq!(*c1.value(), 30);
283        assert_eq!(c1.kids(), 1);
284
285        let c1_0 = c1.child(0);
286        assert_eq!(*c1_0.value(), 40);
287        assert_eq!(c1_0.kids(), 0);
288    }
289
290    #[test]
291    fn multiple_trees() {
292        let mut trees: Trees<Vec<u64>> = Default::default();
293        trees.push_tree(branch(1u64, vec![leaf(2), leaf(3)]));
294        trees.push_tree(leaf(100u64));
295        trees.push_tree(branch(200u64, vec![leaf(300)]));
296
297        assert_eq!(trees.len(), 3);
298
299        let borrowed = trees.borrow();
300
301        let t0 = borrowed.get(0);
302        assert_eq!(*t0.value(), 1);
303        assert_eq!(t0.kids(), 2);
304        assert_eq!(*t0.child(0).value(), 2);
305        assert_eq!(*t0.child(1).value(), 3);
306
307        let t1 = borrowed.get(1);
308        assert_eq!(*t1.value(), 100);
309        assert_eq!(t1.kids(), 0);
310
311        let t2 = borrowed.get(2);
312        assert_eq!(*t2.value(), 200);
313        assert_eq!(t2.kids(), 1);
314        assert_eq!(*t2.child(0).value(), 300);
315    }
316
317    #[test]
318    fn ref_index() {
319        let mut trees: Trees<Vec<u64>> = Default::default();
320        trees.push_tree(branch(1u64, vec![leaf(2)]));
321
322        let root = (&trees).get(0);
323        assert_eq!(*root.value(), 1);
324        assert_eq!(*root.child(0).value(), 2);
325    }
326
327    #[test]
328    fn clear_and_reuse() {
329        let mut trees: Trees<Vec<u64>> = Default::default();
330        trees.push_tree(leaf(42u64));
331        assert_eq!(trees.len(), 1);
332
333        trees.clear();
334        assert_eq!(trees.len(), 0);
335
336        trees.push_tree(leaf(99u64));
337        assert_eq!(trees.len(), 1);
338        assert_eq!(*trees.borrow().get(0).value(), 99);
339    }
340
341    #[test]
342    fn as_from_bytes() {
343        let mut trees: Trees<Vec<u64>> = Default::default();
344        trees.push_tree(branch(10u64, vec![leaf(20), branch(30, vec![leaf(40)])]));
345        trees.push_tree(leaf(100u64));
346
347        let borrowed = trees.borrow();
348        let rebuilt = Trees::<&[u64], &[u64]>::from_bytes(
349            &mut borrowed.as_bytes().map(|(_, bytes)| bytes)
350        );
351        assert_eq!(rebuilt.len(), 2);
352
353        let root = rebuilt.get(0);
354        assert_eq!(*root.value(), 10);
355        assert_eq!(root.kids(), 2);
356        assert_eq!(*root.child(0).value(), 20);
357        assert_eq!(*root.child(1).value(), 30);
358        assert_eq!(*root.child(1).child(0).value(), 40);
359
360        let t1 = rebuilt.get(1);
361        assert_eq!(*t1.value(), 100);
362        assert_eq!(t1.kids(), 0);
363    }
364
365    #[test]
366    fn columnar_strings() {
367        use crate::Strings;
368
369        let mut trees: Trees<Strings> = Default::default();
370        trees.push_tree(branch(
371            "root".to_string(),
372            vec![leaf("left".to_string()), leaf("right".to_string())],
373        ));
374
375        let borrowed = trees.borrow();
376        let root = borrowed.get(0);
377        assert_eq!(root.value(), b"root");
378        assert_eq!(root.kids(), 2);
379        assert_eq!(root.child(0).value(), b"left");
380        assert_eq!(root.child(1).value(), b"right");
381    }
382
383    #[test]
384    fn deep_tree() {
385        let mut trees: Trees<Vec<u64>> = Default::default();
386        // Build a chain: 0 -> 1 -> 2 -> 3 -> 4
387        let mut tree = leaf(4u64);
388        for i in (0..4).rev() {
389            tree = branch(i, vec![tree]);
390        }
391        trees.push_tree(tree);
392
393        let borrowed = trees.borrow();
394        let mut node = borrowed.get(0);
395        for i in 0..5u64 {
396            assert_eq!(*node.value(), i);
397            if i < 4 {
398                assert_eq!(node.kids(), 1);
399                node = node.child(0);
400            } else {
401                assert_eq!(node.kids(), 0);
402            }
403        }
404    }
405}
columnar/adts/tree.rs

columnar/adts/
tree.rs
