astreum 0.2.61__py3-none-any.whl → 0.3.9__py3-none-any.whl

astreum/{_storage/patricia.py → storage/models/trie.py}

@@ -1,14 +1,13 @@
-import blake3
-from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
-
-from .atom import Atom, AtomKind, ZERO32
-
-if TYPE_CHECKING:
-    from .._node import Node
+from typing import Dict, List, Optional, Set, Tuple, TYPE_CHECKING
+
+from .atom import Atom, AtomKind, ZERO32
 
-class PatriciaNode:
+if TYPE_CHECKING:
+    from .._node import Node
+
+class TrieNode:
     """
-    A node in a compressed-key Patricia trie.
+    A node in a compressed-key Binary Radix Tree.
 
     Attributes:
         key_len (int): Number of bits in the `key` prefix that are meaningful.
@@ -31,29 +30,38 @@ class PatriciaNode:
         self.value = value
         self.child_0 = child_0
         self.child_1 = child_1
-        self._hash: Optional[bytes] = None
-
-    def hash(self) -> bytes:
-        """
-        Compute and cache the BLAKE3 hash of this node's serialized form.
-        """
-        if self._hash is None:
-            self._hash = blake3.blake3(self.to_bytes()).digest()
-        return self._hash
+        self._hash: Optional[bytes] = None
+
+    def hash(self) -> bytes:
+        """
+        Compute and cache the canonical hash for this node (its type-atom id).
+        """
+        if self._hash is None:
+            head_hash, _ = self._render_atoms()
+            self._hash = head_hash
+        return self._hash
+
+    def to_bytes(self) -> bytes:
+        """
+        Serialize for hashing: key_len (u16 big-endian) + key payload +
+        child_0 (or ZERO32) + child_1 (or ZERO32) + value.
+        """
+        key_len_bytes = self.key_len.to_bytes(2, "big", signed=False)
+        child0 = self.child_0 or ZERO32
+        child1 = self.child_1 or ZERO32
+        value = self.value or b""
+        return key_len_bytes + self.key + child0 + child1 + value
     
-    def to_atoms(self) -> Tuple[bytes, List[Atom]]:
+    def _render_atoms(self) -> Tuple[bytes, List[Atom]]:
         """
         Materialise this node with the canonical atom layout used by the
