diff --git a/src/seriousdb/serializer.py b/src/seriousdb/serializer.py index 2aa49c3..1d02b70 100644 --- a/src/seriousdb/serializer.py +++ b/src/seriousdb/serializer.py @@ -24,6 +24,8 @@ PAGE_SIZE, SLOT_FORMAT, SLOT_SIZE, + InternalNode, + LeafNode, Node, SdbMetadata, Slot, @@ -70,40 +72,37 @@ def serialize(cls, node: Node) -> bytes: if node.page_id == 0: raise SerializationError("page 0 is reserved for database metadata") - if node.leaf: - if node.children_ids or node.leftmost_child_id != 0: - raise SerializationError("leaf nodes cannot have child IDs") - - if not 0 <= node.next_page_id <= 0xFFFFFFFF: - raise SerializationError("invalid next page ID") + records: list[bytes] = [] + page_type: bytes = b"" + next_page_id: UInt32 = UInt32(0) + leftmost_child_id: UInt32 = UInt32(0) + if isinstance(node, LeafNode): if len(node.keys) != len(node.values): raise SerializationError( "leaf keys and values must have the same length" ) - else: - if node.values or node.next_page_id != 0: - raise SerializationError( - "internal nodes cannot have values or a next-page ID" - ) + for key, value in zip(node.keys, node.values): + records.append(cls._serialize_leaf_record(key, value)) - if not 0 <= node.leftmost_child_id <= 0xFFFFFFFF: - raise SerializationError("invalid leftmost child ID") + page_type = LEAF_PAGE + next_page_id = node.next_page_id + leftmost_child_id = UInt32(0) - if len(node.children_ids) != len(node.keys): + elif isinstance(node, InternalNode): + if len(node.keys) != len(node.children_ids): raise SerializationError( "internal nodes must have one child ID for each key" ) - - records: list[bytes] = [] - - if node.leaf: - for key, value in zip(node.keys, node.values): - records.append(cls._serialize_leaf_record(key, value)) - else: for key, child_id in zip(node.keys, node.children_ids): records.append(cls._serialize_internal_record(key, child_id)) + page_type = INTERNAL_PAGE + next_page_id = UInt32(0) + leftmost_child_id = node.leftmost_child_id + else: + raise SerializationError(f"unsupported node type: {type(node).__name__}") + slot_count = len(records) free_start = PAGE_HEADER_SIZE + slot_count * SLOT_SIZE free_end = PAGE_SIZE @@ -141,16 +140,14 @@ def serialize(cls, node: Node) -> bytes: slot.length, ) - page_type: bytes = LEAF_PAGE if node.leaf else INTERNAL_PAGE - header = struct.pack( PAGE_HEADER_FORMAT, page_type, slot_count, free_start, free_end, - node.next_page_id if node.leaf else 0, - node.leftmost_child_id if not node.leaf else 0, + next_page_id, + leftmost_child_id, ) page[:PAGE_HEADER_SIZE] = header @@ -234,12 +231,10 @@ def deserialize(cls, page_id: int, data: bytes) -> Node: keys.append(key) values.append(value) - return Node( + return LeafNode( page_id=UInt32(page_id), - leaf=True, keys=keys, values=values, - children_ids=[], next_page_id=UInt32(next_page_id), ) @@ -253,11 +248,9 @@ def deserialize(cls, page_id: int, data: bytes) -> Node: keys.append(key) children_ids.append(child_id) - return Node( + return InternalNode( page_id=UInt32(page_id), - leaf=False, keys=keys, - values=[], children_ids=[UInt32(child_id) for child_id in children_ids], leftmost_child_id=UInt32(leftmost_child_id), ) diff --git a/src/seriousdb/storage_format.py b/src/seriousdb/storage_format.py index 49767ad..379520e 100644 --- a/src/seriousdb/storage_format.py +++ b/src/seriousdb/storage_format.py @@ -4,7 +4,7 @@ import struct from dataclasses import dataclass, field -from typing import Final +from typing import Final, TypeAlias from seriousdb.types import UInt8, UInt16, UInt32 @@ -143,18 +143,32 @@ def _zero_uint32() -> UInt32: @dataclass(slots=True) -class Node: - """In-memory representation of a B+ tree page.""" +class LeafNode: + """Base class for a leaf node.""" page_id: UInt32 - leaf: bool keys: list[bytes] values: list[bytes] - children_ids: list[UInt32] # hacky way to give default value with "function call" to make ruff happpy next_page_id: UInt32 = field(default_factory=_zero_uint32) + + +@dataclass(slots=True) +class InternalNode: + """Base class for Internal node.""" + + page_id: UInt32 + keys: list[bytes] + children_ids: list[UInt32] leftmost_child_id: UInt32 = field(default_factory=_zero_uint32) + def all_children_ids(self) -> list[UInt32]: + """Return all child IDs in key order.""" + return [self.leftmost_child_id, *self.children_ids] + + +Node: TypeAlias = LeafNode | InternalNode + @dataclass(slots=True) class Slot: diff --git a/tests/test_serializer.py b/tests/test_serializer.py index 9b41d37..ab5a52e 100644 --- a/tests/test_serializer.py +++ b/tests/test_serializer.py @@ -15,7 +15,8 @@ PAGE_SIZE, SLOT_FORMAT, SLOT_SIZE, - Node, + InternalNode, + LeafNode, SdbMetadata, ) from seriousdb.types import UInt8, UInt16, UInt32 @@ -23,12 +24,10 @@ class TestPageSerializer: def test_leaf_round_trip(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[b"foo", b"bar"], values=[b"Alice", b"Bob"], - children_ids=[], ) data = PageSerializer.serialize(node) @@ -37,11 +36,9 @@ def test_leaf_round_trip(self) -> None: assert result == node def test_internal_round_trip(self) -> None: - node = Node( + node = InternalNode( page_id=UInt32(2), - leaf=False, keys=[b"foo", b"bar"], - values=[], children_ids=[UInt32(3), UInt32(4)], leftmost_child_id=UInt32(5), ) @@ -52,12 +49,10 @@ def test_internal_round_trip(self) -> None: assert result == node def test_empty_leaf_round_trip(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[], values=[], - children_ids=[], ) data = PageSerializer.serialize(node) @@ -66,11 +61,9 @@ def test_empty_leaf_round_trip(self) -> None: assert result == node def test_empty_internal_round_trip(self) -> None: - node = Node( + node = InternalNode( page_id=UInt32(1), - leaf=False, keys=[], - values=[], children_ids=[], ) @@ -80,12 +73,10 @@ def test_empty_internal_round_trip(self) -> None: assert result == node def test_serialized_page_has_correct_size(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[b"key"], values=[b"value"], - children_ids=[], ) data = PageSerializer.serialize(node) @@ -93,12 +84,10 @@ def test_serialized_page_has_correct_size(self) -> None: assert len(data) == PAGE_SIZE def test_page_zero_cannot_be_serialized(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(0), - leaf=True, keys=[], values=[], - children_ids=[], ) with pytest.raises(SerializationError): @@ -117,62 +106,10 @@ def test_page_larger_than_page_size_is_rejected(self) -> None: PageSerializer.deserialize(UInt32(1), bytes(PAGE_SIZE + 1)) def test_leaf_rejects_mismatched_keys_and_values(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[b"key"], values=[], - children_ids=[], - ) - - with pytest.raises(SerializationError): - PageSerializer.serialize(node) - - def test_leaf_rejects_children(self) -> None: - node = Node( - page_id=UInt32(1), - leaf=True, - keys=[], - values=[], - children_ids=[UInt32(2)], - ) - - with pytest.raises(SerializationError): - PageSerializer.serialize(node) - - def test_leaf_rejects_leftmost_child(self) -> None: - node = Node( - page_id=UInt32(1), - leaf=True, - keys=[], - values=[], - children_ids=[], - leftmost_child_id=UInt32(2), - ) - - with pytest.raises(SerializationError): - PageSerializer.serialize(node) - - def test_internal_rejects_values(self) -> None: - node = Node( - page_id=UInt32(1), - leaf=False, - keys=[], - values=[b"value"], - children_ids=[], - ) - - with pytest.raises(SerializationError): - PageSerializer.serialize(node) - - def test_internal_rejects_next_page_id(self) -> None: - node = Node( - page_id=UInt32(1), - leaf=False, - keys=[], - values=[], - children_ids=[], - next_page_id=UInt32(2), ) with pytest.raises(SerializationError): @@ -222,12 +159,10 @@ def test_internal_with_next_page_id_is_rejected(self) -> None: PageSerializer.deserialize(UInt32(1), bytes(data)) def test_non_zero_reserved_header_bytes_are_rejected(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[], values=[], - children_ids=[], ) data = bytearray(PageSerializer.serialize(node)) @@ -239,12 +174,10 @@ def test_non_zero_reserved_header_bytes_are_rejected(self) -> None: PageSerializer.deserialize(UInt32(1), bytes(data)) def test_leaf_next_page_id_round_trip(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[b"a"], values=[b"b"], - children_ids=[], next_page_id=UInt32(7), ) @@ -252,12 +185,10 @@ def test_leaf_next_page_id_round_trip(self) -> None: assert PageSerializer.deserialize(UInt32(1), data) == node def test_slot_pointing_into_slot_dirctory_is_rejected(self) -> None: - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[b"k"], values=[b"v"], - children_ids=[], ) data = bytearray(PageSerializer.serialize(node)) @@ -270,12 +201,10 @@ def test_exactly_full_page_round_trips(self) -> None: overhead = ( PAGE_HEADER_SIZE + SLOT_SIZE + struct.calcsize(LEAF_RECORD_HEADER_FORMAT) ) - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[b"k"], values=[b"v" * (PAGE_SIZE - overhead - 1)], - children_ids=[], ) data = PageSerializer.serialize(node) assert PageSerializer.deserialize(UInt32(1), data) == node @@ -284,20 +213,16 @@ def test_one_byte_too_large_is_rejected(self) -> None: overhead = ( PAGE_HEADER_SIZE + SLOT_SIZE + struct.calcsize(LEAF_RECORD_HEADER_FORMAT) ) - node = Node( + node = LeafNode( page_id=UInt32(1), - leaf=True, keys=[b"k"], values=[b"v" * (PAGE_SIZE - overhead)], - children_ids=[], ) with pytest.raises(SerializationError): PageSerializer.serialize(node) def test_internal_rejects_mismatched_keys_and_children(self) -> None: - node = Node( - page_id=UInt32(1), leaf=False, keys=[b"k"], values=[], children_ids=[] - ) + node = InternalNode(page_id=UInt32(1), keys=[b"k"], children_ids=[]) with pytest.raises(SerializationError): PageSerializer.serialize(node)