From d2eaf7809b92a623c07ec8f0858891221efbbad0 Mon Sep 17 00:00:00 2001 From: Gabriel Levcovitz Date: Fri, 6 Jun 2025 16:54:37 -0300 Subject: [PATCH 1/2] feat(nano): implement storage module --- hathor/nanocontracts/storage/__init__.py | 29 ++ hathor/nanocontracts/storage/backends.py | 100 +++++ hathor/nanocontracts/storage/block_storage.py | 131 ++++++ .../nanocontracts/storage/changes_tracker.py | 269 ++++++++++++ .../nanocontracts/storage/contract_storage.py | 368 +++++++++++++++++ hathor/nanocontracts/storage/factory.py | 77 ++++ .../storage/maybedeleted_nc_type.py | 81 ++++ hathor/nanocontracts/storage/node_nc_type.py | 81 ++++ hathor/nanocontracts/storage/patricia_trie.py | 389 ++++++++++++++++++ hathor/nanocontracts/storage/token_proxy.py | 36 ++ hathor/nanocontracts/storage/types.py | 27 ++ hathor/transaction/token_creation_tx.py | 8 + 12 files changed, 1596 insertions(+) create mode 100644 hathor/nanocontracts/storage/__init__.py create mode 100644 hathor/nanocontracts/storage/backends.py create mode 100644 hathor/nanocontracts/storage/block_storage.py create mode 100644 hathor/nanocontracts/storage/changes_tracker.py create mode 100644 hathor/nanocontracts/storage/contract_storage.py create mode 100644 hathor/nanocontracts/storage/factory.py create mode 100644 hathor/nanocontracts/storage/maybedeleted_nc_type.py create mode 100644 hathor/nanocontracts/storage/node_nc_type.py create mode 100644 hathor/nanocontracts/storage/patricia_trie.py create mode 100644 hathor/nanocontracts/storage/token_proxy.py create mode 100644 hathor/nanocontracts/storage/types.py diff --git a/hathor/nanocontracts/storage/__init__.py b/hathor/nanocontracts/storage/__init__.py new file mode 100644 index 000000000..37f274af2 --- /dev/null +++ b/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', +] diff --git a/hathor/nanocontracts/storage/backends.py b/hathor/nanocontracts/storage/backends.py new file mode 100644 index 000000000..d331c2831 --- /dev/null +++ b/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) diff --git a/hathor/nanocontracts/storage/block_storage.py b/hathor/nanocontracts/storage/block_storage.py new file mode 100644 index 000000000..d5edc2896 --- /dev/null +++ b/hathor/nanocontracts/storage/block_storage.py @@ -0,0 +1,131 @@ +# 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 + +import pickle +from enum import Enum +from typing import TYPE_CHECKING, NamedTuple, Optional + +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 + +if TYPE_CHECKING: + from hathor.transaction.token_creation_tx import TokenDescription + + +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.""" + + 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 = pickle.loads(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 = pickle.dumps(token_description) + self._block_trie.update(bytes(key), token_description_bytes) diff --git a/hathor/nanocontracts/storage/changes_tracker.py b/hathor/nanocontracts/storage/changes_tracker.py new file mode 100644 index 000000000..c7deb34d1 --- /dev/null +++ b/hathor/nanocontracts/storage/changes_tracker.py @@ -0,0 +1,269 @@ +# 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. + +import itertools +from dataclasses import dataclass +from enum import Enum +from types import MappingProxyType +from typing import Any, TypeVar + +from typing_extensions import override + +from hathor.conf.settings import HATHOR_TOKEN_UID +from hathor.nanocontracts.exception import NCInsufficientFunds, NCTokenAlreadyExists +from hathor.nanocontracts.nc_types import NCType +from hathor.nanocontracts.storage.contract_storage import ( + AttrKey, + Balance, + BalanceKey, + MutableBalance, + NCContractStorage, +) +from hathor.nanocontracts.storage.types import _NOT_PROVIDED, DeletedKey, DeletedKeyType +from hathor.nanocontracts.types import ContractId, TokenUid +from hathor.transaction.token_creation_tx import TokenDescription + +T = TypeVar('T') +D = TypeVar('D') + + +class _NCAuthorityState(Enum): + """The tri-state of an authority during execution.""" + NONE = 'none' + GRANTED = 'granted' + REVOKED = 'revoked' + + +@dataclass(slots=True, kw_only=True) +class _NCAuthorityDiff: + """Track the tri-state diff of each authority.""" + mint: _NCAuthorityState = _NCAuthorityState.NONE + melt: _NCAuthorityState = _NCAuthorityState.NONE + + def grant_mint(self) -> bool: + """Return whether the final mint state of this diff in granted.""" + return self.mint is _NCAuthorityState.GRANTED + + def grant_melt(self) -> bool: + """Return whether the final melt state of this diff in granted.""" + return self.melt is _NCAuthorityState.GRANTED + + def revoke_mint(self) -> bool: + """Return whether the final mint state of this diff in revoked.""" + return self.mint is _NCAuthorityState.REVOKED + + def revoke_melt(self) -> bool: + """Return whether the final melt state of this diff in revoked.""" + return self.melt is _NCAuthorityState.REVOKED + + +class NCChangesTracker(NCContractStorage): + """Keep track of changes during the execution of a contract's method. + + These changes are not committed to the storage.""" + + def __init__(self, nc_id: ContractId, storage: NCContractStorage): + self.storage = storage + self.nc_id = nc_id + + self.data: dict[AttrKey, tuple[Any, NCType | None]] = {} + self._balance_diff: dict[BalanceKey, int] = {} + self._authorities_diff: dict[BalanceKey, _NCAuthorityDiff] = {} + self._created_tokens: dict[TokenUid, TokenDescription] = {} + + self.has_been_commited = False + self.has_been_blocked = False + + def create_token(self, token_id: TokenUid, token_name: str, token_symbol: str) -> None: + """Create a new token in this changes tracker.""" + if self.has_token(token_id): + raise NCTokenAlreadyExists + self._created_tokens[token_id] = TokenDescription( + token_id=token_id, + token_name=token_name, + token_symbol=token_symbol, + ) + + def has_token(self, token_id: TokenUid) -> bool: + """Return True if a given token_id already exists.""" + if token_id in self._created_tokens: + return True + return self.storage.has_token(token_id) + + def get_balance_diff(self) -> MappingProxyType[BalanceKey, int]: + """Return the balance diff of this change tracker.""" + return MappingProxyType(self._balance_diff) + + @override + def check_if_locked(self) -> None: + if self.has_been_commited: + raise RuntimeError('you cannot change any value after the commit has been executed') + elif self.has_been_blocked: + raise RuntimeError('you cannot change any value after the changes have been blocked') + + def block(self) -> None: + """Block the changes and prevent them from being committed.""" + self.check_if_locked() + self.has_been_blocked = True + + @override + def get_obj(self, key: bytes, nc_type: NCType[T], *, default: D = _NOT_PROVIDED) -> T | D: + obj_key = self._to_attr_key(key) + obj: T | D | DeletedKeyType + if obj_key in self.data: + obj, _ = self.data[obj_key] + else: + # XXX: extra variable used so mypy can infer the correct type + obj_td = self.storage.get_obj(key, nc_type, default=default) + obj = obj_td + if obj is DeletedKey: + raise KeyError(key) + assert not isinstance(obj, DeletedKeyType) + return obj + + @override + def put_obj(self, key: bytes, nc_type: NCType[T], data: T) -> None: + self.check_if_locked() + obj_key = self._to_attr_key(key) + self.data[obj_key] = (data, nc_type) + + @override + def del_obj(self, key: bytes) -> None: + self.check_if_locked() + obj_key = self._to_attr_key(key) + self.data[obj_key] = (DeletedKey, None) + + @override + def has_obj(self, key: bytes) -> bool: + obj_key = self._to_attr_key(key) + if obj_key in self.data: + obj, _ = self.data[obj_key] + return obj is not DeletedKey + else: + return self.storage.has_obj(key) + + @override + def commit(self) -> None: + """Save the changes in the storage.""" + self.check_if_locked() + for attr_key, (obj, nc_type) in self.data.items(): + if obj is not DeletedKey: + assert nc_type is not None + assert not isinstance(obj, DeletedKeyType) + self.storage.put_obj(attr_key.key, nc_type, obj) + else: + self.storage.del_obj(attr_key.key) + + for balance_key, amount in self._balance_diff.items(): + self.storage.add_balance(balance_key.token_uid, amount) + + for balance_key, diff in self._authorities_diff.items(): + self.storage.grant_authorities( + balance_key.token_uid, + grant_mint=diff.grant_mint(), + grant_melt=diff.grant_melt(), + ) + self.storage.revoke_authorities( + balance_key.token_uid, + revoke_mint=diff.revoke_mint(), + revoke_melt=diff.revoke_melt(), + ) + + for td in self._created_tokens.values(): + self.storage.create_token(TokenUid(td.token_id), td.token_name, td.token_symbol) + + self.has_been_commited = True + + def reset(self) -> None: + """Discard all local changes without persisting.""" + self.data = {} + self._balance_diff = {} + + @override + def _get_mutable_balance(self, token_uid: bytes) -> MutableBalance: + internal_key = BalanceKey(self.nc_id, token_uid) + balance = self.storage._get_mutable_balance(token_uid) + balance_diff = self._balance_diff.get(internal_key, 0) + authorities_diff = self._authorities_diff.get(internal_key, _NCAuthorityDiff()) + + balance.value += balance_diff + balance.grant_authorities( + grant_mint=authorities_diff.grant_mint(), + grant_melt=authorities_diff.grant_melt(), + ) + balance.revoke_authorities( + revoke_mint=authorities_diff.revoke_mint(), + revoke_melt=authorities_diff.revoke_melt(), + ) + + return balance + + def validate_balances_are_positive(self) -> None: + """Check that all final balances are positive. If not, it raises NCInsufficientFunds.""" + for balance_key in self._balance_diff.keys(): + balance = self.get_balance(balance_key.token_uid) + if balance.value < 0: + raise NCInsufficientFunds( + f'negative balance for contract {self.nc_id.hex()} ' + f'(balance={balance} token_uid={balance_key.token_uid.hex()})' + ) + + @override + def get_all_balances(self) -> dict[BalanceKey, Balance]: + all_balance_keys: itertools.chain[BalanceKey] = itertools.chain( + self.storage.get_all_balances().keys(), + # There might be tokens in the change tracker that are still + # not on storage, so we must check and add them as well + self._balance_diff.keys(), + self._authorities_diff.keys(), + ) + + return {key: self.get_balance(key.token_uid) for key in set(all_balance_keys)} + + @override + def add_balance(self, token_uid: bytes, amount: int) -> None: + self.check_if_locked() + internal_key = BalanceKey(self.nc_id, token_uid) + old = self._balance_diff.get(internal_key, 0) + new = old + amount + self._balance_diff[internal_key] = new + + @override + def grant_authorities(self, token_uid: bytes, *, grant_mint: bool, grant_melt: bool) -> None: + assert token_uid != HATHOR_TOKEN_UID + self.check_if_locked() + internal_key = BalanceKey(self.nc_id, token_uid) + diff = self._authorities_diff.get(internal_key, _NCAuthorityDiff()) + diff.mint = _NCAuthorityState.GRANTED if grant_mint else diff.mint + diff.melt = _NCAuthorityState.GRANTED if grant_melt else diff.melt + self._authorities_diff[internal_key] = diff + + @override + def revoke_authorities(self, token_uid: bytes, *, revoke_mint: bool, revoke_melt: bool) -> None: + assert token_uid != HATHOR_TOKEN_UID + self.check_if_locked() + internal_key = BalanceKey(self.nc_id, token_uid) + diff = self._authorities_diff.get(internal_key, _NCAuthorityDiff()) + diff.mint = _NCAuthorityState.REVOKED if revoke_mint else diff.mint + diff.melt = _NCAuthorityState.REVOKED if revoke_melt else diff.melt + self._authorities_diff[internal_key] = diff + + def is_empty(self) -> bool: + # this method is only called in view contexts, so it's impossible for the balance to have changed. + assert not bool(self._balance_diff) + return not bool(self.data) + + @override + def get_root_id(self) -> bytes: + raise NotImplementedError diff --git a/hathor/nanocontracts/storage/contract_storage.py b/hathor/nanocontracts/storage/contract_storage.py new file mode 100644 index 000000000..eedd73012 --- /dev/null +++ b/hathor/nanocontracts/storage/contract_storage.py @@ -0,0 +1,368 @@ +# 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. + +# XXX: avoid using `from __future__ import annotations` here because `make_dataclass_nc_type` doesn't support it + +import hashlib +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum +from typing import TypeVar + +from hathor.conf.settings import HATHOR_TOKEN_UID +from hathor.nanocontracts.nc_types import BytesNCType, NCType +from hathor.nanocontracts.nc_types.dataclass_nc_type import make_dataclass_nc_type +from hathor.nanocontracts.storage.maybedeleted_nc_type import MaybeDeletedNCType +from hathor.nanocontracts.storage.patricia_trie import PatriciaTrie +from hathor.nanocontracts.storage.token_proxy import TokenProxy +from hathor.nanocontracts.storage.types import _NOT_PROVIDED, DeletedKey, DeletedKeyType +from hathor.nanocontracts.types import BlueprintId, TokenUid, VertexId +from hathor.serialization import Deserializer, Serializer + +T = TypeVar('T') +D = TypeVar('D') + +_BYTES_NC_TYPE: NCType[bytes] = BytesNCType() + + +class _Tag(Enum): + ATTR = b'\0' + BALANCE = b'\1' + METADATA = b'\2' + + +class TrieKey(ABC): + @abstractmethod + def __bytes__(self) -> bytes: + raise NotImplementedError + + +@dataclass(frozen=True, slots=True) +class AttrKey(TrieKey): + nc_id: bytes + key: bytes + + def __bytes__(self) -> bytes: + return _Tag.ATTR.value + hashlib.sha1(self.key).digest() + + +@dataclass(frozen=True, slots=True) +class BalanceKey(TrieKey): + nc_id: bytes + token_uid: bytes + + def __bytes__(self) -> bytes: + return _Tag.BALANCE.value + self.token_uid + + +@dataclass(slots=True, frozen=True, kw_only=True) +class Balance: + """ + The balance of a token in the storage, which includes its value (amount of tokens), and the + stored authorities. This class is immutable and therefore suitable to be used externally. + """ + value: int + can_mint: bool + can_melt: bool + + def to_mutable(self) -> 'MutableBalance': + return MutableBalance( + value=self.value, + can_mint=self.can_mint, + can_melt=self.can_melt, + ) + + +@dataclass(slots=True, kw_only=True) +class MutableBalance: + """ + The balance of a token in the storage, which includes its value (amount of tokens), + and the stored authorities. This is a mutable version of the `Balance` class and + therefore only suitable to be used in NCContractStorage and its subclasses. + """ + value: int + can_mint: bool + can_melt: bool + + def grant_authorities(self, *, grant_mint: bool, grant_melt: bool) -> None: + """Grant authorities to this balance, returning a new updated one.""" + self.can_mint = self.can_mint or grant_mint + self.can_melt = self.can_melt or grant_melt + + def revoke_authorities(self, *, revoke_mint: bool, revoke_melt: bool) -> None: + """Revoke authorities from this balance, returning a new updated one.""" + self.can_mint = self.can_mint and not revoke_mint + self.can_melt = self.can_melt and not revoke_melt + + @staticmethod + def get_default() -> 'MutableBalance': + """Get the default empty balance.""" + return MutableBalance(value=0, can_mint=False, can_melt=False) + + def to_immutable(self) -> Balance: + return Balance( + value=self.value, + can_mint=self.can_mint, + can_melt=self.can_melt, + ) + + +_BALANCE_NC_TYPE: NCType[MutableBalance] = make_dataclass_nc_type(MutableBalance) + + +@dataclass(frozen=True, slots=True) +class MetadataKey(TrieKey): + nc_id: bytes + key: bytes + + def __bytes__(self) -> bytes: + return _Tag.METADATA.value + hashlib.sha1(self.key).digest() + + +_BLUEPRINT_ID_KEY = b'blueprint_id' + + +class NCContractStorage: + """This is the storage used by NanoContracts. + + This implementation works for both memory and rocksdb backends.""" + + def __init__(self, *, trie: PatriciaTrie, nc_id: VertexId, token_proxy: TokenProxy) -> None: + # State (balances, metadata and attributes) + self._trie: PatriciaTrie = trie + + # Nano contract id + self.nc_id = nc_id + + # Flag to check whether any change or commit can be executed. + self.is_locked = False + + self._token_proxy = token_proxy + + def has_token(self, token_id: TokenUid) -> bool: + """Return True if token_id exists in the current block.""" + return self._token_proxy.has_token(token_id) + + def create_token(self, token_id: TokenUid, token_name: str, token_symbol: str) -> None: + """Create a new token in the current block.""" + self._token_proxy.create_token(token_id, token_name, token_symbol) + + def lock(self) -> None: + """Lock the storage for changes or commits.""" + self.is_locked = True + + def unlock(self) -> None: + """Unlock the storage.""" + self.is_locked = False + + def check_if_locked(self) -> None: + """Raise a runtime error if the wallet is locked.""" + if self.is_locked: + raise RuntimeError('you cannot modify or commit if the storage is locked') + + def _serialize(self, obj: T | DeletedKeyType, nc_type: NCType[T] | None) -> bytes: + """Serialize a obj to be stored on the trie.""" + serializer = Serializer.build_bytes_serializer() + if nc_type is None: + assert obj is DeletedKey, 'nc_type=None must only be used when obj=DeletedKey' + assert not isinstance(nc_type, MaybeDeletedNCType), 'nested MaybeDeletedNCType' + MaybeDeletedNCType(nc_type).serialize(serializer, obj) + return bytes(serializer.finalize()) + + def _deserialize(self, content: bytes, nc_type: NCType[T]) -> T | DeletedKeyType: + """Deserialize a obj stored on the trie.""" + deserializer = Deserializer.build_bytes_deserializer(content) + assert not isinstance(nc_type, MaybeDeletedNCType), 'nested MaybeDeletedNCType' + obj = MaybeDeletedNCType(nc_type).deserialize(deserializer) + if isinstance(obj, DeletedKeyType): + return DeletedKey + return obj + + def _trie_has_key(self, trie_key: TrieKey) -> bool: + """Returns True if trie-key exists and is not deleted.""" + try: + value_bytes = self._trie.get(bytes(trie_key)) + except KeyError: + return False + if MaybeDeletedNCType.is_deleted_key(value_bytes): + return False + return True + + def _trie_get_obj(self, trie_key: TrieKey, nc_type: NCType[T], *, default: D = _NOT_PROVIDED) -> T | D: + """Internal method that gets the object stored at a given trie-key.""" + obj: T | DeletedKeyType + key_bytes = bytes(trie_key) + try: + content = self._trie.get(key_bytes) + except KeyError: + obj = DeletedKey + else: + # XXX: extra variable used so mypy can infer the correct type + obj_t = self._deserialize(content, nc_type) + obj = obj_t + if obj is DeletedKey: + if default is _NOT_PROVIDED: + raise KeyError(f'trie_key={key_bytes!r}') + return default + assert not isinstance(obj, DeletedKeyType) + return obj + + def _trie_update(self, trie_key: TrieKey, nc_type: NCType[T] | None, obj: T | DeletedKeyType) -> None: + """Internal method that updates the object stored at a given trie-key + + For convenience `nc_type=None` is accepted when `obj=DeletedKey`, since it doesn't affect the serialization, so + knowing the actual NCType isn't needed. + """ + content = self._serialize(obj, nc_type) + self._trie.update(bytes(trie_key), content) + + def _to_attr_key(self, key: bytes) -> AttrKey: + """Return the actual key used in the storage.""" + assert isinstance(key, bytes) + return AttrKey(self.nc_id, key) + + def get_obj(self, key: bytes, nc_type: NCType[T], *, default: D = _NOT_PROVIDED) -> T | D: + """Return the object stored at the given `key`, deserialized with the given NCType. + + XXX: using a different NCType to deserialize than was used to serialize can result in successful + deserialization and cause silent errors. + + It raises KeyError if key is not found and a default is not provided. + """ + obj_key = self._to_attr_key(key) + try: + obj = self._trie_get_obj(obj_key, nc_type, default=default) + except KeyError as e: + raise KeyError(f'key={key!r} key_bytes={bytes(obj_key)!r}') from e + return obj + + def put_obj(self, key: bytes, nc_type: NCType[T], obj: T) -> None: + """Store the `object` for the provided `key` serialized with the given NCType. + """ + self.check_if_locked() + obj_key = self._to_attr_key(key) + self._trie_update(obj_key, nc_type, obj) + + def del_obj(self, key: bytes) -> None: + """Delete `key` from storage. + """ + self.check_if_locked() + obj_key = self._to_attr_key(key) + self._trie_update(obj_key, None, DeletedKey) + + def has_obj(self, key: bytes) -> bool: + """whether an object with the given `key` exists in the storage, also False if the object was deleted.""" + obj_key = self._to_attr_key(key) + return self._trie_has_key(obj_key) + + def _get_metadata(self, key: bytes) -> bytes: + """Return the metadata stored at the given key.""" + metadata_key = MetadataKey(self.nc_id, key) + return self._trie_get_obj(metadata_key, _BYTES_NC_TYPE) + + def _put_metadata(self, key: bytes, metadata_bytes: bytes) -> None: + """Store a new metadata at the given key.""" + metadata_key = MetadataKey(self.nc_id, key) + self._trie_update(metadata_key, _BYTES_NC_TYPE, metadata_bytes) + + def get_blueprint_id(self) -> BlueprintId: + """Return the blueprint id of the contract.""" + return BlueprintId(VertexId(self._get_metadata(_BLUEPRINT_ID_KEY))) + + def set_blueprint_id(self, blueprint_id: BlueprintId, /) -> None: + """Set a new blueprint id for the contract.""" + return self._put_metadata(_BLUEPRINT_ID_KEY, blueprint_id) + + def get_balance(self, token_uid: bytes) -> Balance: + """Return the contract balance for a token.""" + return self._get_mutable_balance(token_uid).to_immutable() + + def _get_mutable_balance(self, token_uid: bytes) -> MutableBalance: + """Return the mutable balance for a token. For internal use only.""" + balance_key = BalanceKey(self.nc_id, TokenUid(token_uid)) + balance = self._trie_get_obj(balance_key, _BALANCE_NC_TYPE, default=MutableBalance.get_default()) + assert isinstance(balance, MutableBalance) + return balance + + def get_all_balances(self) -> dict[BalanceKey, Balance]: + """Return the contract balances of all tokens.""" + balances: dict[BalanceKey, Balance] = {} + balance_tag = self._trie._encode_key(_Tag.BALANCE.value) + + node = self._trie._find_nearest_node(balance_tag) + if node.key.startswith(balance_tag): + balance_root = node + else: + for prefix, child_id in node.children.items(): + child = self._trie.get_node(child_id) + if child.key.startswith(balance_tag): + balance_root = child + break + else: + # No balance found. + return balances + + for node, _, is_leaf in self._trie.iter_dfs(node=balance_root): + if node.content is None: + # Skip all nodes with no content. + continue + # Found a token. + assert node.content is not None + balance = self._deserialize(node.content, _BALANCE_NC_TYPE) + assert isinstance(balance, MutableBalance) + token_uid = TokenUid(self._trie._decode_key(node.key)[1:]) + key = BalanceKey(self.nc_id, token_uid) + balances[key] = balance.to_immutable() + return balances + + def add_balance(self, token_uid: bytes, amount: int) -> None: + """Change the contract balance value for a token. The amount will be added to the previous balance value. + + Note that the provided `amount` might be negative, but not the result.""" + self.check_if_locked() + balance_key = BalanceKey(self.nc_id, TokenUid(token_uid)) + balance = self._trie_get_obj(balance_key, _BALANCE_NC_TYPE, default=MutableBalance.get_default()) + assert isinstance(balance, MutableBalance) + balance.value += amount + assert balance.value >= 0, f'balance cannot be negative: {balance.value}' + self._trie_update(balance_key, _BALANCE_NC_TYPE, balance) + + def grant_authorities(self, token_uid: bytes, *, grant_mint: bool, grant_melt: bool) -> None: + """Grant authorities to the contract for a token.""" + assert token_uid != HATHOR_TOKEN_UID + self.check_if_locked() + balance_key = BalanceKey(self.nc_id, TokenUid(token_uid)) + balance = self._trie_get_obj(balance_key, _BALANCE_NC_TYPE, default=MutableBalance.get_default()) + assert isinstance(balance, MutableBalance) + balance.grant_authorities(grant_mint=grant_mint, grant_melt=grant_melt) + self._trie_update(balance_key, _BALANCE_NC_TYPE, balance) + + def revoke_authorities(self, token_uid: bytes, *, revoke_mint: bool, revoke_melt: bool) -> None: + """Revoke authorities from the contract for a token.""" + assert token_uid != HATHOR_TOKEN_UID + self.check_if_locked() + balance_key = BalanceKey(self.nc_id, TokenUid(token_uid)) + balance = self._trie_get_obj(balance_key, _BALANCE_NC_TYPE, default=MutableBalance.get_default()) + assert isinstance(balance, MutableBalance) + balance.revoke_authorities(revoke_mint=revoke_mint, revoke_melt=revoke_melt) + self._trie_update(balance_key, _BALANCE_NC_TYPE, balance) + + def commit(self) -> None: + """Flush all local changes to the storage.""" + self.check_if_locked() + self._trie.commit() + + def get_root_id(self) -> bytes: + """Return the current merkle root id of the trie.""" + return self._trie.root.id diff --git a/hathor/nanocontracts/storage/factory.py b/hathor/nanocontracts/storage/factory.py new file mode 100644 index 000000000..db746f55c --- /dev/null +++ b/hathor/nanocontracts/storage/factory.py @@ -0,0 +1,77 @@ +# 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 +from typing import TYPE_CHECKING, Optional + +from hathor.nanocontracts.storage.backends import MemoryNodeTrieStore, NodeTrieStore, RocksDBNodeTrieStore +from hathor.nanocontracts.storage.block_storage import NCBlockStorage + +if TYPE_CHECKING: + from hathor.nanocontracts.storage.patricia_trie import NodeId, PatriciaTrie + from hathor.storage import RocksDBStorage + from hathor.transaction.block import Block + + +class NCStorageFactory(ABC): + _store: 'NodeTrieStore' + + @staticmethod + def bytes_to_node_id(node_id: Optional[bytes]) -> Optional['NodeId']: + from hathor.nanocontracts.storage.patricia_trie import 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 + trie = PatriciaTrie(self._store, root_id=self.bytes_to_node_id(root_id)) + return trie + + def get_block_storage_from_block(self, block: Block) -> NCBlockStorage: + raise NotImplementedError('temporarily removed during nano merge') + + def get_block_storage(self, block_root_id: bytes) -> NCBlockStorage: + """Return a non-empty block storage.""" + trie = self._get_trie(block_root_id) + return NCBlockStorage(trie) + + def get_empty_block_storage(self) -> NCBlockStorage: + """Create an empty block storage.""" + trie = self._get_trie(None) + return NCBlockStorage(trie) + + +class NCMemoryStorageFactory(NCStorageFactory): + """Factory to create a memory storage for a contract. + + As it is a memory storage, the factory keeps all contract stored data on + its attribute `self.data`. + """ + + def __init__(self) -> None: + # This attribute stores data from all contracts. + self._store = MemoryNodeTrieStore() + + +class NCRocksDBStorageFactory(NCStorageFactory): + """Factory to create a RocksDB storage for a contract. + """ + + def __init__(self, rocksdb_storage: 'RocksDBStorage') -> None: + # This store keeps data from all contracts. + self._store = RocksDBNodeTrieStore(rocksdb_storage) diff --git a/hathor/nanocontracts/storage/maybedeleted_nc_type.py b/hathor/nanocontracts/storage/maybedeleted_nc_type.py new file mode 100644 index 000000000..efcaf1676 --- /dev/null +++ b/hathor/nanocontracts/storage/maybedeleted_nc_type.py @@ -0,0 +1,81 @@ +# Copyright 2025 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 typing import TypeVar + +from typing_extensions import override + +from hathor.nanocontracts.nc_types import NCType +from hathor.nanocontracts.storage.types import DeletedKey, DeletedKeyType +from hathor.serialization import Deserializer, Serializer +from hathor.serialization.encoding.bool import decode_bool + +T = TypeVar('T') + + +class MaybeDeletedNCType(NCType[T | DeletedKeyType]): + """ Used internally to wrap a NCType or Delete + """ + + __slots__ = ('_value',) + _value: NCType[T] | None + + def __init__(self, wrapped_value: NCType[T] | None) -> None: + self._value = wrapped_value + + @classmethod + def is_deleted_key(cls, data: bytes) -> bool: + """ Shortcut to check if serializing data would result in a `DeletedKey`. + + It is possible to do that because of the serialization layout, it basically boils down to checking the first + byte of data, this is done indirectly but using the same implementation that `MaybeDeletedNCType.deserialize` + uses. + """ + deserializer = Deserializer.build_bytes_deserializer(data) + has_value = decode_bool(deserializer) + return not has_value + + @override + def _check_value(self, value: T | DeletedKeyType, /, *, deep: bool) -> None: + if isinstance(value, DeletedKeyType): + assert value is DeletedKey + return + if deep: + if self._value is None: + raise ValueError('missing inner NCType') + self._value._check_value(value, deep=deep) + + @override + def _serialize(self, serializer: Serializer, value: T | DeletedKeyType, /) -> None: + from hathor.serialization.encoding.bool import encode_bool + if value is DeletedKey: + encode_bool(serializer, False) + else: + if self._value is None: + raise ValueError('missing inner NCType') + assert not isinstance(value, DeletedKeyType) + encode_bool(serializer, True) + self._value.serialize(serializer, value) + + @override + def _deserialize(self, deserializer: Deserializer, /) -> T | DeletedKeyType: + has_value = decode_bool(deserializer) + if has_value: + if self._value is None: + raise ValueError('missing inner NCType') + return self._value.deserialize(deserializer) + else: + return DeletedKey diff --git a/hathor/nanocontracts/storage/node_nc_type.py b/hathor/nanocontracts/storage/node_nc_type.py new file mode 100644 index 000000000..dd9a74ee3 --- /dev/null +++ b/hathor/nanocontracts/storage/node_nc_type.py @@ -0,0 +1,81 @@ +# Copyright 2025 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 typing import TYPE_CHECKING + +from typing_extensions import override + +from hathor.nanocontracts.nc_types import ( + BytesLikeNCType, + BytesNCType, + DictNCType, + NCType, + OptionalNCType, + VarUint32NCType, +) +from hathor.serialization import Deserializer, Serializer + +if TYPE_CHECKING: + from hathor.nanocontracts.storage.patricia_trie import Node, NodeId + + +class NodeNCType(NCType['Node']): + """ Used internally to (de)serialize a Node into/from the database. + """ + + __slots__ = ('_key', '_length', '_content', '_children', '_id') + _key: NCType[bytes] + _length: NCType[int] + _content: NCType[bytes | None] + _children: NCType[dict[bytes, NodeId]] + # XXX: id is not optional, we're indicating that only nodes with id can be stored + _id: NCType[NodeId] + + def __init__(self) -> None: + from hathor.nanocontracts.storage.patricia_trie import NodeId + self._key = BytesNCType() + self._length = VarUint32NCType() + self._content = OptionalNCType(BytesNCType()) + # XXX: ignores because mypy can't figure out that BytesLikeNCType[NodeId] provides a NCType[NodeId] + self._children = DictNCType(BytesNCType(), BytesLikeNCType(NodeId)) # type: ignore[assignment] + self._id = BytesLikeNCType(NodeId) + + @override + def _check_value(self, value: Node, /, *, deep: bool) -> None: + from hathor.nanocontracts.storage.patricia_trie import Node + if not isinstance(value, Node): + raise TypeError('expected Node class') + + @override + def _serialize(self, serializer: Serializer, node: Node, /) -> None: + # XXX: the order is important, must be the same between de/serialization + self._key.serialize(serializer, node.key) + self._length.serialize(serializer, node.length) + self._content.serialize(serializer, node.content) + self._children.serialize(serializer, node.children) + self._id.serialize(serializer, node.id) + + @override + def _deserialize(self, deserializer: Deserializer, /) -> Node: + from hathor.nanocontracts.storage.patricia_trie import DictChildren, Node + + # XXX: the order is important, must be the same between de/serialization + key = self._key.deserialize(deserializer) + length = self._length.deserialize(deserializer) + content = self._content.deserialize(deserializer) + children = DictChildren(self._children.deserialize(deserializer)) + id_ = self._id.deserialize(deserializer) + return Node(key=key, length=length, content=content, children=children, _id=id_) diff --git a/hathor/nanocontracts/storage/patricia_trie.py b/hathor/nanocontracts/storage/patricia_trie.py new file mode 100644 index 000000000..edb1e0b64 --- /dev/null +++ b/hathor/nanocontracts/storage/patricia_trie.py @@ -0,0 +1,389 @@ +# 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. + +import hashlib +from dataclasses import dataclass, field +from itertools import chain +from typing import Iterable, NamedTuple, NewType, Optional + +from hathor.nanocontracts.storage.backends import NodeTrieStore + +NodeId = NewType('NodeId', bytes) + + +class DictChildren(dict[bytes, NodeId]): + """Data structure to store children of tree nodes.""" + def find_prefix(self, a: bytes) -> Optional[tuple[bytes, NodeId]]: + """Find the key that is a prefix of `a`.""" + # TODO Optimize search. + for key, node_id in self.items(): + if a.startswith(key): + return key, node_id + return None + + def copy(self): + """Return a copy of itself.""" + return DictChildren(self) + + +@dataclass(kw_only=True, slots=True) +class Node: + """This is a node in the Patricia trie. + + Each node can carry an object or not. If a node does not carry an object, its key has never been directly added to + the trie but it was created because some keys have the same prefix. + + Note: We might be able to remove the length. + """ + + key: bytes + length: int + content: Optional[bytes] = None + children: DictChildren = field(default_factory=DictChildren) + _id: Optional[NodeId] = None + + @property + def id(self) -> NodeId: + assert self._id is not None + return self._id + + def copy(self, content: Optional[bytes] = None, children: Optional[DictChildren] = None) -> 'Node': + """Generate a copy of this node except by the id field.""" + content = content if content is not None else self.content + children = children if children is not None else self.children.copy() + return Node(key=self.key, length=self.length, content=content, children=children) + + def calculate_id(self) -> NodeId: + """Calculate a merkle hash to serve as node id. + + This method assumes that all children already have their ids calculated. + """ + h = hashlib.sha1() + h.update(self.key) + if self.content is not None: + h.update(self.content) + sorted_child_ids = sorted(list(self.children.values())) + for child_id in sorted_child_ids: + h.update(child_id) + return NodeId(h.digest()) + + def update_id(self) -> None: + """Update node id.""" + assert self._id is None + self._id = self.calculate_id() + + +class IterDFSNode(NamedTuple): + """Item yielded by `PatriciaTrie.iter_dfs()`.""" + node: Node + height: int + is_leaf: bool + + +class PatriciaTrie: + """This object manages one or more Patricia tries; each Patricia trie is a compressed radix trie. + + All nodes are immutable. So every update will create a new path of nodes from leaves to a new root. + + - The tree structure must be the same regardless of the order the items are added. + """ + + __slots__ = ('_local_changes', '_db', 'root') + + def __init__(self, store: NodeTrieStore, *, root_id: Optional[NodeId] = None) -> None: + self._local_changes: dict[NodeId, Node] = {} + self._db = store + if root_id is None: + self.root: Node = Node(key=b'', length=0) + self.root.update_id() + self._db[self.root.id] = self.root + else: + self.root = self._db[root_id] + assert self.root.id == root_id + + def get_store(self) -> NodeTrieStore: + return self._db + + def commit(self) -> None: + """Flush all local changes from self.root to the database. All other nodes not accessed from self.root + will be discarded. + + This method should be called after all changes have been made to reduce the total number of nodes. + """ + self._commit_dfs(self.root) + self._local_changes = {} + + def _commit_dfs(self, node: Node) -> None: + """Auxiliary method to run a dfs from self.root and flush local changes to the database.""" + self._add_to_db_or_assert(node) + for child_id in node.children.values(): + child = self._local_changes.get(child_id, None) + if child is not None: + self._commit_dfs(child) + else: + assert child_id in self._db + + def _add_to_db_or_assert(self, node: Node) -> None: + """Auxiliary method to either add to the database or check consistency.""" + if node.id in self._db: + assert self._db[node.id] == node + else: + self._db[node.id] = node + + def rollback(self) -> None: + """Discard all local changes.""" + self._local_changes = {} + + def is_dirty(self) -> bool: + """Check if there is any pending local change.""" + return bool(self._local_changes) + + def get_node(self, node_id: NodeId) -> Node: + """Return a node from local changes or the database.""" + if node_id in self._local_changes: + return self._local_changes[node_id] + return self._db[node_id] + + def iter_dfs(self, *, node: Optional[Node] = None) -> Iterable[IterDFSNode]: + """Iterate from a node in a depth-first search.""" + if node is None: + node = self.root + assert node is not None + yield from self._iter_dfs(node=node, depth=0) + + def _iter_dfs(self, *, node: Node, depth: int) -> Iterable[IterDFSNode]: + """Iterate from a node in a depth-first search.""" + is_leaf = bool(not node.children) + yield IterDFSNode(node, depth, is_leaf) + for _, child_id in node.children.items(): + child = self.get_node(child_id) + yield from self._iter_dfs(node=child, depth=depth + 1) + + def _find_nearest_node(self, + key: bytes, + *, + root_id: Optional[NodeId] = None, + log_path: Optional[list[tuple[bytes, Node]]] = None) -> Node: + """Find the nearest node in the trie starting from root_id. + + Notice that it does not have to be a match. The nearest node will share the longest common + prefix with the provided key. + """ + + node: Node + if root_id is None: + node = self.root + else: + node = self.get_node(root_id) + + last_match: bytes = b'' + + while True: + if log_path is not None: + log_path.append((last_match, node)) + + if node.key == key: + return node + + suffix = key[node.length:] + match = node.children.find_prefix(suffix) + if match is not None: + last_match, next_node_id = match + else: + return node + + node = self.get_node(next_node_id) + + @staticmethod + def _find_longest_common_prefix(a: bytes, b: bytes) -> int: + """Return the index of the longest common prefix between `a` and `b`. + + If a and b does not share any prefix, returns -1. + Otherwise, return an integer in the range [0, min(|a|, |b|) - 1]. + """ + n = min(len(a), len(b)) + for i in range(n): + if a[i] != b[i]: + return i - 1 + return n - 1 + + def print_dfs(self, node: Optional[Node] = None, *, depth: int = 0) -> None: + if node is None: + node = self.root + + prefix = ' ' * depth + print(f'{prefix}key: {node.key!r}') + print(f'{prefix}length: {node.length}') + print(f'{prefix}content: {node.content!r}') + print(f'{prefix}n_children: {len(node.children)}') + print(f'{prefix}id: {node.id.hex()}') + print() + for k, child_id in node.children.items(): + print(f' {prefix}--- {k!r} ---') + child = self.get_node(child_id) + self.print_dfs(child, depth=depth + 1) + + def _build_path(self, log_path: list[tuple[bytes, Node]], new_nodes: list[tuple[bytes, Node]]) -> None: + """Build a new path of nodes from the new nodes being added and the current nodes at the trie.""" + prev_suffix: bytes | None = None + + prev_suffix, _ = new_nodes[0] + log_path_copy: list[tuple[bytes, Node]] = [] + for suffix, node in log_path[::-1]: + new_node = node.copy() + assert prev_suffix is not None + del new_node.children[prev_suffix] + log_path_copy.append((suffix, new_node)) + prev_suffix = suffix + + prev: Node | None = None + prev_suffix = None + for suffix, node in chain(new_nodes[::-1], log_path_copy): + if prev is not None: + assert prev.id is not None + assert prev_suffix is not None + node.children[prev_suffix] = prev.id + node.update_id() + self._local_changes[node.id] = node + prev = node + prev_suffix = suffix + + assert prev is not None + self.root = prev + + def _encode_key(self, key: bytes) -> bytes: + """Encode key for internal use. + + This encoding mechanism is utilized to limit the maximum number of children a node can have.""" + return key.hex().encode('ascii') + + def _decode_key(self, key: bytes) -> bytes: + """Decode key from internal format to the provided one. + + During the trie operation, keys are split and they might not be a valid hex string. + In this cases, we append a '0' at the end. + """ + if len(key) % 2 == 1: + key += b'0' + return bytes.fromhex(key.decode('ascii')) + + def _update(self, key: bytes, content: bytes) -> None: + """Internal method to update a key. + + This method never updates a node. It actually copies the node and creates a new path + from that node to the root. + """ + # The new_nodes carries the nodes that currently do not exist in the store. + # These nodes still do not have an id. Their ids will be calculated in the _build_path() method. + new_nodes: list[tuple[bytes, Node]] = [] + + # The log_path is used to backtrack the nearest node to the root. These nodes will be copied in + # the _build_path() method. + log_path: list[tuple[bytes, Node]] = [] + + # First, search for the nearest node to `key`. It either matches the key or is a prefix of the key. + parent = self._find_nearest_node(key, log_path=log_path) + # The last item in the log_path is equal to the returned node. We discard it because the parent + # will be added to the `new_nodes` later. + parent_match, _ = log_path.pop() + + if parent.key == key: + # If the nearest node stores `key`, then we will just copy it and build a new path up to the root. + new_nodes.append((parent_match, parent.copy(content=content))) + self._build_path(log_path, new_nodes) + return + + # If this point is reached, then `parent.key` is a prefix of `key`. So we have to check whether + # any of parent's children shares a prefix with `key` too. Notice that at most one children can + # share a prefix with `key`. + # TODO Optimize this search. + suffix = key[parent.length:] + for k, _v in parent.children.items(): + idx = self._find_longest_common_prefix(suffix, k) + if idx < 0: + # No share with `key`. So skip it. + continue + + # Found the child the shares a prefix with `key`. So we can stop the search. + # Now we have to add a "split node" between the parent and its child. + # + # Before: parent -> child + # After: parent -> split -> child + common_key = key[:parent.length + idx + 1] + common_key_suffix = suffix[:idx + 1] + + split = Node( + key=common_key, + length=len(common_key), + ) + split.children[k[idx + 1:]] = _v + + parent_children_copy = parent.children.copy() + del parent_children_copy[k] + new_nodes.append((parent_match, parent.copy(children=parent_children_copy))) + + # Either the split node's key equals to `key` or not. + if split.key == key: + # If they are equal, the split node will store the object and we are done. + split.content = content + new_nodes.append((common_key_suffix, split)) + self._build_path(log_path, new_nodes) + return + + # Otherwise, the split node will be the parent of the new node that will be created + # to store the object. + parent = split + parent_match = common_key_suffix + break + + # Finally, create the new node that will store the object. + assert parent.key != key + suffix = key[parent.length:] + child = Node( + key=key, + length=len(key), + content=content, + ) + new_nodes.append((parent_match, parent.copy())) + new_nodes.append((suffix, child)) + self._build_path(log_path, new_nodes) + + def _get(self, key: bytes, *, root_id: Optional[NodeId] = None) -> bytes: + """Internal method to get the object-bytes of a key.""" + if key == b'': + raise KeyError('key cannot be empty') + node = self._find_nearest_node(key, root_id=root_id) + if node.key != key: + raise KeyError + if node.content is None: + raise KeyError + return node.content + + def update(self, key: bytes, content: bytes) -> None: + """Update the object of a key. This method might change the root of the trie.""" + real_key = self._encode_key(key) + return self._update(real_key, content) + + def get(self, key: bytes, *, root_id: Optional[NodeId] = None) -> bytes: + """Return the object of a key.""" + real_key = self._encode_key(key) + return self._get(real_key, root_id=root_id) + + def has_key(self, key: bytes, *, root_id: Optional[NodeId] = None) -> bool: + """Return true if the key exists.""" + try: + self.get(key, root_id=root_id) + except KeyError: + return False + return True diff --git a/hathor/nanocontracts/storage/token_proxy.py b/hathor/nanocontracts/storage/token_proxy.py new file mode 100644 index 000000000..107362e3a --- /dev/null +++ b/hathor/nanocontracts/storage/token_proxy.py @@ -0,0 +1,36 @@ +# 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 typing import TYPE_CHECKING + +if TYPE_CHECKING: + from hathor.nanocontracts.storage.block_storage import NCBlockStorage + from hathor.nanocontracts.types import TokenUid + + +class TokenProxy: + """A proxy used to limit access to only the tokens method of a block storage. + """ + def __init__(self, block_storage: NCBlockStorage) -> None: + self.__block_storage = block_storage + + def has_token(self, token_id: TokenUid) -> bool: + """Proxy to block_storage.has_token().""" + return self.__block_storage.has_token(token_id) + + def create_token(self, token_id: TokenUid, token_name: str, token_symbol: str) -> None: + """Proxy to block_storage.create_token().""" + self.__block_storage.create_token(token_id, token_name, token_symbol) diff --git a/hathor/nanocontracts/storage/types.py b/hathor/nanocontracts/storage/types.py new file mode 100644 index 000000000..4df166c2b --- /dev/null +++ b/hathor/nanocontracts/storage/types.py @@ -0,0 +1,27 @@ +# 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 typing import Any + + +class DeletedKeyType: + pass + + +# Placeholder to mark a key as deleted in a dict. +DeletedKey = DeletedKeyType() + +# Sentinel value to differentiate where a user has provided a default value or not. +# Since _NOT_PROVIDED is a unique object, it is guaranteed not to be equal to any other value. +_NOT_PROVIDED: Any = object() diff --git a/hathor/transaction/token_creation_tx.py b/hathor/transaction/token_creation_tx.py index 629050197..9a246d116 100644 --- a/hathor/transaction/token_creation_tx.py +++ b/hathor/transaction/token_creation_tx.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import dataclass from struct import error as StructError, pack from typing import Any, Optional @@ -35,6 +36,13 @@ TOKEN_INFO_VERSION = 1 +@dataclass(slots=True, frozen=True, kw_only=True) +class TokenDescription: + token_id: bytes + token_name: str + token_symbol: str + + class TokenCreationTransaction(Transaction): def __init__( self, From 24fa8832ef8979922686e1a36d3d4f93555432a8 Mon Sep 17 00:00:00 2001 From: Gabriel Levcovitz Date: Fri, 6 Jun 2025 16:53:34 -0300 Subject: [PATCH 2/2] fix(nano): address part 4 issues (#258) --- hathor/nanocontracts/storage/block_storage.py | 13 ++++++------- hathor/nanocontracts/storage/changes_tracker.py | 2 ++ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/hathor/nanocontracts/storage/block_storage.py b/hathor/nanocontracts/storage/block_storage.py index d5edc2896..c36cf4c44 100644 --- a/hathor/nanocontracts/storage/block_storage.py +++ b/hathor/nanocontracts/storage/block_storage.py @@ -14,18 +14,15 @@ from __future__ import annotations -import pickle from enum import Enum -from typing import TYPE_CHECKING, NamedTuple, Optional +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 -if TYPE_CHECKING: - from hathor.transaction.token_creation_tx import TokenDescription - class _Tag(Enum): CONTRACT = b'\0' @@ -50,6 +47,8 @@ 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 @@ -108,7 +107,7 @@ 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 = pickle.loads(token_description_bytes) + token_description = self._TOKEN_DESCRIPTION_NC_TYPE.from_bytes(token_description_bytes) return token_description def has_token(self, token_id: TokenUid) -> bool: @@ -127,5 +126,5 @@ def create_token(self, token_id: TokenUid, token_name: str, token_symbol: str) - key = TokenKey(token_id) token_description = TokenDescription(token_id=token_id, token_name=token_name, token_symbol=token_symbol) - token_description_bytes = pickle.dumps(token_description) + token_description_bytes = self._TOKEN_DESCRIPTION_NC_TYPE.to_bytes(token_description) self._block_trie.update(bytes(key), token_description_bytes) diff --git a/hathor/nanocontracts/storage/changes_tracker.py b/hathor/nanocontracts/storage/changes_tracker.py index c7deb34d1..3d289d01d 100644 --- a/hathor/nanocontracts/storage/changes_tracker.py +++ b/hathor/nanocontracts/storage/changes_tracker.py @@ -262,6 +262,8 @@ def revoke_authorities(self, token_uid: bytes, *, revoke_mint: bool, revoke_melt def is_empty(self) -> bool: # this method is only called in view contexts, so it's impossible for the balance to have changed. assert not bool(self._balance_diff) + assert not bool(self._authorities_diff) + assert not bool(self._created_tokens) return not bool(self.data) @override