scorex_crypto_avltree/
batch_node.rs

1use crate::operation::*;
2use blake2::digest::{Update, VariableOutput};
3use blake2::VarBlake2b;
4use bytes::{Buf, BufMut, Bytes, BytesMut};
5use std::cell::RefCell;
6//use debug_cell::RefCell;
7use std::cmp::Ordering;
8use std::fmt;
9use std::rc::Rc;
10
11// Do not use bytes -1, 0, or 1 -- these are for balance
12pub(crate) const LEAF_IN_PACKAGED_PROOF: u8 = 2;
13pub(crate) const LABEL_IN_PACKAGED_PROOF: u8 = 3;
14pub(crate) const END_OF_TREE_IN_PACKAGED_PROOF: u8 = 4;
15
16pub type Balance = i8;
17pub type SerializedAdProof = Bytes;
18pub type NodeId = Rc<RefCell<Node>>;
19pub type Resolver = fn(&Digest32) -> Node;
20
21#[derive(Debug, Clone)]
22pub struct NodeHeader {
23    pub visited: bool,
24    pub is_new: bool,
25    pub label: Option<Digest32>,
26    pub key: Option<ADKey>,
27}
28
29#[derive(Debug, Clone)]
30pub struct InternalNode {
31    pub hdr: NodeHeader,
32    pub balance: Balance,
33    pub left: NodeId,
34    pub right: NodeId,
35}
36
37#[derive(Debug, Clone)]
38pub struct LeafNode {
39    pub hdr: NodeHeader,
40    pub value: ADValue,
41    pub next_node_key: ADKey,
42}
43
44#[derive(Debug, Clone)]
45pub enum Node {
46    LabelOnly(NodeHeader),
47    Internal(InternalNode),
48    Leaf(LeafNode),
49}
50
51const INTERNAL_NODE_PREFIX: u8 = 0;
52const LEAF_NODE_PREFIX: u8 = 1;
53
54impl Node {
55    pub fn visited(&self) -> bool {
56        self.hdr().visited
57    }
58
59    pub fn reset(&mut self) -> bool {
60        let hdr = self.hdr_mut();
61        let was_new = hdr.is_new;
62        hdr.is_new = false;
63        hdr.visited = false;
64        was_new
65    }
66
67    pub fn mark_visited(&mut self, visited: bool) {
68        let hdr = self.hdr_mut();
69        hdr.visited = visited;
70    }
71
72    pub fn is_new(&self) -> bool {
73        self.hdr().is_new
74    }
75
76    pub fn get_label(&self) -> Digest32 {
77        self.hdr().label.unwrap()
78    }
79
80    pub fn label(&mut self) -> Digest32 {
81        if let Some(label) = self.hdr().label {
82            return label;
83        }
84        match self {
85            Node::LabelOnly(hdr) => hdr.label.unwrap(),
86            Node::Leaf(node) => {
87                let mut hasher = VarBlake2b::new(32).unwrap();
88                hasher.update(&[0u8; 1]);
89                hasher.update(&node.hdr.key.as_ref().unwrap()[..]);
90                hasher.update(&node.value[..]);
91                hasher.update(&node.next_node_key[..]);
92                let mut label: Digest32 = Default::default();
93                label.copy_from_slice(&hasher.finalize_boxed());
94                node.hdr.label = Some(label);
95                label
96            }
97            Node::Internal(node) => {
98                let mut hasher = VarBlake2b::new(32).unwrap();
99                hasher.update(&[1u8; 1]);
100                hasher.update(&[node.balance as u8; 1]);
101                hasher.update(node.left.borrow_mut().label());
102                hasher.update(node.right.borrow_mut().label());
103                let mut label: Digest32 = Default::default();
104                label.copy_from_slice(&hasher.finalize_boxed());
105                node.hdr.label = Some(label);
106                label
107            }
108        }
109    }
110
111    pub fn left(&self) -> NodeId {
112        if let Node::Internal(node) = self {
113            node.left.clone()
114        } else {
115            panic!("not internal node");
116        }
117    }
118
119    pub fn right(&self) -> NodeId {
120        if let Node::Internal(node) = self {
121            node.right.clone()
122        } else {
123            panic!("not internal node");
124        }
125    }
126    pub fn balance(&self) -> Balance {
127        if let Node::Internal(node) = self {
128            node.balance
129        } else {
130            panic!("not internal node");
131        }
132    }
133    pub fn value(&self) -> ADValue {
134        if let Node::Leaf(node) = self {
135            node.value.clone()
136        } else {
137            panic!("not leaf node");
138        }
139    }
140    pub fn key(&self) -> ADValue {
141        let hdr = self.hdr();
142        hdr.key.as_ref().unwrap().clone()
143    }
144    pub fn next_node_key(&self) -> ADKey {
145        if let Node::Leaf(node) = self {
146            node.next_node_key.clone()
147        } else {
148            panic!("not leaf node");
149        }
150    }
151    pub fn is_leaf(&self) -> bool {
152        match self {
153            Node::Leaf(_) => true,
154            _ => false,
155        }
156    }
157    pub fn is_internal(&self) -> bool {
158        match self {
159            Node::Internal(_) => true,
160            _ => false,
161        }
162    }
163    pub fn new_label(label: &Digest32) -> NodeId {
164        Rc::new(RefCell::new(Node::LabelOnly(NodeHeader::new(
165            Some(*label),
166            None,
167        ))))
168    }
169
170    // Private methods
171    fn hdr(&'_ self) -> &'_ NodeHeader {
172        match self {
173            Node::LabelOnly(hdr) => &hdr,
174            Node::Internal(node) => &node.hdr,
175            Node::Leaf(node) => &node.hdr,
176        }
177    }
178
179    fn hdr_mut(&'_ mut self) -> &'_ mut NodeHeader {
180        match self {
181            Node::LabelOnly(ref mut hdr) => hdr,
182            Node::Internal(ref mut node) => &mut node.hdr,
183            Node::Leaf(ref mut node) => &mut node.hdr,
184        }
185    }
186
187    fn reset_recursive(node: &NodeId) {
188        if node.borrow_mut().reset() {
189            if let Node::Internal(r) = &*node.borrow() {
190                Self::reset_recursive(&r.left);
191                Self::reset_recursive(&r.right);
192            }
193        }
194    }
195}
196
197impl NodeHeader {
198    pub fn new(label: Option<Digest32>, key: Option<ADKey>) -> NodeHeader {
199        NodeHeader {
200            visited: false,
201            is_new: true,
202            key,
203            label,
204        }
205    }
206}
207
208impl InternalNode {
209    pub fn new(key: Option<ADKey>, left: &NodeId, right: &NodeId, balance: Balance) -> NodeId {
210        Rc::new(RefCell::new(Node::Internal(InternalNode {
211            hdr: NodeHeader::new(None, key),
212            left: left.clone(),
213            right: right.clone(),
214            balance,
215        })))
216    }
217
218    pub fn update_key(node: &NodeId, key: &ADKey) -> NodeId {
219        if let Node::Internal(this) = &mut *node.borrow_mut() {
220            if this.hdr.is_new {
221                this.hdr.key = Some(key.clone());
222            } else {
223                return Self::new(Some(key.clone()), &this.left, &this.right, this.balance);
224            }
225        } else {
226            panic!("Not internal node");
227        }
228        node.clone()
229    }
230
231    pub fn update(node: &NodeId, left: &NodeId, right: &NodeId, balance: Balance) -> NodeId {
232        if let Node::Internal(this) = &mut *node.borrow_mut() {
233            if this.hdr.is_new {
234                this.left = left.clone();
235                this.right = right.clone();
236                this.balance = balance;
237                this.hdr.label = None;
238            } else {
239                return Self::new(this.hdr.key.clone(), left, right, balance);
240            }
241        } else {
242            panic!("Not internal node");
243        }
244        node.clone()
245    }
246}
247
248impl LeafNode {
249    pub fn update(node: &NodeId, key: &ADKey, value: &ADValue, next_node_key: &ADKey) -> NodeId {
250        if let Node::Leaf(this) = &mut *node.borrow_mut() {
251            if this.hdr.is_new {
252                this.hdr.key = Some(key.clone());
253                this.value = value.clone();
254                this.next_node_key = next_node_key.clone();
255                this.hdr.label = None;
256            } else {
257                return Self::new(key, value, next_node_key);
258            }
259        } else {
260            panic!("Not leaf node");
261        }
262        node.clone()
263    }
264
265    pub fn new(key: &ADKey, value: &ADValue, next_node_key: &ADKey) -> NodeId {
266        Rc::new(RefCell::new(Node::Leaf(LeafNode {
267            hdr: NodeHeader::new(None, Some(key.clone())),
268            value: value.clone(),
269            next_node_key: next_node_key.clone(),
270        })))
271    }
272}
273
274#[derive(Clone)]
275pub struct AVLTree {
276    pub root: Option<NodeId>,
277    pub height: usize,
278    pub key_length: usize,
279    pub value_length: Option<usize>,
280    pub resolver: Resolver,
281}
282
283impl AVLTree {
284    pub fn new(resolver: Resolver, key_length: usize, value_length: Option<usize>) -> AVLTree {
285        AVLTree {
286            key_length,
287            value_length,
288            resolver,
289            height: 0,
290            root: None,
291        }
292    }
293
294    pub fn left(&self, node: &NodeId) -> NodeId {
295        if let Node::Internal(r) = &mut *node.borrow_mut() {
296            self.resolve(&mut r.left)
297        } else {
298            panic!("Not internal node");
299        }
300    }
301
302    pub fn right(&self, node: &NodeId) -> NodeId {
303        if let Node::Internal(r) = &mut *node.borrow_mut() {
304            self.resolve(&mut r.right)
305        } else {
306            panic!("Not internal node");
307        }
308    }
309
310    pub fn balance(&self, node: &NodeId) -> Balance {
311        node.borrow().balance()
312    }
313
314    pub fn label(&self, node: &NodeId) -> Digest32 {
315        node.borrow_mut().label()
316    }
317
318    pub fn key(&self, node: &NodeId) -> ADKey {
319        node.borrow().key()
320    }
321
322    pub fn value(&self, node: &NodeId) -> ADValue {
323        node.borrow().value()
324    }
325
326    pub fn next_node_key(&self, node: &NodeId) -> ADKey {
327        node.borrow().next_node_key()
328    }
329
330    pub fn visited(&self, node: &NodeId) -> bool {
331        node.borrow().visited()
332    }
333
334    pub fn is_new(&self, node: &NodeId) -> bool {
335        node.borrow().is_new()
336    }
337
338    pub fn mark_visited(&self, node: &NodeId, visited: bool) {
339        node.borrow_mut().mark_visited(visited)
340    }
341
342    pub fn resolve(&self, child: &mut NodeId) -> NodeId {
343        let mut resolved_node: Option<NodeId> = None;
344        if let Node::LabelOnly(hdr) = &*child.borrow() {
345            resolved_node = Some(Rc::new(RefCell::new((self.resolver)(&hdr.label.unwrap()))));
346        }
347        if let Some(node) = resolved_node {
348            *child = node
349        }
350        child.clone()
351    }
352
353    pub fn copy(&self, node: &NodeId) -> Node {
354        let n = &mut *node.borrow_mut();
355        if let Node::Internal(r) = n {
356            let _ = self.resolve(&mut r.left);
357            let _ = self.resolve(&mut r.right);
358        }
359        n.clone()
360    }
361
362    pub fn extract_nodes(
363        &self,
364        extractor: &mut dyn FnMut(&mut Node) -> bool,
365    ) -> Option<Vec<NodeId>> {
366        if let Some(root) = &self.root {
367            let mut set = Vec::new();
368            self.extract_nodes_recursive(extractor, root, &mut set);
369            Some(set)
370        } else {
371            None
372        }
373    }
374
375    pub fn extract_first_node(
376        &self,
377        extractor: &mut dyn FnMut(&mut Node) -> bool,
378    ) -> Option<NodeId> {
379        if let Some(root) = &self.root {
380            self.extract_first_node_recursive(extractor, root)
381        } else {
382            None
383        }
384    }
385
386    fn extract_first_node_recursive(
387        &self,
388        extractor: &mut dyn FnMut(&mut Node) -> bool,
389        node: &NodeId,
390    ) -> Option<NodeId> {
391        let nr = &mut *node.borrow_mut();
392        if let Node::Internal(r) = nr {
393            self.extract_first_node_recursive(extractor, &self.resolve(&mut r.left))
394                .or(self.extract_first_node_recursive(extractor, &self.resolve(&mut r.right)))
395        } else if extractor(nr) {
396            Some(node.clone())
397        } else {
398            None
399        }
400    }
401
402    pub fn contains(&self, node: &NodeId) -> bool {
403        if let Some(root) = &self.root {
404            self.contains_recursive(root, &self.key(node), &self.label(node), false)
405        } else {
406            false
407        }
408    }
409
410    pub fn contains_key(&self, key: &ADKey, label: &Digest32) -> bool {
411        if let Some(root) = &self.root {
412            self.contains_recursive(root, key, label, false)
413        } else {
414            false
415        }
416    }
417
418    fn contains_recursive(
419        &self,
420        node: &NodeId,
421        key: &ADKey,
422        label: &Digest32,
423        key_found: bool,
424    ) -> bool {
425        if &self.label(node) == label {
426            true
427        } else {
428            if let Node::Internal(r) = &mut *node.borrow_mut() {
429                if key_found {
430                    self.contains_recursive(&self.resolve(&mut r.left), key, label, true)
431                } else {
432                    match (*key).cmp(r.hdr.key.as_ref().unwrap()) {
433                        Ordering::Equal =>
434                        // found in the tree -- go one step right, then left to the leaf
435                        {
436                            self.contains_recursive(&self.resolve(&mut r.right), key, label, true)
437                        }
438                        Ordering::Less =>
439                        // going left, not yet found
440                        {
441                            self.contains_recursive(&self.resolve(&mut r.left), key, label, false)
442                        }
443                        Ordering::Greater => {
444                            self.contains_recursive(&self.resolve(&mut r.right), key, label, false)
445                        }
446                    }
447                }
448            } else {
449                false
450            }
451        }
452    }
453
454    fn extract_nodes_recursive(
455        &self,
456        extractor: &mut dyn FnMut(&mut Node) -> bool,
457        node: &NodeId,
458        set: &mut Vec<NodeId>,
459    ) {
460        let nr = &mut *node.borrow_mut();
461        if let Node::Internal(r) = nr {
462            self.extract_nodes_recursive(extractor, &self.resolve(&mut r.left), set);
463            self.extract_nodes_recursive(extractor, &self.resolve(&mut r.right), set);
464        } else if extractor(nr) {
465            set.push(node.clone())
466        }
467    }
468
469    fn fmt_recursive(&self, f: &mut fmt::Formatter, node: &NodeId, depth: usize) -> fmt::Result {
470        write!(f, "{:1$}", "  ", depth)?;
471        match &*node.borrow() {
472            Node::Leaf(leaf) => {
473                writeln!(
474                    f,
475                    "At leaf label={:?}, key={:?} next_node_key={:?}, value={:?}",
476                    leaf.hdr.label, leaf.hdr.key, leaf.next_node_key, leaf.value
477                )
478            }
479            Node::Internal(r) => {
480                writeln!(
481                    f,
482                    "Internal node label={:?}, balance={:?}",
483                    r.hdr.label, r.balance
484                )?;
485                self.fmt_recursive(f, &r.left, depth + 1)?;
486                self.fmt_recursive(f, &r.right, depth + 1)
487            }
488            Node::LabelOnly(hdr) => {
489                writeln!(f, "Label-only node label={:?}", hdr.label)
490            }
491        }
492    }
493
494    pub fn reset(&self) {
495        if let Some(root) = &self.root {
496            Node::reset_recursive(root);
497        }
498    }
499
500    pub fn pack(&self, node: NodeId) -> Bytes {
501        let mut buf = BytesMut::new();
502        match &*node.borrow() {
503            Node::Internal(node) => {
504                buf.put_u8(INTERNAL_NODE_PREFIX);
505                buf.put_i8(node.balance);
506                buf.extend_from_slice(&node.hdr.key.as_ref().unwrap());
507                buf.extend_from_slice(&node.left.borrow_mut().label());
508                buf.extend_from_slice(&node.right.borrow_mut().label());
509            }
510            Node::Leaf(leaf) => {
511                buf.put_u8(LEAF_NODE_PREFIX);
512                buf.extend_from_slice(&leaf.hdr.key.as_ref().unwrap());
513                if let Some(value_length) = self.value_length {
514                    assert!(leaf.value.len() == value_length as usize);
515                    buf.extend_from_slice(&leaf.value);
516                } else {
517                    buf.put_u32(leaf.value.len() as u32);
518                    buf.extend_from_slice(&leaf.value);
519                }
520                buf.extend_from_slice(&leaf.next_node_key);
521            }
522            Node::LabelOnly(_) => panic!("LabelOnly nodes should not be serialized"),
523        }
524        buf.freeze()
525    }
526
527    pub fn unpack(&self, bytes: &Bytes) -> NodeId {
528        let mut buf = BytesMut::from(&bytes[..]);
529        match buf.get_u8() {
530            INTERNAL_NODE_PREFIX => {
531                let balance = buf.get_i8();
532                let key = Some(buf.copy_to_bytes(self.key_length));
533                let mut left: Digest32 = Default::default();
534                buf.copy_to_slice(&mut left);
535                let mut right: Digest32 = Default::default();
536                buf.copy_to_slice(&mut right);
537                InternalNode::new(
538                    key,
539                    &Node::new_label(&left),
540                    &Node::new_label(&right),
541                    balance,
542                )
543            }
544            LEAF_NODE_PREFIX => {
545                let key = buf.copy_to_bytes(self.key_length);
546                let value: Bytes;
547                if let Some(value_length) = self.value_length {
548                    value = buf.copy_to_bytes(value_length);
549                } else {
550                    let value_length = buf.get_u32() as usize;
551                    value = buf.copy_to_bytes(value_length);
552                }
553                let next_node_key = buf.copy_to_bytes(self.key_length);
554                LeafNode::new(&key, &value, &next_node_key)
555            }
556            _ => {
557                panic!("Unexpected node prefix");
558            }
559        }
560    }
561
562    pub fn positive_infinity_key(&self) -> ADKey {
563        Bytes::from(vec![0xFFu8; self.key_length])
564    }
565
566    pub fn negative_infinity_key(&self) -> ADKey {
567        Bytes::from(vec![0u8; self.key_length])
568    }
569}
570
571impl fmt::Display for AVLTree {
572    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
573        if let Some(root) = &self.root {
574            self.fmt_recursive(f, &root, 0)
575        } else {
576            writeln!(f, "Empty tree")
577        }
578    }
579}
scorex_crypto_avltree/batch_node.rs

scorex_crypto_avltree/
batch_node.rs
