feat(nano): implement new metadata [part 10]

hathor/indexes/manager.py

@@ -20,6 +20,7 @@
 from typing import TYPE_CHECKING, Iterator, Optional
 
 from structlog import get_logger
+from typing_extensions import assert_never
 
 from hathor.indexes.address_index import AddressIndex
 from hathor.indexes.base_index import BaseIndex
@@ -35,6 +36,7 @@
 from hathor.indexes.tokens_index import TokensIndex
 from hathor.indexes.utxo_index import UtxoIndex
 from hathor.transaction import BaseTransaction
+from hathor.transaction.nc_execution_state import NCExecutionState
 from hathor.util import tx_progress
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -204,11 +206,141 @@ def _manually_initialize(self, tx_storage: 'TransactionStorage') -> None:
     def update(self, tx: BaseTransaction) -> None:
         """ This is the new update method that indexes should use instead of add_tx/del_tx
         """
+        self.nc_update_add(tx)
+
         # XXX: this _should_ be here, but it breaks some tests, for now this is done explicitly in hathor.manager
         # self.mempool_tips.update(tx)
         if self.utxo:
             self.utxo.update(tx)
 
+    def nc_update_add(self, tx: BaseTransaction) -> None:
+        from hathor.conf.settings import HATHOR_TOKEN_UID
+        from hathor.nanocontracts.runner.types import (
+            NCSyscallRecord,
+            SyscallCreateContractRecord,
+            SyscallUpdateTokensRecord,
+        )
+        from hathor.nanocontracts.types import ContractId
+        from hathor.transaction.nc_execution_state import NCExecutionState
+
+        if not tx.is_nano_contract():
+            return
+
+        meta = tx.get_metadata()
+        if meta.nc_execution != NCExecutionState.SUCCESS:
+            return
+
+        assert meta.nc_calls
+        first_call = meta.nc_calls[0]
+        nc_syscalls: list[NCSyscallRecord] = []
+
+        # Add to indexes.
+        for call in meta.nc_calls:
+            # Txs that call other contracts are added to those contracts' history. This includes calls to `initialize`.
+            if self.nc_history:
+                self.nc_history.add_single_key(call.contract_id, tx)
+
+            # Accumulate all syscalls.
+            nc_syscalls.extend(call.index_updates)
+
+        created_contracts: set[ContractId] = set()
+        for syscall in nc_syscalls:
+            match syscall:
+                case SyscallCreateContractRecord(blueprint_id=blueprint_id, contract_id=contract_id):
+                    assert contract_id not in created_contracts, f'contract {contract_id.hex()} created multiple times'
+                    assert contract_id != first_call.contract_id, (
+                        f'contract {contract_id.hex()} cannot make a syscall to create itself'
+                    )
+                    created_contracts.add(contract_id)
+
+                    # Txs that create other contracts are added to the NC creation index and blueprint index.
+                    # They're already added to the NC history index, above.
+                    if self.nc_creation:
+                        self.nc_creation.manually_add_tx(tx)
+
+                    if self.blueprint_history:
+                        self.blueprint_history.add_single_key(blueprint_id, tx)
+
+                case SyscallUpdateTokensRecord():
+                    # Minted/melted tokens are added/removed to/from the tokens index,
+                    # and the respective destroyed/created HTR too.
+                    if self.tokens:
+                        try:
+                            self.tokens.get_token_info(syscall.token_uid)
+                        except KeyError:
+                            # If the token doesn't exist in the index yet, it must be a token creation syscall.
+                            from hathor.nanocontracts.runner.types import SyscallRecordType
+                            assert syscall.type is SyscallRecordType.CREATE_TOKEN, syscall.type
+                            assert syscall.token_name is not None and syscall.token_symbol is not None
+                            self.tokens.create_token_info(syscall.token_uid, syscall.token_name, syscall.token_symbol)
+
+                        self.tokens.add_to_total(syscall.token_uid, syscall.token_amount)
+                        self.tokens.add_to_total(HATHOR_TOKEN_UID, syscall.htr_amount)
+
+                case _:
+                    assert_never(syscall)
+
+    def nc_update_remove(self, tx: BaseTransaction) -> None:
+        from hathor.conf.settings import HATHOR_TOKEN_UID
+        from hathor.nanocontracts.runner.types import (
+            NCSyscallRecord,
+            SyscallCreateContractRecord,
+            SyscallUpdateTokensRecord,
+        )
+        from hathor.nanocontracts.types import NC_INITIALIZE_METHOD, ContractId
+
+        if not tx.is_nano_contract():
+            return
+
+        meta = tx.get_metadata()
+        assert meta.nc_execution is NCExecutionState.SUCCESS
+        assert meta.nc_calls
+        first_call = meta.nc_calls[0]
+        nc_syscalls: list[NCSyscallRecord] = []
+
+        # Remove from indexes, but we must keep the first call's contract still in the indexes.
+        for call in meta.nc_calls:
+            # Remove from nc_history except where it's the same contract as the first call.
+            if self.nc_history and call.contract_id != first_call.contract_id:
+                self.nc_history.remove_single_key(call.contract_id, tx)
+
+            # Accumulate all syscalls.
+            nc_syscalls.extend(call.index_updates)
+
+        created_contracts: set[ContractId] = set()
+        for syscall in nc_syscalls:
+            match syscall:
+                case SyscallCreateContractRecord(blueprint_id=blueprint_id, contract_id=contract_id):
+                    assert contract_id not in created_contracts, f'contract {contract_id.hex()} created multiple times'
+                    assert contract_id != first_call.contract_id, (
+                        f'contract {contract_id.hex()} cannot make a syscall to create itself'
+                    )
+                    created_contracts.add(contract_id)
+
+                    # Remove only when the first call is not creating a contract, that is,
+                    # if the tx itself is a nc creation, it must be kept in the indexes.
+                    if first_call.method_name != NC_INITIALIZE_METHOD:
+                        # Remove from nc_creation.
+                        if self.nc_creation:
+                            self.nc_creation.del_tx(tx)
+
+                        # Remove from blueprint_history.
+                        if self.blueprint_history:
+                            self.blueprint_history.remove_single_key(blueprint_id, tx)
+
+                case SyscallUpdateTokensRecord():
+                    # Undo the tokens update.
+                    if self.tokens:
+                        self.tokens.add_to_total(syscall.token_uid, -syscall.token_amount)
+                        self.tokens.add_to_total(HATHOR_TOKEN_UID, -syscall.htr_amount)
+
+                        from hathor.nanocontracts.runner.types import SyscallRecordType
+                        if syscall.type is SyscallRecordType.CREATE_TOKEN:
+                            self.tokens.destroy_token(syscall.token_uid)
+
+                case _:
+                    assert_never(syscall)
+
     def add_tx(self, tx: BaseTransaction) -> bool:
         """ Add a transaction to the indexes
 

hathor/indexes/rocksdb_tokens_index.py

@@ -17,6 +17,7 @@
 from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, TypedDict, cast
 
 from structlog import get_logger
+from typing_extensions import assert_never, override
 
 from hathor.conf.settings import HathorSettings
 from hathor.indexes.rocksdb_utils import (
@@ -27,6 +28,12 @@
     to_internal_token_uid,
 )
 from hathor.indexes.tokens_index import TokenIndexInfo, TokensIndex, TokenUtxoInfo
+from hathor.nanocontracts.types import (
+    NCAcquireAuthorityAction,
+    NCDepositAction,
+    NCGrantAuthorityAction,
+    NCWithdrawalAction,
+)
 from hathor.transaction import BaseTransaction, Transaction
 from hathor.transaction.base_transaction import TxVersion
 from hathor.util import collect_n, json_dumpb, json_loadb
@@ -168,7 +175,7 @@ def _to_value_info(self, info: _InfoDict) -> bytes:
     def _from_value_info(self, value: bytes) -> _InfoDict:
         return cast(_InfoDict, json_loadb(value))
 
-    def _create_token_info(self, token_uid: bytes, name: str, symbol: str, total: int = 0) -> None:
+    def create_token_info(self, token_uid: bytes, name: str, symbol: str, total: int = 0) -> None:
         key = self._to_key_info(token_uid)
         old_value = self._db.get((self._cf, key))
         assert old_value is None
@@ -179,7 +186,7 @@ def _create_token_info(self, token_uid: bytes, name: str, symbol: str, total: in
         })
         self._db.put((self._cf, key), value)
 
-    def _destroy_token(self, token_uid: bytes) -> None:
+    def destroy_token(self, token_uid: bytes) -> None:
         import rocksdb
 
         # a writebatch works similar to a "SQL transaction" in that if it fails, either all persist or none
@@ -218,14 +225,15 @@ def _remove_authority_utxo(self, token_uid: bytes, tx_hash: bytes, index: int, *
         self._db.delete((self._cf, self._to_key_authority(token_uid, TokenUtxoInfo(tx_hash, index), is_mint=is_mint)))
 
     def _create_genesis_info(self) -> None:
-        self._create_token_info(
+        self.create_token_info(
             self._settings.HATHOR_TOKEN_UID,
             self._settings.HATHOR_TOKEN_NAME,
             self._settings.HATHOR_TOKEN_SYMBOL,
             self._settings.GENESIS_TOKENS,
         )
 
-    def _add_to_total(self, token_uid: bytes, amount: int) -> None:
+    @override
+    def add_to_total(self, token_uid: bytes, amount: int) -> None:
         key_info = self._to_key_info(token_uid)
         old_value_info = self._db.get((self._cf, key_info))
         if token_uid == self._settings.HATHOR_TOKEN_UID and old_value_info is None:
@@ -237,18 +245,6 @@ def _add_to_total(self, token_uid: bytes, amount: int) -> None:
         new_value_info = self._to_value_info(dict_info)
         self._db.put((self._cf, key_info), new_value_info)
 
-    def _subtract_from_total(self, token_uid: bytes, amount: int) -> None:
-        key_info = self._to_key_info(token_uid)
-        old_value_info = self._db.get((self._cf, key_info))
-        if token_uid == self._settings.HATHOR_TOKEN_UID and old_value_info is None:
-            self._create_genesis_info()
-            old_value_info = self._db.get((self._cf, key_info))
-        assert old_value_info is not None
-        dict_info = self._from_value_info(old_value_info)
-        dict_info['total'] -= amount
-        new_value_info = self._to_value_info(dict_info)
-        self._db.put((self._cf, key_info), new_value_info)
-
     def _add_utxo(self, tx: BaseTransaction, index: int) -> None:
         """ Add tx to mint/melt indexes and total amount
         """
@@ -263,7 +259,7 @@ def _add_utxo(self, tx: BaseTransaction, index: int) -> None:
                 # add to melt index
                 self._add_authority_utxo(token_uid, tx.hash, index, is_mint=False)
         else:
-            self._add_to_total(token_uid, tx_output.value)
+            self.add_to_total(token_uid, tx_output.value)
 
     def _remove_utxo(self, tx: BaseTransaction, index: int) -> None:
         """ Remove tx from mint/melt indexes and total amount
