ethereum.shanghai.trie

  1"""
  2State Trie
  3^^^^^^^^^^
  4
  5.. contents:: Table of Contents
  6    :backlinks: none
  7    :local:
  8
  9Introduction
 10------------
 11
 12The state trie is the structure responsible for storing
 13`.fork_types.Account` objects.
 14"""
 15
 16import copy
 17from dataclasses import dataclass, field
 18from typing import (
 19    Callable,
 20    Dict,
 21    Generic,
 22    List,
 23    Mapping,
 24    MutableMapping,
 25    Optional,
 26    Sequence,
 27    TypeVar,
 28    Union,
 29    cast,
 30)
 31
 32from ethereum.crypto.hash import keccak256
 33from ethereum.paris import trie as previous_trie
 34from ethereum.utils.ensure import ensure
 35from ethereum.utils.hexadecimal import hex_to_bytes
 36
 37from .. import rlp
 38from ..base_types import U256, Bytes, Uint, slotted_freezable
 39from .fork_types import (
 40    Account,
 41    Address,
 42    LegacyTransaction,
 43    Receipt,
 44    Root,
 45    Withdrawal,
 46    encode_account,
 47)
 48
 49# note: an empty trie (regardless of whether it is secured) has root:
 50#
 51#   keccak256(RLP(b''))
 52#       ==
 53#   56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421 # noqa: E501,SC10
 54#
 55# also:
 56#
 57#   keccak256(RLP(()))
 58#       ==
 59#   1dcc4de8dec75d7aab85b567b6ccd41ad312451b948a7413f0a142fd40d49347 # noqa: E501,SC10
 60#
 61# which is the sha3Uncles hash in block header with no uncles
 62EMPTY_TRIE_ROOT = Root(
 63    hex_to_bytes(
 64        "56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421"
 65    )
 66)
 67
 68Node = Union[
 69    Account, Bytes, LegacyTransaction, Receipt, Uint, U256, Withdrawal, None
 70]
 71K = TypeVar("K", bound=Bytes)
 72V = TypeVar(
 73    "V",
 74    Optional[Account],
 75    Optional[Bytes],
 76    Bytes,
 77    Optional[Union[LegacyTransaction, Bytes]],
 78    Optional[Union[Receipt, Bytes]],
 79    Optional[Union[Withdrawal, Bytes]],
 80    Uint,
 81    U256,
 82)
 83
 84
 85@slotted_freezable
 86@dataclass
 87class LeafNode:
 88    """Leaf node in the Merkle Trie"""
 89
 90    rest_of_key: Bytes
 91    value: rlp.RLP
 92
 93
 94@slotted_freezable
 95@dataclass
 96class ExtensionNode:
 97    """Extension node in the Merkle Trie"""
 98
 99    key_segment: Bytes
