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 _get_key_after(self, node: HexaryTrieNode, key: Nibbles,
traversed: Nibbles) -> Optional[Nibbles]:
"""
Find the next key in the trie after key
:param node: the source node to search for the next key after `key`
:param key: the starting key used to seek the nearest key on the right
:param traversed: the nibbles already traversed to get down to `node`
:return: the complete key that is immediately to the right of `key` or None,
if no key is immediately to the right (under `node`)
"""
for next_segment in node.sub_segments:
if key[:len(next_segment)] > next_segment:
# This segment is to the left of the key, keep looking...
continue
else:
# Either: found the exact match, or the next result to the right
# Either way, we'll want to take a look
next_node = self._trie.traverse_from(node, next_segment)
common, key_remaining, segment_remaining = consume_common_prefix(
key, next_segment)
if len(segment_remaining) == 0:
# Found a perfect match! Keep looking for keys to the right of the target
next_key = self._get_key_after(
next_node,
key_remaining,
traversed + next_segment,
)
if next_key is None:
# Could not find a key to the right in any sub-node.
# In other words, *only* the target key is in this sub-trie
# So keep looking to the right...
continue
else:
# We successfully found a key to the right in a subtree, return it up
return next_key
else:
# Found no exact match, and are now looking for the next possible key
return self._get_next_key(next_node,
traversed + next_segment)
if node.suffix > key:
# This leaf node is to the right of the target key
return traversed + node.suffix
else:
# Nothing found in any sub-segments
return None
def should_continue(node_path, end_path):
common, node_rest, end_rest = consume_common_prefix(node_path, end_path)
if len(node_rest) == 0:
# e.g. node_path, end_path = ffa, ffaf
# everything from ffa{0-f} is yet to be explored
return True
if len(node_rest) and len(end_rest):
# e.g. node_path, end_path = 0, 11
# everything is left to be explored!
return node_rest[0] < end_rest[0]
if len(end_rest) == 0:
# e.g. node_path, end_path = 110, 11
# everything which remains will lie after {end_path}
return False
print(node_path, end_path, common, node_rest, end_rest)
raise Exception('oh no')
def _traverse_extension(self, node, trie_key):
current_key = extract_key(node)
common_prefix, current_key_remainder, trie_key_remainder = consume_common_prefix(
current_key,
trie_key,
)
if len(current_key_remainder) == 0:
# The full extension node's key was consumed
return node[1], trie_key_remainder
elif len(trie_key_remainder) == 0:
# The trie key was consumed before reaching the end of the extension node's key
raise _PartialTraversal
else:
# The trie key and extension node key branch away from each other, so there
# is no node at the specified key.
return BLANK_NODE, ()
def traverse_prefix(db, root, prefix):
"""
Return all nodes with the given prefix
"""
for node in traverse_node(db, tuple(), root, prefix):
common, node_rest, prefix_rest = consume_common_prefix(
node.path, prefix)
if len(node_rest) == 0 and len(prefix_rest) != 0:
# we haven't reached the starting node yet
continue
if len(node_rest) == 0 and len(prefix_rest) == 0:
# this is the starting node!
yield node
continue
if len(prefix_rest) == 0 and len(node_rest):
# this node is part of our prefix!
yield node
continue
if len(prefix_rest) and len(node_rest):
# this node is past our prefix
break
def test_consume_common_prefix(left, right, expected):
actual_a = consume_common_prefix(left, right)
actual_b = consume_common_prefix(right, left)
expected_b = (expected[0], expected[2], expected[1])
assert actual_a == expected
assert actual_b == expected_b