@@ -280,7 +276,7 @@ def _remove_utxo(self, tx: BaseTransaction, index: int) -> None:
                 # remove from melt index
                 self._remove_authority_utxo(token_uid, tx.hash, index, is_mint=False)
         else:
-            self._subtract_from_total(token_uid, tx_output.value)
+            self.add_to_total(token_uid, -tx_output.value)
 
     def add_tx(self, tx: BaseTransaction) -> None:
         # if it's a TokenCreationTransaction, update name and symbol
@@ -292,7 +288,7 @@ def add_tx(self, tx: BaseTransaction) -> None:
             key_info = self._to_key_info(tx.hash)
             token_info = self._db.get((self._cf, key_info))
             if token_info is None:
-                self._create_token_info(tx.hash, tx.token_name, tx.token_symbol)
+                self.create_token_info(tx.hash, tx.token_name, tx.token_symbol)
 
         if tx.is_transaction:
             # Adding this tx to the transactions key list
@@ -308,6 +304,24 @@ def add_tx(self, tx: BaseTransaction) -> None:
             self.log.debug('add utxo', tx=tx.hash_hex, index=index)
             self._add_utxo(tx, index)
 
+        # Handle actions from Nano Contracts.
+        if tx.is_nano_contract():
+            assert isinstance(tx, Transaction)
+            nano_header = tx.get_nano_header()
+            ctx = nano_header.get_context()
+            for action in ctx.__all_actions__:
+                match action:
+                    case NCDepositAction():
+                        self.add_to_total(action.token_uid, action.amount)
+                    case NCWithdrawalAction():
+                        self.add_to_total(action.token_uid, -action.amount)
+                    case NCGrantAuthorityAction() | NCAcquireAuthorityAction():
+                        # These actions don't affect the nc token balance,
+                        # so no need for any special handling on the index.
+                        pass
+                    case _:
+                        assert_never(action)
+
     def remove_tx(self, tx: BaseTransaction) -> None:
         for tx_input in tx.inputs:
             spent_tx = tx.get_spent_tx(tx_input)
