fuel_merkle/common/
path_iterator.rs

1use crate::common::{
2    node::{
3        ChildKeyResult,
4        ChildResult,
5        ParentNode,
6    },
7    path::{
8        Path,
9        Side,
10    },
11};
12
13/// # Path Iterator
14///
15/// A naturally arising property of binary trees is that a leaf index encodes
16/// the unique path needed to traverse from the root of the tree to that leaf.
17/// The index's binary representation can be read left to right as a sequence of
18/// traversal instructions: a `0` bit means "descend left" and a `1` bit means
19/// "descend right". By following the `x` bits composing the index, starting at
20/// the root, descending to the left child at each `0`, descending to the right
21/// child at each `1`, we arrive at the leaf position, having touched every node
22/// position along the path formed by this index. Note that this algorithm does
23/// not prescribe how to descend from one node to the next; it describes merely
24/// the direction in which to descend at each step.
25///
26/// Alternatively, this can be interpreted as reading the index's most
27/// significant bit (MSB) at an offset `n`: read the `n`th bit to the right of
28/// the MSB. Here, `n` is a given step in the tree traversal, starting at 0, and
29/// incrementing by 1 at each depth until the leaf is reached. The
30/// traversal path is then the list of nodes calculated by traversing the tree
31/// using the instruction (`0` or `1`) indicated at `x`<sub>`n`</sub>, where `x`
32/// is the index in binary representation, and `n` is the offset for each digit
33/// in `x` from the MSB.
34///
35/// Reversing this path gives us the path from the leaf to the root.
36///
37/// Imagine a 3-bit integer type `u3` underpinning a tree's leaf indices. 3 bits
38/// give our tree a maximum height of 3, and a maximum number of leaf nodes
39/// 2<sup>3</sup> = 8. For demonstration, internal nodes are numbered using
40/// in-order indices (note that this would require an integer type with 4 bits
41/// or more in practice). In-order indexing provides a deterministic way to
42/// descend from one node to the next (see [Position](crate::common::Position)).
43///
44/// ```text
45///                             07
46///                            /  \
47///                           /    \
48///                          /      \
49///                         /        \
50///                        /          \
51///                       /            \
52///                     03              11
53///                    /  \            /  \
54///                   /    \          /    \
55///                 01      05      09      13
56///                /  \    /  \    /  \    /  \
57/// In-order idx: 00  02  04  06  08  10  12  14
58///     Leaf idx:  0   1   2   3   4   5   6   7
59/// ```
60///
61/// Let us now find the path to leaf with index `6`. In the above diagram, this
62/// is the seventh leaf in the leaf layer. A priori, we can see that the path
63/// from the root to this leaf is represented by the following list of in-order
64/// indices: `07, 11, 13, 12` (note that the leaf index that corresponds to the
65/// in-order index `12` is `6`).
66///
67/// ```text
68/// 0d6: u3 = 0b110
69///         = Right, Right, Left
70/// ```
71///
72/// Starting at the tree's root at index `07`, we can follow the instructions
73/// encoded by the binary representation of leaf `6` (`0b110`). In combination
74/// with our in-order index rules for descending nodes, we evaluate the
75/// following:
76///
77/// 1. The first bit is `1`; move right from `07` to `11`.
78/// 2. The next bit is `1`; move right from `11` to `13`.
79/// 3. The next and final bit is `0`; move left from `13` to `12`.
80///
81/// We have arrived at the desired leaf position with in-order index `12` and
82/// leaf index `6`. Indeed, following the instructions at each bit has produced
83/// the same list of positional indices that we observed earlier: `07, 11, 13,
84/// 12`.
85pub struct PathIter<T: ParentNode> {
86    leaf_key: T::Key,
87    current: Option<(ChildResult<T>, ChildKeyResult<T>)>,
88    current_offset: u32,
89}
90
91impl<T> PathIter<T>
92where
93    T: ParentNode + Clone,
94    T::Key: Clone,
95{
96    pub fn new(root: &T, leaf_key: &T::Key) -> Self {
97        let initial = (Ok(root.clone()), Ok(root.key()));
98
99        #[rustfmt::skip]
100        // The initial offset from the most significant bit (MSB).
101        //
102        // The offset from the MSB indicates which bit to read when deducing the
103        // path from the root to the leaf. As we descend down the tree,
104        // increasing the traversal depth, we increment this offset and read the
105        // corresponding bit to get the next traversal instruction.
106        //
107        // In the general case, we start by reading the first bit of the path at
108        // offset 0. This happens when the path fills its allocated memory;
109        // e.g., a path of 256 instructions is encoded within a 256 bit
110        // allocation for the leaf key. This also means that the key size in
111        // bits is equal to the maximum height of the tree.
112        //
113        // In the case that the length of the path is less than the number of
114        // bits in the key, the initial offset from the MSB must be augmented to
115        // accommodate the shortened path. This occurs when the key is allocated
116        // with a larger address space to reduce collisions of node addresses.
117        //
118        // E.g,
119        // With an 8-bit key and heights 1 through 7:
120        //
121        // Height Depth
122        // 7      0                        127                    Offset = Bits - Height = 8 - 7 = 1
123        //                                 / \
124        //                                /   \
125        // ...                          ...   ...
126        //                              /       \
127        //                             /         \
128        // 3      4                   07         247              Offset = Bits - Height = 8 - 3 = 5
129        //                           /  \        / \
130        //                          /    \     ...  \
131        //                         /      \          \
132        //                        /        \          \
133        //                       /          \          \
134        //                      /            \          \
135        // 2      5           03              11        251       Offset = Bits - Height = 8 - 2 = 6
136        //                   /  \            /  \       / \
137        //                  /    \          /    \    ...  \
138        // 1      6       01      05      09      13       253    Offset = Bits - Height = 8 - 1 = 7
139        //               /  \    /  \    /  \    /  \      / \
140        // 0      7     00  02  04  06  08  10  12  14   252 254
141        //              00  01  02  03  04  05  06  07   126 127
142        //
143        let initial_offset = T::key_size_bits().checked_sub(root.height())
144        .expect("Root height more than key size allows, ParentNode impl is incorrect");
145        Self {
146            leaf_key: leaf_key.clone(),
147            current: Some(initial),
148            current_offset: initial_offset,
149        }
150    }
151}
152
153impl<T> Iterator for PathIter<T>
154where
155    T: ParentNode,
156    T::Key: Path,
157{
158    type Item = (ChildResult<T>, ChildKeyResult<T>);
159
160    fn next(&mut self) -> Option<Self::Item> {
161        let value = self.current.take();
162
163        if let Some((ref path_node, _)) = value {
164            match path_node {
165                Ok(path_node) if path_node.is_node() => {
166                    let path = &self.leaf_key;
167                    let instruction = path.get_instruction(self.current_offset);
168                    self.current = instruction.map(|instruction| {
169                        // get_instruction ensures current_offset is ok
170                        #[allow(clippy::arithmetic_side_effects)]
171                        {
172                            self.current_offset += 1;
173                        }
174                        match instruction {
175                            Side::Left => {
176                                (path_node.left_child(), path_node.right_child_key())
177                            }
178                            Side::Right => {
179                                (path_node.right_child(), path_node.left_child_key())
180                            }
181                        }
182                    });
183                }
184                // Terminate the iterator if any of the following are true:
185                //    - The path node is a leaf (traversal is complete)
186                //    - The left or right child was not found and returned a ChildNotFound
187                //      error
188                //    - The left or right child returned any other error
189                _ => self.current = None,
190            }
191        }
192
193        value
194    }
195}
196
197pub trait AsPathIterator<T: ParentNode> {
198    fn as_path_iter(&self, leaf_key: &T::Key) -> PathIter<T>;
199}
200
201impl<T> AsPathIterator<T> for T
202where
203    T: ParentNode + Clone,
204    T::Key: Clone,
205{
206    fn as_path_iter(&self, leaf_key: &T::Key) -> PathIter<T> {
207        PathIter::new(self, leaf_key)
208    }
209}
210
211#[cfg(test)]
212#[allow(
213    clippy::restriction,
214    clippy::cast_possible_wrap,
215    clippy::cast_sign_loss
216)]
217mod test {
218    use crate::common::{
219        node::{
220            ChildKeyResult,
221            ChildResult,
222            Node,
223            ParentNode,
224        },
225        AsPathIterator,
226        Bytes8,
227    };
228    use alloc::vec::Vec;
229    use core::convert::Infallible;
230
231    #[derive(Debug, Clone, PartialEq)]
232    struct TestNode {
233        value: u64,
234    }
235
236    impl TestNode {
237        pub fn in_order_index(&self) -> u64 {
238            self.value
239        }
240
241        pub fn leaf_index(&self) -> u64 {
242            assert!(self.is_leaf());
243            self.value / 2
244        }
245
246        pub fn from_in_order_index(index: u64) -> Self {
247            Self { value: index }
248        }
249
250        pub fn from_leaf_index(index: u64) -> Self {
251            Self { value: index * 2 }
252        }
253
254        pub fn height(&self) -> u32 {
255            (!self.in_order_index()).trailing_zeros()
256        }
257
258        pub fn is_leaf(&self) -> bool {
259            self.in_order_index() % 2 == 0
260        }
261
262        fn child(&self, direction: i64) -> Self {
263            assert!(!self.is_leaf());
264            let shift = 1 << (self.height() - 1);
265            let index = self.in_order_index() as i64 + shift * direction;
266            Self::from_in_order_index(index as u64)
267        }
268    }
269
270    impl Node for TestNode {
271        type Key = Bytes8;
272
273        fn height(&self) -> u32 {
274            TestNode::height(self)
275        }
276
277        #[allow(clippy::arithmetic_side_effects, clippy::cast_possible_truncation)] // const
278        fn key_size_bits() -> u32 {
279            core::mem::size_of::<Self::Key>() as u32 * 8
280        }
281
282        fn leaf_key(&self) -> Self::Key {
283            TestNode::leaf_index(self).to_be_bytes()
284        }
285
286        fn is_leaf(&self) -> bool {
287            TestNode::is_leaf(self)
288        }
289
290        fn is_node(&self) -> bool {
291            !TestNode::is_leaf(self)
292        }
293    }
294
295    impl ParentNode for TestNode {
296        type ChildKey = u64;
297        type Error = Infallible;
298
299        fn key(&self) -> Self::ChildKey {
300            self.in_order_index()
301        }
302
303        fn left_child(&self) -> ChildResult<Self> {
304            Ok(TestNode::child(self, -1))
305        }
306
307        fn left_child_key(&self) -> ChildKeyResult<Self> {
308            self.left_child().map(|node| node.in_order_index())
309        }
310
311        fn right_child(&self) -> ChildResult<Self> {
312            Ok(TestNode::child(self, 1))
313        }
314
315        fn right_child_key(&self) -> ChildKeyResult<Self> {
316            self.right_child().map(|node| node.in_order_index())
317        }
318    }
319
320    #[test]
321    fn test_path_iter_returns_path() {
322        //               07
323        //              /  \
324        //             /    \
325        //            /      \
326        //           /        \
327        //          /          \
328        //         /            \
329        //       03              11
330        //      /  \            /  \
331        //     /    \          /    \
332        //   01      05      09      13
333        //  /  \    /  \    /  \    /  \
334        // 00  02  04  06  08  10  12  14
335        // 00  01  02  03  04  05  06  07
336        //
337        type Node = TestNode;
338        let root = Node::from_in_order_index(7);
339
340        {
341            let leaf = Node::from_leaf_index(0);
342            let path: Vec<TestNode> = root
343                .as_path_iter(&leaf.leaf_key())
344                .map(|(path, _)| path.unwrap())
345                .collect();
346            let expected_path = vec![
347                Node::from_in_order_index(7),
348                Node::from_in_order_index(3),
349                Node::from_in_order_index(1),
350                Node::from_leaf_index(0),
351            ];
352            assert_eq!(path, expected_path);
353        }
354
355        {
356            let leaf = Node::from_leaf_index(1);
357            let path: Vec<TestNode> = root
358                .as_path_iter(&leaf.leaf_key())
359                .map(|(path, _)| path.unwrap())
360                .collect();
361            let expected_path = vec![
362                Node::from_in_order_index(7),
363                Node::from_in_order_index(3),
364                Node::from_in_order_index(1),
365                Node::from_leaf_index(1),
366            ];
367            assert_eq!(path, expected_path);
368        }
369
370        {
371            let leaf = Node::from_leaf_index(2);
372            let path: Vec<TestNode> = root
373                .as_path_iter(&leaf.leaf_key())
374                .map(|(path, _)| path.unwrap())
375                .collect();
376            let expected_path = vec![
377                Node::from_in_order_index(7),
378                Node::from_in_order_index(3),
379                Node::from_in_order_index(5),
380                Node::from_leaf_index(2),
381            ];
382            assert_eq!(path, expected_path);
383        }
384
385        {
386            let leaf = Node::from_leaf_index(3);
387            let path: Vec<TestNode> = root
388                .as_path_iter(&leaf.leaf_key())
389                .map(|(path, _)| path.unwrap())
390                .collect();
391            let expected_path = vec![
392                Node::from_in_order_index(7),
393                Node::from_in_order_index(3),
394                Node::from_in_order_index(5),
395                Node::from_leaf_index(3),
396            ];
397            assert_eq!(path, expected_path);
398        }
399
400        {
401            let leaf = Node::from_leaf_index(4);
402            let path: Vec<TestNode> = root
403                .as_path_iter(&leaf.leaf_key())
404                .map(|(path, _)| path.unwrap())
405                .collect();
406            let expected_path = vec![
407                Node::from_in_order_index(7),
408                Node::from_in_order_index(11),
409                Node::from_in_order_index(9),
410                Node::from_leaf_index(4),
411            ];
412            assert_eq!(path, expected_path);
413        }
414
415        {
416            let leaf = Node::from_leaf_index(5);
417            let path: Vec<TestNode> = root
418                .as_path_iter(&leaf.leaf_key())
419                .map(|(path, _)| path.unwrap())
420                .collect();
421            let expected_path = vec![
422                Node::from_in_order_index(7),
423                Node::from_in_order_index(11),
424                Node::from_in_order_index(9),
425                Node::from_leaf_index(5),
426            ];
427            assert_eq!(path, expected_path);
428        }
429
430        {
431            let leaf = Node::from_leaf_index(6);
432            let path: Vec<TestNode> = root
433                .as_path_iter(&leaf.leaf_key())
434                .map(|(path, _)| path.unwrap())
435                .collect();
436            let expected_path = vec![
437                Node::from_in_order_index(7),
438                Node::from_in_order_index(11),
439                Node::from_in_order_index(13),
440                Node::from_leaf_index(6),
441            ];
442            assert_eq!(path, expected_path);
443        }
444
445        {
446            let leaf = Node::from_leaf_index(7);
447            let path: Vec<TestNode> = root
448                .as_path_iter(&leaf.leaf_key())
449                .map(|(path, _)| path.unwrap())
450                .collect();
451            let expected_path = vec![
452                Node::from_in_order_index(7),
453                Node::from_in_order_index(11),
454                Node::from_in_order_index(13),
455                Node::from_leaf_index(7),
456            ];
457            assert_eq!(path, expected_path);
458        }
459    }
460
461    #[test]
462    fn test_path_iter_returns_side_nodes() {
463        //               07
464        //              /  \
465        //             /    \
466        //            /      \
467        //           /        \
468        //          /          \
469        //         /            \
470        //       03              11
471        //      /  \            /  \
472        //     /    \          /    \
473        //   01      05      09      13
474        //  /  \    /  \    /  \    /  \
475        // 00  02  04  06  08  10  12  14
476        // 00  01  02  03  04  05  06  07
477        //
478        type Node = TestNode;
479        let root = Node::from_in_order_index(7); // 2^3 - 1
480
481        {
482            let leaf = Node::from_leaf_index(0);
483            let side: Vec<_> = root
484                .as_path_iter(&leaf.leaf_key())
485                .map(|(_, side)| side.unwrap())
486                .collect();
487            let expected_side = vec![
488                Node::from_in_order_index(7).value,  // Root's dummy side
489                Node::from_in_order_index(11).value, // Sibling of node 3
490                Node::from_in_order_index(5).value,  // Sibling of node 1
491                Node::from_leaf_index(1).value,      // Sibling of leaf 0
492            ];
493            assert_eq!(side, expected_side);
494        }
495
496        {
497            let leaf = Node::from_leaf_index(1);
498            let side: Vec<_> = root
499                .as_path_iter(&leaf.leaf_key())
500                .map(|(_, side)| side.unwrap())
501                .collect();
502            let expected_side = vec![
503                Node::from_in_order_index(7).value,  // Root's dummy side
504                Node::from_in_order_index(11).value, // Sibling of node 3
505                Node::from_in_order_index(5).value,  // Sibling of node 1
506                Node::from_leaf_index(0).value,      // Sibling of leaf 1
507            ];
508            assert_eq!(side, expected_side);
509        }
510
511        {
512            let leaf = Node::from_leaf_index(2);
513            let side: Vec<_> = root
514                .as_path_iter(&leaf.leaf_key())
515                .map(|(_, side)| side.unwrap())
516                .collect();
517            let expected_side = vec![
518                Node::from_in_order_index(7).value,  // Root's dummy side
519                Node::from_in_order_index(11).value, // Sibling of node 3
520                Node::from_in_order_index(1).value,  // Sibling of node 5
521                Node::from_leaf_index(3).value,      // Sibling of leaf 2
522            ];
523            assert_eq!(side, expected_side);
524        }
525
526        {
527            let leaf = Node::from_leaf_index(3);
528            let side: Vec<_> = root
529                .as_path_iter(&leaf.leaf_key())
530                .map(|(_, side)| side.unwrap())
531                .collect();
532            let expected_side = vec![
533                Node::from_in_order_index(7).value,  // Root's dummy side
534                Node::from_in_order_index(11).value, // Sibling of node 3
535                Node::from_in_order_index(1).value,  // Sibling of node 5
536                Node::from_leaf_index(2).value,      // Sibling of leaf 3
537            ];
538            assert_eq!(side, expected_side);
539        }
540
541        {
542            let leaf = Node::from_leaf_index(4);
543            let side: Vec<_> = root
544                .as_path_iter(&leaf.leaf_key())
545                .map(|(_, side)| side.unwrap())
546                .collect();
547            let expected_side = vec![
548                Node::from_in_order_index(7).value,  // Root's dummy side
549                Node::from_in_order_index(3).value,  // Sibling of node 11
550                Node::from_in_order_index(13).value, // Sibling of node 9
551                Node::from_leaf_index(5).value,      // Sibling of leaf 4
552            ];
553            assert_eq!(side, expected_side);
554        }
555
556        {
557            let leaf = Node::from_leaf_index(5);
558            let side: Vec<_> = root
559                .as_path_iter(&leaf.leaf_key())
560                .map(|(_, side)| side.unwrap())
561                .collect();
562            let expected_side = vec![
563                Node::from_in_order_index(7).value,  // Root's dummy side
564                Node::from_in_order_index(3).value,  // Sibling of node 11
565                Node::from_in_order_index(13).value, // Sibling of node 9
566                Node::from_leaf_index(4).value,      // Sibling of leaf 5
567            ];
568            assert_eq!(side, expected_side);
569        }
570
571        {
572            let leaf = Node::from_leaf_index(6);
573            let side: Vec<_> = root
574                .as_path_iter(&leaf.leaf_key())
575                .map(|(_, side)| side.unwrap())
576                .collect();
577            let expected_side = vec![
578                Node::from_in_order_index(7).value, // Root's dummy side
579                Node::from_in_order_index(3).value, // Sibling of node 11
580                Node::from_in_order_index(9).value, // Sibling of node 13
581                Node::from_leaf_index(7).value,     // Sibling of leaf 6
582            ];
583            assert_eq!(side, expected_side);
584        }
585
586        {
587            let leaf = Node::from_leaf_index(7);
588            let side: Vec<_> = root
589                .as_path_iter(&leaf.leaf_key())
590                .map(|(_, side)| side.unwrap())
591                .collect();
592            let expected_side = vec![
593                Node::from_in_order_index(7).value, // Root's dummy side
594                Node::from_in_order_index(3).value, // Sibling of node 11
595                Node::from_in_order_index(9).value, // Sibling of node 13
596                Node::from_leaf_index(6).value,     // Sibling of leaf 7
597            ];
598            assert_eq!(side, expected_side);
599        }
600    }
601
602    #[test]
603    fn test_path_iter_height_4() {
604        type Node = TestNode;
605        let root = Node::from_in_order_index(15); // 2^4 - 1
606        let leaf = Node::from_leaf_index(4); // 0b0100
607
608        let path: Vec<TestNode> = root
609            .as_path_iter(&leaf.leaf_key())
610            .map(|(path, _)| path.unwrap())
611            .collect();
612
613        let expected_path = vec![
614            Node::from_in_order_index(15),
615            Node::from_in_order_index(7),
616            Node::from_in_order_index(11),
617            Node::from_in_order_index(9),
618            Node::from_in_order_index(8),
619        ];
620        assert_eq!(path, expected_path);
621    }
622
623    #[test]
624    fn test_path_iter_height_8() {
625        type Node = TestNode;
626        let root = Node::from_in_order_index(255); // 2^8 - 1
627        let leaf = Node::from_leaf_index(61); // 0b00111101
628
629        let path: Vec<TestNode> = root
630            .as_path_iter(&leaf.leaf_key())
631            .map(|(path, _)| path.unwrap())
632            .collect();
633
634        let expected_path = vec![
635            Node::from_in_order_index(255),
636            Node::from_in_order_index(127),
637            Node::from_in_order_index(63),
638            Node::from_in_order_index(95),
639            Node::from_in_order_index(111),
640            Node::from_in_order_index(119),
641            Node::from_in_order_index(123),
642            Node::from_in_order_index(121),
643            Node::from_leaf_index(61),
644        ];
645        assert_eq!(path, expected_path);
646    }
647
648    #[test]
649    fn test_path_iter_returns_root_root_when_root_is_leaf() {
650        type Node = TestNode;
651        let root = Node::from_in_order_index(0);
652        let leaf = Node::from_leaf_index(0);
653
654        let (path, side): (Vec<TestNode>, Vec<_>) = root
655            .as_path_iter(&leaf.leaf_key())
656            .map(|(path, side)| (path.unwrap(), side.unwrap()))
657            .unzip();
658
659        let expected_path = vec![Node::from_in_order_index(0)];
660        let expected_side = vec![0];
661        assert_eq!(path, expected_path);
662        assert_eq!(side, expected_side);
663    }
664}