Skip to main content

ab_core_primitives/
trees.rs

1//! Merkle Trees used by the protocol.
2//!
3//! Every tree hashes its inner nodes with its own key, so a node of one tree can't be interpreted
4//! as a node of another tree. Trees with a number of leaves that isn't known at compile time also
5//! commit to it.
6
7mod alias;
8mod hashers;
9#[cfg(test)]
10mod tests;
11
12use crate::address::Address;
13use ab_blake3::{BLOCK_LEN, OUT_LEN};
14use ab_io_type::trivial_type::TrivialType;
15use ab_merkle_tree::MerkleHasher;
16use ab_merkle_tree::balanced::BalancedMerkleTree;
17use ab_merkle_tree::blake3::Blake3MerkleHasher;
18use ab_merkle_tree::mmr::MerkleMountainRange;
19use ab_merkle_tree::sparse::SparseMerkleTree;
20use alias::UnbalancedMerkleTreeAlias;
21use array_reshape::FlattenEach;
22use hashers::{Blake3SizedMerkleHasher, Blake3StaticMerkleHasher};
23
24type SegmentRootsMerkleHasher =
25    Blake3SizedMerkleHasher<{ ab_blake3::const_hash(b"segment-roots-merkle-tree") }>;
26type OwnSegmentsMerkleHasher =
27    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"own-segments-merkle-tree") }>;
28type ShardOwnSegmentsMerkleHasher =
29    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"shard-own-segments-merkle-tree") }>;
30type LeafShardBodyMerkleHasher =
31    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"leaf-shard-body-merkle-tree") }>;
32type TransactionsMerkleHasher =
33    Blake3SizedMerkleHasher<{ ab_blake3::const_hash(b"transactions-merkle-tree") }>;
34type IntermediateShardBodyMerkleHasher =
35    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"intermediate-shard-body-merkle-tree") }>;
36type LeafShardsSegmentsMerkleHasher =
37    Blake3SizedMerkleHasher<{ ab_blake3::const_hash(b"leaf-shards-segments-merkle-tree") }>;
38type ChildLeafShardBlocksMerkleHasher =
39    Blake3SizedMerkleHasher<{ ab_blake3::const_hash(b"child-leaf-shard-blocks-merkle-tree") }>;
40type ChildIntermediateShardBlocksMerkleHasher = Blake3SizedMerkleHasher<
41    { ab_blake3::const_hash(b"child-intermediate-shard-blocks-merkle-tree") },
42>;
43type BeaconChainBodyMerkleHasher =
44    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"beacon-chain-body-merkle-tree") }>;
45type IntermediateShardsSegmentsMerkleHasher =
46    Blake3SizedMerkleHasher<{ ab_blake3::const_hash(b"intermediate-shards-segments-merkle-tree") }>;
47type BeaconChainHeaderMerkleHasher =
48    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"beacon-chain-header-merkle-tree") }>;
49type IntermediateShardHeaderMerkleHasher =
50    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"intermediate-shard-header-merkle-tree") }>;
51type LeafShardHeaderMerkleHasher =
52    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"leaf-shard-header-merkle-tree") }>;
53type SuperSegmentMerkleHasher =
54    Blake3SizedMerkleHasher<{ ab_blake3::const_hash(b"super-segment-merkle-tree") }>;
55type StateOwnersSparseMerkleHasher =
56    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"state-owners-sparse-merkle-tree") }>;
57type StateOwnerContractsSparseMerkleHasher = Blake3StaticMerkleHasher<
58    { ab_blake3::const_hash(b"state-owner-contracts-sparse-merkle-tree") },
59>;
60type BlockMerkleMountainRangeHasher =
61    Blake3StaticMerkleHasher<{ ab_blake3::const_hash(b"block-merkle-mountain-range") }>;
62
63const {
64    let keys = [
65        Blake3MerkleHasher::KEY,
66        SegmentRootsMerkleHasher::BASE_KEY,
67        OwnSegmentsMerkleHasher::KEY,
68        ShardOwnSegmentsMerkleHasher::KEY,
69        LeafShardBodyMerkleHasher::KEY,
70        TransactionsMerkleHasher::BASE_KEY,
71        IntermediateShardBodyMerkleHasher::KEY,
72        LeafShardsSegmentsMerkleHasher::BASE_KEY,
73        ChildLeafShardBlocksMerkleHasher::BASE_KEY,
74        ChildIntermediateShardBlocksMerkleHasher::BASE_KEY,
75        BeaconChainBodyMerkleHasher::KEY,
76        IntermediateShardsSegmentsMerkleHasher::BASE_KEY,
77        BeaconChainHeaderMerkleHasher::KEY,
78        IntermediateShardHeaderMerkleHasher::KEY,
79        LeafShardHeaderMerkleHasher::KEY,
80        SuperSegmentMerkleHasher::BASE_KEY,
81        StateOwnersSparseMerkleHasher::KEY,
82        StateOwnerContractsSparseMerkleHasher::KEY,
83        BlockMerkleMountainRangeHasher::KEY,
84    ];
85
86    let mut i = 0;
87    while i < keys.len() {
88        let mut j = i + 1;
89        while j < keys.len() {
90            assert!(keys[i] != keys[j], "Keys of Merkle Trees must be unique");
91            j += 1;
92        }
93        i += 1;
94    }
95}
96
97/// Number of bits in [`Address`], which is the number of levels of state sparse Merkle Trees
98const ADDRESS_BITS: u8 = Address::SIZE as u8 * u8::BITS as u8;
99
100/// Segment roots of own segments of a shard
101pub type SegmentRootsMerkleTree = <SegmentRootsMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
102/// Own segments of a shard
103pub type OwnSegmentsMerkleTree<'a> = BalancedMerkleTree<'a, 2, OwnSegmentsMerkleHasher>;
104/// Own segments of a leaf shard
105pub type ShardOwnSegmentsMerkleTree<'a> = BalancedMerkleTree<'a, 2, ShardOwnSegmentsMerkleHasher>;
106/// Leaf shard block body
107pub type LeafShardBodyMerkleTree<'a> = BalancedMerkleTree<'a, 2, LeafShardBodyMerkleHasher>;
108/// Transactions of a leaf shard block
109pub type TransactionsMerkleTree = <TransactionsMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
110/// Intermediate shard block body
111pub type IntermediateShardBodyMerkleTree =
112    <IntermediateShardBodyMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
113/// Own segments of child leaf shard blocks of an intermediate shard block
114pub type LeafShardsSegmentsMerkleTree =
115    <LeafShardsSegmentsMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
116/// Child leaf shard block roots of an intermediate shard block
117pub type ChildLeafShardBlocksMerkleTree =
118    <ChildLeafShardBlocksMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
119/// Child intermediate shard block roots of a beacon chain block
120pub type ChildIntermediateShardBlocksMerkleTree =
121    <ChildIntermediateShardBlocksMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
122/// Beacon chain block body
123pub type BeaconChainBodyMerkleTree<'a> = BalancedMerkleTree<'a, 4, BeaconChainBodyMerkleHasher>;
124/// Segments of child intermediate shard blocks of a beacon chain block
125pub type IntermediateShardsSegmentsMerkleTree =
126    <IntermediateShardsSegmentsMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
127/// Beacon chain block header fields, its root is the block root
128pub type BeaconChainHeaderMerkleTree =
129    <BeaconChainHeaderMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
130/// Intermediate shard block header fields, its root is the block root
131pub type IntermediateShardHeaderMerkleTree =
132    <IntermediateShardHeaderMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
133/// Leaf shard block header fields, its root is the block root
134pub type LeafShardHeaderMerkleTree =
135    <LeafShardHeaderMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
136/// Segment roots with their positions in a super segment
137pub type SuperSegmentMerkleTree = <SuperSegmentMerkleHasher as UnbalancedMerkleTreeAlias>::Tree;
138/// Global state: state roots of owners by owner address
139pub type StateOwnersSparseMerkleTree =
140    SparseMerkleTree<ADDRESS_BITS, StateOwnersSparseMerkleHasher>;
141/// State of one owner: hashes of contents of slots by the address of the managing contract
142pub type StateOwnerContractsSparseMerkleTree =
143    SparseMerkleTree<ADDRESS_BITS, StateOwnerContractsSparseMerkleHasher>;
144/// Beacon chain block roots.
145///
146/// NOTE: `u32` is smaller than `BlockNumber`'s internal `u64` but will be sufficient for a long
147/// time and substantially decrease the size of the in-memory data structure.
148pub type BlockMerkleMountainRange =
149    MerkleMountainRange<{ u64::from(u32::MAX) }, BlockMerkleMountainRangeHasher>;
150
151// `MerkleHasher` is implemented for each hasher separately, a generic implementation would need
152// `<Self as MerkleHasher>::OUT_LEN` with generic `Self` in method signatures, which requires
153// `gca_macroless_args` that this crate doesn't use
154// TODO: Replace the macro with generic implementations once `gca_macroless_args` is
155//  enabled in this crate:
156// // `<Self as MerkleHasher>::OUT_LEN` doesn't unify with `OUT_LEN` in generic code, hence
157// // reshaping
158// impl<const KEY: [u8; KEY_LEN]> MerkleHasher for Blake3StaticMerkleHasher<KEY> {
159//     const OUT_LEN: usize = OUT_LEN;
160//
161//     #[inline(always)]
162//     fn hash_pair(
163//         &self,
164//         pair: &[[u8; <Self as MerkleHasher>::OUT_LEN]; 2],
165//     ) -> [u8; <Self as MerkleHasher>::OUT_LEN] {
166//         let hash =
167//             ab_blake3::single_block_keyed_hash(self.node_key(), pair.flatten_ref::<BLOCK_LEN>());
168//         // Reshaping inside `Option` results in fewer copies of the hash than after `expect()`
169//         hash.map(|hash| [hash].flatten())
170//             .expect("Exactly one block worth of data; qed")
171//     }
172//
173//     #[inline(always)]
174//     fn hash_pairs<const NUM_PAIRS: usize>(
175//         &self,
176//         pairs: &[[[u8; <Self as MerkleHasher>::OUT_LEN]; 2]; NUM_PAIRS],
177//     ) -> [[u8; <Self as MerkleHasher>::OUT_LEN]; NUM_PAIRS] {
178//         let pairs = pairs
179//             .as_slice()
180//             .flatten_each_ref::<BLOCK_LEN>()
181//             .as_array::<NUM_PAIRS>()
182//             .expect("Same number of pairs as in the input; qed");
183//         let mut hashes = [[0; <Self as MerkleHasher>::OUT_LEN]; NUM_PAIRS];
184//         ab_blake3::single_block_keyed_hash_many_exact(
185//             self.node_key(),
186//             pairs,
187//             hashes.rechunk_mut(),
188//         );
189//         hashes
190//     }
191//
192//     #[inline(always)]
193//     fn hash_single_leaf(
194//         &self,
195//         leaf: &[u8; <Self as MerkleHasher>::OUT_LEN],
196//     ) -> [u8; <Self as MerkleHasher>::OUT_LEN] {
197//         [self.single_leaf_root(array::from_ref(leaf).flatten_ref::<OUT_LEN>())].flatten()
198//     }
199// }
200//
201// // Same body as above
202// impl<const BASE_KEY: [u8; KEY_LEN]> MerkleHasher for Blake3SizedMerkleHasher<BASE_KEY> { ... }
203macro_rules! impl_merkle_hasher {
204    ($($hasher:ty),+ $(,)?) => {
205        $(
206            impl MerkleHasher for $hasher {
207                const OUT_LEN: usize = OUT_LEN;
208
209                #[inline(always)]
210                fn hash_pair(&self, pair: &[[u8; OUT_LEN]; 2]) -> [u8; OUT_LEN] {
211                    ab_blake3::single_block_keyed_hash(self.node_key(), pair.as_flattened())
212                        .expect("Exactly one block worth of data; qed")
213                }
214
215                #[inline(always)]
216                fn hash_pairs<const NUM_PAIRS: usize>(
217                    &self,
218                    pairs: &[[[u8; OUT_LEN]; 2]; NUM_PAIRS],
219                ) -> [[u8; OUT_LEN]; NUM_PAIRS] {
220                    let pairs = pairs
221                        .as_slice()
222                        .flatten_each_ref::<BLOCK_LEN>()
223                        .as_array::<NUM_PAIRS>()
224                        .expect("Same number of pairs as in the input; qed");
225                    let mut hashes = [[0; OUT_LEN]; NUM_PAIRS];
226                    ab_blake3::single_block_keyed_hash_many_exact(
227                        self.node_key(),
228                        pairs,
229                        &mut hashes,
230                    );
231                    hashes
232                }
233
234                #[inline(always)]
235                fn hash_single_leaf(&self, leaf: &[u8; OUT_LEN]) -> [u8; OUT_LEN] {
236                    self.single_leaf_root(leaf)
237                }
238            }
239        )+
240    };
241}
242
243impl_merkle_hasher!(
244    SegmentRootsMerkleHasher,
245    OwnSegmentsMerkleHasher,
246    ShardOwnSegmentsMerkleHasher,
247    LeafShardBodyMerkleHasher,
248    TransactionsMerkleHasher,
249    IntermediateShardBodyMerkleHasher,
250    LeafShardsSegmentsMerkleHasher,
251    ChildLeafShardBlocksMerkleHasher,
252    ChildIntermediateShardBlocksMerkleHasher,
253    BeaconChainBodyMerkleHasher,
254    IntermediateShardsSegmentsMerkleHasher,
255    BeaconChainHeaderMerkleHasher,
256    IntermediateShardHeaderMerkleHasher,
257    LeafShardHeaderMerkleHasher,
258    SuperSegmentMerkleHasher,
259    StateOwnersSparseMerkleHasher,
260    StateOwnerContractsSparseMerkleHasher,
261    BlockMerkleMountainRangeHasher,
262);