def _set_kv_node(self, node, trie_key, value):
current_key = extract_key(node)
common_prefix, current_key_remainder, trie_key_remainder = consume_common_prefix(
current_key,
trie_key,
)
is_extension = is_extension_node(node)
if not current_key_remainder and not trie_key_remainder:
if is_leaf_node(node):
return [node[0], value]
else:
sub_node = self._get_node(node[1])
# TODO: this needs to cleanup old storage.
new_node = self._set(sub_node, trie_key_remainder, value)
elif not current_key_remainder:
if is_extension:
sub_node = self._get_node(node[1])
# TODO: this needs to cleanup old storage.
new_node = self._set(sub_node, trie_key_remainder, value)
else:
subnode_position = trie_key_remainder[0]
subnode_key = compute_leaf_key(trie_key_remainder[1:])
sub_node = [subnode_key, value]
new_node = [BLANK_NODE] * 16 + [node[1]]
new_node[subnode_position] = self._persist_node(sub_node)
else:
new_node = [BLANK_NODE] * 17
if len(current_key_remainder) == 1 and is_extension:
new_node[current_key_remainder[0]] = node[1]
else:
if is_extension:
compute_key_fn = compute_extension_key
else:
compute_key_fn = compute_leaf_key
new_node[current_key_remainder[0]] = self._persist_node([
compute_key_fn(current_key_remainder[1:]),
node[1],
])
if trie_key_remainder:
new_node[trie_key_remainder[0]] = self._persist_node([
compute_leaf_key(trie_key_remainder[1:]),
value,
])
else:
new_node[-1] = value
if common_prefix:
new_node_key = self._persist_node(new_node)
return [compute_extension_key(common_prefix), new_node_key]
else:
return new_node
def _normalize_branch_node(self, node):
"""
A branch node which is left with only a single non-blank item should be
turned into either a leaf or extension node.
"""
iter_node = iter(node)
if any(iter_node) and any(iter_node):
return node
if node[16]:
return [compute_leaf_key([]), node[16]]
sub_node_idx, sub_node_hash = next(
(idx, v) for idx, v in enumerate(node[:16]) if v)
sub_node = self.get_node(sub_node_hash)
sub_node_type = get_node_type(sub_node)
self._prune_node(sub_node)
if sub_node_type in {NODE_TYPE_LEAF, NODE_TYPE_EXTENSION}:
new_subnode_key = encode_nibbles(
tuple(
itertools.chain(
[sub_node_idx],
decode_nibbles(sub_node[0]),
)))
return [new_subnode_key, sub_node[1]]
elif sub_node_type == NODE_TYPE_BRANCH:
subnode_hash = self._persist_node(sub_node)
return [encode_nibbles([sub_node_idx]), subnode_hash]
else:
raise Exception("Invariant: this code block should be unreachable")
def _make_simulated_node(self) -> HexaryTrieNode:
from trie.utils.nodes import (
key_starts_with,
compute_extension_key,
compute_leaf_key,
)
actual_node = self.node
key_tail = self.untraversed_tail
actual_sub_segments = actual_node.sub_segments
if len(key_tail) == 0:
raise ValueError(
"Can only raise a TraversedPartialPath when some series of nibbles was untraversed"
)
if len(actual_sub_segments) == 0:
if not key_starts_with(actual_node.suffix, key_tail):
raise ValidationError(
f"Internal traverse bug: {actual_node.suffix} does not start with {key_tail}"
)
else:
trimmed_suffix = Nibbles(actual_node.suffix[len(key_tail):])
return HexaryTrieNode(
(),
actual_node.value,
trimmed_suffix,
[compute_leaf_key(trimmed_suffix), actual_node.raw[1]],
NodeType(NODE_TYPE_LEAF),
)
elif len(actual_sub_segments) == 1:
extension = actual_sub_segments[0]
if not key_starts_with(extension, key_tail):
raise ValidationError(
f"Internal traverse bug: extension {extension} does not start with {key_tail}"
)
elif len(key_tail) == len(extension):
raise ValidationError(
f"Internal traverse bug: {key_tail} should not equal {extension}"
)
else:
trimmed_extension = Nibbles(extension[len(key_tail):])
return HexaryTrieNode(
(trimmed_extension, ),
actual_node.value,
actual_node.suffix,
[compute_extension_key(trimmed_extension), actual_node.raw[1]],
NodeType(NODE_TYPE_EXTENSION),
)
else:
raise ValidationError(
f"Can only partially traverse into leaf or extension, got {actual_node}"
)
def _set(self, node, trie_key, value):
node_type = get_node_type(node)
if node_type == NODE_TYPE_BLANK:
return [
compute_leaf_key(trie_key),
value,
]
elif node_type in {NODE_TYPE_LEAF, NODE_TYPE_EXTENSION}:
return self._set_kv_node(node, trie_key, value)
elif node_type == NODE_TYPE_BRANCH:
return self._set_branch_node(node, trie_key, value)
else:
raise Exception("Invariant: This shouldn't ever happen")
def test_TraversedPartialPath_keeps_node_value(key_encoding):
node_key = (0, 0xf, 9)
untraversed_tail = node_key[:1]
remaining_key = node_key[1:]
node_value = b'unicorns'
node = annotate_node([key_encoding(node_key), node_value])
tpp = TraversedPartialPath(node_key, node, untraversed_tail)
simulated_node = tpp.simulated_node
assert simulated_node.raw[1] == node_value
if key_encoding is compute_leaf_key:
assert simulated_node.sub_segments == ()
assert simulated_node.suffix == remaining_key
assert simulated_node.raw[0] == compute_leaf_key(remaining_key)
assert simulated_node.value == node_value
elif key_encoding is compute_extension_key:
assert simulated_node.sub_segments == (remaining_key, )
assert simulated_node.suffix == ()
assert simulated_node.raw[0] == compute_extension_key(remaining_key)
else:
raise Exception("Unsupported way to encode keys: {key_encoding}")
