diem_types/proof/accumulator/mod.rs
1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright 2021 Conflux Foundation. All rights reserved.
5// Conflux is free software and distributed under GNU General Public License.
6// See http://www.gnu.org/licenses/
7
8//! This module implements an in-memory Merkle Accumulator that is similar to
9//! what we use in storage. This accumulator will only store a small portion of
10//! the tree -- for any subtree that is full, we store only the root. Also we
11//! only store the frozen nodes, therefore this structure will always store up
12//! to `Log(n)` number of nodes, where `n` is the total number of leaves in
13//! the tree.
14//!
15//! This accumulator is immutable once constructed. If we append new leaves to
16//! the tree we will obtain a new accumulator instance and the old one remains
17//! unchanged.
18
19#[cfg(test)]
20mod accumulator_test;
21
22use super::MerkleTreeInternalNode;
23use crate::proof::definition::{LeafCount, MAX_ACCUMULATOR_LEAVES};
24use anyhow::{ensure, format_err, Result};
25use diem_crypto::{
26 hash::{CryptoHash, CryptoHasher, ACCUMULATOR_PLACEHOLDER_HASH},
27 HashValue,
28};
29use std::marker::PhantomData;
30
31/// The Accumulator implementation.
32pub struct InMemoryAccumulator<H> {
33 /// Represents the roots of all the full subtrees from left to right in
34 /// this accumulator. For example, if we have the following
35 /// accumulator, this vector will have two hashes that correspond to
36 /// `X` and `e`.
37 ///
38 /// ```text
39 /// root
40 /// / \
41 /// / \
42 /// / \
43 /// X o
44 /// / \ / \
45 /// / \ / \
46 /// o o o placeholder
47 /// / \ / \ / \
48 /// a b c d e placeholder
49 /// ```
50 frozen_subtree_roots: Vec<HashValue>,
51
52 /// The total number of leaves in this accumulator.
53 num_leaves: LeafCount,
54
55 /// The root hash of this accumulator.
56 root_hash: HashValue,
57
58 phantom: PhantomData<H>,
59}
60
61impl<H> InMemoryAccumulator<H>
62where H: CryptoHasher
63{
64 /// Constructs a new accumulator with roots of existing frozen subtrees.
65 /// Returns error if the number of frozen subtree roots does not match
66 /// the number of leaves.
67 pub fn new(
68 frozen_subtree_roots: Vec<HashValue>, num_leaves: LeafCount,
69 ) -> Result<Self> {
70 ensure!(
71 frozen_subtree_roots.len() == num_leaves.count_ones() as usize,
72 "The number of frozen subtrees does not match the number of leaves. \
73 frozen_subtree_roots.len(): {}. num_leaves: {}.",
74 frozen_subtree_roots.len(),
75 num_leaves,
76 );
77
78 let root_hash =
79 Self::compute_root_hash(&frozen_subtree_roots, num_leaves);
80
81 Ok(Self {
82 frozen_subtree_roots,
83 num_leaves,
84 root_hash,
85 phantom: PhantomData,
86 })
87 }
88
89 /// Constructs a new accumulator with given leaves.
90 pub fn from_leaves(leaves: &[HashValue]) -> Self {
91 Self::default().append(leaves)
92 }
93
94 /// Appends a list of new leaves to an existing accumulator. Since the
95 /// accumulator is immutable, the existing one remains unchanged and a
96 /// new one representing the result is returned.
97 pub fn append(&self, leaves: &[HashValue]) -> Self {
98 let mut frozen_subtree_roots = self.frozen_subtree_roots.clone();
99 let mut num_leaves = self.num_leaves;
100 for leaf in leaves {
101 Self::append_one(&mut frozen_subtree_roots, num_leaves, *leaf);
102 num_leaves += 1;
103 }
104
105 Self::new(frozen_subtree_roots, num_leaves).expect(
106 "Appending leaves to a valid accumulator should create another valid accumulator.",
107 )
108 }
109
110 /// Appends one leaf. This will update `frozen_subtree_roots` to store new
111 /// frozen root nodes and remove old nodes if they are now part of a
112 /// larger frozen subtree.
113 fn append_one(
114 frozen_subtree_roots: &mut Vec<HashValue>,
115 num_existing_leaves: LeafCount, leaf: HashValue,
116 ) {
117 // For example, this accumulator originally had N = 7 leaves. Appending
118 // a leaf is like adding one to this number N: 0b0111 + 1 =
119 // 0b1000. Every time we carry a bit to the left we merge the
120 // rightmost two subtrees and compute their parent.
121 //
122 // ```text
123 // A
124 // / \
125 // / \
126 // o o B
127 // / \ / \ / \
128 // o o o o o o o
129 // ```
130
131 // First just append the leaf.
132 frozen_subtree_roots.push(leaf);
133
134 // Next, merge the last two subtrees into one. If `num_existing_leaves`
135 // has N trailing ones, the carry will happen N times.
136 let num_trailing_ones = (!num_existing_leaves).trailing_zeros();
137
138 for _i in 0..num_trailing_ones {
139 let right_hash =
140 frozen_subtree_roots.pop().expect("Invalid accumulator.");
141 let left_hash =
142 frozen_subtree_roots.pop().expect("Invalid accumulator.");
143 let parent_hash =
144 MerkleTreeInternalNode::<H>::new(left_hash, right_hash).hash();
145 frozen_subtree_roots.push(parent_hash);
146 }
147 }
148
149 /// Appends a list of new subtrees to the existing accumulator. This is
150 /// similar to \[`append`\](Accumulator::append) except that the new
151 /// leaves themselves are not known and they are represented by
152 /// `subtrees`. As an example, given the following accumulator that
153 /// currently has 10 leaves, the frozen subtree roots and the new subtrees
154 /// are annotated below. Note that in this case `subtrees[0]` represents
155 /// two new leaves `A` and `B`, `subtrees[1]` represents four new leaves
156 /// `C`, `D`, `E` and `F`, `subtrees[2]` represents four new leaves `G`,
157 /// `H`, `I` and `J`, and the last `subtrees[3]` represents one new leaf
158 /// `K`.
159 ///
160 /// ```text
161 /// new_root
162 /// / \
163 /// / \
164 /// / \
165 /// / \
166 /// / \
167 /// / \
168 /// / \
169 /// / \
170 /// / \
171 /// / \
172 /// / \
173 /// / \
174 /// / \
175 /// old_root o
176 /// / \ / \
177 /// / \ / placeholder
178 /// / \ /
179 /// / \ /
180 /// / \ /
181 /// / \ o
182 /// / \ / \
183 /// / \ / \
184 /// / o / \
185 /// frozen_subtree_roots[0] / \ / \
186 /// / \ / \ / \
187 /// / \ / \ / \
188 /// o o o subtrees[1] subtrees[2] o
189 /// / \ / \ / \ / \ / \ / \
190 /// o o o o frozen_subtree_roots[1] subtrees[0] o o o o o placeholder
191 /// / \ / \ / \ / \ / \ / \ / \ / \ / \ / \ / \
192 /// o o o o o o o o o o A B C D E F G H I J K (subtrees[3]) placeholder
193 /// ```
194 pub fn append_subtrees(
195 &self, subtrees: &[HashValue], num_new_leaves: LeafCount,
196 ) -> Result<Self> {
197 ensure!(
198 num_new_leaves <= MAX_ACCUMULATOR_LEAVES - self.num_leaves,
199 "Too many new leaves. self.num_leaves: {}. num_new_leaves: {}.",
200 self.num_leaves,
201 num_new_leaves,
202 );
203
204 if self.num_leaves == 0 {
205 return Self::new(subtrees.to_vec(), num_new_leaves);
206 }
207
208 let mut current_subtree_roots = self.frozen_subtree_roots.clone();
209 let mut current_num_leaves = self.num_leaves;
210 let mut remaining_new_leaves = num_new_leaves;
211 let mut subtree_iter = subtrees.iter();
212
213 // Check if we want to combine a new subtree with the rightmost frozen
214 // subtree. To do that this new subtree needs to represent
215 // `rightmost_frozen_subtree_size` leaves, so we need to have at
216 // least this many new leaves remaining.
217 let mut rightmost_frozen_subtree_size =
218 1 << current_num_leaves.trailing_zeros();
219 while remaining_new_leaves >= rightmost_frozen_subtree_size {
220 // Note that after combining the rightmost frozen subtree of size X
221 // with a new subtree, we obtain a subtree of size 2X.
222 // If there was already a frozen subtree of size 2X, we
223 // need to carry this process further.
224 let mut mask = rightmost_frozen_subtree_size;
225 let mut current_hash = *subtree_iter
226 .next()
227 .ok_or_else(|| format_err!("Too few subtrees."))?;
228 while current_num_leaves & mask != 0 {
229 let left_hash = current_subtree_roots
230 .pop()
231 .expect("This frozen subtree must exist.");
232 current_hash =
233 MerkleTreeInternalNode::<H>::new(left_hash, current_hash)
234 .hash();
235 mask <<= 1;
236 }
237 current_subtree_roots.push(current_hash);
238
239 current_num_leaves += rightmost_frozen_subtree_size;
240 remaining_new_leaves -= rightmost_frozen_subtree_size;
241 rightmost_frozen_subtree_size = mask;
242 }
243
244 // Now all the new subtrees are smaller than the rightmost frozen
245 // subtree. We just append all of them. Note that if the number
246 // of new subtrees does not actually match the number
247 // of new leaves, `Self::new` below will raise an error.
248 current_num_leaves += remaining_new_leaves;
249 current_subtree_roots.extend(subtree_iter);
250
251 Self::new(current_subtree_roots, current_num_leaves)
252 }
253
254 /// Returns the root hash of the accumulator.
255 pub fn root_hash(&self) -> HashValue { self.root_hash }
256
257 pub fn version(&self) -> u64 {
258 if self.num_leaves() == 0 {
259 0
260 } else {
261 self.num_leaves() - 1
262 }
263 }
264
265 /// Computes the root hash of an accumulator given the frozen subtree roots
266 /// and the number of leaves in this accumulator.
267 fn compute_root_hash(
268 frozen_subtree_roots: &[HashValue], num_leaves: LeafCount,
269 ) -> HashValue {
270 match frozen_subtree_roots.len() {
271 0 => return *ACCUMULATOR_PLACEHOLDER_HASH,
272 1 => return frozen_subtree_roots[0],
273 _ => (),
274 }
275
276 // The trailing zeros do not matter since anything below the lowest
277 // frozen subtree is already represented by the subtree roots.
278 let mut bitmap = num_leaves >> num_leaves.trailing_zeros();
279 let mut current_hash = *ACCUMULATOR_PLACEHOLDER_HASH;
280 let mut frozen_subtree_iter = frozen_subtree_roots.iter().rev();
281
282 while bitmap > 0 {
283 current_hash = if bitmap & 1 != 0 {
284 MerkleTreeInternalNode::<H>::new(
285 *frozen_subtree_iter
286 .next()
287 .expect("This frozen subtree should exist."),
288 current_hash,
289 )
290 } else {
291 MerkleTreeInternalNode::<H>::new(
292 current_hash,
293 *ACCUMULATOR_PLACEHOLDER_HASH,
294 )
295 }
296 .hash();
297 bitmap >>= 1;
298 }
299
300 current_hash
301 }
302
303 /// Returns the set of frozen subtree roots in this accumulator
304 pub fn frozen_subtree_roots(&self) -> &Vec<HashValue> {
305 &self.frozen_subtree_roots
306 }
307
308 /// Returns the total number of leaves in this accumulator.
309 pub fn num_leaves(&self) -> LeafCount { self.num_leaves }
310}
311
312// We manually implement Debug because H (CryptoHasher) does not implement
313// Debug.
314impl<H> std::fmt::Debug for InMemoryAccumulator<H> {
315 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
316 write!(
317 f,
318 "Accumulator {{ frozen_subtree_roots: {:?}, num_leaves: {:?} }}",
319 self.frozen_subtree_roots, self.num_leaves
320 )
321 }
322}
323
324impl<H> Default for InMemoryAccumulator<H>
325where H: CryptoHasher
326{
327 fn default() -> Self {
328 Self::new(vec![], 0)
329 .expect("Constructing empty accumulator should work.")
330 }
331}