feat(nano): implement storage module [part 4]

hathor/nanocontracts/storage/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2023 Hathor Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from hathor.nanocontracts.storage.block_storage import NCBlockStorage
+from hathor.nanocontracts.storage.changes_tracker import NCChangesTracker
+from hathor.nanocontracts.storage.contract_storage import NCContractStorage
+from hathor.nanocontracts.storage.factory import NCMemoryStorageFactory, NCRocksDBStorageFactory, NCStorageFactory
+from hathor.nanocontracts.storage.types import DeletedKey
+
+__all__ = [
+    'NCBlockStorage',
+    'NCContractStorage',
+    'NCChangesTracker',
+    'NCMemoryStorageFactory',
+    'NCRocksDBStorageFactory',
+    'NCStorageFactory',
+    'DeletedKey',
+]

hathor/nanocontracts/storage/backends.py

@@ -0,0 +1,100 @@
+# Copyright 2023 Hathor Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from hathor.nanocontracts.storage.node_nc_type import NodeNCType
+from hathor.serialization import Deserializer, Serializer
+from hathor.storage.rocksdb_storage import RocksDBStorage
+
+if TYPE_CHECKING:
+    from hathor.nanocontracts.storage.patricia_trie import Node
+
+
+class NodeTrieStore(ABC):
+    @abstractmethod
+    def __getitem__(self, key: bytes) -> Node:
+        raise NotImplementedError
+
+    @abstractmethod
+    def __setitem__(self, key: bytes, item: Node) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def __contains__(self, key: bytes) -> bool:
+        raise NotImplementedError
+
+
+class MemoryNodeTrieStore(NodeTrieStore):
+    def __init__(self) -> None:
+        self._db: dict[bytes, Node] = {}
+
+    def __getitem__(self, key: bytes) -> Node:
+        return self._db[key]
+
+    def __setitem__(self, key: bytes, item: Node) -> None:
+        self._db[key] = item
+
+    def __len__(self) -> int:
+        return len(self._db)
+
+    def __contains__(self, key: bytes) -> bool:
+        return key in self._db
+
+
+class RocksDBNodeTrieStore(NodeTrieStore):
+    _CF_NAME = b'nc-state'
+    _KEY_LENGTH = b'length'
+
+    def __init__(self, rocksdb_storage: RocksDBStorage) -> None:
+        self._rocksdb_storage = rocksdb_storage
+        self._db = self._rocksdb_storage.get_db()
+        self._cf_key = self._rocksdb_storage.get_or_create_column_family(self._CF_NAME)
+        self._node_nc_type = NodeNCType()
+
+    def _serialize_node(self, node: Node, /) -> bytes:
+        serializer = Serializer.build_bytes_serializer()
+        self._node_nc_type.serialize(serializer, node)
+        return bytes(serializer.finalize())
+
+    def _deserialize_node(self, node_bytes: bytes, /) -> Node:
+        deserializer = Deserializer.build_bytes_deserializer(node_bytes)
+        node = self._node_nc_type.deserialize(deserializer)
+        deserializer.finalize()
+        return node
+
+    def __getitem__(self, key: bytes) -> Node:
+        item_bytes = self._db.get((self._cf_key, key))
+        if item_bytes is None:
+            raise KeyError(key.hex())
+        return self._deserialize_node(item_bytes)
+
+    def __setitem__(self, key: bytes, item: Node) -> None:
+        item_bytes = self._serialize_node(item)
+        self._db.put((self._cf_key, key), item_bytes)
+
+    def __len__(self) -> int:
+        it = self._db.iterkeys()
+        it.seek_to_first()
+        return sum(1 for _ in it)
+
+    def __contains__(self, key: bytes) -> bool:
+        return bool(self._db.get((self._cf_key, key)) is not None)

hathor/nanocontracts/storage/block_storage.py

