| // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| // Copyright by contributors to this project. |
| // SPDX-License-Identifier: (Apache-2.0 OR MIT) |
| |
| use alloc::vec::Vec; |
| use core::{fmt::Debug, hash::Hash}; |
| use mls_rs_codec::{MlsDecode, MlsEncode}; |
| |
| use super::node::LeafIndex; |
| |
| pub trait TreeIndex: |
| Send + Sync + Eq + Clone + Debug + Default + MlsEncode + MlsDecode + Hash + Ord |
| { |
| fn root(&self) -> Self; |
| |
| fn left_unchecked(&self) -> Self; |
| fn right_unchecked(&self) -> Self; |
| |
| fn parent_sibling(&self, leaf_count: &Self) -> Option<ParentSibling<Self>>; |
| fn is_leaf(&self) -> bool; |
| fn is_in_tree(&self, root: &Self) -> bool; |
| |
| #[cfg(any(feature = "secret_tree_access", feature = "private_message"))] |
| fn zero() -> Self; |
| |
| #[cfg(any(feature = "secret_tree_access", feature = "private_message", test))] |
| fn left(&self) -> Option<Self> { |
| (!self.is_leaf()).then(|| self.left_unchecked()) |
| } |
| |
| #[cfg(any(feature = "secret_tree_access", feature = "private_message", test))] |
| fn right(&self) -> Option<Self> { |
| (!self.is_leaf()).then(|| self.right_unchecked()) |
| } |
| |
| fn direct_copath(&self, leaf_count: &Self) -> Vec<CopathNode<Self>> { |
| let root = leaf_count.root(); |
| |
| if !self.is_in_tree(&root) { |
| return Vec::new(); |
| } |
| |
| let mut path = Vec::new(); |
| let mut parent = self.clone(); |
| |
| while let Some(ps) = parent.parent_sibling(leaf_count) { |
| path.push(CopathNode::new(ps.parent.clone(), ps.sibling)); |
| parent = ps.parent; |
| } |
| |
| path |
| } |
| } |
| |
| #[derive(Clone, PartialEq, Eq, Debug)] |
| pub struct CopathNode<T> { |
| pub path: T, |
| pub copath: T, |
| } |
| |
| impl<T: Clone + PartialEq + Eq + core::fmt::Debug> CopathNode<T> { |
| pub fn new(path: T, copath: T) -> CopathNode<T> { |
| CopathNode { path, copath } |
| } |
| } |
| |
| #[derive(Clone, PartialEq, Eq, Debug)] |
| pub struct ParentSibling<T> { |
| pub parent: T, |
| pub sibling: T, |
| } |
| |
| impl<T: Clone + PartialEq + Eq + core::fmt::Debug> ParentSibling<T> { |
| pub fn new(parent: T, sibling: T) -> ParentSibling<T> { |
| ParentSibling { parent, sibling } |
| } |
| } |
| |
| macro_rules! impl_tree_stdint { |
| ($t:ty) => { |
| impl TreeIndex for $t { |
| fn root(&self) -> $t { |
| *self - 1 |
| } |
| |
| /// Panicks if `x` is even in debug, overflows in release. |
| fn left_unchecked(&self) -> Self { |
| *self ^ (0x01 << (self.trailing_ones() - 1)) |
| } |
| |
| /// Panicks if `x` is even in debug, overflows in release. |
| fn right_unchecked(&self) -> Self { |
| *self ^ (0x03 << (self.trailing_ones() - 1)) |
| } |
| |
| fn parent_sibling(&self, leaf_count: &Self) -> Option<ParentSibling<Self>> { |
| if self == &leaf_count.root() { |
| return None; |
| } |
| |
| let lvl = self.trailing_ones(); |
| let p = (self & !(1 << (lvl + 1))) | (1 << lvl); |
| |
| let s = if *self < p { |
| p.right_unchecked() |
| } else { |
| p.left_unchecked() |
| }; |
| |
| Some(ParentSibling::new(p, s)) |
| } |
| |
| fn is_leaf(&self) -> bool { |
| self & 1 == 0 |
| } |
| |
| fn is_in_tree(&self, root: &Self) -> bool { |
| *self <= 2 * root |
| } |
| |
| #[cfg(any(feature = "secret_tree_access", feature = "private_message"))] |
| fn zero() -> Self { |
| 0 |
| } |
| } |
| }; |
| } |
| |
| impl_tree_stdint!(u32); |
| |
| #[cfg(test)] |
| impl_tree_stdint!(u64); |
| |
| pub fn leaf_lca_level(x: u32, y: u32) -> u32 { |
| let mut xn = x; |
| let mut yn = y; |
| let mut k = 0; |
| |
| while xn != yn { |
| xn >>= 1; |
| yn >>= 1; |
| k += 1; |
| } |
| |
| k |
| } |
| |
| pub fn subtree(x: u32) -> (LeafIndex, LeafIndex) { |
| let breadth = 1 << x.trailing_ones(); |
| ( |
| LeafIndex((x + 1 - breadth) >> 1), |
| LeafIndex(((x + breadth) >> 1) + 1), |
| ) |
| } |
| |
| pub struct BfsIterTopDown { |
| level: usize, |
| mask: usize, |
| level_end: usize, |
| ctr: usize, |
| } |
| |
| impl BfsIterTopDown { |
| pub fn new(num_leaves: usize) -> Self { |
| let depth = num_leaves.trailing_zeros() as usize; |
| Self { |
| level: depth + 1, |
| mask: (1 << depth) - 1, |
| level_end: 1, |
| ctr: 0, |
| } |
| } |
| } |
| |
| impl Iterator for BfsIterTopDown { |
| type Item = usize; |
| |
| fn next(&mut self) -> Option<Self::Item> { |
| if self.ctr == self.level_end { |
| if self.level == 1 { |
| return None; |
| } |
| self.level_end = (((self.level_end - 1) << 1) | 1) + 1; |
| self.level -= 1; |
| self.ctr = 0; |
| self.mask >>= 1; |
| } |
| let res = Some((self.ctr << self.level) | self.mask); |
| self.ctr += 1; |
| res |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| use itertools::Itertools; |
| use serde::{Deserialize, Serialize}; |
| |
| #[cfg(target_arch = "wasm32")] |
| use wasm_bindgen_test::wasm_bindgen_test as test; |
| |
| #[derive(Serialize, Deserialize)] |
| struct TestCase { |
| n_leaves: u32, |
| n_nodes: u32, |
| root: u32, |
| left: Vec<Option<u32>>, |
| right: Vec<Option<u32>>, |
| parent: Vec<Option<u32>>, |
| sibling: Vec<Option<u32>>, |
| } |
| |
| pub fn node_width(n: u32) -> u32 { |
| if n == 0 { |
| 0 |
| } else { |
| 2 * (n - 1) + 1 |
| } |
| } |
| |
| #[test] |
| fn test_bfs_iterator() { |
| let expected = [7, 3, 11, 1, 5, 9, 13, 0, 2, 4, 6, 8, 10, 12, 14]; |
| let bfs = BfsIterTopDown::new(8); |
| assert_eq!(bfs.collect::<Vec<_>>(), expected); |
| } |
| |
| #[cfg_attr(coverage_nightly, coverage(off))] |
| fn generate_tree_math_test_cases() -> Vec<TestCase> { |
| let mut test_cases = Vec::new(); |
| |
| for log_n_leaves in 0..8 { |
| let n_leaves = 1 << log_n_leaves; |
| let n_nodes = node_width(n_leaves); |
| let left = (0..n_nodes).map(|x| x.left()).collect::<Vec<_>>(); |
| let right = (0..n_nodes).map(|x| x.right()).collect::<Vec<_>>(); |
| |
| let (parent, sibling) = (0..n_nodes) |
| .map(|x| { |
| x.parent_sibling(&n_leaves) |
| .map(|ps| (ps.parent, ps.sibling)) |
| .unzip() |
| }) |
| .unzip(); |
| |
| test_cases.push(TestCase { |
| n_leaves, |
| n_nodes, |
| root: n_leaves.root(), |
| left, |
| right, |
| parent, |
| sibling, |
| }) |
| } |
| |
| test_cases |
| } |
| |
| fn load_test_cases() -> Vec<TestCase> { |
| load_test_case_json!(tree_math, generate_tree_math_test_cases()) |
| } |
| |
| #[test] |
| fn test_tree_math() { |
| let test_cases = load_test_cases(); |
| |
| for case in test_cases { |
| assert_eq!(node_width(case.n_leaves), case.n_nodes); |
| assert_eq!(case.n_leaves.root(), case.root); |
| |
| for x in 0..case.n_nodes { |
| assert_eq!(x.left(), case.left[x as usize]); |
| assert_eq!(x.right(), case.right[x as usize]); |
| |
| let (p, s) = x |
| .parent_sibling(&case.n_leaves) |
| .map(|ps| (ps.parent, ps.sibling)) |
| .unzip(); |
| |
| assert_eq!(p, case.parent[x as usize]); |
| assert_eq!(s, case.sibling[x as usize]); |
| } |
| } |
| } |
| |
| #[test] |
| fn test_direct_path() { |
| let expected: Vec<Vec<u32>> = [ |
| [0x01, 0x03, 0x07, 0x0f].to_vec(), |
| [0x03, 0x07, 0x0f].to_vec(), |
| [0x01, 0x03, 0x07, 0x0f].to_vec(), |
| [0x07, 0x0f].to_vec(), |
| [0x05, 0x03, 0x07, 0x0f].to_vec(), |
| [0x03, 0x07, 0x0f].to_vec(), |
| [0x05, 0x03, 0x07, 0x0f].to_vec(), |
| [0x0f].to_vec(), |
| [0x09, 0x0b, 0x07, 0x0f].to_vec(), |
| [0x0b, 0x07, 0x0f].to_vec(), |
| [0x09, 0x0b, 0x07, 0x0f].to_vec(), |
| [0x07, 0x0f].to_vec(), |
| [0x0d, 0x0b, 0x07, 0x0f].to_vec(), |
| [0x0b, 0x07, 0x0f].to_vec(), |
| [0x0d, 0x0b, 0x07, 0x0f].to_vec(), |
| [].to_vec(), |
| [0x11, 0x13, 0x17, 0x0f].to_vec(), |
| [0x13, 0x17, 0x0f].to_vec(), |
| [0x11, 0x13, 0x17, 0x0f].to_vec(), |
| [0x17, 0x0f].to_vec(), |
| [0x15, 0x13, 0x17, 0x0f].to_vec(), |
| [0x13, 0x17, 0x0f].to_vec(), |
| [0x15, 0x13, 0x17, 0x0f].to_vec(), |
| [0x0f].to_vec(), |
| [0x19, 0x1b, 0x17, 0x0f].to_vec(), |
| [0x1b, 0x17, 0x0f].to_vec(), |
| [0x19, 0x1b, 0x17, 0x0f].to_vec(), |
| [0x17, 0x0f].to_vec(), |
| [0x1d, 0x1b, 0x17, 0x0f].to_vec(), |
| [0x1b, 0x17, 0x0f].to_vec(), |
| [0x1d, 0x1b, 0x17, 0x0f].to_vec(), |
| ] |
| .to_vec(); |
| |
| for (i, item) in expected.iter().enumerate() { |
| let path = (i as u32) |
| .direct_copath(&16) |
| .into_iter() |
| .map(|cp| cp.path) |
| .collect_vec(); |
| |
| assert_eq!(item, &path) |
| } |
| } |
| |
| #[test] |
| fn test_copath_path() { |
| let expected: Vec<Vec<u32>> = [ |
| [0x02, 0x05, 0x0b, 0x17].to_vec(), |
| [0x05, 0x0b, 0x17].to_vec(), |
| [0x00, 0x05, 0x0b, 0x17].to_vec(), |
| [0x0b, 0x17].to_vec(), |
| [0x06, 0x01, 0x0b, 0x17].to_vec(), |
| [0x01, 0x0b, 0x17].to_vec(), |
| [0x04, 0x01, 0x0b, 0x17].to_vec(), |
| [0x17].to_vec(), |
| [0x0a, 0x0d, 0x03, 0x17].to_vec(), |
| [0x0d, 0x03, 0x17].to_vec(), |
| [0x08, 0x0d, 0x03, 0x17].to_vec(), |
| [0x03, 0x17].to_vec(), |
| [0x0e, 0x09, 0x03, 0x17].to_vec(), |
| [0x09, 0x03, 0x17].to_vec(), |
| [0x0c, 0x09, 0x03, 0x17].to_vec(), |
| [].to_vec(), |
| [0x12, 0x15, 0x1b, 0x07].to_vec(), |
| [0x15, 0x1b, 0x07].to_vec(), |
| [0x10, 0x15, 0x1b, 0x07].to_vec(), |
| [0x1b, 0x07].to_vec(), |
| [0x16, 0x11, 0x1b, 0x07].to_vec(), |
| [0x11, 0x1b, 0x07].to_vec(), |
| [0x14, 0x11, 0x1b, 0x07].to_vec(), |
| [0x07].to_vec(), |
| [0x1a, 0x1d, 0x13, 0x07].to_vec(), |
| [0x1d, 0x13, 0x07].to_vec(), |
| [0x18, 0x1d, 0x13, 0x07].to_vec(), |
| [0x13, 0x07].to_vec(), |
| [0x1e, 0x19, 0x13, 0x07].to_vec(), |
| [0x19, 0x13, 0x07].to_vec(), |
| [0x1c, 0x19, 0x13, 0x07].to_vec(), |
| ] |
| .to_vec(); |
| |
| for (i, item) in expected.iter().enumerate() { |
| let copath = (i as u32) |
| .direct_copath(&16) |
| .into_iter() |
| .map(|cp| cp.copath) |
| .collect_vec(); |
| |
| assert_eq!(item, &copath) |
| } |
| } |
| } |
