openmls/binary_tree/array_representation/
treemath.rs

1use std::cmp::Ordering;
2
3use serde::{Deserialize, Serialize};
4use tls_codec::{TlsDeserialize, TlsDeserializeBytes, TlsSerialize, TlsSize};
5
6pub(crate) const MAX_TREE_SIZE: u32 = (1 << 30) - 1;
7pub(crate) const MIN_TREE_SIZE: u32 = 1;
8
9/// LeafNodeIndex references a leaf node in a tree.
10#[derive(
11    Debug,
12    Clone,
13    Copy,
14    PartialEq,
15    Eq,
16    PartialOrd,
17    Ord,
18    Hash,
19    Serialize,
20    Deserialize,
21    TlsDeserialize,
22    TlsDeserializeBytes,
23    TlsSerialize,
24    TlsSize,
25)]
26pub struct LeafNodeIndex(u32);
27
28impl std::fmt::Display for LeafNodeIndex {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        f.write_fmt(format_args!("{:?}", self.0))
31    }
32}
33
34impl LeafNodeIndex {
35    /// Create a new `LeafNodeIndex` from a `u32`.
36    pub fn new(index: u32) -> Self {
37        LeafNodeIndex(index)
38    }
39
40    /// Return the inner value as `u32`.
41    pub fn u32(&self) -> u32 {
42        self.0
43    }
44
45    /// Return the inner value as `usize`.
46    pub fn usize(&self) -> usize {
47        self.u32() as usize
48    }
49
50    /// Return the index as a TreeNodeIndex value.
51    fn to_tree_index(self) -> u32 {
52        self.0 * 2
53    }
54
55    /// Warning: Only use when the node index represents a leaf node
56    fn from_tree_index(node_index: u32) -> Self {
57        debug_assert!(node_index.is_multiple_of(2));
58        LeafNodeIndex(node_index / 2)
59    }
60}
61
62/// ParentNodeIndex references a parent node in a tree.
63#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
64pub struct ParentNodeIndex(u32);
65
66impl ParentNodeIndex {
67    /// Create a new `ParentNodeIndex` from a `u32`.
68    pub(crate) fn new(index: u32) -> Self {
69        ParentNodeIndex(index)
70    }
71
72    /// Return the inner value as `u32`.
73    pub fn u32(&self) -> u32 {
74        self.0
75    }
76
77    pub(crate) fn usize(&self) -> usize {
78        self.0 as usize
79    }
80
81    /// Return the index as a TreeNodeIndex value.
82    fn to_tree_index(self) -> u32 {
83        self.0 * 2 + 1
84    }
85
86    /// Warning: Only use when the node index represents a parent node
87    fn from_tree_index(node_index: u32) -> Self {
88        debug_assert!(node_index > 0);
89        debug_assert!(node_index % 2 == 1);
90        ParentNodeIndex((node_index - 1) / 2)
91    }
92}
93
94#[cfg(test)]
95impl ParentNodeIndex {
96    /// Re-exported for testing.
97    pub(crate) fn test_from_tree_index(node_index: u32) -> Self {
98        Self::from_tree_index(node_index)
99    }
100}
101
102#[cfg(any(feature = "test-utils", test))]
103impl ParentNodeIndex {
104    /// Re-exported for testing.
105    pub(crate) fn test_to_tree_index(self) -> u32 {
106        self.to_tree_index()
107    }
108}
109
110impl From<LeafNodeIndex> for TreeNodeIndex {
111    fn from(leaf_index: LeafNodeIndex) -> Self {
112        TreeNodeIndex::Leaf(leaf_index)
113    }
114}
115
116impl From<ParentNodeIndex> for TreeNodeIndex {
117    fn from(parent_index: ParentNodeIndex) -> Self {
118        TreeNodeIndex::Parent(parent_index)
119    }
120}
121
122/// TreeNodeIndex references a node in a tree.
123#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
124pub enum TreeNodeIndex {
125    Leaf(LeafNodeIndex),
126    Parent(ParentNodeIndex),
127}
128
129impl TreeNodeIndex {
130    /// Create a new `TreeNodeIndex` from a `u32`.
131    fn new(index: u32) -> Self {
132        if index.is_multiple_of(2) {
133            TreeNodeIndex::Leaf(LeafNodeIndex::from_tree_index(index))
134        } else {
135            TreeNodeIndex::Parent(ParentNodeIndex::from_tree_index(index))
136        }
137    }
138
139    /// Re-exported for testing.
140    #[cfg(any(feature = "test-utils", test))]
141    pub(crate) fn test_new(index: u32) -> Self {
142        Self::new(index)
143    }
144
145    /// Return the inner value as `u32`.
146    fn u32(&self) -> u32 {
147        match self {
148            TreeNodeIndex::Leaf(index) => index.to_tree_index(),
149            TreeNodeIndex::Parent(index) => index.to_tree_index(),
150        }
151    }
152
153    /// Re-exported for testing.
154    #[cfg(any(feature = "test-utils", feature = "crypto-debug", test))]
155    pub(crate) fn test_u32(&self) -> u32 {
156        self.u32()
157    }
158
159    /// Return the inner value as `usize`.
160    #[cfg(any(feature = "test-utils", test))]
161    fn usize(&self) -> usize {
162        self.u32() as usize
163    }
164
165    /// Re-exported for testing.
166    #[cfg(any(feature = "test-utils", test))]
167    pub(crate) fn test_usize(&self) -> usize {
168        self.usize()
169    }
170}
171
172impl Ord for TreeNodeIndex {
173    fn cmp(&self, other: &TreeNodeIndex) -> Ordering {
174        self.u32().cmp(&other.u32())
175    }
176}
177
178impl PartialOrd for TreeNodeIndex {
179    fn partial_cmp(&self, other: &TreeNodeIndex) -> Option<Ordering> {
180        Some(self.cmp(other))
181    }
182}
183
184#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
185pub(crate) struct TreeSize(u32);
186
187impl TreeSize {
188    /// Create a new `TreeSize` from `nodes`, which will be rounded up to the
189    /// next power of 2. The tree size then reflects the smallest tree that can
190    /// contain the number of nodes.
191    pub(crate) fn new(nodes: u32) -> Self {
192        let k = log2(nodes);
193        TreeSize((1 << (k + 1)) - 1)
194    }
195
196    /// Creates a new `TreeSize` from a specific leaf count
197    #[cfg(any(feature = "test-utils", feature = "extensions-draft", test))]
198    pub(crate) fn from_leaf_count(leaf_count: u32) -> Self {
199        TreeSize::new(leaf_count * 2)
200    }
201
202    /// Return the number of leaf nodes in the tree.
203    pub(crate) fn leaf_count(&self) -> u32 {
204        (self.0 / 2) + 1
205    }
206
207    /// Return the number of parent nodes in the tree.
208    pub(crate) fn parent_count(&self) -> u32 {
209        self.0 / 2
210    }
211
212    /// Return the inner value as `u32`.
213    pub(crate) fn u32(&self) -> u32 {
214        self.0
215    }
216
217    /// Returns `true` if the leaf is in the left subtree and `false` otherwise.
218    /// If there is only one leaf in the tree, it returns `false`.
219    pub(crate) fn leaf_is_left(&self, leaf_index: LeafNodeIndex) -> bool {
220        leaf_index.u32() < self.leaf_count() / 2
221    }
222
223    /// Increase the size.
224    pub(super) fn inc(&mut self) {
225        self.0 = self.0 * 2 + 1;
226    }
227
228    /// Decrease the size.
229    pub(super) fn dec(&mut self) {
230        debug_assert!(self.0 >= 2);
231        if self.0 >= 2 {
232            self.0 = self.0.div_ceil(2) - 1;
233        } else {
234            self.0 = 0;
235        }
236    }
237}
238
239#[test]
240fn tree_size() {
241    assert_eq!(TreeSize::new(1).u32(), 1);
242    assert_eq!(TreeSize::new(3).u32(), 3);
243    assert_eq!(TreeSize::new(5).u32(), 7);
244    assert_eq!(TreeSize::new(7).u32(), 7);
245    assert_eq!(TreeSize::new(9).u32(), 15);
246    assert_eq!(TreeSize::new(11).u32(), 15);
247    assert_eq!(TreeSize::new(13).u32(), 15);
248    assert_eq!(TreeSize::new(15).u32(), 15);
249    assert_eq!(TreeSize::new(17).u32(), 31);
250}
251
252/// Test if the leaf is in the left subtree.
253#[test]
254fn test_leaf_is_left() {
255    assert!(!TreeSize::new(1).leaf_is_left(LeafNodeIndex::new(0)));
256
257    assert!(TreeSize::new(3).leaf_is_left(LeafNodeIndex::new(0)));
258    assert!(!TreeSize::new(3).leaf_is_left(LeafNodeIndex::new(1)));
259
260    assert!(TreeSize::new(5).leaf_is_left(LeafNodeIndex::new(0)));
261    assert!(TreeSize::new(5).leaf_is_left(LeafNodeIndex::new(1)));
262    assert!(!TreeSize::new(5).leaf_is_left(LeafNodeIndex::new(2)));
263    assert!(!TreeSize::new(5).leaf_is_left(LeafNodeIndex::new(3)));
264
265    assert!(TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(0)));
266    assert!(TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(1)));
267    assert!(TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(2)));
268    assert!(TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(3)));
269    assert!(!TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(4)));
270    assert!(!TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(5)));
271    assert!(!TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(6)));
272    assert!(!TreeSize::new(15).leaf_is_left(LeafNodeIndex::new(7)));
273}
274
275fn log2(x: u32) -> usize {
276    if x == 0 {
277        return 0;
278    }
279    (31 - x.leading_zeros()) as usize
280}
281
282pub fn level(index: u32) -> usize {
283    // The cast is always valid, as there is at most 32 trailing ones
284    index.trailing_ones() as usize
285}
286
287pub(crate) fn root(size: TreeSize) -> TreeNodeIndex {
288    let size = size.u32();
289    debug_assert!(size > 0);
290    TreeNodeIndex::new((1 << log2(size)) - 1)
291}
292
293pub(crate) fn left(index: ParentNodeIndex) -> TreeNodeIndex {
294    let x = index.to_tree_index();
295    let k = level(x);
296    debug_assert!(k > 0);
297    let index = x ^ (0x01 << (k - 1));
298    TreeNodeIndex::new(index)
299}
300
301pub(crate) fn right(index: ParentNodeIndex) -> TreeNodeIndex {
302    let x = index.to_tree_index();
303    let k = level(x);
304    debug_assert!(k > 0);
305    let index = x ^ (0x03 << (k - 1));
306    TreeNodeIndex::new(index)
307}
308
309/// Warning: There is no check about the tree size and whether the parent is
310/// beyond the root
311fn parent(x: TreeNodeIndex) -> ParentNodeIndex {
312    let x = x.u32();
313    let k = level(x);
314    let b = (x >> (k + 1)) & 0x01;
315    let index = (x | (1 << k)) ^ (b << (k + 1));
316    ParentNodeIndex::from_tree_index(index)
317}
318
319/// Re-exported for testing.
320#[cfg(any(feature = "test-utils", test))]
321pub(crate) fn test_parent(index: TreeNodeIndex) -> ParentNodeIndex {
322    parent(index)
323}
324
325fn sibling(index: TreeNodeIndex) -> TreeNodeIndex {
326    let p = parent(index);
327    match index.u32().cmp(&p.to_tree_index()) {
328        Ordering::Less => right(p),
329        Ordering::Greater => left(p),
330        Ordering::Equal => left(p),
331    }
332}
333
334/// Re-exported for testing.
335#[cfg(any(feature = "test-utils", test))]
336pub(crate) fn test_sibling(index: TreeNodeIndex) -> TreeNodeIndex {
337    sibling(index)
338}
339
340/// Direct path from a node to the root.
341/// Does not include the node itself.
342pub(crate) fn direct_path(node_index: LeafNodeIndex, size: TreeSize) -> Vec<ParentNodeIndex> {
343    let r = root(size).u32();
344
345    let mut d = vec![];
346    let mut x = node_index.to_tree_index();
347    while x != r {
348        let parent = parent(TreeNodeIndex::new(x));
349        d.push(parent);
350        x = parent.to_tree_index();
351    }
352    d
353}
354
355/// Copath of a leaf node.
356pub(crate) fn copath(leaf_index: LeafNodeIndex, size: TreeSize) -> Vec<TreeNodeIndex> {
357    // Start with leaf
358    let mut full_path = vec![TreeNodeIndex::Leaf(leaf_index)];
359    let mut direct_path = direct_path(leaf_index, size);
360    if !direct_path.is_empty() {
361        // Remove root
362        direct_path.pop();
363    }
364    full_path.append(
365        &mut direct_path
366            .iter()
367            .map(|i| TreeNodeIndex::Parent(*i))
368            .collect(),
369    );
370
371    full_path.into_iter().map(sibling).collect()
372}
373
374/// Common ancestor of two leaf nodes, aka the node where their direct paths
375/// intersect.
376pub(super) fn lowest_common_ancestor(x: LeafNodeIndex, y: LeafNodeIndex) -> ParentNodeIndex {
377    let x = x.to_tree_index();
378    let y = y.to_tree_index();
379    let (lx, ly) = (level(x) + 1, level(y) + 1);
380    if (lx <= ly) && (x >> ly == y >> ly) {
381        return ParentNodeIndex::from_tree_index(y);
382    } else if (ly <= lx) && (x >> lx == y >> lx) {
383        return ParentNodeIndex::from_tree_index(x);
384    }
385
386    let (mut xn, mut yn) = (x, y);
387    let mut k = 0;
388    while xn != yn {
389        xn >>= 1;
390        yn >>= 1;
391        k += 1;
392    }
393    ParentNodeIndex::from_tree_index((xn << k) + (1 << (k - 1)) - 1)
394}
395
396/// The common direct path of two leaf nodes, i.e. the path from their common
397/// ancestor to the root.
398pub(crate) fn common_direct_path(
399    x: LeafNodeIndex,
400    y: LeafNodeIndex,
401    size: TreeSize,
402) -> Vec<ParentNodeIndex> {
403    let mut x_path = direct_path(x, size);
404    let mut y_path = direct_path(y, size);
405    x_path.reverse();
406    y_path.reverse();
407
408    let mut common_path = vec![];
409
410    for (x, y) in x_path.iter().zip(y_path.iter()) {
411        if x == y {
412            common_path.push(*x);
413        } else {
414            break;
415        }
416    }
417
418    common_path.reverse();
419    common_path
420}
421
422#[cfg(any(feature = "test-utils", test))]
423pub(crate) fn node_width(n: usize) -> usize {
424    if n == 0 {
425        0
426    } else {
427        2 * (n - 1) + 1
428    }
429}
430
431pub(crate) fn is_node_in_tree(node_index: TreeNodeIndex, size: TreeSize) -> bool {
432    node_index.u32() < size.u32()
433}
434
435#[test]
436fn test_node_in_tree() {
437    let tests = [(0u32, 3u32), (1, 3), (2, 5), (5, 7), (2, 11)];
438    for test in tests.iter() {
439        assert!(is_node_in_tree(
440            TreeNodeIndex::new(test.0),
441            TreeSize::new(test.1)
442        ));
443    }
444}
445
446#[test]
447fn test_node_not_in_tree() {
448    let tests = [(3u32, 1u32), (13, 7)];
449    for test in tests.iter() {
450        assert!(!is_node_in_tree(
451            TreeNodeIndex::new(test.0),
452            TreeSize::new(test.1)
453        ));
454    }
455}
openmls/binary_tree/array_representation/treemath.rs

openmls/binary_tree/array_representation/
treemath.rs
