1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    proof::{AccumulatorConsistencyProof, MerkleTreeInternalNode, TransactionAccumulatorSummary},
    transaction::Version,
};
use diem_crypto::hash::{
    CryptoHash, HashValue, TransactionAccumulatorHasher, ACCUMULATOR_PLACEHOLDER_HASH,
};
use std::{cell::RefCell, collections::HashMap};

type Children = (HashValue, Option<HashValue>);

/// An immutable transaction accumulator (not a summary, since it stores all leaf
/// nodes and caches internal nodes) that allows for easily computing root hashes,
/// consistency proofs, etc... at any versions in the range `0..=self.version()`.
///
/// This implementation intentionally eschews the existing storage implementation
/// and `Position` APIs to be as different as possible so we can better guarantee
/// correctness across implementations. Here, we mostly use "naive" recursive
/// algorithms (with internal caches).
///
/// Note: intended for test-only code; panics whenever something goes wrong.
#[derive(Clone, Debug)]
pub struct MockTransactionAccumulator {
    /// The list of leaf nodes.
    leaves: Vec<HashValue>,
    /// Cache that maps child nodes to their parent node.
    c2p: RefCell<HashMap<Children, HashValue>>,
    /// Cache that maps from a parent node to its child nodes.
    p2c: RefCell<HashMap<HashValue, Children>>,
}

impl MockTransactionAccumulator {
    /// Create a full transaction accumulator from a list of leaf node hashes.
    pub fn from_leaves(leaves: Vec<HashValue>) -> Self {
        assert!(!leaves.is_empty());
        Self {
            leaves,
            c2p: RefCell::new(HashMap::new()),
            p2c: RefCell::new(HashMap::new()),
        }
    }

    /// Create an accumulator with some mock leaf hash values at the given version.
    pub fn with_version(version: Version) -> Self {
        Self::from_leaves(mock_txn_hashes(version))
    }

    pub fn version(&self) -> Version {
        self.leaves.len() as u64 - 1
    }

    fn hash_internal_inner(pair: Children) -> HashValue {
        let (left, maybe_right) = pair;
        let right = maybe_right.unwrap_or(*ACCUMULATOR_PLACEHOLDER_HASH);
        MerkleTreeInternalNode::<TransactionAccumulatorHasher>::new(left, right).hash()
    }

    /// Compute the parent node hash from a pair of siblings while also updating
    /// internal caches.
    fn hash_internal(&self, pair: Children) -> HashValue {
        let parent = *self
            .c2p
            .borrow_mut()
            .entry(pair)
            .or_insert_with(|| Self::hash_internal_inner(pair));
        self.p2c.borrow_mut().insert(parent, pair);
        parent
    }

    /// Get the accumulator root hash at a specific version. Note that this method
    /// has the side effect of seeding the parent<->child caches.
    pub fn get_root_hash(&self, version: Version) -> HashValue {
        assert!(version <= self.version());
        let leaves_view = &self.leaves[0..=version as usize];

        let mut level = leaves_view.to_vec();
        while level.len() != 1 {
            let siblings_iter = level.chunks(2);
            level = siblings_iter
                .map(|pair| self.hash_internal((pair[0], pair.get(1).copied())))
                .collect();
        }

        level.into_iter().next().unwrap()
    }

    fn children(&self, parent: HashValue) -> Option<Children> {
        // precondition: parent in internal node set
        self.p2c.borrow().get(&parent).copied()
    }

    fn height(&self, subtree_root: HashValue) -> u64 {
        // precondition: subtree_root in internal node set
        match self.children(subtree_root) {
            // leaves have no children
            None => 0,
            Some((left, _)) => self.height(left) + 1,
        }
    }