@@ -0,0 +1,130 @@
+# Copyright 2023 Hathor Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import NamedTuple, Optional
+
+from hathor.nanocontracts.nc_types.dataclass_nc_type import make_dataclass_nc_type
+from hathor.nanocontracts.storage.contract_storage import NCContractStorage
+from hathor.nanocontracts.storage.patricia_trie import NodeId, PatriciaTrie
+from hathor.nanocontracts.storage.token_proxy import TokenProxy
+from hathor.nanocontracts.types import ContractId, TokenUid
+
+
+class _Tag(Enum):
+    CONTRACT = b'\0'
+    TOKEN = b'\1'
+
+
+class ContractKey(NamedTuple):
+    nc_id: bytes
+
+    def __bytes__(self):
+        return _Tag.CONTRACT.value + self.nc_id
+
+
+class TokenKey(NamedTuple):
+    token_id: bytes
+
+    def __bytes__(self):
+        return _Tag.TOKEN.value + self.token_id
+
+
+class NCBlockStorage:
+    """This is the storage used by NanoContracts.
+
+    This implementation works for both memory and rocksdb backends."""
+    from hathor.transaction.token_creation_tx import TokenDescription
+    _TOKEN_DESCRIPTION_NC_TYPE = make_dataclass_nc_type(TokenDescription)
+
+    def __init__(self, block_trie: PatriciaTrie) -> None:
+        self._block_trie: PatriciaTrie = block_trie
+
+    def has_contract(self, contract_id: ContractId) -> bool:
+        try:
+            self.get_contract_root_id(contract_id)
+        except KeyError:
+            return False
+        else:
+            return True
+
+    def get_contract_root_id(self, contract_id: ContractId) -> bytes:
+        """Return the root id of a contract's storage."""
+        key = ContractKey(contract_id)
+        return self._block_trie.get(bytes(key))
+
+    def update_contract_trie(self, nc_id: ContractId, root_id: bytes) -> None:
+        key = ContractKey(nc_id)
+        self._block_trie.update(bytes(key), root_id)
+
+    def commit(self) -> None:
+        """Flush all local changes to the storage."""
+        self._block_trie.commit()
+
+    def get_root_id(self) -> bytes:
+        """Return the current merkle root id of the trie."""
+        return self._block_trie.root.id
+
+    @staticmethod
+    def bytes_to_node_id(node_id: Optional[bytes]) -> Optional[NodeId]:
+        if node_id is None:
+            return node_id
+        return NodeId(node_id)
+
+    def _get_trie(self, root_id: Optional[bytes]) -> 'PatriciaTrie':
+        """Return a PatriciaTrie object with a given root."""
+        from hathor.nanocontracts.storage.patricia_trie import PatriciaTrie
+        store = self._block_trie.get_store()
+        trie = PatriciaTrie(store, root_id=self.bytes_to_node_id(root_id))
+        return trie
+
+    def get_contract_storage(self, contract_id: ContractId) -> NCContractStorage:
+        nc_root_id = self.get_contract_root_id(contract_id)
+        trie = self._get_trie(nc_root_id)
+        token_proxy = TokenProxy(self)
+        return NCContractStorage(trie=trie, nc_id=contract_id, token_proxy=token_proxy)
+
+    def get_empty_contract_storage(self, contract_id: ContractId) -> NCContractStorage:
+        """Create a new contract storage instance for a given contract."""
+        trie = self._get_trie(None)
+        token_proxy = TokenProxy(self)
+        return NCContractStorage(trie=trie, nc_id=contract_id, token_proxy=token_proxy)
+
+    def get_token_description(self, token_id: TokenUid) -> TokenDescription:
+        """Return the token description for a given token_id."""
+        key = TokenKey(token_id)
+        token_description_bytes = self._block_trie.get(bytes(key))
+        token_description = self._TOKEN_DESCRIPTION_NC_TYPE.from_bytes(token_description_bytes)
+        return token_description
+
+    def has_token(self, token_id: TokenUid) -> bool:
+        """Return True if the token_id already exists in this block's nano state."""
+        key = TokenKey(token_id)
+        try:
+            self._block_trie.get(bytes(key))
+        except KeyError:
+            return False
+        else:
+            return True
+
+    def create_token(self, token_id: TokenUid, token_name: str, token_symbol: str) -> None:
+        """Create a new token in this block's nano state."""
+        from hathor.transaction.token_creation_tx import TokenDescription
+
+        key = TokenKey(token_id)
+        token_description = TokenDescription(token_id=token_id, token_name=token_name, token_symbol=token_symbol)
+        token_description_bytes = self._TOKEN_DESCRIPTION_NC_TYPE.to_bytes(token_description)
+        self._block_trie.update(bytes(key), token_description_bytes)