snarkvm_circuit_collections/kary_merkle_tree/
mod.rs

1// Copyright 2024-2025 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
16mod helpers;
17pub use helpers::{BooleanHash, LeafHash, PathHash};
18
19mod verify;
20
21#[cfg(all(test, feature = "console"))]
22use snarkvm_circuit_types::environment::assert_scope;
23
24use snarkvm_circuit_types::{Boolean, Field, U16, U64, environment::prelude::*};
25
26pub struct KaryMerklePath<E: Environment, PH: PathHash<E>, const DEPTH: u8, const ARITY: u8> {
27    /// The leaf index for the path.
28    leaf_index: U64<E>,
29    /// The `siblings` contains a list of sibling hashes from the leaf to the root.
30    siblings: Vec<Vec<PH::Hash>>,
31}
32
33#[cfg(feature = "console")]
34impl<E: Environment, PH: PathHash<E>, const DEPTH: u8, const ARITY: u8> Inject for KaryMerklePath<E, PH, DEPTH, ARITY> {
35    type Primitive = console::kary_merkle_tree::KaryMerklePath<PH::Primitive, DEPTH, ARITY>;
36
37    /// Initializes a Merkle path from the given mode and native Merkle path.
38    fn new(mode: Mode, merkle_path: Self::Primitive) -> Self {
39        // Initialize the leaf index.
40        let leaf_index = U64::new(mode, console::U64::new(merkle_path.leaf_index()));
41        // Initialize the Merkle path siblings.
42        let siblings: Vec<Vec<_>> = merkle_path
43            .siblings()
44            .iter()
45            .map(|nodes| nodes.iter().map(|node| Inject::new(mode, *node)).collect())
46            .collect();
47
48        // Ensure the Merkle path has the correct arity.
49        for sibling in &siblings {
50            if sibling.len() != ARITY.saturating_sub(1) as usize {
51                return E::halt("Merkle path is not the correct depth");
52            }
53        }
54        // Ensure the Merkle path is the correct depth.
55        match siblings.len() == DEPTH as usize {
56            // Return the Merkle path.
57            true => Self { leaf_index, siblings },
58            false => E::halt("Merkle path is not the correct depth"),
59        }
60    }
61}
62
63#[cfg(feature = "console")]
64impl<E: Environment, PH: PathHash<E>, const DEPTH: u8, const ARITY: u8> Eject for KaryMerklePath<E, PH, DEPTH, ARITY> {
65    type Primitive = console::kary_merkle_tree::KaryMerklePath<PH::Primitive, DEPTH, ARITY>;
66
67    /// Ejects the mode of the Merkle path.
68    fn eject_mode(&self) -> Mode {
69        (&self.leaf_index, &self.siblings).eject_mode()
70    }
71
72    /// Ejects the Merkle path.
73    fn eject_value(&self) -> Self::Primitive {
74        match Self::Primitive::try_from((*self.leaf_index.eject_value(), self.siblings.eject_value())) {
75            Ok(merkle_path) => merkle_path,
76            Err(error) => E::halt(format!("Failed to eject the Merkle path: {error}")),
77        }
78    }
79}
80
81#[cfg(all(test, feature = "console"))]
82mod tests {
83    use super::*;
84    use console::{
85        algorithms::{BHP512 as NativeBHP512, BHP1024 as NativeBHP1024},
86        kary_merkle_tree::KaryMerkleTree,
87    };
88    use snarkvm_circuit_algorithms::BHP512;
89    use snarkvm_circuit_network::AleoV0 as Circuit;
90    use snarkvm_utilities::{TestRng, Uniform};
91
92    use anyhow::Result;
93
94    const ITERATIONS: u128 = 100;
95
96    fn check_new<const DEPTH: u8, const ARITY: u8>(
97        mode: Mode,
98        num_constants: u64,
99        num_public: u64,
100        num_private: u64,
101        num_constraints: u64,
102    ) -> Result<()> {
103        let mut rng = TestRng::default();
104
105        type PH = BHP512<Circuit>;
106
107        type NativeLH = NativeBHP1024<<Circuit as Environment>::Network>;
108        type NativePH = NativeBHP512<<Circuit as Environment>::Network>;
109
110        let leaf_hasher = NativeLH::setup("AleoMerklePathTest0")?;
111        let path_hasher = NativePH::setup("AleoMerklePathTest1")?;
112
113        let mut create_leaves = |num_leaves| {
114            (0..num_leaves)
115                .map(|_| console::Field::<<Circuit as Environment>::Network>::rand(&mut rng).to_bits_le())
116                .collect::<Vec<_>>()
117        };
118
119        for i in 0..ITERATIONS {
120            // Determine the number of leaves.
121            let num_leaves = core::cmp::min((ARITY as u128).pow(DEPTH as u32), i);
122            // Compute the leaves.
123            let leaves = create_leaves(num_leaves);
124            // Compute the Merkle tree.
125            let merkle_tree =
126                KaryMerkleTree::<NativeLH, NativePH, DEPTH, ARITY>::new(&leaf_hasher, &path_hasher, &leaves)?;
127
128            for (index, leaf) in leaves.iter().enumerate() {
129                // Compute the Merkle path.
130                let merkle_path = merkle_tree.prove(index, leaf)?;
131
132                // // Initialize the Merkle leaf.
133                // let leaf: Vec<Boolean<_>> = Inject::new(mode, leaf.clone());
134
135                Circuit::scope(format!("New {mode}"), || {
136                    let candidate = KaryMerklePath::<Circuit, PH, DEPTH, ARITY>::new(mode, merkle_path.clone());
137                    assert_eq!(merkle_path, candidate.eject_value());
138                    assert_scope!(num_constants, num_public, num_private, num_constraints);
139                });
140                Circuit::reset();
141            }
142        }
143        Ok(())
144    }
145
146    #[test]
147    fn test_new_constant() -> Result<()> {
148        check_new::<32, 3>(Mode::Constant, 128, 0, 0, 0)
149    }
150
151    #[test]
152    fn test_new_public() -> Result<()> {
153        check_new::<32, 3>(Mode::Public, 0, 128, 0, 64)
154    }
155
156    #[test]
157    fn test_new_private() -> Result<()> {
158        check_new::<32, 3>(Mode::Private, 0, 0, 128, 64)
159    }
160}