@@ -324,7 +338,25 @@ def remove_tx(self, tx: BaseTransaction) -> None:
 
         # if it's a TokenCreationTransaction, remove it from index
         if tx.version == TxVersion.TOKEN_CREATION_TRANSACTION:
-            self._destroy_token(tx.hash)
+            self.destroy_token(tx.hash)
+
+        # Handle actions from Nano Contracts.
+        if tx.is_nano_contract():
+            assert isinstance(tx, Transaction)
+            nano_header = tx.get_nano_header()
+            ctx = nano_header.get_context()
+            for action in ctx.__all_actions__:
+                match action:
+                    case NCDepositAction():
+                        self.add_to_total(action.token_uid, -action.amount)
+                    case NCWithdrawalAction():
+                        self.add_to_total(action.token_uid, action.amount)
+                    case NCGrantAuthorityAction() | NCAcquireAuthorityAction():
+                        # These actions don't affect the nc token balance,
+                        # so no need for any special handling on the index.
+                        pass
+                    case _:
+                        assert_never(action)
 
     def iter_all_tokens(self) -> Iterator[tuple[bytes, TokenIndexInfo]]:
         self.log.debug('seek to start')

hathor/indexes/tokens_index.py

@@ -114,6 +114,16 @@ def get_token_info(self, token_uid: bytes) -> TokenIndexInfo:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def create_token_info(self, token_uid: bytes, name: str, symbol: str, total: int = 0) -> None:
+        """Create a token info for a new token."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def destroy_token(self, token_uid: bytes) -> None:
+        """Destroy a token."""
+        raise NotImplementedError
+
     @abstractmethod
     def get_transactions_count(self, token_uid: bytes) -> int:
         """ Get quantity of transactions from requested token
