From 4e78fbe86ef4d030968d00f461926abbcb813943 Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Mon, 2 Oct 2023 17:42:35 -0400
Subject: [PATCH] docs: document non-zero amounts in State::increment_balances
 (#760)

* feat: document non-zero amounts in State::increment_balances

* handle zero balance increment

---------

Co-authored-by: rakita <dragan0rakita@gmail.com>
---
 crates/revm/src/db/states/cache_account.rs | 16 +++++++++-------
 crates/revm/src/db/states/state.rs         | 14 +++++++++++++-
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/crates/revm/src/db/states/cache_account.rs b/crates/revm/src/db/states/cache_account.rs
index f6b85436c6..c9b29b803a 100644
--- a/crates/revm/src/db/states/cache_account.rs
+++ b/crates/revm/src/db/states/cache_account.rs
@@ -293,13 +293,15 @@ impl CacheAccount {
     /// Increment balance by `balance` amount. Assume that balance will not
     /// overflow or be zero.
     ///
-    /// Note: to skip some edge cases we assume that additional balance is never zero.
-    /// And as increment is always related to block fee/reward and withdrawals this is correct.
-    pub fn increment_balance(&mut self, balance: u128) -> TransitionAccount {
-        self.account_info_change(|info| {
-            info.balance += U256::from(balance);
-        })
-        .1
+    /// Note: only if balance is zero we would return None as no transition would be made.
+    pub fn increment_balance(&mut self, balance: u128) -> Option<TransitionAccount> {
+        if balance == 0 {
+            return None;
+        }
+        let (_, transition) = self.account_info_change(|info| {
+            info.balance = info.balance.saturating_add(U256::from(balance));
+        });
+        Some(transition)
     }
 
     fn account_info_change<T, F: FnOnce(&mut AccountInfo) -> T>(
diff --git a/crates/revm/src/db/states/state.rs b/crates/revm/src/db/states/state.rs
index 5166a0a82b..293b49e420 100644
--- a/crates/revm/src/db/states/state.rs
+++ b/crates/revm/src/db/states/state.rs
@@ -82,6 +82,10 @@ impl<DB: Database> State<DB> {
     /// If account is not found inside cache state it will be loaded from database.
     ///
     /// Update will create transitions for all accounts that are updated.
+    ///
+    /// Like [CacheAccount::increment_balance], this assumes that incremented balances are not
+    /// zero, and will not overflow once incremented. If using this to implement withdrawals, zero
+    /// balances must be filtered out before calling this function.
     pub fn increment_balances(
         &mut self,
         balances: impl IntoIterator<Item = (Address, u128)>,
@@ -89,8 +93,16 @@ impl<DB: Database> State<DB> {
         // make transition and update cache state
         let mut transitions = Vec::new();
         for (address, balance) in balances {
+            if balance == 0 {
+                continue;
+            }
             let original_account = self.load_cache_account(address)?;
-            transitions.push((address, original_account.increment_balance(balance)))
+            transitions.push((
+                address,
+                original_account
+                    .increment_balance(balance)
+                    .expect("Balance is not zero"),
+            ))
         }
         // append transition
         if let Some(s) = self.transition_state.as_mut() {