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 _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 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}")