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