Fix migration of dense diffs

Author: michaelsproul

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -6783,6 +6783,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     }
 
     /// As for `chain_dump` but dumping only the portion of the chain newer than `from_slot`.
+    #[allow(clippy::type_complexity)]
     pub fn chain_dump_from_slot(
         &self,
         from_slot: Slot,

beacon_node/beacon_chain/src/schema_change/migration_schema_v24.rs

@@ -36,6 +36,32 @@ pub fn store_full_state_v22<E: EthSpec>(
     Ok(())
 }
 
+/// Fetch a V22 state from the database either as a full state or using block replay.
+pub fn get_state_v22<T: BeaconChainTypes>(
+    db: &Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+    state_root: &Hash256,
+    spec: &ChainSpec,
+) -> Result<Option<BeaconState<T::EthSpec>>, Error> {
+    let Some(summary) = db.get_item::<HotStateSummaryV22>(state_root)? else {
+        return Ok(None);
+    };
+    let Some(base_state) =
+        get_full_state_v22(&db.hot_db, &summary.epoch_boundary_state_root, spec)?
+    else {
+        return Ok(None);
+    };
+    // Loading hot states via block replay doesn't care about the schema version, so we can use
+    // the DB's current method for this.
+    let update_cache = false;
+    db.load_hot_state_using_replay(
+        base_state,
+        summary.slot,
+        summary.latest_block_root,
+        update_cache,
+    )
+    .map(Some)
+}
+
 pub fn get_full_state_v22<KV: KeyValueStore<E>, E: EthSpec>(
     db: &KV,
     state_root: &Hash256,
@@ -139,6 +165,20 @@ pub struct HotStateSummaryV22 {
     epoch_boundary_state_root: Hash256,
 }
 
+impl StoreItem for HotStateSummaryV22 {
+    fn db_column() -> DBColumn {
+        DBColumn::BeaconStateSummary
+    }
+
+    fn as_store_bytes(&self) -> Vec<u8> {
+        self.as_ssz_bytes()
+    }
+
+    fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
+        Ok(Self::from_ssz_bytes(bytes)?)
+    }
+}
+
 pub fn upgrade_to_v24<T: BeaconChainTypes>(
     db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
 ) -> Result<Vec<KeyValueStoreOp>, Error> {
@@ -319,19 +359,8 @@ pub fn upgrade_to_v24<T: BeaconChainTypes>(
 
                 match storage_strategy {
                     StorageStrategy::DiffFrom(_) | StorageStrategy::Snapshot => {
-                        // We choose to not compute states during the epoch with block replay for
-                        // simplicity. Therefore we can only support a hot heriarchy config where the
-                        // lowest layer is >= log2(SLOTS_PER_EPOCH).
-                        // FIXME(tree-states): we need to remove this restriction.
-                        if slot % T::EthSpec::slots_per_epoch() != 0 {
-                            return Err(Error::MigrationError(format!(
-                                "Hot hierarchy config lowest value must be >= {}",
-                                T::EthSpec::slots_per_epoch().trailing_zeros()
-                            )));
-                        }
-
-                        // Load the full state and re-store it as a snapshot or diff.
-                        let state = get_full_state_v22(&db.hot_db, &state_root, &db.spec)?
+                        // Load the state and re-store it as a snapshot or diff.
+                        let state = get_state_v22::<T>(&db, &state_root, &db.spec)?
                             .ok_or(Error::MissingState(state_root))?;
 
                         // Store immediately so that future diffs can load and diff from it.

beacon_node/beacon_chain/tests/store_tests.rs

@@ -3681,7 +3681,7 @@ async fn ancestor_state_root_prior_to_split() {
                     store
                         .load_hot_state_summary(&ancestor_state_root)
                         .unwrap()
-                        .expect(&format!(
+                        .unwrap_or_else(|| panic!(
                             "no summary found for {ancestor_state_root:?} (slot {ancestor_slot})"
                         ))
                         .slot,

beacon_node/store/src/hot_cold_store.rs

@@ -1791,7 +1791,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
         }
     }
 
-    fn load_hot_state_using_replay(
+    pub fn load_hot_state_using_replay(
         &self,
         base_state: BeaconState<E>,
         slot: Slot,