    /// A node is frozen if it is complete. A node is complete if it's a leaf node
    /// or both of its children are complete. Note that a complete right child
    /// also implies that the left child is complete (since this is an accumulator).
    fn is_frozen(&self, subtree_root: HashValue) -> bool {
        // precondition: subtree_root in internal node set
        match self.children(subtree_root) {
            // leaf node ==> frozen
            None => true,
            // right subtree complete ==> left subtree also complete ==> frozen
            Some((_, Some(right))) => self.is_frozen(right),
            // not complete ==> not frozen
            Some((_, None)) => false,
        }
    }

    /// f(n) := if n.is_frozen => [n]
    ///         else           => f(n.left) || f(n.right)
    fn frozen_subtrees(&self, subtree_root: HashValue) -> Vec<HashValue> {
        // precondition: subtree_root in internal node set
        if self.is_frozen(subtree_root) {
            vec![subtree_root]
        } else {
            let (left, maybe_right) = self.children(subtree_root).unwrap();
            let mut left_subtrees = self.frozen_subtrees(left);
            let right_subtrees = maybe_right
                .map(|right| self.frozen_subtrees(right))
                .unwrap_or_else(Vec::new);
            left_subtrees.extend(right_subtrees);
            left_subtrees
        }
    }

    pub fn get_accumulator_summary(&self, version: Version) -> TransactionAccumulatorSummary {
        assert!(version <= self.version());

        let genesis_consistency_proof = self.get_consistency_proof(None, version);
        TransactionAccumulatorSummary::try_from_genesis_proof(genesis_consistency_proof, version)
            .unwrap()
    }

    pub fn get_consistency_proof(
        &self,
        start_version: Option<Version>,
        end_version: Version,
    ) -> AccumulatorConsistencyProof {
        assert!(start_version <= Some(end_version));
        assert!(end_version <= self.version());

        let maybe_old_root = start_version.map(|v| self.get_root_hash(v));
        let new_root = self.get_root_hash(end_version);

        let height_diff = maybe_old_root
            .map(|old_root| self.height(new_root) - self.height(old_root))
            .unwrap_or(0);

        let subtrees = self.frozen_subtree_diff(maybe_old_root, new_root, height_diff);
        AccumulatorConsistencyProof::new(subtrees)
    }

    /// Given a subtree root at an older version and a subtree root at a newer
    /// version, find the frozen subtree siblings needed to bring the frozen
    /// subtrees at the old subtree to the frozen subtrees at the new subtree.
    ///
    /// This method is recursive, typically starting at the root hashes for the
    /// old and new accumulators. Note that the new accumulator root can start at
    /// a greater height. After recursing left `height_diff` times, the old and
    /// new subtree roots will be tracking the same position. When old and new
    /// subtrees are tracking the same position, `height_diff == 0`.
    ///
    /// The old subtree parameter is an Option; when None, it means there is no
    /// subtree at this position in the old accumulator.
    ///
    /// ```not_rust
    /// // base case:
    /// // position does not exist in old subtree
    /// f(None, new, _) := frozen_subtrees(new)
    /// // node is frozen in both accumulators
    /// f(Some(same), same, 0) := []
    ///
    /// // inductive step:
    /// // same position
    /// f(Some(old), new, 0)
    ///     := f(Some(old.left), new.left, 0) || new.maybe_right.map(|new_right| f(old.maybe_right, new_right, 0))
    /// // new.height > old.height ==> old is in new's left descendents
    /// f(Some(old), new, dh > 0)
    ///     := f(Some(old), new.left, dh - 1) || f(None, new.right, _)
    /// ```
    fn frozen_subtree_diff(
        &self,
        maybe_old_subtree: Option<HashValue>,
        new_subtree: HashValue,
        height_diff: u64,
    ) -> Vec<HashValue> {
        // precondition: Some(old_subtree) and new_subtree in internal node set

        let old_subtree = match maybe_old_subtree {
            Some(old_subtree) => old_subtree,
            // When we are in a disjoint subtree from the old accumulator, we just
            // need the frozen subtrees.
            None => return self.frozen_subtrees(new_subtree),
        };

        if height_diff == 0 {
            // the old subtree and new subtree are at the same position

            if old_subtree == new_subtree {
                // the same subtree is frozen in both accumulators, don't need
                // anything here.
                vec![]
            } else {
                // this subtree is different between the old and new accumulator

                let (old_left, maybe_old_right) = self
                    .children(old_subtree)
                    .expect("cannot have two different leaf nodes");
                let (new_left, maybe_new_right) = self
                    .children(new_subtree)
                    .expect("cannot have two different leaf nodes");

                // recurse left and right
                let mut left_subtrees = self.frozen_subtree_diff(Some(old_left), new_left, 0);
                let right_subtrees = maybe_new_right
                    .map(|new_right| self.frozen_subtree_diff(maybe_old_right, new_right, 0))
                    .unwrap_or_else(Vec::new);
                left_subtrees.extend(right_subtrees);
                left_subtrees
            }
        } else {
            // the new subtree is still above the old subtree (which is somewhere
            // in the new subtree's left descendents).

            let (left, maybe_right) = self
                .children(new_subtree)
                .expect("non-zero height_diff implies new_subtree is an internal node, which must have children");
            let right = maybe_right
                .expect("non-zero height_diff implies new subtree must have a right child");
            // recurse to the left
            let mut left_subtrees =
                self.frozen_subtree_diff(Some(old_subtree), left, height_diff - 1);
            // this is effectively just frozen_subtrees(new.right), since the right
            // subtree will always be disjoint
            let right_subtrees = self.frozen_subtree_diff(None, right, 0);
            left_subtrees.extend(right_subtrees);
            left_subtrees
        }
    }
}

