snarkvm_circuit_collections/kary_merkle_tree/
mod.rs

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