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

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Merkle proof types.

pub(crate) mod definition;
#[cfg(all(test, feature = "std"))]
pub(crate) mod proptest_proof;

use crate::{
    proof::SparseMerkleNode::{Internal, Leaf},
    SimpleHasher,
};

#[cfg(all(test, feature = "std"))]
use proptest_derive::Arbitrary;

pub use self::definition::{SparseMerkleProof, SparseMerkleRangeProof, UpdateMerkleProof};
use crate::{KeyHash, ValueHash, SPARSE_MERKLE_PLACEHOLDER_HASH};
use borsh::{BorshDeserialize, BorshSerialize};
use serde::{Deserialize, Serialize};

pub const LEAF_DOMAIN_SEPARATOR: &[u8] = b"JMT::LeafNode";
pub const INTERNAL_DOMAIN_SEPARATOR: &[u8] = b"JMT::IntrnalNode";

#[cfg_attr(all(test, feature = "std"), derive(Arbitrary))]
#[derive(
    Serialize, Deserialize, Clone, Copy, Eq, PartialEq, BorshSerialize, BorshDeserialize, Debug,
)]
/// A [`SparseMerkleNode`] is either a null node, an internal sparse node or a leaf node.
/// This is useful in the delete case to know if we need to coalesce the leaves on deletion.
/// The [`SparseMerkleNode`] needs to store either a [`SparseMerkleInternalNode`] or a [`SparseMerkleLeafNode`]
/// to be able to safely assert that the node is either a leaf or an internal node. Indeed,
/// if one stores the node/leaf hash directly into the structure, any malicious prover would
/// be able to forge the node/leaf type, as this assertion wouldn't be checked.
/// Providing a [`SparseMerkleInternalNode`] or a [`SparseMerkleLeafNode`] structure is sufficient to
/// prove the node type as one would need to reverse the hash function to forge them.
pub(crate) enum SparseMerkleNode {
    // The default sparse node
    Null,
    // The internal sparse merkle tree node
    Internal(SparseMerkleInternalNode),
    // The leaf sparse merkle tree node
    Leaf(SparseMerkleLeafNode),
}

impl SparseMerkleNode {
    pub(crate) fn hash<H: SimpleHasher>(&self) -> [u8; 32] {
        match self {
            SparseMerkleNode::Null => SPARSE_MERKLE_PLACEHOLDER_HASH,
            Internal(node) => node.hash::<H>(),
            Leaf(node) => node.hash::<H>(),
        }
    }
}

#[derive(
    Serialize, Deserialize, Clone, Copy, Eq, PartialEq, BorshSerialize, BorshDeserialize, Debug,
)]
#[cfg_attr(all(test, feature = "std"), derive(Arbitrary))]
pub(crate) struct SparseMerkleInternalNode {
    left_child: [u8; 32],
    right_child: [u8; 32],
}

impl SparseMerkleInternalNode {
    pub fn new(left_child: [u8; 32], right_child: [u8; 32]) -> Self {
        Self {
            left_child,
            right_child,
        }
    }

    pub fn hash<H: SimpleHasher>(&self) -> [u8; 32] {
        let mut hasher = H::new();
        // chop a vowel to fit in 16 bytes
        hasher.update(INTERNAL_DOMAIN_SEPARATOR);
        hasher.update(&self.left_child);
        hasher.update(&self.right_child);
        hasher.finalize()
    }
}

#[derive(Eq, Copy, Serialize, Deserialize, borsh::BorshSerialize, borsh::BorshDeserialize)]
pub struct SparseMerkleLeafNode {
    key_hash: KeyHash,
    value_hash: ValueHash,
}

// Manually implement Arbitrary to get the correct bounds. The derived Arbitrary impl adds a spurious
// H: Debug bound even with the proptest(no_bound) annotation
#[cfg(any(test))]
impl proptest::arbitrary::Arbitrary for SparseMerkleLeafNode {
    type Parameters = ();
    type Strategy = proptest::strategy::BoxedStrategy<Self>;

    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
        use proptest::{arbitrary::any, strategy::Strategy};
        (any::<KeyHash>(), any::<ValueHash>())
            .prop_map(|(key_hash, value_hash)| Self {
                key_hash,
                value_hash,
            })
            .boxed()
    }
}

// Manually implement Clone to circumvent [incorrect auto-bounds](https://github.com/rust-lang/rust/issues/26925)
// TODO: Switch back to #[derive] once the perfect_derive feature lands
impl Clone for SparseMerkleLeafNode {
    fn clone(&self) -> Self {
        Self {
            key_hash: self.key_hash.clone(),
            value_hash: self.value_hash.clone(),
        }
    }
}

// Manually implement Debug to circumvent [incorrect auto-bounds](https://github.com/rust-lang/rust/issues/26925)
// TODO: Switch back to #[derive] once the perfect_derive feature lands
impl core::fmt::Debug for SparseMerkleLeafNode {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("SparseMerkleLeafNode")
            .field("key_hash", &self.key_hash)
            .field("value_hash", &self.value_hash)
            .finish()
    }
}

// Manually implement PartialEq to circumvent [incorrect auto-bounds](https://github.com/rust-lang/rust/issues/26925)
// TODO: Switch back to #[derive] once the perfect_derive feature lands
impl PartialEq for SparseMerkleLeafNode {
    fn eq(&self, other: &Self) -> bool {
        self.key_hash == other.key_hash && self.value_hash == other.value_hash
    }
}

impl SparseMerkleLeafNode {
    pub(crate) fn new(key_hash: KeyHash, value_hash: ValueHash) -> Self {
        SparseMerkleLeafNode {
            key_hash,
            value_hash,
        }
    }

    pub(crate) fn key_hash(&self) -> KeyHash {
        self.key_hash
    }

    pub(crate) fn hash<H: SimpleHasher>(&self) -> [u8; 32] {
        let mut hasher = H::new();
        hasher.update(LEAF_DOMAIN_SEPARATOR);
        hasher.update(&self.key_hash.0);
        hasher.update(&self.value_hash.0);
        hasher.finalize()
    }
}