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