100    subnode: rlp.RLP
101
102
103@slotted_freezable
104@dataclass
105class BranchNode:
106    """Branch node in the Merkle Trie"""
107
108    subnodes: List[rlp.RLP]
109    value: rlp.RLP
110
111
112InternalNode = Union[LeafNode, ExtensionNode, BranchNode]
113
114
115def encode_internal_node(node: Optional[InternalNode]) -> rlp.RLP:
116    """
117    Encodes a Merkle Trie node into its RLP form. The RLP will then be
118    serialized into a `Bytes` and hashed unless it is less that 32 bytes
119    when serialized.
120
121    This function also accepts `None`, representing the absence of a node,
122    which is encoded to `b""`.
123
124    Parameters
125    ----------
126    node : Optional[InternalNode]
127        The node to encode.
128
129    Returns
130    -------
131    encoded : `rlp.RLP`
132        The node encoded as RLP.
133    """
134    unencoded: rlp.RLP
135    if node is None:
136        unencoded = b""
137    elif isinstance(node, LeafNode):
138        unencoded = (
139            nibble_list_to_compact(node.rest_of_key, True),
140            node.value,
141        )
142    elif isinstance(node, ExtensionNode):
143        unencoded = (
144            nibble_list_to_compact(node.key_segment, False),
145            node.subnode,
146        )
147    elif isinstance(node, BranchNode):
148        unencoded = node.subnodes + [node.value]
149    else:
150        raise AssertionError(f"Invalid internal node type {type(node)}!")
151
152    encoded = rlp.encode(unencoded)
153    if len(encoded) < 32:
154        return unencoded
155    else:
156        return keccak256(encoded)
157
158
159def encode_node(node: Node, storage_root: Optional[Bytes] = None) -> Bytes:
160    """
161    Encode a Node for storage in the Merkle Trie.
162
163    Currently mostly an unimplemented stub.
164    """
165    if isinstance(node, Account):
166        assert storage_root is not None
167        return encode_account(node, storage_root)
168    elif isinstance(node, (LegacyTransaction, Receipt, Withdrawal, U256)):
169        return rlp.encode(cast(rlp.RLP, node))
170    elif isinstance(node, Bytes):
171        return node
172    else:
173        return previous_trie.encode_node(node, storage_root)
174
175
176@dataclass
177class Trie(Generic[K, V]):
178    """
179    The Merkle Trie.
180    """
181
182    secured: bool
183    default: V
184    _data: Dict[K, V] = field(default_factory=dict)
185
186
187def copy_trie(trie: Trie[K, V]) -> Trie[K, V]:
188    """
189    Create a copy of `trie`. Since only frozen objects may be stored in tries,
190    the contents are reused.
191
192    Parameters
193    ----------
194    trie: `Trie`
195        Trie to copy.
196
197    Returns
198    -------
199    new_trie : `Trie[K, V]`
200        A copy of the trie.
201    """
202    return Trie(trie.secured, trie.default, copy.copy(trie._data))
203
204
205def trie_set(trie: Trie[K, V], key: K, value: V) -> None:
206    """
207    Stores an item in a Merkle Trie.
208
209    This method deletes the key if `value == trie.default`, because the Merkle
210    Trie represents the default value by omitting it from the trie.
211
212    Parameters
213    ----------
214    trie: `Trie`
215        Trie to store in.
216    key : `Bytes`
217        Key to lookup.
218    value : `V`
219        Node to insert at `key`.
220    """
221    if value == trie.default:
222        if key in trie._data:
223            del trie._data[key]
224    else:
225        trie._data[key] = value
226
227
228def trie_get(trie: Trie[K, V], key: K) -> V:
229    """
230    Gets an item from the Merkle Trie.
231
232    This method returns `trie.default` if the key is missing.
233
234    Parameters
235    ----------
236    trie:
237        Trie to lookup in.
238    key :
239        Key to lookup.
240
241    Returns
242    -------
243    node : `V`
244        Node at `key` in the trie.
245    """
246    return trie._data.get(key, trie.default)
247
248
249def common_prefix_length(a: Sequence, b: Sequence) -> int:
250    """
251    Find the longest common prefix of two sequences.
252    """
253    for i in range(len(a)):
254        if i >= len(b) or a[i] != b[i]:
255            return i
256    return len(a)
257
258
259def nibble_list_to_compact(x: Bytes, is_leaf: bool) -> Bytes:
260    """
261    Compresses nibble-list into a standard byte array with a flag.
262
263    A nibble-list is a list of byte values no greater than `15`. The flag is
264    encoded in high nibble of the highest byte. The flag nibble can be broken
265    down into two two-bit flags.
266
267    Highest nibble::
268
269        +---+---+----------+--------+
270        | _ | _ | is_leaf | parity |
271        +---+---+----------+--------+
272          3   2      1         0
273
274
275    The lowest bit of the nibble encodes the parity of the length of the
276    remaining nibbles -- `0` when even and `1` when odd. The second lowest bit
277    is used to distinguish leaf and extension nodes. The other two bits are not
278    used.
279
280    Parameters
281    ----------
282    x :
283        Array of nibbles.
284    is_leaf :
285        True if this is part of a leaf node, or false if it is an extension
286        node.
287
288    Returns
289    -------
290    compressed : `bytearray`
291        Compact byte array.
292    """
293    compact = bytearray()
294
295    if len(x) % 2 == 0:  # ie even length
296        compact.append(16 * (2 * is_leaf))
297        for i in range(0, len(x), 2):
298            compact.append(16 * x[i] + x[i + 1])
299    else:
300        compact.append(16 * ((2 * is_leaf) + 1) + x[0])
301        for i in range(1, len(x), 2):
302            compact.append(16 * x[i] + x[i + 1])
303
304    return Bytes(compact)
305
306
307def bytes_to_nibble_list(bytes_: Bytes) -> Bytes:
308    """
309    Converts a `Bytes` into to a sequence of nibbles (bytes with value < 16).
310
311    Parameters
312    ----------
313    bytes_:
314        The `Bytes` to convert.
315
316    Returns
317    -------
318    nibble_list : `Bytes`
319        The `Bytes` in nibble-list format.
320    """
321    nibble_list = bytearray(2 * len(bytes_))
322    for byte_index, byte in enumerate(bytes_):
323        nibble_list[byte_index * 2] = (byte & 0xF0) >> 4
324        nibble_list[byte_index * 2 + 1] = byte & 0x0F
325    return Bytes(nibble_list)
326
327
328def _prepare_trie(
329    trie: Trie[K, V],
330    get_storage_root: Optional[Callable[[Address], Root]] = None,
331) -> Mapping[Bytes, Bytes]:
332    """
333    Prepares the trie for root calculation. Removes values that are empty,
334    hashes the keys (if `secured == True`) and encodes all the nodes.
335
336    Parameters
337    ----------
338    trie :
339        The `Trie` to prepare.
340    get_storage_root :
341        Function to get the storage root of an account. Needed to encode
342        `Account` objects.
343
344    Returns
345    -------
346    out : `Mapping[ethereum.base_types.Bytes, Node]`
347        Object with keys mapped to nibble-byte form.
348    """
349    mapped: MutableMapping[Bytes, Bytes] = {}
350
351    for preimage, value in trie._data.items():
352        if isinstance(value, Account):
353            assert get_storage_root is not None
354            address = Address(preimage)
355            encoded_value = encode_node(value, get_storage_root(address))
356        else:
357            encoded_value = encode_node(value)
358        # Empty values are represented by their absence
359        ensure(encoded_value != b"", AssertionError)
360        key: Bytes
361        if trie.secured:
362            # "secure" tries hash keys once before construction
363            key = keccak256(preimage)
364        else:
365            key = preimage
366        mapped[bytes_to_nibble_list(key)] = encoded_value
367
368    return mapped
369
370
371def root(
372    trie: Trie[K, V],
373    get_storage_root: Optional[Callable[[Address], Root]] = None,
374) -> Root:
375    """
376    Computes the root of a modified merkle patricia trie (MPT).
377
378    Parameters
379    ----------
380    trie :
381        `Trie` to get the root of.
382    get_storage_root :
383        Function to get the storage root of an account. Needed to encode
384        `Account` objects.
385
386
387    Returns
388    -------
389    root : `.fork_types.Root`
390        MPT root of the underlying key-value pairs.
391    """
392    obj = _prepare_trie(trie, get_storage_root)
393
394    root_node = encode_internal_node(patricialize(obj, Uint(0)))
395    if len(rlp.encode(root_node)) < 32:
396        return keccak256(rlp.encode(root_node))
397    else:
398        assert isinstance(root_node, Bytes)
399        return Root(root_node)
400
401
402def patricialize(
403    obj: Mapping[Bytes, Bytes], level: Uint
404) -> Optional[InternalNode]:
405    """
406    Structural composition function.
407
408    Used to recursively patricialize and merkleize a dictionary. Includes
409    memoization of the tree structure and hashes.
410
411    Parameters
412    ----------
413    obj :
414        Underlying trie key-value pairs, with keys in nibble-list format.
415    level :
416        Current trie level.
417
418    Returns
419    -------
420    node : `ethereum.base_types.Bytes`
421        Root node of `obj`.
422    """
423    if len(obj) == 0:
424        return None
425
426    arbitrary_key = next(iter(obj))
427
428    # if leaf node
429    if len(obj) == 1:
430        leaf = LeafNode(arbitrary_key[level:], obj[arbitrary_key])
431        return leaf
432
433    # prepare for extension node check by finding max j such that all keys in
434    # obj have the same key[i:j]
435    substring = arbitrary_key[level:]
436    prefix_length = len(substring)
437    for key in obj:
438        prefix_length = min(
439            prefix_length, common_prefix_length(substring, key[level:])
440        )
441
442        # finished searching, found another key at the current level
443        if prefix_length == 0:
444            break
445
446    # if extension node
447    if prefix_length > 0:
448        prefix = arbitrary_key[level : level + prefix_length]
449        return ExtensionNode(
450            prefix,
451            encode_internal_node(patricialize(obj, level + prefix_length)),
452        )
453
454    branches: List[MutableMapping[Bytes, Bytes]] = []
455    for _ in range(16):
456        branches.append({})
457    value = b""
458    for key in obj:
459        if len(key) == level:
460            # shouldn't ever have an account or receipt in an internal node
461            if isinstance(obj[key], (Account, Receipt, Uint)):
462                raise AssertionError
463            value = obj[key]
464        else:
465            branches[key[level]][key] = obj[key]
466
467    return BranchNode(
468        [
469            encode_internal_node(patricialize(branches[k], level + 1))
470            for k in range(16)
471        ],
472        value,
473    )