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

use crate::{HashBuilder, EMPTY_ROOT_HASH};
use alloc::vec::Vec;
use alloy_primitives::B256;
use alloy_rlp::Encodable;
use nybbles::Nibbles;

/// Adjust the index of an item for rlp encoding.
pub const fn adjust_index_for_rlp(i: usize, len: usize) -> usize {
    if i > 0x7f {
        i
    } else if i == 0x7f || i + 1 == len {
        0
    } else {
        i + 1
    }
}

/// Compute a trie root of the collection of rlp encodable items.
pub fn ordered_trie_root<T: Encodable>(items: &[T]) -> B256 {
    ordered_trie_root_with_encoder(items, |item, buf| item.encode(buf))
}

/// Compute a trie root of the collection of items with a custom encoder.
pub fn ordered_trie_root_with_encoder<T, F>(items: &[T], mut encode: F) -> B256
where
    F: FnMut(&T, &mut Vec<u8>),
{
    if items.is_empty() {
        return EMPTY_ROOT_HASH;
    }

    let mut value_buffer = Vec::new();

    let mut hb = HashBuilder::default();
    let items_len = items.len();
    for i in 0..items_len {
        let index = adjust_index_for_rlp(i, items_len);

        let index_buffer = alloy_rlp::encode_fixed_size(&index);

        value_buffer.clear();
        encode(&items[index], &mut value_buffer);

        hb.add_leaf(Nibbles::unpack(&index_buffer), &value_buffer);
    }

    hb.root()
}

/// Ethereum specific trie root functions.
#[cfg(feature = "ethereum")]
pub use ethereum::*;
#[cfg(feature = "ethereum")]
mod ethereum {
    use super::*;
    use crate::TrieAccount;
    use alloy_primitives::{keccak256, Address, U256};

    /// Hashes storage keys, sorts them and them calculates the root hash of the storage trie.
    /// See [`storage_root_unsorted`] for more info.
    pub fn storage_root_unhashed(storage: impl IntoIterator<Item = (B256, U256)>) -> B256 {
        storage_root_unsorted(storage.into_iter().map(|(slot, value)| (keccak256(slot), value)))
    }

    /// Sorts and calculates the root hash of account storage trie.
    /// See [`storage_root`] for more info.
    pub fn storage_root_unsorted(storage: impl IntoIterator<Item = (B256, U256)>) -> B256 {
        let mut v = Vec::from_iter(storage);
        v.sort_unstable_by_key(|(key, _)| *key);
        storage_root(v)
    }

    /// Calculates the root hash of account storage trie.
    ///
    /// # Panics
    ///
    /// If the items are not in sorted order.
    pub fn storage_root(storage: impl IntoIterator<Item = (B256, U256)>) -> B256 {
        let mut hb = HashBuilder::default();
        for (hashed_slot, value) in storage {
            hb.add_leaf(
                Nibbles::unpack(hashed_slot),
                alloy_rlp::encode_fixed_size(&value).as_ref(),
            );
        }
        hb.root()
    }

    /// Hashes and sorts account keys, then proceeds to calculating the root hash of the state
    /// represented as MPT.
    /// See [`state_root_unsorted`] for more info.
    pub fn state_root_ref_unhashed<'a, A: Into<TrieAccount> + Clone + 'a>(
        state: impl IntoIterator<Item = (&'a Address, &'a A)>,
    ) -> B256 {
        state_root_unsorted(
            state.into_iter().map(|(address, account)| (keccak256(address), account.clone())),
        )
    }

    /// Hashes and sorts account keys, then proceeds to calculating the root hash of the state
    /// represented as MPT.
    /// See [`state_root_unsorted`] for more info.
    pub fn state_root_unhashed<A: Into<TrieAccount>>(
        state: impl IntoIterator<Item = (Address, A)>,
    ) -> B256 {
        state_root_unsorted(
            state.into_iter().map(|(address, account)| (keccak256(address), account)),
        )
    }

    /// Sorts the hashed account keys and calculates the root hash of the state represented as MPT.
    /// See [`state_root`] for more info.
    pub fn state_root_unsorted<A: Into<TrieAccount>>(
        state: impl IntoIterator<Item = (B256, A)>,
    ) -> B256 {
        let mut vec = Vec::from_iter(state);
        vec.sort_unstable_by_key(|(key, _)| *key);
        state_root(vec)
    }

    /// Calculates the root hash of the state represented as MPT.
    ///
    /// Corresponds to [geth's `deriveHash`](https://github.com/ethereum/go-ethereum/blob/6c149fd4ad063f7c24d726a73bc0546badd1bc73/core/genesis.go#L119).
    ///
    /// # Panics
    ///
    /// If the items are not in sorted order.
    pub fn state_root<A: Into<TrieAccount>>(state: impl IntoIterator<Item = (B256, A)>) -> B256 {
        let mut hb = HashBuilder::default();
        let mut account_rlp_buf = Vec::new();
        for (hashed_key, account) in state {
            account_rlp_buf.clear();
            account.into().encode(&mut account_rlp_buf);
            hb.add_leaf(Nibbles::unpack(hashed_key), &account_rlp_buf);
        }
        hb.root()
    }
}