fn mock_txn_hashes(version: Version) -> Vec<HashValue> {
    (0..=version).map(HashValue::from_u64).collect::<Vec<_>>()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{block_info::BlockInfo, ledger_info::LedgerInfo};
    use proptest::prelude::*;

    fn mock_ledger_info(version: Version, root_hash: HashValue) -> LedgerInfo {
        LedgerInfo::new(
            BlockInfo::new(0, 0, HashValue::zero(), root_hash, version, 0, None),
            HashValue::zero(),
        )
    }

    #[test]
    fn test_mock_accumulator() {
        let end = 255;

        // A
        let accumulator = MockTransactionAccumulator::with_version(end);
        assert_eq!(accumulator.version(), end);

        let end_li = mock_ledger_info(end, accumulator.get_root_hash(end));

        // check for { m1, m2 : 0 <= m1 <= m2 <= end } that accumulators A and B have the same root hash
        // |------------------------------------------->|        A: accumulator
        // |--------------->|------------>|------------>|        B: B_{None,m1}.extend(B_{m1,m2}).extend(B_{m2,end})
        // 0               m1            m2            end
        //    B_{None,m1}      B_{m1,m2}     B_{m2,end}    B_{i,j}: accumulator.consistency_proof(i, j)
        proptest!(|((m1, m2) in (0..=end).prop_flat_map(|m1| (Just(m1), m1..=end)))| {
            // B_{None,m1}
            let accumulator_summary = accumulator.get_accumulator_summary(m1);
            assert_eq!(accumulator_summary.version(), m1);
            assert_eq!(accumulator_summary.root_hash(), accumulator.get_root_hash(m1));

            // B_{m1,m2}
            let mid1_to_mid2 = accumulator.get_consistency_proof(Some(m1), m2);
            let mid2_li = mock_ledger_info(m2, accumulator.get_root_hash(m2));

            // B_{None,m1}.extend(B_{m1,m2})
            let accumulator_summary = accumulator_summary
                .try_extend_with_proof(&mid1_to_mid2, &mid2_li)
                .unwrap();

            // B_{m2,end}
            let mid2_to_end = accumulator.get_consistency_proof(Some(m2), end);

            // B_{None,m1}.extend(B_{m1,m2}).extend(B_{m2,end})
            accumulator_summary.try_extend_with_proof(&mid2_to_end, &end_li).unwrap();
        });
    }
}