feat: track balance increases in state diffs

bolt-sidecar/src/builder/template.rs

@@ -7,12 +7,14 @@ use ethereum_consensus::{
     deneb::mainnet::{Blob, BlobsBundle},
 };
 use reth_primitives::TransactionSigned;
-use std::collections::HashMap;
 use tracing::warn;
 
 use crate::{
     common::transactions::max_transaction_cost,
-    primitives::{AccountState, FullTransaction, SignedConstraints, TransactionExt},
+    primitives::{
+        diffs::{AccountDiff, BalanceDiff, StateDiff},
+        AccountState, FullTransaction, SignedConstraints, TransactionExt,
+    },
 };
 
 /// A block template that serves as a fallback block, but is also used
@@ -34,7 +36,7 @@ pub struct BlockTemplate {
 
 impl BlockTemplate {
     /// Return the state diff of the block template.
-    pub fn get_diff(&self, address: &Address) -> Option<(u64, U256)> {
+    pub fn get_diff(&self, address: &Address) -> Option<AccountDiff> {
         self.state_diff.get_diff(address)
     }
 
@@ -123,14 +125,46 @@ impl BlockTemplate {
     pub fn add_constraints(&mut self, constraints: SignedConstraints) {
         for constraint in &constraints.message.transactions {
             let max_cost = max_transaction_cost(constraint);
+
+            // Increase the nonce and decrease the balance of the sender
             self.state_diff
                 .diffs
                 .entry(*constraint.sender().expect("recovered sender"))
-                .and_modify(|(nonce, balance)| {
-                    *nonce += 1;
-                    *balance += max_cost;
+                .and_modify(|diff| {
+                    *diff = AccountDiff::new(
+                        diff.nonce().saturating_add(1),
+                        BalanceDiff::new(
+                            diff.balance().increase(),
+                            diff.balance().decrease().saturating_add(max_cost),
+                        ),
+                    )
                 })
-                .or_insert((1, max_cost));
+                .or_insert(AccountDiff::new(1, BalanceDiff::new(U256::ZERO, max_cost)));
+
+            // If there is an ETH transfer and it's not a contract creation, increase the balance
+            // of the recipient so that it can send inclusion requests on this preconfirmed state.
+            let value = constraint.tx.value();
+            if value.is_zero() {
+                continue;
+            }
+            let Some(recipient) = constraint.to() else { continue };
+
+            self.state_diff
+                .diffs
+                .entry(recipient)
+                .and_modify(|diff| {
+                    *diff = AccountDiff::new(
+                        diff.nonce().saturating_add(1),
+                        BalanceDiff::new(
+                            diff.balance().increase().saturating_add(constraint.tx.value()),
+                            diff.balance().decrease(),
+                        ),
+                    )
+                })
+                .or_insert(AccountDiff::new(
+                    0,
+                    BalanceDiff::new(constraint.tx.value(), U256::ZERO),
+                ));
         }
 
         self.signed_constraints_list.push(constraints);
@@ -141,13 +175,38 @@ impl BlockTemplate {
         let constraints = self.signed_constraints_list.remove(index);
 
         for constraint in &constraints.message.transactions {
+            let max_cost = max_transaction_cost(constraint);
+
             self.state_diff
                 .diffs
                 .entry(*constraint.sender().expect("recovered sender"))
-                .and_modify(|(nonce, balance)| {
-                    *nonce = nonce.saturating_sub(1);
-                    *balance -= max_transaction_cost(constraint);
+                .and_modify(|diff| {
+                    *diff = AccountDiff::new(
+                        diff.nonce().saturating_sub(1),
+                        BalanceDiff::new(
+                            diff.balance().increase(),
+                            diff.balance().decrease().saturating_sub(max_cost),
+                        ),
+                    )
                 });
+
+            // If there is an ETH transfer and it's not a contract creation, remove the balance
+            // increase of the recipient.
+            let value = constraint.tx.value();
+            if value.is_zero() {
+                continue;
+            }
+            let Some(recipient) = constraint.to() else { continue };
+
+            self.state_diff.diffs.entry(recipient).and_modify(|diff| {
+                *diff = AccountDiff::new(
+                    diff.nonce().saturating_sub(1),
+                    BalanceDiff::new(
+                        diff.balance().increase().saturating_sub(constraint.tx.value()),
+                        diff.balance().decrease(),
+                    ),
+                )
+            });
         }
     }
 
@@ -196,20 +255,3 @@ impl BlockTemplate {
         }
     }
 }
-
-/// StateDiff tracks the intermediate changes to the state according to the block template.
-#[derive(Debug, Default)]
-pub struct StateDiff {
-    /// Map of diffs per address. Each diff is a tuple of the nonce and balance diff
-    /// that should be applied to the current state.
-    pub(crate) diffs: HashMap<Address, (u64, U256)>,
-}
-
-impl StateDiff {
-    /// Returns a tuple of the nonce and balance diff for the given address.
-    /// The nonce diff should be added to the current nonce, the balance diff should be subtracted
-    /// from the current balance.
-    pub fn get_diff(&self, address: &Address) -> Option<(u64, U256)> {
-        self.diffs.get(address).copied()
-    }
-}

bolt-sidecar/src/chain_io/utils.rs

@@ -48,7 +48,7 @@ fn pubkey_hash_digest(key: &BlsPublicKey) -> B512 {
 ///
 /// Example usage:
 ///
-/// ```rust no_run
+/// ```rust ignore
 /// sol! {
 ///    library ErrorLib {
 ///       error SomeError(uint256 code);

bolt-sidecar/src/primitives/diffs.rs

@@ -0,0 +1,98 @@
+use std::collections::HashMap;
+
+use alloy::primitives::{Address, U256};
+
+/// StateDiff tracks the intermediate changes to the state according to the block template.
+#[derive(Debug, Default)]
+pub struct StateDiff {
+    /// Map of diffs per address. Each diff is a tuple of the nonce and balance diff
+    /// that should be applied to the current state.
+    pub(crate) diffs: HashMap<Address, AccountDiff>,
+}
+
+impl StateDiff {
+    /// Returns a tuple of the nonce and balance diff for the given address.
+    /// The nonce diff should be added to the current nonce, the balance diff should be subtracted
+    /// from the current balance.
+    pub fn get_diff(&self, address: &Address) -> Option<AccountDiff> {
+        self.diffs.get(address).copied()
+    }
+}
+
+/// AccountDiff tracks the changes to an account's nonce and balance.
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub struct AccountDiff {
+    /// The nonce of the account.
+    nonce: u64,
+    /// The balance diff of the account.
+    balance: BalanceDiff,
+}
+
+impl AccountDiff {
+    /// Creates a new account diff with the given nonce and balance diff.
+    pub fn new(nonce: u64, balance: BalanceDiff) -> Self {
+        Self { nonce, balance }
+    }
+
+    /// Returns the nonce diff of the account.
+    pub fn nonce(&self) -> u64 {
+        self.nonce
+    }
+
+    /// Returns the balance diff of the account.
+    pub fn balance(&self) -> BalanceDiff {
+        self.balance
+    }
+}
+
+/// A balance diff is a tuple of consisting of a balance increase and a balance decrease.
+///
+/// An `increase` should be _added_ to the current balance, while a `decrease` should be _subtracted_.
+///
+/// Example:
+/// ```rs
+/// let balance = U256::from(100);
+/// let balance_diff = BalanceDiff::new(U256::from(50), U256::from(10));
+/// assert_eq!(balance_diff.apply(balance), U256::from(140));
+/// ```
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub struct BalanceDiff {
+    /// The balance increase.
+    increase: U256,
+    /// The balance decrease.
+    decrease: U256,
+}
+
+impl BalanceDiff {
+    /// Creates a new balance diff with the given increase and decrease.
+    pub fn new(increase: U256, decrease: U256) -> Self {
+        Self { increase, decrease }
+    }
+
+    /// Returns the increase of the balance diff.
+    pub fn increase(&self) -> U256 {
+        self.increase
+    }
+
+    /// Returns the decrease of the balance diff.
+    pub fn decrease(&self) -> U256 {
+        self.decrease
+    }
+
+    /// Applies the balance diff to the given balance.
+    pub fn apply(&self, balance: U256) -> U256 {
+        balance.saturating_add(self.increase).saturating_sub(self.decrease)
+    }
+}
+
+/// A trait for applying a balance diff to a U256 balance.
+pub trait BalanceDiffApplier {
+    /// Applies the balance diff to the given balance.
+    fn apply_diff(&self, diff: BalanceDiff) -> U256;
+}
+
+impl BalanceDiffApplier for U256 {
+    fn apply_diff(&self, diff: BalanceDiff) -> U256 {
+        self.saturating_add(diff.increase).saturating_sub(diff.decrease)
+    }
+}

bolt-sidecar/src/primitives/mod.rs

@@ -45,6 +45,9 @@ pub mod signature;
 /// JSON-RPC helper types and functions.
 pub mod jsonrpc;
 
+/// Types and utilties relates to calculating account diffs.
+pub mod diffs;
+
 /// An alias for a Beacon Chain slot number
 pub type Slot = u64;