From 9a61ac9795915ea1e25a69e3cf9a50ecdbf8e6d9 Mon Sep 17 00:00:00 2001
From: Sarah <SarahADavisA@gmail.com>
Date: Thu, 28 Aug 2025 23:19:41 +0800
Subject: [PATCH] Implements TODO

---
 crates/stages/stages/src/stages/merkle.rs | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/crates/stages/stages/src/stages/merkle.rs b/crates/stages/stages/src/stages/merkle.rs
index 00e1177ed02..eac56fea9c2 100644
--- a/crates/stages/stages/src/stages/merkle.rs
+++ b/crates/stages/stages/src/stages/merkle.rs
@@ -401,10 +401,19 @@ where
             // Validation passed, apply unwind changes to the database.
             provider.write_trie_updates(&updates)?;
 
-            // TODO(alexey): update entities checkpoint
+            // Update entities checkpoint to reflect the unwind operation
+            // Since we're unwinding, we need to recalculate the total entities at the target block
+            let accounts = tx.entries::<tables::HashedAccounts>()?;
+            let storages = tx.entries::<tables::HashedStorages>()?;
+            let total = (accounts + storages) as u64;
+            entities_checkpoint.total = total;
+            entities_checkpoint.processed = total;
         }
 
-        Ok(UnwindOutput { checkpoint: StageCheckpoint::new(input.unwind_to) })
+        Ok(UnwindOutput {
+            checkpoint: StageCheckpoint::new(input.unwind_to)
+                .with_entities_stage_checkpoint(entities_checkpoint),
+        })
     }
 }