-        storage layer: a leading SYMBOL atom with payload ``b"radix"`` whose
+        storage layer: a leading SYMBOL atom with payload ``b"trie"`` whose
         ``next`` pointer links to four BYTES atoms containing, in order:
         key (len byte + key payload), child_0 hash, child_1 hash, value bytes.
         Returns the top atom hash and the emitted atoms.
         """
-        if self.key_len > 255:
-            raise ValueError("Patricia key length > 255 bits cannot be encoded in a single atom field")
-
         entries: List[bytes] = [
-            bytes([self.key_len]) + self.key,
+            self.key_len.to_bytes(2, "big", signed=False) + self.key,
             self.child_0 or ZERO32,
             self.child_1 or ZERO32,
             self.value or b"",
@@ -62,104 +70,75 @@ class PatriciaNode:
         data_atoms: List[Atom] = []
         next_hash = ZERO32
         for payload in reversed(entries):
-            atom = Atom.from_data(data=payload, next_hash=next_hash, kind=AtomKind.BYTES)
+            atom = Atom(data=payload, next_id=next_hash, kind=AtomKind.BYTES)
             data_atoms.append(atom)
             next_hash = atom.object_id()
 
         data_atoms.reverse()
 
-        type_atom = Atom.from_data(
-            data=b"radix",
-            next_hash=next_hash,
-            kind=AtomKind.SYMBOL,
-        )
+        type_atom = Atom(data=b"trie", next_id=next_hash, kind=AtomKind.SYMBOL)
 
         atoms = data_atoms + [type_atom]
         return type_atom.object_id(), atoms
 
-    @classmethod
-    def from_atoms(
-        cls,
-        node: "Node",
-        head_hash: bytes,
-    ) -> "PatriciaNode":
-        """
-        Reconstruct a node from the atom chain rooted at `head_hash`, using the
-        supplied `node` instance to resolve atom object ids.
-        """
-        if head_hash == ZERO32:
-            raise ValueError("empty atom chain for Patricia node")
-
-        def _atom_kind(atom: Optional[Atom]) -> Optional[AtomKind]:
-            kind_value = getattr(atom, "kind", None)
-            if isinstance(kind_value, AtomKind):
-                return kind_value
-            if isinstance(kind_value, int):
-                try:
-                    return AtomKind(kind_value)
-                except ValueError:
-                    return None
-            return None
-
-        def _require_atom(atom_hash: Optional[bytes], context: str) -> Atom:
-            if not atom_hash or atom_hash == ZERO32:
-                raise ValueError(f"missing {context}")
-            atom = node.storage_get(atom_hash)
-            if atom is None:
-                raise ValueError(f"missing {context}")
-            return atom
-
-        type_atom = _require_atom(head_hash, "Patricia type atom")
-        if _atom_kind(type_atom) is not AtomKind.SYMBOL:
-            raise ValueError("malformed Patricia node (type atom kind)")
-        if type_atom.data != b"radix":
-            raise ValueError("not a Patricia node (type mismatch)")
-
-        entries: List[bytes] = []
-        current = type_atom.next
-        hops = 0
-
-        while current and current != ZERO32 and hops < 4:
-            atom = node.storage_get(current)
-            if atom is None:
-                raise ValueError("missing atom while decoding Patricia node")
-            if _atom_kind(atom) is not AtomKind.BYTES:
-                raise ValueError("Patricia node detail atoms must be bytes")
-            entries.append(atom.data)
-            current = atom.next
-            hops += 1
-
-        if current and current != ZERO32:
-            raise ValueError("too many fields while decoding Patricia node")
-
-        if len(entries) != 4:
-            raise ValueError("incomplete atom sequence for Patricia node")
-
-        key_entry = entries[0]
-        if not key_entry:
+    def to_atoms(self) -> Tuple[bytes, List[Atom]]:
+        head_hash, atoms = self._render_atoms()
+        self._hash = head_hash
+        return head_hash, atoms
+
+    @classmethod
+    def from_atoms(
+        cls,
+        node: "Node",
+        head_hash: bytes,
+    ) -> "TrieNode":
+        """
+        Reconstruct a node from the atom chain rooted at `head_hash`, using the
+        supplied `node` instance to resolve atom object ids.
+        """
+        if head_hash == ZERO32:
+            raise ValueError("empty atom chain for Patricia node")
+
+        atom_chain = node.get_atom_list_from_storage(head_hash)
+        if atom_chain is None or len(atom_chain) != 5:
+            raise ValueError("malformed Patricia atom chain")
+
+        type_atom, key_atom, child0_atom, child1_atom, value_atom = atom_chain
+
+        if type_atom.kind is not AtomKind.SYMBOL:
+            raise ValueError("malformed Patricia node (type atom kind)")
+        if type_atom.data != b"trie":
+            raise ValueError("not a Patricia node (type mismatch)")
+
+        for detail in (key_atom, child0_atom, child1_atom, value_atom):
+            if detail.kind is not AtomKind.BYTES:
+                raise ValueError("Patricia node detail atoms must be bytes")
+
+        key_entry = key_atom.data
+        if len(key_entry) < 2:
             raise ValueError("missing key entry while decoding Patricia node")
-        key_len = key_entry[0]
-        key = key_entry[1:]
-        child_0 = entries[1] if entries[1] != ZERO32 else None
-        child_1 = entries[2] if entries[2] != ZERO32 else None
-        value = entries[3]
-
-        return cls(key_len=key_len, key=key, value=value, child_0=child_0, child_1=child_1)
+        key_len = int.from_bytes(key_entry[:2], "big", signed=False)
+        key = key_entry[2:]
+        child_0 = child0_atom.data if child0_atom.data != ZERO32 else None
+        child_1 = child1_atom.data if child1_atom.data != ZERO32 else None
+        value = value_atom.data
+
+        return cls(key_len=key_len, key=key, value=value, child_0=child_0, child_1=child_1)
 
-class PatriciaTrie:
+class Trie:
     """
