From bf3d6b690e868f21cdb94eb9a1e6443baf4bb975 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Wed, 18 Dec 2024 18:48:00 +0800
Subject: [PATCH] Complete migration

---
 beacon_node/beacon_chain/src/migrate.rs | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/beacon_node/beacon_chain/src/migrate.rs b/beacon_node/beacon_chain/src/migrate.rs
index 6ba9e78ab1..f435d6adb6 100644
--- a/beacon_node/beacon_chain/src/migrate.rs
+++ b/beacon_node/beacon_chain/src/migrate.rs
@@ -11,6 +11,7 @@ use std::sync::{mpsc, Arc};
 use std::thread;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use store::hot_cold_store::{migrate_database, HotColdDBError};
+use store::iter::StateRootsIterator;
 use store::{DBColumn, Error, HotStateSummary, ItemStore, StoreItem, StoreOp};
 pub use store::{HotColdDB, MemoryStore};
 use types::{
@@ -354,6 +355,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
 
         let old_finalized_checkpoint = match Self::prune_hot_db(
             db.clone(),
+            finalized_state_root.into(),
             &finalized_state,
             notif.finalized_checkpoint,
             notif.genesis_block_root,
@@ -464,6 +466,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
     #[allow(clippy::too_many_arguments)]
     fn prune_hot_db(
         store: Arc<HotColdDB<E, Hot, Cold>>,
+        new_finalized_state_hash: Hash256,
         new_finalized_state: &BeaconState<E>,
         new_finalized_checkpoint: Checkpoint,
         genesis_block_root: Hash256,
@@ -564,6 +567,20 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
         let block_roots_descendant_of_finalized =
             summaries_dag.descendant_block_roots_of(&new_finalized_checkpoint.root)?;
 
+        // TODO(hdiff): Could re-use the summaries dag
+        let newly_finalized_chain =
+            std::iter::once(Ok((new_finalized_slot, new_finalized_state_hash)))
+                .chain(
+                    StateRootsIterator::new(&store, new_finalized_state)
+                        .map(|res| res.map(|(state_root, slot)| (slot, state_root))),
+                )
+                .take_while(|res| {
+                    res.as_ref()
+                        .map_or(true, |(slot, _)| *slot >= old_finalized_slot)
+                })
+                .map(|res| res.map(|(_, state_root)| state_root))
+                .collect::<Result<Vec<Hash256>, _>>()?;
+
         // Compute the set of finalized state roots that we must keep to make the dynamic HDiff system
         // work.
         // TODO(hdiff): must not delete *all* finalized hdiffs. Instead, keep
@@ -584,6 +601,11 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
                 if block_roots_descendant_of_finalized.contains(&summary.latest_block_root) {
                     // Keep this state is the post state of a viable head, or a state advance from a
                     // viable head.
+                } else if newly_finalized_chain.contains(&state_root)
+                    && required_finalized_diff_state_slots.contains(&slot)
+                {
+                    // Keep this state and diff as it's necessary for the finalized portion of the
+                    // HDiff links
                 } else {
                     // TODO(hdiff) compute the root of the hdiffs to keep, else prune
                     abandoned_blocks.insert(SignedBeaconBlockHash::from(summary.latest_block_root));