@@ -139,3 +149,8 @@ def get_newer_transactions(self, token_uid: bytes, timestamp: int, hash_bytes: b
         """ Get transactions from the timestamp/hash_bytes reference to the newest
         """
         raise NotImplementedError
+
+    @abstractmethod
+    def add_to_total(self, token_uid: bytes, amount: int) -> None:
+        """Add an amount to the total of `token_uid`. The amount may be negative."""
+        raise NotImplementedError

hathor/nanocontracts/balance_rules.py

@@ -101,7 +101,9 @@ class _DepositRules(BalanceRules[NCDepositAction]):
 
     @override
     def verification_rule(self, token_dict: dict[TokenUid, TokenInfo]) -> None:
-        raise NotImplementedError('temporarily removed during nano merge')
+        token_info = token_dict.get(self.action.token_uid, TokenInfo.get_default())
+        token_info.amount = token_info.amount + self.action.amount
+        token_dict[self.action.token_uid] = token_info
 
     @override
     def nc_callee_execution_rule(self, callee_changes_tracker: NCChangesTracker) -> None:
@@ -123,7 +125,9 @@ class _WithdrawalRules(BalanceRules[NCWithdrawalAction]):
 
     @override
     def verification_rule(self, token_dict: dict[TokenUid, TokenInfo]) -> None:
-        raise NotImplementedError('temporarily removed during nano merge')
+        token_info = token_dict.get(self.action.token_uid, TokenInfo.get_default())
+        token_info.amount = token_info.amount - self.action.amount
+        token_dict[self.action.token_uid] = token_info
 
     @override
     def nc_callee_execution_rule(self, callee_changes_tracker: NCChangesTracker) -> None:
@@ -145,7 +149,17 @@ class _GrantAuthorityRules(BalanceRules[NCGrantAuthorityAction]):
 
     @override
     def verification_rule(self, token_dict: dict[TokenUid, TokenInfo]) -> None:
-        raise NotImplementedError('temporarily removed during nano merge')
+        assert self.action.token_uid != HATHOR_TOKEN_UID
+        token_info = token_dict.get(self.action.token_uid, TokenInfo.get_default())
+        if self.action.mint and not token_info.can_mint:
+            raise NCInvalidAction(
+                f'{self.action.name} token {self.action.token_uid.hex()} requires mint, but no input has it'
+            )
+
+        if self.action.melt and not token_info.can_melt:
+            raise NCInvalidAction(
+                f'{self.action.name} token {self.action.token_uid.hex()} requires melt, but no input has it'
+            )
 
     @override
     def nc_callee_execution_rule(self, callee_changes_tracker: NCChangesTracker) -> None:
@@ -187,7 +201,11 @@ class _AcquireAuthorityRules(BalanceRules[NCAcquireAuthorityAction]):
 
     @override
     def verification_rule(self, token_dict: dict[TokenUid, TokenInfo]) -> None:
-        raise NotImplementedError('temporarily removed during nano merge')
+        assert self.action.token_uid != HATHOR_TOKEN_UID
+        token_info = token_dict.get(self.action.token_uid, TokenInfo.get_default())
+        token_info.can_mint = token_info.can_mint or self.action.mint
+        token_info.can_melt = token_info.can_melt or self.action.melt
+        token_dict[self.action.token_uid] = token_info
 
     @override
     def nc_callee_execution_rule(self, callee_changes_tracker: NCChangesTracker) -> None: