Merge pull request #1523 from mintlayer/fix/delegation_inputs_spending

Keep delegations in the db

chainstate/constraints-value-accumulator/src/accounts_balances_tracker.rs

@@ -0,0 +1,88 @@
+// Copyright (c) 2024 RBB S.r.l
+// [email protected]
+// SPDX-License-Identifier: MIT
+// Licensed under the MIT License;
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://github.com/mintlayer/mintlayer-core/blob/master/LICENSE
+//
+// 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.
+
+use std::collections::{btree_map::Entry, BTreeMap};
+
+use common::{
+    chain::{AccountSpending, AccountType, DelegationId},
+    primitives::Amount,
+};
+
+use crate::Error;
+
+#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
+enum GenericAccountId {
+    Delegation(DelegationId),
+}
+
+impl From<AccountSpending> for GenericAccountId {
+    fn from(account: AccountSpending) -> Self {
+        match account {
+            AccountSpending::DelegationBalance(id, _) => Self::Delegation(id),
+        }
+    }
+}
+
+impl From<GenericAccountId> for AccountType {
+    fn from(value: GenericAccountId) -> Self {
+        match value {
+            GenericAccountId::Delegation(id) => AccountType::Delegation(id),
+        }
+    }
+}
+
+pub struct AccountsBalancesTracker<'a, DelegationBalanceGetterFn> {
+    balances: BTreeMap<GenericAccountId, Amount>,
+
+    delegation_balance_getter: &'a DelegationBalanceGetterFn,
+}
+
+impl<'a, DelegationBalanceGetterFn> AccountsBalancesTracker<'a, DelegationBalanceGetterFn>
+where
+    DelegationBalanceGetterFn: Fn(DelegationId) -> Result<Option<Amount>, Error>,
+{
+    pub fn new(delegation_balance_getter: &'a DelegationBalanceGetterFn) -> Self {
+        Self {
+            balances: BTreeMap::new(),
+            delegation_balance_getter,
+        }
+    }
+
+    pub fn spend_from_account(&mut self, account: AccountSpending) -> Result<(), Error> {
+        match self.balances.entry(account.clone().into()) {
+            Entry::Vacant(e) => {
+                let (balance, spending) = match account {
+                    AccountSpending::DelegationBalance(id, spending) => {
+                        let balance = (self.delegation_balance_getter)(id)?
+                            .ok_or(Error::AccountBalanceNotFound(account.clone().into()))?;
+                        (balance, spending)
+                    }
+                };
+                let new_balance = (balance - spending)
+                    .ok_or(Error::NegativeAccountBalance(account.clone().into()))?;
+                e.insert(new_balance);
+            }
+            Entry::Occupied(mut e) => {
+                let balance = e.get_mut();
+                let spending = match account {
+                    AccountSpending::DelegationBalance(_, spending) => spending,
+                };
+                *balance = (*balance - spending)
+                    .ok_or(Error::NegativeAccountBalance(account.clone().into()))?;
+            }
+        };
+        Ok(())
+    }
+}

chainstate/constraints-value-accumulator/src/constraints_accumulator.rs

@@ -18,12 +18,15 @@ use std::{collections::BTreeMap, num::NonZeroU64};
 use common::{
     chain::{
         output_value::OutputValue, timelock::OutputTimeLock, AccountCommand, AccountSpending,
-        ChainConfig, DelegationId, PoolId, TxInput, TxOutput, UtxoOutPoint,
+        AccountsBalancesCheckVersion, ChainConfig, DelegationId, PoolId, TxInput, TxOutput,
+        UtxoOutPoint,
     },
     primitives::{Amount, BlockHeight, CoinOrTokenId, Fee, Subsidy},
 };
 use utils::ensure;
 
+use crate::accounts_balances_tracker::AccountsBalancesTracker;
+
 use super::{accumulated_fee::AccumulatedFee, insert_or_increase, Error};
 
 /// `ConstrainedValueAccumulator` helps avoiding messy inputs/outputs combinations analysis by
@@ -72,6 +75,8 @@ impl ConstrainedValueAccumulator {
 
         let mut accumulator = Self::new();
         let mut total_fee_deducted = Amount::ZERO;
+        let mut accounts_balances_tracker =
+            AccountsBalancesTracker::new(&delegation_balance_getter);
 
         for (input, input_utxo) in inputs.iter().zip(inputs_utxos.iter()) {
             match input {
@@ -91,6 +96,7 @@ impl ConstrainedValueAccumulator {
                         chain_config,
                         block_height,
                         outpoint.account(),
+                        &mut accounts_balances_tracker,
                         &delegation_balance_getter,
                     )?;
                 }
@@ -194,19 +200,32 @@ impl ConstrainedValueAccumulator {
         chain_config: &ChainConfig,
         block_height: BlockHeight,
         account: &AccountSpending,
+        accounts_balances_tracker: &mut AccountsBalancesTracker<DelegationBalanceGetterFn>,
         delegation_balance_getter: &DelegationBalanceGetterFn,
     ) -> Result<(), Error>
     where
         DelegationBalanceGetterFn: Fn(DelegationId) -> Result<Option<Amount>, Error>,
     {
         match account {
             AccountSpending::DelegationBalance(delegation_id, spend_amount) => {
-                let delegation_balance = delegation_balance_getter(*delegation_id)?
-                    .ok_or(Error::DelegationBalanceNotFound(*delegation_id))?;
-                ensure!(
-                    *spend_amount <= delegation_balance,
-                    Error::AttemptToPrintMoney(CoinOrTokenId::Coin)
-                );
+                match chain_config
+                    .chainstate_upgrades()
+                    .version_at_height(block_height)
+                    .1
+                    .accounts_balances_version()
+                {
+                    AccountsBalancesCheckVersion::V0 => {
+                        let delegation_balance = delegation_balance_getter(*delegation_id)?
+                            .ok_or(Error::DelegationBalanceNotFound(*delegation_id))?;
+                        ensure!(
+                            *spend_amount <= delegation_balance,
+                            Error::AttemptToPrintMoney(CoinOrTokenId::Coin)
+                        );
+                    }
+                    AccountsBalancesCheckVersion::V1 => {
+                        accounts_balances_tracker.spend_from_account(account.clone())?;
+                    }
+                }
 
                 let maturity_distance =
                     chain_config.staking_pool_spend_maturity_block_count(block_height);

chainstate/constraints-value-accumulator/src/error.rs

@@ -14,7 +14,7 @@
 // limitations under the License.
 
 use common::{
-    chain::{DelegationId, PoolId, UtxoOutPoint},
+    chain::{AccountType, DelegationId, PoolId, UtxoOutPoint},
     primitives::CoinOrTokenId,
 };
 
@@ -44,4 +44,8 @@ pub enum Error {
     SpendingNonSpendableOutput(UtxoOutPoint),
     #[error("Balance not found for delegation `{0}`")]
     DelegationBalanceNotFound(DelegationId),
+    #[error("Account balance not found for `{0:?}`")]
+    AccountBalanceNotFound(AccountType),
+    #[error("Negative account balance for `{0:?}`")]
+    NegativeAccountBalance(AccountType),
 }

chainstate/constraints-value-accumulator/src/lib.rs

@@ -17,6 +17,7 @@ use std::collections::BTreeMap;
 
 use common::primitives::Amount;
 
+mod accounts_balances_tracker;
 mod accumulated_fee;
 mod constraints_accumulator;
 mod error;

chainstate/src/detail/ban_score.rs

@@ -503,6 +503,8 @@ impl BanScore for constraints_value_accumulator::Error {
             constraints_value_accumulator::Error::SpendingNonSpendableOutput(_) => 100,
             constraints_value_accumulator::Error::AttemptToViolateFeeRequirements => 100,
             constraints_value_accumulator::Error::DelegationBalanceNotFound(_) => 0,
+            constraints_value_accumulator::Error::AccountBalanceNotFound(_) => 0,
+            constraints_value_accumulator::Error::NegativeAccountBalance(_) => 100,
         }
     }
 }

chainstate/storage/src/internal/version.rs

@@ -19,7 +19,7 @@ use serialization::{Decode, Encode};
 pub struct ChainstateStorageVersion(u32);
 
 impl ChainstateStorageVersion {
-    pub const CURRENT: Self = Self(5);
+    pub const CURRENT: Self = Self(6);
 
     pub fn new(value: u32) -> Self {
         Self(value)

chainstate/test-suite/src/tests/chainstate_storage_tests.rs

@@ -24,9 +24,10 @@ use common::{
     chain::{
         output_value::OutputValue,
         tokens::{make_token_id, NftIssuance, TokenAuxiliaryData, TokenIssuanceV0},
-        ChainstateUpgrade, Destination, NetUpgrades, NftIdMismatchCheck, OutPointSourceId,
-        RewardDistributionVersion, TokenIssuanceVersion, TokensFeeVersion,
-        TokensTickerMaxLengthVersion, Transaction, TxInput, TxOutput, UtxoOutPoint,
+        AccountsBalancesCheckVersion, ChainstateUpgrade, Destination, NetUpgrades,
+        NftIdMismatchCheck, OutPointSourceId, RewardDistributionVersion, TokenIssuanceVersion,
+        TokensFeeVersion, TokensTickerMaxLengthVersion, Transaction, TxInput, TxOutput,
+        UtxoOutPoint,
     },
     primitives::{Amount, Id, Idable},
 };
@@ -121,6 +122,7 @@ fn store_fungible_token_v0(#[case] seed: Seed) {
                                 TokensFeeVersion::V1,
                                 TokensTickerMaxLengthVersion::V1,
                                 NftIdMismatchCheck::Yes,
+                                AccountsBalancesCheckVersion::V1,
                             ),
                         )])
                         .unwrap(),
@@ -201,6 +203,7 @@ fn store_nft_v0(#[case] seed: Seed) {
                                 TokensFeeVersion::V1,
                                 TokensTickerMaxLengthVersion::V1,
                                 NftIdMismatchCheck::Yes,
+                                AccountsBalancesCheckVersion::V1,
                             ),
                         )])
                         .unwrap(),
@@ -386,6 +389,7 @@ fn store_aux_data_from_issue_nft(#[case] seed: Seed) {
                                 TokensFeeVersion::V1,
                                 TokensTickerMaxLengthVersion::V1,
                                 NftIdMismatchCheck::Yes,
+                                AccountsBalancesCheckVersion::V1,
                             ),
                         )])
                         .unwrap(),