-    A compressed-key Patricia trie supporting get and put.
+    A compressed-key Binary Radix Tree supporting get and put.
     """
 
-    def __init__(
-        self,
-        root_hash: Optional[bytes] = None,
-    ) -> None:
-        """
-        :param root_hash: optional hash of existing root node
-        """
-        self.nodes: Dict[bytes, PatriciaNode] = {}
-        self.root_hash = root_hash
+    def __init__(
+        self,
+        root_hash: Optional[bytes] = None,
+    ) -> None:
+        """
+        :param root_hash: optional hash of existing root node
+        """
+        self.nodes: Dict[bytes, TrieNode] = {}
+        self.root_hash = root_hash
 
     @staticmethod
     def _bit(buf: bytes, idx: int) -> bool:
@@ -189,67 +168,113 @@ class PatriciaTrie:
                 return False
         return True
 
-    def _fetch(self, storage_node: "Node", h: bytes) -> Optional[PatriciaNode]:
-        """
-        Fetch a node by hash, consulting the in-memory cache first and falling
-        back to the atom storage provided by `storage_node`.
-        """
-        cached = self.nodes.get(h)
-        if cached is not None:
-            return cached
-
-        if storage_node.storage_get(h) is None:
-            return None
-
-        pat_node = PatriciaNode.from_atoms(storage_node, h)
-        self.nodes[h] = pat_node
-        return pat_node
+    def _fetch(self, storage_node: "Node", h: bytes) -> Optional[TrieNode]:
+        """
+        Fetch a node by hash, consulting the in-memory cache first and falling
+        back to the atom storage provided by `storage_node`.
+        """
+        cached = self.nodes.get(h)
+        if cached is not None:
+            return cached
+
+        if storage_node.storage_get(h) is None:
+            return None
+
+        pat_node = TrieNode.from_atoms(storage_node, h)
+        self.nodes[h] = pat_node
+        return pat_node
 
     def get(self, storage_node: "Node", key: bytes) -> Optional[bytes]:
         """
         Return the stored value for `key`, or None if absent.
         """
-        # Empty trie?
-        if self.root_hash is None:
-            return None
+        # Empty trie?
+        if self.root_hash is None:
+            return None
+
+        current = self._fetch(storage_node, self.root_hash)
+        if current is None:
+            return None
+
+        key_pos = 0  # bit offset into key
+
+        while current is not None:
+            # 1) Check that this node's prefix matches the key here
+            if not self._match_prefix(current.key, current.key_len, key, key_pos):
+                return None
+            key_pos += current.key_len
+
+            # 2) If we've consumed all bits of the search key:
+            if key_pos == len(key) * 8:
+                # Return value only if this node actually stores one
+                return current.value
+
+            # 3) Decide which branch to follow via next bit
+            try:
+                next_bit = self._bit(key, key_pos)
+            except IndexError:
+                return None
+
+            child_hash = current.child_1 if next_bit else current.child_0
+            if child_hash is None:
+                return None  # dead end
+
+            # 4) Fetch child and continue descent
+            current = self._fetch(storage_node, child_hash)
+            if current is None:
+                return None  # dangling pointer
+
+            key_pos += 1  # consumed routing bit
+
+        return None
 
-        current = self._fetch(storage_node, self.root_hash)
-        if current is None:
-            return None
+    def get_all(self, storage_node: "Node") -> Dict[bytes, bytes]:
+        """
+        Return a mapping of every key/value pair stored in the trie.
+        """
+        if self.root_hash is None or self.root_hash == ZERO32:
+            return {}
 
-        key_pos = 0  # bit offset into key
+        def _bits_from_payload(payload: bytes, bit_length: int) -> str:
+            if bit_length <= 0 or not payload:
+                return ""
+            bit_stream = "".join(f"{byte:08b}" for byte in payload)
+            return bit_stream[:bit_length]
 
-        while current is not None:
-            # 1) Check that this node's prefix matches the key here
-            if not self._match_prefix(current.key, current.key_len, key, key_pos):
-                return None
-            key_pos += current.key_len
+        def _bits_to_bytes(bit_string: str) -> bytes:
+            if not bit_string:
+                return b""
+            pad = (8 - (len(bit_string) % 8)) % 8
+            bit_string = bit_string + ("0" * pad)
+            byte_len = len(bit_string) // 8
+            return int(bit_string, 2).to_bytes(byte_len, "big")
 
-            # 2) If we've consumed all bits of the search key:
-            if key_pos == len(key) * 8:
-                # Return value only if this node actually stores one
-                return current.value
+        results: Dict[bytes, bytes] = {}
+        stack: List[Tuple[bytes, str]] = [(self.root_hash, "")]
+        visited: Set[bytes] = set()
 
-            # 3) Decide which branch to follow via next bit
-            try:
-                next_bit = self._bit(key, key_pos)
-            except IndexError:
-                return None
+        while stack:
+            node_hash, prefix_bits = stack.pop()
+            if not node_hash or node_hash == ZERO32 or node_hash in visited:
+                continue
+            visited.add(node_hash)
 
-            child_hash = current.child_1 if next_bit else current.child_0
-            if child_hash is None:
-                return None  # dead end
+            pat_node = TrieNode.from_atoms(storage_node, node_hash)
+            node_bits = _bits_from_payload(pat_node.key, pat_node.key_len)
+            combined_bits = prefix_bits + node_bits
 
-            # 4) Fetch child and continue descent
-            current = self._fetch(storage_node, child_hash)
-            if current is None:
-                return None  # dangling pointer
+            if pat_node.value is not None:
+                key_bytes = _bits_to_bytes(combined_bits)
+                results[key_bytes] = pat_node.value
 
-            key_pos += 1  # consumed routing bit
+            if pat_node.child_0:
+                stack.append((pat_node.child_0, combined_bits + "0"))
+            if pat_node.child_1:
+                stack.append((pat_node.child_1, combined_bits + "1"))
 
-        return None
+        return results
 
-    def put(self, storage_node: "Node", key: bytes, value: bytes) -> None:
+    def put(self, storage_node: "Node", key: bytes, value: bytes) -> None:
         """
         Insert or update `key` with `value` in-place.
         """
@@ -262,64 +287,64 @@ class PatriciaTrie:
             return
 
         # S2 – traversal bookkeeping
-        stack: List[Tuple[PatriciaNode, bytes, int]] = []  # (parent, parent_hash, dir_bit)
-        current = self._fetch(storage_node, self.root_hash)
-        assert current is not None
+        stack: List[Tuple[TrieNode, bytes, int]] = []  # (parent, parent_hash, dir_bit)
+        current = self._fetch(storage_node, self.root_hash)
+        assert current is not None
         key_pos = 0
 
         # S4 – main descent loop
         while True:
             # 4.1 – prefix mismatch? → split
-            if not self._match_prefix(current.key, current.key_len, key, key_pos):
-                self._split_and_insert(current, stack, key, key_pos, value)
+            if not self._match_prefix(current.key, current.key_len, key, key_pos):
+                self._split_and_insert(current, stack, key, key_pos, value)
                 return
 
             # 4.2 – consume this prefix
-            key_pos += current.key_len
+            key_pos += current.key_len
 
             # 4.3 – matched entire key → update value
             if key_pos == total_bits:
-                old_hash = current.hash()
-                current.value = value
-                self._invalidate_hash(current)
-                new_hash = current.hash()
-                if new_hash != old_hash:
-                    self.nodes.pop(old_hash, None)
-                self.nodes[new_hash] = current
+                old_hash = current.hash()
+                current.value = value
+                self._invalidate_hash(current)
+                new_hash = current.hash()
+                if new_hash != old_hash:
+                    self.nodes.pop(old_hash, None)
+                self.nodes[new_hash] = current
                 self._bubble(stack, new_hash)
                 return
 
             # 4.4 – routing bit
             next_bit = self._bit(key, key_pos)
-            child_hash = current.child_1 if next_bit else current.child_0
+            child_hash = current.child_1 if next_bit else current.child_0
 
             # 4.6 – no child → easy append leaf
             if child_hash is None:
-                self._append_leaf(current, next_bit, key, key_pos, value, stack)
+                self._append_leaf(current, next_bit, key, key_pos, value, stack)
                 return
 
             # 4.7 – push current node onto stack
-            stack.append((current, current.hash(), int(next_bit)))
+            stack.append((current, current.hash(), int(next_bit)))
 
             # 4.8 – fetch child and continue
-            child = self._fetch(storage_node, child_hash)
-            if child is None:
-                # Dangling pointer: treat as missing child
-                parent, _, _ = stack[-1]
-                self._append_leaf(parent, next_bit, key, key_pos, value, stack[:-1])
-                return
-
-            current = child
-            key_pos += 1  # consumed routing bit
+            child = self._fetch(storage_node, child_hash)
+            if child is None:
+                # Dangling pointer: treat as missing child
+                parent, _, _ = stack[-1]
+                self._append_leaf(parent, next_bit, key, key_pos, value, stack[:-1])
+                return
+
+            current = child
+            key_pos += 1  # consumed routing bit
 
     def _append_leaf(
         self,
-        parent: PatriciaNode,
+        parent: TrieNode,
         dir_bit: bool,
         key: bytes,
         key_pos: int,
         value: bytes,
-        stack: List[Tuple[PatriciaNode, bytes, int]],
+        stack: List[Tuple[TrieNode, bytes, int]],
     ) -> None:
         tail_len = len(key) * 8 - (key_pos + 1)
         tail_bits, tail_len = self._bit_slice(key, key_pos + 1, tail_len)
@@ -342,8 +367,8 @@ class PatriciaTrie:
 
     def _split_and_insert(
         self,
-        node: PatriciaNode,
-        stack: List[Tuple[PatriciaNode, bytes, int]],
+        node: TrieNode,
+        stack: List[Tuple[TrieNode, bytes, int]],
         key: bytes,
         key_pos: int,
         value: bytes,
@@ -416,46 +441,45 @@ class PatriciaTrie:
         value: Optional[bytes],
         child0: Optional[bytes],
         child1: Optional[bytes],
-    ) -> PatriciaNode:
-        node = PatriciaNode(prefix_len, prefix_bits, value, child0, child1)
+    ) -> TrieNode:
+        node = TrieNode(prefix_len, prefix_bits, value, child0, child1)
         self.nodes[node.hash()] = node
         return node
 
-    def _invalidate_hash(self, node: PatriciaNode) -> None:
+    def _invalidate_hash(self, node: TrieNode) -> None:
         """Clear cached hash so next .hash() recomputes."""
         node._hash = None  # type: ignore
 
     def _bubble(
         self,
-        stack: List[Tuple[PatriciaNode, bytes, int]],
+        stack: List[Tuple[TrieNode, bytes, int]],
         new_hash: bytes
     ) -> None:
         """
         Propagate updated child-hash `new_hash` up the ancestor stack,
         rebasing each parent's pointer, invalidating and re-hashing.
         """
-        while stack:
-            parent, old_hash, dir_bit = stack.pop()
-
-            if dir_bit == 0:
-                parent.child_0 = new_hash
+        while stack:
+            parent, old_hash, dir_bit = stack.pop()
+
+            if dir_bit == 0:
+                parent.child_0 = new_hash
             else:
                 parent.child_1 = new_hash
 
             self._invalidate_hash(parent)
             new_hash = parent.hash()
             if new_hash != old_hash:
-                self.nodes.pop(old_hash, None)
-            self.nodes[new_hash] = parent
-            
-        self.root_hash = new_hash
-
-
-    def _bit_slice(
-        self,
-        buf: bytes,
-        start_bit: int,
-        length: int
+                self.nodes.pop(old_hash, None)
+            self.nodes[new_hash] = parent
+            
+        self.root_hash = new_hash
+
+    def _bit_slice(
+        self,
+        buf: bytes,
+        start_bit: int,
+        length: int
     ) -> tuple[bytes, int]:
         """
         Extract `length` bits from `buf` starting at `start_bit` (MSB-first),

