def set(self, key: bytes, value: bytes) -> Tuple[Hash32]:
"""
Returns all updated hashes in root->leaf order
"""
validate_is_bytes(key)
validate_length(key, self._key_size)
validate_is_bytes(value)
path = to_int(key)
node = value
_, branch = self._get(key)
proof_update = [] # Keep track of proof updates
target_bit = 1
# branch is in root->leaf order, so flip
for sibling_node in reversed(branch):
# Set
node_hash = keccak(node)
proof_update.append(node_hash)
self.db[node_hash] = node
# Update
if (path & target_bit):
node = sibling_node + node_hash
else:
node = node_hash + sibling_node
target_bit <<= 1
# Finally, update root hash
self.root_hash = keccak(node)
self.db[self.root_hash] = node
# updates need to be in root->leaf order, so flip back
return tuple(reversed(proof_update))
def _get(self, key: bytes) -> Tuple[bytes, Tuple[Hash32]]:
"""
Returns db value and branch in root->leaf order
"""
validate_is_bytes(key)
validate_length(key, self._key_size)
branch = []
target_bit = 1 << (self.depth - 1)
path = to_int(key)
node_hash = self.root_hash
# Append the sibling node to the branch
# Iterate on the parent
for _ in range(self.depth):
node = self.db[node_hash]
left, right = node[:32], node[32:]
if path & target_bit:
branch.append(left)
node_hash = right
else:
branch.append(right)
node_hash = left
target_bit >>= 1
# Value is the last hash in the chain
# NOTE: Didn't do exception here for testing purposes
return self.db[node_hash], tuple(branch)
def encode_branch_node(left_child_node_hash, right_child_node_hash):
"""
Serializes a branch node
"""
validate_is_bytes(left_child_node_hash)
validate_length(left_child_node_hash, 32)
validate_is_bytes(right_child_node_hash)
validate_length(right_child_node_hash, 32)
return BRANCH_TYPE_PREFIX + left_child_node_hash + right_child_node_hash
def delete(self, key: bytes) -> Tuple[Hash32]:
"""
Equals to setting the value to None
Returns all updated hashes in root->leaf order
"""
validate_is_bytes(key)
validate_length(key, self._key_size)
return self.set(key, self._default)
def exists(self, key: bytes) -> bool:
validate_is_bytes(key)
validate_length(key, self._key_size)
try:
self.get(key)
return True
except KeyError:
return False
def encode_kv_node(keypath, child_node_hash):
"""
Serializes a key/value node
"""
if keypath is None or keypath == b'':
raise ValidationError("Key path can not be empty")
validate_is_bytes(keypath)
validate_is_bytes(child_node_hash)
validate_length(child_node_hash, 32)
return KV_TYPE_PREFIX + encode_from_bin_keypath(keypath) + child_node_hash
def __init__(self, key: bytes, value: bytes, branch: Sequence[Hash32]):
validate_is_bytes(key)
validate_is_bytes(value)
validate_length(branch, len(key) * 8)
self._key = key
self._key_size = len(key)
self._value = value
self._branch = list(branch) # Avoid issues with mutable lists
self._branch_size = len(branch)
def from_db(cls,
db: Dict[bytes, bytes],
root_hash: Hash32,
key_size: int = 32,
default: bytes = BLANK_NODE):
smt = cls(key_size=key_size, default=default)
# If db is provided, and is not consistent,
# there may be a silent error. Can't solve that easily.
smt.db = db
# Set root_hash, so we know where to start
validate_is_bytes(root_hash)
validate_length(root_hash, 32) # Must be a bytes32 hash
smt.root_hash = root_hash
return smt
def update(self, key: bytes, value: bytes, node_updates: Sequence[Hash32]):
"""
Merge an update for another key with the one we are tracking internally.
:param key: keypath of the update we are processing
:param value: value of the update we are processing
:param node_updates: sequence of sibling nodes (in root->leaf order)
must be at least as large as the first diverging
key in the keypath
"""
validate_is_bytes(key)
validate_length(key, self._key_size)
# Path diff is the logical XOR of the updated key and this account
path_diff = (to_int(self.key) ^ to_int(key))
# Same key (diff of 0), update the tracked value
if path_diff == 0:
self._value = value
# No need to update branch
else:
# Find the first mismatched bit between keypaths. This is
# where the branch point occurs, and we should update the
# sibling node in the source branch at the branch point.
# NOTE: Keys are in MSB->LSB (root->leaf) order.
# Node lists are in root->leaf order.
# Be sure to convert between them effectively.
for bit in reversed(range(self._branch_size)):
if path_diff & (1 << bit) > 0:
branch_point = (self._branch_size - 1) - bit
break
# NOTE: node_updates only has to be as long as necessary
# to obtain the update. This allows an optimization
# of pruning updates to the maximum possible depth
# that would be required to update, which may be
# significantly smaller than the tree depth.
if len(node_updates) <= branch_point:
raise ValidationError("Updated node list is not deep enough")
# Update sibling node in the branch where our key differs from the update
self._branch[branch_point] = node_updates[branch_point]
def calc_root(key: bytes, value: bytes, branch: Sequence[Hash32]) -> Hash32:
"""
Obtain the merkle root of a given key/value/branch set.
Can be used to validate a merkle proof or compute it's value from data.
:param key: the keypath to decide the ordering of the sibling nodes in the branch
:param value: the value (or leaf) that starts the merkle proof computation
:param branch: the sequence of sibling nodes used to recursively perform the computation
:return: the root hash of the merkle proof computation
.. doctest::
>>> key = b'\x02' # Keypath
>>> value = b'' # Value (or leaf)
>>> branch = tuple([b'\x00'] * 8) # Any list of hashes
>>> calc_root(key, value, branch)
b'.+4IKt[\xd2\x14\xe4).\xf5\xc6\n\x11=\x01\xe89\xa1Z\x07#\xfd~(;\xfb\xb8\x8a\x0e'
"""
validate_is_bytes(key)
validate_is_bytes(value)
validate_length(branch, len(key) * 8)
path = to_int(key)
target_bit = 1
# traverse the path in leaf->root order
# branch is in root->leaf order (key is in MSB to LSB order)
node_hash = keccak(value)
for sibling_node in reversed(branch):
if path & target_bit:
node_hash = keccak(sibling_node + node_hash)
else:
node_hash = keccak(node_hash + sibling_node)
target_bit <<= 1
return node_hash
