arborist/
lib.rs

1pub mod hash;
2
3use bonsai::{
4    children, expand, first_leaf, last_leaf, log2, relative_depth, subtree_index_to_general,
5};
6use hash::{hash, zero_hash};
7use std::cmp::min;
8use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
9use std::convert::From;
10
11pub type K = u128;
12pub type V = [u8; 32];
13
14#[derive(Debug, Default, PartialEq)]
15pub struct Tree {
16    map: BTreeMap<K, V>,
17}
18
19impl Tree {
20    pub fn new() -> Self {
21        Self {
22            map: BTreeMap::new(),
23        }
24    }
25
26    pub fn to_subtree(mut self, root: K) -> Self {
27        self.map = self
28            .map
29            .into_iter()
30            .map(|(k, v)| (subtree_index_to_general(root, k), v))
31            .collect();
32
33        self
34    }
35
36    pub fn get(&self, key: &K) -> Option<&V> {
37        self.map.get(key)
38    }
39
40    pub fn insert(&mut self, key: K, val: V) -> Option<V> {
41        self.map.insert(key, val)
42    }
43
44    pub fn keys(&self) -> BTreeSet<K> {
45        self.map.keys().cloned().collect()
46    }
47
48    pub fn fill_subtree(&mut self, root: K, depth: u32, default: &V) {
49        let mut keys: BinaryHeap<u128> = self
50            .keys()
51            .intersection(&self._leaf_keys(root, depth))
52            .cloned()
53            .collect();
54
55        while let Some(key) = keys.pop() {
56            if key <= root {
57                break;
58            }
59
60            let (left, right, parent) = expand(key);
61
62            if !self.map.contains_key(&parent) {
63                let mut get_or_insert = |n: K| -> V {
64                    *self.map.entry(n).or_insert(zero_hash(
65                        default,
66                        relative_depth(n, first_leaf(root, depth as u128)),
67                    ))
68                };
69
70                let left = get_or_insert(left);
71                let right = get_or_insert(right);
72
73                self.map.insert(parent, hash(&left, &right));
74                keys.push(parent);
75            }
76        }
77    }
78
79    pub fn trim(mut self) -> Self {
80        self._trim();
81        self
82    }
83
84    pub fn raw_insert_bytes(&mut self, rooted_at: K, bytes: Vec<u8>) {
85        let len = bytes.len() as K;
86        let padded_len = len
87            .checked_next_power_of_two()
88            .expect("compiled code to fit in tree");
89        let depth = log2(padded_len / 32);
90        let first: K = first_leaf(rooted_at, depth);
91
92        for i in (0..len).step_by(32) {
93            let begin = i as usize;
94            let end = min(i + 32, len) as usize;
95
96            let chunk_len = if i + 32 < len {
97                32
98            } else {
99                if end % 32 != 0 {
100                    end % 32
101                } else {
102                    32
103                }
104            };
105
106            let mut buf = [0u8; 32];
107            buf[0..chunk_len].copy_from_slice(&bytes[begin..end]);
108
109            self.map.insert(first + (i / 32), buf);
110        }
111    }
112
113    pub fn insert_bytes(&mut self, rooted_at: K, bytes: Vec<u8>) {
114        let len = bytes.len() as K;
115        let padded_len = len
116            .checked_next_power_of_two()
117            .expect("compiled code to fit in tree");
118        let depth = log2(padded_len / 32);
119
120        self.raw_insert_bytes(rooted_at, bytes);
121        self.fill_subtree(rooted_at, depth as u32, &[0; 32]);
122        self._trim();
123    }
124
125    pub fn insert_subtree(&mut self, rooted_at: K, tree: Tree) {
126        for (k, v) in tree.to_subtree(rooted_at).map {
127            self.insert(k, v);
128        }
129    }
130
131    fn _leaf_keys(&self, root: K, depth: u32) -> BTreeSet<K> {
132        (first_leaf(root, depth as u128)..=last_leaf(root, depth as u128)).collect()
133    }
134
135    fn _trim(&mut self) {
136        for key in self.keys() {
137            let (left, right) = children(key);
138
139            if self.map.contains_key(&left) || self.map.contains_key(&right) {
140                self.map.remove(&key);
141            }
142        }
143    }
144}
145
146impl From<Tree> for BTreeMap<K, V> {
147    fn from(tree: Tree) -> BTreeMap<K, V> {
148        tree.map
149    }
150}
151
152impl From<BTreeMap<K, V>> for Tree {
153    fn from(map: BTreeMap<K, V>) -> Tree {
154        Tree { map }
155    }
156}
arborist/lib.rs

arborist/
lib.rs
