some rs datalayer via store

chia/_tests/core/data_layer/test_data_rpc.py

@@ -417,20 +417,26 @@ async def test_keys_values_ancestors(
         assert res is not None
         store_id = bytes32.from_hexstr(res["id"])
         await farm_block_check_singleton(data_layer, full_node_api, ph, store_id, wallet=wallet_rpc_api.service)
+        reference_hashes = []
         key1 = b"a"
         value1 = b"\x01\x02"
+        reference_hashes.append(leaf_hash(key=key1, value=value1))
         changelist: list[dict[str, str]] = [{"action": "insert", "key": key1.hex(), "value": value1.hex()}]
         key2 = b"b"
         value2 = b"\x03\x02"
+        reference_hashes.append(leaf_hash(key=key2, value=value2))
         changelist.append({"action": "insert", "key": key2.hex(), "value": value2.hex()})
         key3 = b"c"
         value3 = b"\x04\x05"
+        reference_hashes.append(leaf_hash(key=key3, value=value3))
         changelist.append({"action": "insert", "key": key3.hex(), "value": value3.hex()})
         key4 = b"d"
         value4 = b"\x06\x03"
+        reference_hashes.append(leaf_hash(key=key4, value=value4))
         changelist.append({"action": "insert", "key": key4.hex(), "value": value4.hex()})
         key5 = b"e"
         value5 = b"\x07\x01"
+        reference_hashes.append(leaf_hash(key=key5, value=value5))
         changelist.append({"action": "insert", "key": key5.hex(), "value": value5.hex()})
         res = await data_rpc_api.batch_update({"id": store_id.hex(), "changelist": changelist})
         update_tx_rec0 = res["tx_id"]
@@ -448,7 +454,7 @@ async def test_keys_values_ancestors(
         assert len(keys["keys"]) == len(dic)
         for key in keys["keys"]:
             assert key in dic
-        val = await data_rpc_api.get_ancestors({"id": store_id.hex(), "hash": val["keys_values"][4]["hash"]})
+        val = await data_rpc_api.get_ancestors({"id": store_id.hex(), "hash": reference_hashes[3].hex()})
         # todo better assertions for get_ancestors result
         assert len(val["ancestors"]) == 1
         res_before = await data_rpc_api.get_root({"id": store_id.hex()})
@@ -1019,7 +1025,7 @@ async def process_for_data_layer_keys(
             value = await data_layer.get_value(store_id=store_id, key=expected_key)
         except Exception as e:
             # TODO: more specific exceptions...
-            if "Key not found" not in str(e):
+            if "Key not found" not in str(e) and "unknown key" not in str(e):
                 raise  # pragma: no cover
         else:
             if expected_value is None or value == expected_value:
@@ -3437,7 +3443,9 @@ async def test_unsubmitted_batch_update(
             count=NUM_BLOCKS_WITHOUT_SUBMIT, guarantee_transaction_blocks=True
         )
         keys_values = await data_rpc_api.get_keys_values({"id": store_id.hex()})
-        assert keys_values == prev_keys_values
+        assert {item["key"]: item for item in keys_values["keys_values"]} == {
+            item["key"]: item for item in prev_keys_values["keys_values"]
+        }
 
         pending_root = await data_layer.data_store.get_pending_root(store_id=store_id)
         assert pending_root is not None

chia/_tests/core/data_layer/test_data_store.py

@@ -14,6 +14,7 @@
 from typing import Any, BinaryIO, Callable, Optional
 
 import aiohttp
+import chia_rs.datalayer
 import pytest
 from chia_rs.datalayer import TreeIndex
 
@@ -25,8 +26,6 @@
     InternalNode,
     Node,
     OperationType,
-    ProofOfInclusion,
-    ProofOfInclusionLayer,
     Root,
     SerializedNode,
     ServerInfo,
@@ -35,14 +34,14 @@
     Subscription,
     TerminalNode,
     _debug_dump,
+    as_program,
     get_delta_filename_path,
     get_full_tree_filename_path,
     leaf_hash,
 )
-from chia.data_layer.data_store import DataStore
+from chia.data_layer.data_store import DataStore, InternalTypes, LeafTypes, MerkleBlobHint
 from chia.data_layer.download_data import insert_from_delta_file, write_files_for_root
 from chia.data_layer.util.benchmark import generate_datastore
-from chia.data_layer.util.merkle_blob import MerkleBlob, RawInternalMerkleNode, RawLeafMerkleNode
 from chia.types.blockchain_format.program import Program
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.byte_types import hexstr_to_bytes
@@ -492,8 +491,8 @@ async def test_insert_batch_reference_and_side(
     nodes_with_indexes = merkle_blob.get_nodes_with_indexes()
     nodes = [pair[1] for pair in nodes_with_indexes]
     assert len(nodes) == 3
-    assert isinstance(nodes[1], RawLeafMerkleNode)
-    assert isinstance(nodes[2], RawLeafMerkleNode)
+    assert isinstance(nodes[1], LeafTypes)
+    assert isinstance(nodes[2], LeafTypes)
     left_terminal_node = await data_store.get_terminal_node(nodes[1].key, nodes[1].value, store_id)
     right_terminal_node = await data_store.get_terminal_node(nodes[2].key, nodes[2].value, store_id)
     if side == Side.LEFT:
@@ -750,24 +749,24 @@ async def test_proof_of_inclusion_by_hash(data_store: DataStore, store_id: bytes
     await _debug_dump(db=data_store.db_wrapper)
 
     expected_layers = [
-        ProofOfInclusionLayer(
+        chia_rs.datalayer.ProofOfInclusionLayer(
             other_hash_side=Side.RIGHT,
             other_hash=bytes32.fromhex("fb66fe539b3eb2020dfbfadfd601fa318521292b41f04c2057c16fca6b947ca1"),
             combined_hash=bytes32.fromhex("36cb1fc56017944213055da8cb0178fb0938c32df3ec4472f5edf0dff85ba4a3"),
         ),
-        ProofOfInclusionLayer(
+        chia_rs.datalayer.ProofOfInclusionLayer(
             other_hash_side=Side.RIGHT,
             other_hash=bytes32.fromhex("6d3af8d93db948e8b6aa4386958e137c6be8bab726db86789594b3588b35adcd"),
             combined_hash=bytes32.fromhex("5f67a0ab1976e090b834bf70e5ce2a0f0a9cd474e19a905348c44ae12274d30b"),
         ),
-        ProofOfInclusionLayer(
+        chia_rs.datalayer.ProofOfInclusionLayer(
             other_hash_side=Side.LEFT,
             other_hash=bytes32.fromhex("c852ecd8fb61549a0a42f9eb9dde65e6c94a01934dbd9c1d35ab94e2a0ae58e2"),
             combined_hash=bytes32.fromhex("7a5193a4e31a0a72f6623dfeb2876022ab74a48abb5966088a1c6f5451cc5d81"),
         ),
     ]
 
-    assert proof == ProofOfInclusion(node_hash=node.hash, layers=expected_layers)
+    assert proof == chia_rs.datalayer.ProofOfInclusion(node_hash=node.hash, layers=expected_layers)
 
 
 @pytest.mark.anyio
@@ -780,7 +779,7 @@ async def test_proof_of_inclusion_by_hash_no_ancestors(data_store: DataStore, st
 
     proof = await data_store.get_proof_of_inclusion_by_hash(node_hash=node.hash, store_id=store_id)
 
-    assert proof == ProofOfInclusion(node_hash=node.hash, layers=[])
+    assert proof == chia_rs.datalayer.ProofOfInclusion(node_hash=node.hash, layers=[])
 
 
 @pytest.mark.anyio
@@ -792,7 +791,7 @@ async def test_proof_of_inclusion_by_hash_program(data_store: DataStore, store_i
 
     proof = await data_store.get_proof_of_inclusion_by_hash(node_hash=node.hash, store_id=store_id)
 
-    assert proof.as_program() == [
+    assert as_program(proof) == [
         b"\x04",
         [
             bytes32.fromhex("fb66fe539b3eb2020dfbfadfd601fa318521292b41f04c2057c16fca6b947ca1"),
@@ -833,7 +832,7 @@ async def test_proof_of_inclusion_by_hash_bytes(data_store: DataStore, store_id:
         b"\xe2\xa0\xaeX\xe2\x80\x80"
     )
 
-    assert bytes(proof.as_program()) == expected
+    assert bytes(as_program(proof)) == expected
 
 
 # @pytest.mark.anyio
@@ -1286,7 +1285,7 @@ async def write_tree_to_file_old_format(
     node_hash: bytes32,
     store_id: bytes32,
     writer: BinaryIO,
-    merkle_blob: Optional[MerkleBlob] = None,
+    merkle_blob: Optional[MerkleBlobHint] = None,
     hash_to_index: Optional[dict[bytes32, TreeIndex]] = None,
 ) -> None:
     if node_hash == bytes32.zeros:
@@ -1306,13 +1305,13 @@ async def write_tree_to_file_old_format(
     raw_node = merkle_blob.get_raw_node(raw_index)
 
     to_write = b""
-    if isinstance(raw_node, RawInternalMerkleNode):
+    if isinstance(raw_node, InternalTypes):
         left_hash = merkle_blob.get_hash_at_index(raw_node.left)
         right_hash = merkle_blob.get_hash_at_index(raw_node.right)
         await write_tree_to_file_old_format(data_store, root, left_hash, store_id, writer, merkle_blob, hash_to_index)
         await write_tree_to_file_old_format(data_store, root, right_hash, store_id, writer, merkle_blob, hash_to_index)
         to_write = bytes(SerializedNode(False, bytes(left_hash), bytes(right_hash)))
-    elif isinstance(raw_node, RawLeafMerkleNode):
+    elif isinstance(raw_node, LeafTypes):
         node = await data_store.get_terminal_node(raw_node.key, raw_node.value, store_id)
         to_write = bytes(SerializedNode(True, node.key, node.value))
     else:
@@ -1720,7 +1719,7 @@ async def mock_http_download_2(
         filenames = {entry.name for entry in entries}
         assert len(filenames) == num_files + max_full_files - 1
     kv = await data_store.get_keys_values(store_id=store_id)
-    assert kv == kv_before
+    assert set(kv) == set(kv_before)
 
 
 @pytest.mark.anyio
@@ -1758,7 +1757,7 @@ async def test_get_node_by_key_with_overlapping_keys(raw_data_store: DataStore)
                 if random.randint(0, 4) == 0:
                     batch = [{"action": "delete", "key": key}]
                     await raw_data_store.insert_batch(store_id, batch, status=Status.COMMITTED)
-                    with pytest.raises(KeyNotFoundError, match=f"Key not found: {key.hex()}"):
+                    with pytest.raises((KeyNotFoundError, chia_rs.datalayer.UnknownKeyError)):
                         await raw_data_store.get_node_by_key(store_id=store_id, key=key)
 
 
@@ -1827,7 +1826,7 @@ async def test_insert_from_delta_file_correct_file_exists(
         filenames = {entry.name for entry in entries}
         assert len(filenames) == num_files + 2  # 1 full and 6 deltas
     kv = await data_store.get_keys_values(store_id=store_id)
-    assert kv == kv_before
+    assert set(kv) == set(kv_before)
 
 
 @pytest.mark.anyio
@@ -2042,7 +2041,7 @@ async def test_migration(
     data_store.recent_merkle_blobs = LRUCache(capacity=128)
     assert await data_store.get_keys_values(store_id=store_id) == []
     await data_store.migrate_db(tmp_path)
-    assert await data_store.get_keys_values(store_id=store_id) == kv_before
+    assert set(await data_store.get_keys_values(store_id=store_id)) == set(kv_before)
 
 
 @pytest.mark.anyio

chia/data_layer/data_layer.py

@@ -34,6 +34,7 @@
     ProofOfInclusionLayer,
     Root,
     ServerInfo,
+    Side,
     Status,
     StoreProofs,
     Subscription,
@@ -44,6 +45,8 @@
     get_delta_filename_path,
     get_full_tree_filename_path,
     leaf_hash,
+    sibling_hashes,
+    sibling_sides_integer,
     unspecified,
 )
 from chia.data_layer.data_layer_wallet import DataLayerWallet, Mirror, verify_offer
@@ -1065,7 +1068,7 @@ async def process_offered_stores(self, offer_stores: tuple[OfferStore, ...]) ->
                         node_hash=proof_of_inclusion.node_hash,
                         layers=tuple(
                             Layer(
-                                other_hash_side=layer.other_hash_side,
+                                other_hash_side=Side(layer.other_hash_side),
                                 other_hash=layer.other_hash,
                                 combined_hash=layer.combined_hash,
                             )
@@ -1164,12 +1167,12 @@ async def take_offer(
                         for layer in proof.layers
                     ]
                     proof_of_inclusion = ProofOfInclusion(node_hash=proof.node_hash, layers=layers)
-                    sibling_sides_integer = proof_of_inclusion.sibling_sides_integer()
+                    sibling_sides_integer_value = sibling_sides_integer(proof_of_inclusion)
                     proofs_of_inclusion.append(
                         (
                             root.hex(),
-                            str(sibling_sides_integer),
-                            ["0x" + sibling_hash.hex() for sibling_hash in proof_of_inclusion.sibling_hashes()],
+                            str(sibling_sides_integer_value),
+                            ["0x" + sibling_hash.hex() for sibling_hash in sibling_hashes(proof_of_inclusion)],
                         )
                     )
 

chia/data_layer/data_layer_util.py

@@ -5,9 +5,10 @@
 from enum import Enum, IntEnum
 from hashlib import sha256
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 import aiosqlite
+import chia_rs.datalayer
 from typing_extensions import final
 
 from chia.data_layer.data_layer_errors import ProofIntegrityError
@@ -24,6 +25,9 @@
     from chia.data_layer.data_store import DataStore
     from chia.wallet.wallet_node import WalletNode
 
+ProofOfInclusionHint = Union["ProofOfInclusion", chia_rs.datalayer.ProofOfInclusion]
+ProofOfInclusionLayerHint = Union["ProofOfInclusionLayer", chia_rs.datalayer.ProofOfInclusionLayer]
+
 
 def internal_hash(left_hash: bytes32, right_hash: bytes32) -> bytes32:
     # see test for the definition this is optimized from
@@ -187,7 +191,7 @@ class NodeType(IntEnum):
 
 
 @final
-class Side(IntEnum):
+class Side(uint8, Enum):
     LEFT = 0
     RIGHT = 1
 
@@ -278,7 +282,31 @@ def from_hashes(cls, primary_hash: bytes32, other_hash_side: Side, other_hash: b
         return cls(other_hash_side=other_hash_side, other_hash=other_hash, combined_hash=combined_hash)
 
 
-other_side_to_bit = {Side.LEFT: 1, Side.RIGHT: 0}
+def sibling_sides_integer(proof: ProofOfInclusionHint) -> int:
+    # casting to workaround this
+    # class C: ...
+    # class D: ...
+    #
+    # m: list[C | D]
+    # reveal_type(enumerate(m))
+    # # main.py:5: note: Revealed type is "builtins.enumerate[Union[__main__.C, __main__.D]]"
+    #
+    # n: list[C] | list[D]
+    # reveal_type(enumerate(n))
+    # main.py:9: note: Revealed type is "builtins.enumerate[builtins.object]"
+
+    return sum(
+        (1 << index if cast(ProofOfInclusionLayerHint, layer).other_hash_side == Side.LEFT else 0)
+        for index, layer in enumerate(proof.layers)
+    )
+
+
+def sibling_hashes(proof: ProofOfInclusionHint) -> list[bytes32]:
+    return [layer.other_hash for layer in proof.layers]
+
+
+def as_program(proof: ProofOfInclusionHint) -> Program:
+    return Program.to([sibling_sides_integer(proof), sibling_hashes(proof)])
 
 
 @dataclass(frozen=True)
@@ -293,15 +321,6 @@ def root_hash(self) -> bytes32:
 
         return self.layers[-1].combined_hash
 
-    def sibling_sides_integer(self) -> int:
-        return sum(other_side_to_bit[layer.other_hash_side] << index for index, layer in enumerate(self.layers))
-
-    def sibling_hashes(self) -> list[bytes32]:
-        return [layer.other_hash for layer in self.layers]
-
-    def as_program(self) -> Program:
-        return Program.to([self.sibling_sides_integer(), self.sibling_hashes()])
-
     def valid(self) -> bool:
         existing_hash = self.node_hash