astreum/storage/requests.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from threading import RLock
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .. import Node
+
+
+def add_atom_req(node: "Node", atom_id: bytes) -> None:
+    """Mark an atom request as pending."""
+    with node.atom_requests_lock:
+        node.atom_requests.add(atom_id)
+
+
+def has_atom_req(node: "Node", atom_id: bytes) -> bool:
+    """Return True if the atom request is currently tracked."""
+    with node.atom_requests_lock:
+        return atom_id in node.atom_requests
+
+
+def pop_atom_req(node: "Node", atom_id: bytes) -> bool:
+    """Remove the pending request if present. Returns True when removed."""
+    with node.atom_requests_lock:
+        if atom_id in node.atom_requests:
+            node.atom_requests.remove(atom_id)
+            return True
+        return False

astreum/storage/setup.py

@@ -1,15 +1,22 @@
-from pathlib import Path
-from typing import Optional, Dict, Tuple, Any
-
-def storage_setup(config: dict) -> Tuple[Optional[Path], Dict[bytes, Any], int, Dict[bytes, bytes]]:
-    storage_path_str = config.get('storage_path')
-    if storage_path_str is None:
-        storage_path, memory_storage = None, {}
-    else:
-        storage_path = Path(storage_path_str)
-        storage_path.mkdir(parents=True, exist_ok=True)
-        memory_storage = None
-
-    timeout = config.get('storage_get_relay_timeout', 5)
-    storage_index: Dict[bytes, bytes] = {}
-    return storage_path, memory_storage, timeout, storage_index
+from __future__ import annotations
+
+from typing import Any
+
+
+def storage_setup(node: Any, config: dict) -> None:
+    """Initialize hot/cold storage helpers on the node."""
+
+    node.logger.info("Setting up node storage")
+
+    node.hot_storage = {}
+    node.hot_storage_hits = {}
+    node.storage_index = {}
+    node.hot_storage_size = 0
+    node.cold_storage_size = 0
+
+    node.logger.info(
+        "Storage ready (hot_limit=%s bytes, cold_limit=%s bytes, cold_path=%s)",
+        config["hot_storage_default_limit"],
+        config["cold_storage_limit"],
+        config["cold_storage_path"] or "disabled",
+    )