1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use light_bounded_vec::{BoundedVec, Pod};

use crate::errors::ConcurrentMerkleTreeError;

#[derive(Clone, Debug, PartialEq, Eq)]
#[repr(C)]
pub struct ChangelogEntry<const HEIGHT: usize> {
    /// Root.
    pub root: [u8; 32],
    // Path of the changelog.
    pub path: [[u8; 32]; HEIGHT],
    // Index.
    pub index: u64,
}

pub type ChangelogEntry22 = ChangelogEntry<22>;
pub type ChangelogEntry26 = ChangelogEntry<26>;
pub type ChangelogEntry32 = ChangelogEntry<32>;
pub type ChangelogEntry40 = ChangelogEntry<40>;

unsafe impl<const HEIGHT: usize> Pod for ChangelogEntry<HEIGHT> {}

impl<const HEIGHT: usize> ChangelogEntry<HEIGHT> {
    pub fn new(root: [u8; 32], path: [[u8; 32]; HEIGHT], index: usize) -> Self {
        let index = index as u64;
        Self { root, path, index }
    }

    pub fn default_with_index(index: usize) -> Self {
        Self {
            root: [0u8; 32],
            path: [[0u8; 32]; HEIGHT],
            index: index as u64,
        }
    }

    pub fn index(&self) -> usize {
        self.index as usize
    }

    /// Returns an intersection index in the changelog entry which affects the
    /// provided path.
    ///
    /// Determining it can be done by taking a XOR of the leaf index (which was
    /// directly updated in the changelog entry) and the leaf index we are
    /// trying to update.
    ///
    /// The number of bytes in the binary representations of the indexes is
    /// determined by the height of the tree. For example, for the tree with
    /// height 4, update attempt of leaf under index 2 and changelog affecting
    /// index 4, critbit would be:
    ///
    /// 2 ^ 4 = 0b_0010 ^ 0b_0100 = 0b_0110 = 6
    fn intersection_index(&self, leaf_index: usize) -> usize {
        let padding = 64 - HEIGHT;
        let common_path_len = ((leaf_index ^ self.index()) << padding).leading_zeros() as usize;

        (HEIGHT - 1) - common_path_len
    }

    pub fn update_proof(
        &self,
        leaf_index: usize,
        proof: &mut BoundedVec<[u8; 32]>,
        _allow_updates_in_changelog: bool,
    ) -> Result<(), ConcurrentMerkleTreeError> {
        if leaf_index != self.index() {
            let intersection_index = self.intersection_index(leaf_index);
            proof[intersection_index] = self.path[intersection_index];
        } else {
            // This case means that the leaf we are trying to update was
            // already updated. Therefore, updating the proof is impossible.
            // We need to return an error and request the caller
            // to retry the update with a new proof.
            //
            // TODO(vadorovsky): Re-visit optional throwing of this error.
            // if !allow_updates_in_changelog {
            //     return Err(ConcurrentMerkleTreeError::CannotUpdateLeaf);
            // }
        }

        Ok(())
    }

    pub fn update_subtrees(&self, rightmost_index: usize, subtrees: &mut BoundedVec<[u8; 32]>) {
        let (mut current_index, start) = if rightmost_index != self.index() {
            let intersection_index = self.intersection_index(rightmost_index);
            let current_index = rightmost_index + intersection_index;

            subtrees[intersection_index] = self.path[intersection_index];

            (current_index, intersection_index)
        } else {
            (rightmost_index, 0)
        };

        for (i, subtree) in subtrees.iter_mut().enumerate().skip(start) {
            let is_left = current_index % 2 == 0;
            if is_left {
                *subtree = self.path[i];
            }

            current_index /= 2;
        }
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn test_get_rightmost_proof() {}
}