Repository: Chia-Network/chia-blockchain
File: chia/data_layer/data_store.py (Python)

Code before:
        store_id: bytes32,
        root_hash: Union[bytes32, Unspecified] = unspecified,
    ) -> dict[bytes, bytes]:
        pairs = await self.get_keys_values(store_id=store_id, root_hash=root_hash)
        return {node.key: node.value for node in pairs}

    async def get_keys(
        self,
        store_id: bytes32,
        root_hash: Union[bytes32, Unspecified] = unspecified,
    ) -> list[bytes]:
        async with self.db_wrapper.reader():
            if root_hash is unspecified:
                root = await self.get_tree_root(store_id=store_id)
                resolved_root_hash = root.node_hash
            else:
                resolved_root_hash = root_hash

            try:
                merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash)
            except MerkleBlobNotFoundError:
                return []

            kv_ids = merkle_blob.get_keys_values()
            raw_key_ids = [KeyOrValueId(id.raw) for id in kv_ids.keys()]
            table_blobs = await self.get_table_blobs(raw_key_ids, store_id)
            keys: list[bytes] = []
            for kid in kv_ids.keys():
                key = self.get_blob_from_table_blob(table_blobs[KeyOrValueId(kid.raw)], store_id)
                keys.append(key)

        return keys

    def get_reference_kid_side(self, merkle_blob: MerkleBlob, seed: bytes32) -> tuple[KeyId, Side]:
        side_seed = bytes(seed)[0]
        side = Side.LEFT if side_seed < 128 else Side.RIGHT
        reference_node = merkle_blob.get_random_leaf_node(seed)
        kid = reference_node.key
        return (kid, side)

    async def get_terminal_node_from_kid(self, merkle_blob: MerkleBlob, kid: KeyId, store_id: bytes32) -> TerminalNode:
        index = merkle_blob.get_key_index(kid)
        raw_node = merkle_blob.get_raw_node(index)
        assert isinstance(raw_node, chia_rs.datalayer.LeafNode)
        return await self.get_terminal_node(raw_node.key, raw_node.value, store_id)

    async def get_terminal_node_for_seed(self, seed: bytes32, store_id: bytes32) -> Optional[TerminalNode]:
        root = await self.get_tree_root(store_id=store_id)
        if root is None or root.node_hash is None:
            return None


Review comment (suggestion):
maybe worth changing to `async def get_table_blobs(self, kv_ids: list[Union[KeyOrValueId, KeyId, ValueId]]` to avoid the conversions?  though i can't say i know offhand what sqlite does with a `KeyId` directly...  anyways, maybe close this after reading and we'll consider this separately.  just seems like it would be nice to be able to just pass `kv_ids.keys()` and be done.

Code after:
        root_hash: Union[bytes32, Unspecified] = unspecified,
    ) -> dict[bytes, bytes]:
        pairs = await self.get_keys_values(store_id=store_id, root_hash=root_hash)
        return {node.key: node.value for node in pairs}

    async def get_keys(
        self,
        store_id: bytes32,
        root_hash: Union[bytes32, Unspecified] = unspecified,
    ) -> list[bytes]:
        async with self.db_wrapper.reader():
            if root_hash is unspecified:
                root = await self.get_tree_root(store_id=store_id)
                resolved_root_hash = root.node_hash
            else:
                resolved_root_hash = root_hash

            try:
                merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=resolved_root_hash)
            except MerkleBlobNotFoundError:
                return []

            kv_ids = merkle_blob.get_keys_values()
            raw_key_ids = (KeyOrValueId(id.raw) for id in kv_ids.keys())
            table_blobs = await self.get_table_blobs(raw_key_ids, store_id)
            keys: list[bytes] = []
            for kid in kv_ids.keys():
                key = self.get_blob_from_table_blob(table_blobs[KeyOrValueId(kid.raw)], store_id)
                keys.append(key)

        return keys

    def get_reference_kid_side(self, merkle_blob: MerkleBlob, seed: bytes32) -> tuple[KeyId, Side]:
        side_seed = bytes(seed)[0]
        side = Side.LEFT if side_seed < 128 else Side.RIGHT
        reference_node = merkle_blob.get_random_leaf_node(seed)
        kid = reference_node.key
        return (kid, side)

    async def get_terminal_node_from_kid(self, merkle_blob: MerkleBlob, kid: KeyId, store_id: bytes32) -> TerminalNode:
        index = merkle_blob.get_key_index(kid)
        raw_node = merkle_blob.get_raw_node(index)
        assert isinstance(raw_node, chia_rs.datalayer.LeafNode)
        return await self.get_terminal_node(raw_node.key, raw_node.value, store_id)

    async def get_terminal_node_for_seed(self, seed: bytes32, store_id: bytes32) -> Optional[TerminalNode]:
        root = await self.get_tree_root(store_id=store_id)
        if root is None or root.node_hash is None:
            return None

        merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash)