rostl_oram/
heap_tree.rs

1//! Represents a heap tree as an array and provides functions to access it.
2//!
3//! The tree is represented by a reverse lexicographical order binary tree.
4//!                            0
5//!                     1             2
6//!                3         5    4       6
7//! Path:          0         2    1       3
8
9use crate::prelude::PositionType;
10
11/// Represents a heap tree structure.
12#[derive(Debug)]
13pub struct HeapTree<T> {
14  /// Actual storage container
15  pub(crate) tree: Vec<T>,
16  /// Height of the tree, public (tree with a single element has height 1)
17  pub height: usize,
18}
19
20impl<T> HeapTree<T>
21where
22  T: Default + Clone,
23{
24  /// Initialized a new heap tree with a certain height
25  pub fn new(height: usize) -> Self {
26    let tree = vec![T::default(); 2usize.pow(height as u32) - 1];
27    Self { tree, height }
28  }
29}
30
31impl<T> HeapTree<T>
32where
33  T: Clone,
34{
35  /// Initialized a new heap tree with a certain height and a default value
36  pub fn new_with(height: usize, default: T) -> Self {
37    let tree = vec![default; 2usize.pow(height as u32) - 1];
38    Self { tree, height }
39  }
40}
41
42impl<T> HeapTree<T> {
43  /// Get the index of a node at a certain depth and path
44  #[inline]
45  pub fn get_index(&self, depth: usize, path: PositionType) -> usize {
46    debug_assert!(depth < self.height);
47    let level_offset = (1 << depth) - 1;
48    let mask = level_offset as PositionType;
49    level_offset + (path & mask) as usize
50  }
51
52  /// Get a node of a certain path at a certain depth
53  /// Reveals depth and path
54  #[inline]
55  pub fn get_path_at_depth(&self, depth: usize, path: PositionType) -> &T {
56    let index = self.get_index(depth, path);
57    // UNDONE(git-10): Make sure this doesn't have bounds checking and is safe
58    &self.tree[index]
59  }
60
61  /// Get a node of a certain path at a certain depth
62  /// Reveals depth and path
63  #[inline]
64  pub fn get_path_at_depth_mut(&mut self, depth: usize, path: PositionType) -> &mut T {
65    let index = self.get_index(depth, path);
66
67    // UNDONE(git-10): Make sure this doesn't have bounds checking and is safe
68    &mut self.tree[index]
69  }
70
71  /// Given a path and a node at certain depth, return the other child of that node's parent.
72  pub fn get_sibling(&self, depth: usize, path: PositionType) -> &T {
73    let new_path = path ^ (1 << (depth - 1));
74    self.get_path_at_depth(depth, new_path)
75  }
76}
77
78#[cfg(test)]
79mod tests {
80  use crate::prelude::PositionType;
81
82  fn print_depth_pos_index(height: usize, depth: usize, path: PositionType) {
83    debug_assert!(depth < height);
84    let level_offset = (1 << depth) - 1;
85    let mask = level_offset as PositionType;
86    let _ret = level_offset + (path & mask) as usize;
87  }
88  #[test]
89  fn print_heap_tree_info() {
90    for depth in 0..3 {
91      for path in 0..4 {
92        print_depth_pos_index(3, depth, path);
93      }
94    }
95  }
96}