midnight_circuits/map/
cpu.rs

1// This file is part of MIDNIGHT-ZK.
2// Copyright (C) Midnight Foundation
3// SPDX-License-Identifier: Apache-2.0
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// http://www.apache.org/licenses/LICENSE-2.0
8// Unless required by applicable law or agreed to in writing, software
9// distributed under the License is distributed on an "AS IS" BASIS,
10// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11// See the License for the specific language governing permissions and
12// limitations under the License.
13
14//! CPU implementation of Succinct Key-Value Map Representation Using Merkle
15//! Trees
16
17use std::{collections::HashMap, fmt::Debug, marker::PhantomData};
18
19use num_bigint::BigUint;
20
21use crate::{
22    instructions::{hash::HashCPU, map::MapCPU},
23    utils::util::big_to_fe,
24    CircuitField,
25};
26
27/// This constant defines the height of the tree. This is a lower bound
28/// on the security parameter of the primitive, so it needs to be chosen
29/// carefully.
30pub(crate) const TREE_HEIGHT: u8 = 128;
31
32/// A [MapMt] is a succinct key-value map representation using merkle trees. We
33/// do not store all nodes. Instead, we only store the default nodes for each
34/// level, and those that have been modified.
35#[derive(Clone, Debug)]
36pub struct MapMt<F: CircuitField, H: HashCPU<F, F>> {
37    pub(crate) root: F,
38    // We organise nodes by their height and their position in that level.
39    nodes: HashMap<(u8, u128), F>,
40    // Map containing keys and values.
41    map: HashMap<BigUint, F>,
42    // Tree nodes, organised from leaves to root (though the root is treated separately)
43    default_nodes: [F; TREE_HEIGHT as usize],
44    _marker: PhantomData<H>,
45}
46
47impl<F: CircuitField, H: HashCPU<F, F>> PartialEq for MapMt<F, H> {
48    fn eq(&self, other: &Self) -> bool {
49        self.root == other.root
50            && self.nodes == other.nodes
51            && self.default_nodes == other.default_nodes
52            && self.map == other.map
53    }
54}
55
56impl<F: CircuitField, H: HashCPU<F, F>> Eq for MapMt<F, H> {}
57
58impl<F, H> MapCPU<F, F, F> for MapMt<F, H>
59where
60    F: CircuitField,
61    H: HashCPU<F, F>,
62{
63    fn new(default: &F) -> Self {
64        // The set of 'modified' nodes is empty
65        let nodes = HashMap::new();
66        let map = HashMap::new();
67
68        let mut default_nodes = [*default; TREE_HEIGHT as usize];
69
70        for i in 1..TREE_HEIGHT as usize {
71            default_nodes[i] =
72                <H as HashCPU<F, F>>::hash(&[default_nodes[i - 1], default_nodes[i - 1]]);
73        }
74
75        let root = <H as HashCPU<F, F>>::hash(&[default_nodes[127], default_nodes[127]]);
76
77        Self {
78            root,
79            nodes,
80            map,
81            default_nodes,
82            _marker: PhantomData,
83        }
84    }
85
86    fn succinct_repr(&self) -> F {
87        self.root
88    }
89
90    fn insert(&mut self, key: &F, value: &F) {
91        self.map.insert(key.to_biguint(), *value);
92
93        // We initialise the child with the new representation of the element.
94        let mut child = *value;
95        let mut node_index = Self::compute_node_index(key);
96
97        for height in 0..TREE_HEIGHT {
98            self.nodes.insert((height, node_index), child);
99
100            let sibling = self.get_sibling(node_index, height);
101            let (x, y) = conditional_swap(node_index & 1 == 1, &child, &sibling);
102            child = <H as HashCPU<F, F>>::hash(&[x, y]);
103            node_index >>= 1;
104        }
105
106        self.root = child;
107    }
108
109    fn get(&self, key: &F) -> F {
110        self.map.get(&key.to_biguint()).copied().unwrap_or(self.default_nodes[0])
111    }
112}
113
114impl<F, H> IntoIterator for MapMt<F, H>
115where
116    F: CircuitField,
117    H: HashCPU<F, F>,
118{
119    type Item = (F, F);
120    type IntoIter = std::iter::Map<
121        std::collections::hash_map::IntoIter<BigUint, F>,
122        fn((BigUint, F)) -> (F, F),
123    >;
124
125    fn into_iter(self) -> Self::IntoIter {
126        self.map.into_iter().map(|(key, value)| (big_to_fe(key), value))
127    }
128}
129
130impl<F, H> MapMt<F, H>
131where
132    F: CircuitField,
133    H: HashCPU<F, F>,
134{
135    /// Returns the nodes in the path for the given key.
136    pub(crate) fn get_path(&self, key: &F) -> [F; TREE_HEIGHT as usize] {
137        let mut node_index = Self::compute_node_index(key);
138
139        let mut nodes = [F::ZERO; TREE_HEIGHT as usize];
140
141        for (i, val) in nodes.iter_mut().enumerate() {
142            *val = self.get_sibling(node_index, i as u8);
143            node_index >>= 1;
144        }
145
146        nodes
147    }
148
149    /// Verify that the pair (key, value) is part of the existing map.
150    ///
151    /// Used for testing
152    #[cfg(test)]
153    fn verify_mem_proof(root: &F, path: &[F; TREE_HEIGHT as usize], key: &F, value: F) -> bool {
154        let mut node_index = Self::compute_node_index(key);
155        let mut child = value;
156
157        for node in path {
158            let (x, y) = conditional_swap(node_index & 1 == 1, &child, node);
159            child = <H as HashCPU<F, F>>::hash(&[x, y]);
160            node_index >>= 1;
161        }
162
163        *root == child
164    }
165
166    /// Get the sibling of an indexed node at a given height
167    fn get_sibling(&self, node_index: u128, height: u8) -> F {
168        assert!(height == 0 || (node_index < 1 << (TREE_HEIGHT - height)));
169
170        // If index is even, then we need the right sibling (height_index + 1), if
171        // it is odd, then we need the left sibling (height_index - 1).
172        let sibling_index = node_index + 1 - 2 * (node_index & 1);
173
174        // If the sibling does not exist, we use the default node for this height
175        *self
176            .nodes
177            .get(&(height, sibling_index))
178            .unwrap_or(&self.default_nodes[height as usize])
179    }
180
181    /// Get the node index at the leaf level for a given element, represented by
182    /// the first 128 bits of the hash output.
183    fn compute_node_index(element: &F) -> u128 {
184        let hashed_value = <H as HashCPU<F, F>>::hash(&[*element, F::ZERO]);
185        let bytes = hashed_value.to_bytes_le();
186
187        u128::from_le_bytes(bytes[..TREE_HEIGHT as usize / 8].try_into().unwrap())
188    }
189}
190
191// Takes two inputs and conditionally swaps them before hashing.
192fn conditional_swap<F: CircuitField>(cond: bool, left_input: &F, right_input: &F) -> (F, F) {
193    if cond {
194        (*right_input, *left_input)
195    } else {
196        (*left_input, *right_input)
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use rand::SeedableRng;
203    use rand_chacha::ChaCha8Rng;
204
205    use super::*;
206    use crate::hash::poseidon::{constants::PoseidonField, PoseidonChip};
207
208    fn test_map<F, H>()
209    where
210        F: CircuitField,
211        H: HashCPU<F, F> + Debug,
212    {
213        let mut rng = ChaCha8Rng::seed_from_u64(0xc0ffee);
214        let mut mt = MapMt::<F, H>::new(&F::ZERO);
215
216        // Let's add 100 random keys with one as a value.
217        for _ in 0..100 {
218            mt.insert(&F::random(&mut rng), &F::ONE);
219        }
220
221        // Now let's add key one, with one as a value
222        mt.insert(&F::ONE, &F::ONE);
223        assert_eq!(mt.get(&F::ONE), F::ONE);
224
225        // If we insert two times the same element, it should equal the old map
226        let old_mt = mt.clone();
227        mt.insert(&F::ONE, &F::ONE);
228
229        assert_eq!(old_mt, mt);
230
231        // Now we test path generation for proving that a (key, value) pair is part of
232        // the map
233        let one_path = mt.get_path(&F::ONE);
234        let member = MapMt::<F, H>::verify_mem_proof(&mt.root, &one_path, &F::ONE, F::ONE);
235        assert!(member);
236
237        let non_member = MapMt::<F, H>::verify_mem_proof(&mt.root, &one_path, &F::ONE, F::ZERO);
238        assert!(!non_member);
239
240        // Values that have not been explicitly added have a zero as a value
241        let new_value = F::random(&mut rng);
242        let path = mt.get_path(&new_value);
243        let non_member = MapMt::<F, H>::verify_mem_proof(&mt.root, &path, &new_value, F::ZERO);
244        assert!(non_member);
245    }
246
247    fn run_poseidon_test<F: PoseidonField>() {
248        test_map::<F, PoseidonChip<F>>();
249    }
250
251    #[test]
252    fn test_map_poseidon() {
253        run_poseidon_test::<midnight_curves::Fq>();
254    }
255}
midnight_circuits/map/cpu.rs

midnight_circuits/map/
cpu.rs
