snarkvm_console_collections/kary_merkle_tree/path/
mod.rs

1// Copyright 2024 Aleo Network Foundation
2// This file is part of the snarkVM library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use super::*;
17
18#[derive(Clone, Debug, PartialEq, Eq, Hash)]
19pub struct KaryMerklePath<PH: PathHash, const DEPTH: u8, const ARITY: u8> {
20    /// The leaf index for the path.
21    leaf_index: u64,
22    /// The `siblings` contains a list of sibling hashes from the leaf to the root.
23    siblings: Vec<Vec<PH::Hash>>,
24}
25
26impl<PH: PathHash, const DEPTH: u8, const ARITY: u8> KaryMerklePath<PH, DEPTH, ARITY> {
27    /// Returns a new instance of a Merkle path.
28    pub fn try_from((leaf_index, siblings): (u64, Vec<Vec<PH::Hash>>)) -> Result<Self> {
29        // Ensure the Merkle tree depth is greater than 0.
30        ensure!(DEPTH > 0, "Merkle tree depth must be greater than 0");
31        // Ensure the Merkle tree depth is less than or equal to 64.
32        ensure!(DEPTH <= 64u8, "Merkle tree depth must be less than or equal to 64");
33        // Ensure the Merkle tree arity is greater than 1.
34        ensure!(ARITY > 1, "Merkle tree arity must be greater than 1");
35        // Ensure the Merkle tree does not overflow a u128.
36        ensure!((ARITY as u128).checked_pow(DEPTH as u32).is_some(), "Merkle tree size overflowed");
37        // Ensure the leaf index is within the tree depth.
38        ensure!((leaf_index as u128) < (ARITY as u128).saturating_pow(DEPTH as u32), "Out of bounds Merkle leaf index");
39        // Ensure the Merkle path is the correct length.
40        ensure!(siblings.len() == DEPTH as usize, "Found an incorrect Merkle path length");
41        for sibling in &siblings {
42            // Note: The ARITY is guaranteed to be greater than 1 (by the above check).
43            ensure!(sibling.len() == (ARITY - 1) as usize, "Found an incorrect Merkle path arity");
44        }
45        // Return the Merkle path.
46        Ok(Self { leaf_index, siblings })
47    }
48}
49
50impl<PH: PathHash, const DEPTH: u8, const ARITY: u8> KaryMerklePath<PH, DEPTH, ARITY> {
51    /// Returns the leaf index for the path.
52    pub fn leaf_index(&self) -> u64 {
53        self.leaf_index
54    }
55
56    /// Returns the siblings for the path.
57    pub fn siblings(&self) -> &[Vec<PH::Hash>] {
58        &self.siblings
59    }
60
61    /// Returns `true` if the Merkle path is valid for the given root and leaf.
62    pub fn verify<LH: LeafHash<Hash = PH::Hash>>(
63        &self,
64        leaf_hasher: &LH,
65        path_hasher: &PH,
66        root: &PH::Hash,
67        leaf: &LH::Leaf,
68    ) -> bool {
69        // Ensure the leaf index is within the tree depth.
70        if (self.leaf_index as u128) >= (ARITY as u128).saturating_pow(DEPTH as u32) {
71            eprintln!("Found an out of bounds Merkle leaf index");
72            return false;
73        }
74        // Ensure the path length matches the expected depth.
75        if self.siblings.len() != DEPTH as usize {
76            eprintln!("Found an incorrect Merkle path length");
77            return false;
78        }
79
80        // Initialize a tracker for the current hash, by computing the leaf hash to start.
81        let mut current_hash = match leaf_hasher.hash_leaf(leaf) {
82            Ok(candidate_leaf_hash) => candidate_leaf_hash,
83            Err(error) => {
84                eprintln!("Failed to hash the Merkle leaf during verification: {error}");
85                return false;
86            }
87        };
88
89        // Compute the ordering of the current hash and sibling hashes on each level.
90        // The indicator index determines which sibling the current hash is.
91        let Ok(indicator_indexes) = (0..DEPTH)
92            .map(|i| {
93                usize::try_from(self.leaf_index as u128 / (ARITY as u128).saturating_pow(i as u32) % (ARITY as u128))
94            })
95            .collect::<Result<Vec<_>, _>>()
96        else {
97            eprintln!("Found an incorrect Merkle leaf index");
98            return false;
99        };
100
101        // Check levels between leaf level and root.
102        for (indicator_index, sibling_hashes) in indicator_indexes.into_iter().zip_eq(&self.siblings) {
103            // Construct the ordering of sibling hashes for this level.
104            let mut sibling_hashes = sibling_hashes.clone();
105
106            // Insert the current hash into the list of sibling hashes.
107            sibling_hashes.insert(indicator_index, current_hash);
108
109            // Update the current hash for the next level.
110            match path_hasher.hash_children(&sibling_hashes) {
111                Ok(hash) => current_hash = hash,
112                Err(error) => {
113                    eprintln!("Failed to hash the Merkle path during verification: {error}");
114                    return false;
115                }
116            }
117        }
118
119        // Ensure the final hash matches the given root.
120        current_hash == *root
121    }
122}
123
124impl<PH: PathHash, const DEPTH: u8, const ARITY: u8> FromBytes for KaryMerklePath<PH, DEPTH, ARITY> {
125    /// Reads in a Merkle path from a buffer.
126    #[inline]
127    fn read_le<R: Read>(mut reader: R) -> IoResult<Self> {
128        // Read the leaf index.
129        let leaf_index = u64::read_le(&mut reader)?;
130        // Read the Merkle path siblings.
131        let siblings = (0..DEPTH)
132            .map(|_| (0..ARITY).map(|_| FromBytes::read_le(&mut reader)).collect::<IoResult<Vec<_>>>())
133            .collect::<IoResult<Vec<_>>>()?;
134        // Return the Merkle path.
135        Self::try_from((leaf_index, siblings)).map_err(error)
136    }
137}
138
139impl<PH: PathHash, const DEPTH: u8, const ARITY: u8> ToBytes for KaryMerklePath<PH, DEPTH, ARITY> {
140    /// Writes the Merkle path to a buffer.
141    #[inline]
142    fn write_le<W: Write>(&self, mut writer: W) -> IoResult<()> {
143        // Write the leaf index.
144        self.leaf_index.write_le(&mut writer)?;
145        // Write the Merkle path siblings.
146        self.siblings
147            .iter()
148            .try_for_each(|siblings| siblings.iter().try_for_each(|sibling| sibling.write_le(&mut writer)))
149    }
150}
151
152impl<PH: PathHash, const DEPTH: u8, const ARITY: u8> Serialize for KaryMerklePath<PH, DEPTH, ARITY> {
153    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
154        ToBytesSerializer::serialize_with_size_encoding(self, serializer)
155    }
156}
157
158impl<'de, PH: PathHash, const DEPTH: u8, const ARITY: u8> Deserialize<'de> for KaryMerklePath<PH, DEPTH, ARITY> {
159    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
160        FromBytesDeserializer::<Self>::deserialize_with_size_encoding(deserializer, "K-ary Merkle path")
161    }
162}