astreum 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

astreum/node/storage/patricia.py

@@ -0,0 +1,289 @@
+import blake3
+from typing import Optional, List
+from .storage import Storage
+import astreum.utils.bytes_format as bytes_format bytes_format.decode, bytes_format.encode
+
+
+def common_prefix_length(a: bytes, b: bytes) -> int:
+    """Return the number of common prefix bytes between a and b."""
+    i = 0
+    while i < len(a) and i < len(b) and a[i] == b[i]:
+        i += 1
+    return i
+
+
+class PatriciaNode:
+    def __init__(self, key: bytes, value: Optional[bytes], children: Optional[List[bytes]] = None):
+        """
+        Initialize a Patricia node.
+        
+        :param key: A compressed part of the key.
+        :param value: The stored value (if this node represents a complete key) or None.
+        :param children: A list of child node hashes (bytes). The children are ordered by the first
+                         byte of the child's key.
+        """
+        self.key = key
+        self.value = value
+        self.children = children if children is not None else []
+        self._hash: Optional[bytes] = None
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> 'PatriciaNode':
+        """
+        Deserialize a PatriciaNode from its byte representation.
+        
+        Expected format: [key, value, children]
+        where children is a list of child node hashes (bytes).
+        """
+        decoded = bytes_format.decode(data)
+        key, value, children = decoded
+        return cls(key, value, children)
+
+    @classmethod
+    def from_storage(cls, storage: Storage, hash_value: bytes) -> Optional['PatriciaNode']:
+        """
+        Retrieve and deserialize a PatriciaNode from storage using its hash.
+        
+        :param storage: The Storage instance used to retrieve the node.
+        :param hash_value: The hash key under which the node is stored.
+        :return: A PatriciaNode instance if found, otherwise None.
+        """
+        node_bytes = storage.get(hash_value)
+        if node_bytes is None:
+            return None
+        return cls.from_bytes(node_bytes)
+
+    def to_bytes(self) -> bytes:
+        """
+        Serialize the PatriciaNode into bytes using the Astreum format.
+        
+        Structure: [key, value, children]
+        """
+        return bytes_format.encode([self.key, self.value, self.children])
+
+    def hash(self) -> bytes:
+        """
+        Compute (or retrieve a cached) Blake3 hash over the node's serialized bytes.
+        """
+        if self._hash is None:
+            self._hash = blake3.blake3(self.to_bytes()).digest()
+        return self._hash
+
+    def invalidate_hash(self) -> None:
+        """Clear the cached hash so that it is recomputed on next use."""
+        self._hash = None
+
+
+class PatriciaTrie:
+    def __init__(self, storage: Storage, root_hash: Optional[bytes] = None):
+        """
+        Initialize a Patricia Trie.
+        
+        :param storage: A Storage instance for persisting nodes.
+        :param root_hash: Optionally, an existing root hash. If None, the trie is empty.
+        """
+        self.storage = storage
+        self.root_hash = root_hash
+
+    def get(self, key: bytes) -> Optional[bytes]:
+        """
+        Retrieve the value associated with the given key.
+        
+        :param key: The key (as bytes) to search for.
+        :return: The stored value if found, otherwise None.
+        """
+        if self.root_hash is None:
+            return None
+        return self._get(self.root_hash, key)
+
+    def _get(self, node_hash: bytes, key: bytes) -> Optional[bytes]:
+        node = PatriciaNode.from_storage(self.storage, node_hash)
+        if node is None:
+            return None
+
+        cp_len = common_prefix_length(key, node.key)
+        # If node.key completely matches the beginning of key...
+        if cp_len == len(node.key):
+            if cp_len == len(key):
+                return node.value
+            remainder = key[cp_len:]
+            branch = remainder[0]
+            child_hash = self._find_child(node.children, branch)
+            if child_hash is None:
+                return None
+            return self._get(child_hash, remainder)
+        return None
+
+    def put(self, key: bytes, value: bytes) -> None:
+        """
+        Insert or update the key with the provided value.
+        
+        :param key: The key (as bytes) to insert.
+        :param value: The value (as bytes) to associate with the key.
+        """
+        if self.root_hash is None:
+            new_node = PatriciaNode(key, value, [])
+            new_hash = new_node.hash()
+            self.storage.put(new_hash, new_node.to_bytes())
+            self.root_hash = new_hash
+        else:
+            self.root_hash = self._put(self.root_hash, key, value)
+
+    def _put(self, node_hash: bytes, key: bytes, value: bytes) -> bytes:
+        """
+        Recursive helper for inserting or updating a key.
+        
+        Returns the new hash for the node that replaces the current node.
+        """
+        node = PatriciaNode.from_storage(self.storage, node_hash)
+        if node is None:
+            # Node missing: create a new leaf.
+            new_node = PatriciaNode(key, value, [])
+            new_hash = new_node.hash()
+            self.storage.put(new_hash, new_node.to_bytes())
+            return new_hash
+
+        cp_len = common_prefix_length(key, node.key)
+        len_node_key = len(node.key)
+        len_key = len(key)
+
+        # Case 1: Exact match: update the value.
+        if cp_len == len_node_key and cp_len == len_key:
+            node.value = value
+            node.invalidate_hash()
+            new_hash = node.hash()
+            self.storage.put(new_hash, node.to_bytes())
+            self.storage.delete(node_hash)
+            return new_hash
+
+        # Case 2: Node key is a prefix of key (descend to child).
+        if cp_len == len_node_key:
+            remainder = key[cp_len:]
+            branch = remainder[0]
+            child_hash = self._find_child(node.children, branch)
+            if child_hash is not None:
+                new_child_hash = self._put(child_hash, remainder, value)
+                # Update the child pointer in the list.
+                idx = self._find_child_index(node.children, branch)
+                if idx is None:
+                    raise Exception("Child not found during update.")
+                node.children[idx] = new_child_hash
+            else:
+                # Create a new leaf for the remainder.
+                new_leaf = PatriciaNode(remainder, value, [])
+                new_leaf_hash = new_leaf.hash()
+                self.storage.put(new_leaf_hash, new_leaf.to_bytes())
+                self._insert_child(node.children, new_leaf_hash)
+            node.invalidate_hash()
+            new_hash = node.hash()
+            self.storage.put(new_hash, node.to_bytes())
+            self.storage.delete(node_hash)
+            return new_hash
+
+        # Case 3: Key is a prefix of node.key (split node).
+        if cp_len == len_key and cp_len < len_node_key:
+            old_suffix = node.key[cp_len:]
+            node.key = old_suffix  # update node to hold only the suffix
+            node.invalidate_hash()
+
+            # Create a new branch node with the key as prefix and value.
+            branch_node = PatriciaNode(key, value, [])
+            # The existing node becomes a child of the branch.
+            self._insert_child(branch_node.children, node.hash())
+            branch_hash = branch_node.hash()
+            self.storage.put(branch_hash, branch_node.to_bytes())
+            self.storage.put(node.hash(), node.to_bytes())
+            self.storage.delete(node_hash)
+            return branch_hash
+
+        # Case 4: Partial common prefix (split into a branch with two children).
+        if cp_len < len_node_key and cp_len < len_key:
+            common_prefix = key[:cp_len]
+            old_suffix = node.key[cp_len:]
+            new_suffix = key[cp_len:]
+            branch_node = PatriciaNode(common_prefix, None, [])
+            
+            # Adjust the existing node.
+            node.key = old_suffix
+            node.invalidate_hash()
+            self._insert_child(branch_node.children, node.hash())
+
+            # Create a new leaf for the new key’s remaining portion.
+            new_leaf = PatriciaNode(new_suffix, value, [])
+            new_leaf_hash = new_leaf.hash()
+            self.storage.put(new_leaf_hash, new_leaf.to_bytes())
+            self._insert_child(branch_node.children, new_leaf_hash)
+
+            branch_hash = branch_node.hash()
+            self.storage.put(branch_hash, branch_node.to_bytes())
+            self.storage.put(node.hash(), node.to_bytes())
+            self.storage.delete(node_hash)
+            return branch_hash
+
+        raise Exception("Unhandled case in PatriciaTrie.put")
+
+    def _find_child(self, children: List[bytes], branch: int) -> Optional[bytes]:
+        """
+        Perform a binary search over the ordered children list to find the child hash whose
+        branch (first byte of its key) equals the target branch.
+        """
+        lo = 0
+        hi = len(children)
+        while lo < hi:
+            mid = (lo + hi) // 2
+            child_hash = children[mid]
+            child_node = PatriciaNode.from_storage(self.storage, child_hash)
+            if child_node is None or not child_node.key:
+                raise Exception("Child node missing or has empty key.")
+            child_branch = child_node.key[0]
+            if child_branch == branch:
+                return child_hash
+            elif child_branch < branch:
+                lo = mid + 1
+            else:
+                hi = mid
+        return None
+
+    def _find_child_index(self, children: List[bytes], branch: int) -> Optional[int]:
+        """
+        Similar to _find_child but returns the index in the list.
+        """
+        lo = 0
+        hi = len(children)
+        while lo < hi:
+            mid = (lo + hi) // 2
+            child_hash = children[mid]
+            child_node = PatriciaNode.from_storage(self.storage, child_hash)
+            if child_node is None or not child_node.key:
+                raise Exception("Child node missing or has empty key.")
+            child_branch = child_node.key[0]
+            if child_branch == branch:
+                return mid
+            elif child_branch < branch:
+                lo = mid + 1
+            else:
+                hi = mid
+        return None
+
+    def _insert_child(self, children: List[bytes], new_child_hash: bytes) -> None:
+        """
+        Insert a new child hash into the ordered children list.
+        """
+        new_child_node = PatriciaNode.from_storage(self.storage, new_child_hash)
+        if new_child_node is None or not new_child_node.key:
+            raise Exception("New child node missing or has empty key.")
+        new_branch = new_child_node.key[0]
+        lo = 0
+        hi = len(children)
+        while lo < hi:
+            mid = (lo + hi) // 2
+            child_hash = children[mid]
+            child_node = PatriciaNode.from_storage(self.storage, child_hash)
+            if child_node is None or not child_node.key:
+                raise Exception("Child node missing or has empty key.")
+            child_branch = child_node.key[0]
+            if child_branch < new_branch:
+                lo = mid + 1
+            else:
+                hi = mid
+        children.insert(lo, new_child_hash)