From 5ab9e266a2651db9a1e69c76ff03145a0695aba5 Mon Sep 17 00:00:00 2001
From: Brooks <brooks@anza.xyz>
Date: Wed, 16 Oct 2024 13:03:03 -0400
Subject: [PATCH] Supports rehashing with accounts lt hash (#3186)

---
 runtime/src/bank.rs                  | 14 ++++++++------
 runtime/src/bank/accounts_lt_hash.rs |  5 ++---
 runtime/src/snapshot_bank_utils.rs   | 20 ++------------------
 3 files changed, 12 insertions(+), 27 deletions(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 0a4cb785f2c091..2a9a955768177a 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -2928,6 +2928,11 @@ impl Bank {
 
             // freeze is a one-way trip, idempotent
             self.freeze_started.store(true, Relaxed);
+            if self.is_accounts_lt_hash_enabled() {
+                // updating the accounts lt hash must happen *outside* of hash_internal_state() so
+                // that rehash() can be called and *not* modify self.accounts_lt_hash.
+                self.update_accounts_lt_hash();
+            }
             *hash = self.hash_internal_state();
             self.rc.accounts.accounts_db.mark_slot_frozen(self.slot());
         }
@@ -5454,10 +5459,6 @@ impl Bank {
             hash = hashv(&[hash.as_ref(), epoch_accounts_hash.as_ref().as_ref()]);
         };
 
-        let accounts_lt_hash_checksum = self
-            .is_accounts_lt_hash_enabled()
-            .then(|| self.update_accounts_lt_hash());
-
         let buf = self
             .hard_forks
             .read()
@@ -5517,8 +5518,9 @@ impl Bank {
             } else {
                 "".to_string()
             },
-            if let Some(accounts_lt_hash_checksum) = accounts_lt_hash_checksum {
-                format!(", accounts_lt_hash checksum: {accounts_lt_hash_checksum}")
+            if self.is_accounts_lt_hash_enabled() {
+                let checksum = self.accounts_lt_hash.lock().unwrap().0.checksum();
+                format!(", accounts_lt_hash checksum: {checksum}")
             } else {
                 String::new()
             },
diff --git a/runtime/src/bank/accounts_lt_hash.rs b/runtime/src/bank/accounts_lt_hash.rs
index 2fe0fceb3f3072..f32ce16ac4a30d 100644
--- a/runtime/src/bank/accounts_lt_hash.rs
+++ b/runtime/src/bank/accounts_lt_hash.rs
@@ -2,7 +2,7 @@ use {
     super::Bank,
     rayon::prelude::*,
     solana_accounts_db::accounts_db::AccountsDb,
-    solana_lattice_hash::lt_hash::{Checksum as LtChecksum, LtHash},
+    solana_lattice_hash::lt_hash::LtHash,
     solana_measure::{meas_dur, measure::Measure},
     solana_sdk::{
         account::{accounts_equal, AccountSharedData},
@@ -29,12 +29,11 @@ impl Bank {
     /// - mix in its current state
     ///
     /// Since this function is non-idempotent, it should only be called once per bank.
-    pub fn update_accounts_lt_hash(&self) -> LtChecksum {
+    pub fn update_accounts_lt_hash(&self) {
         debug_assert!(self.is_accounts_lt_hash_enabled());
         let delta_lt_hash = self.calculate_delta_lt_hash();
         let mut accounts_lt_hash = self.accounts_lt_hash.lock().unwrap();
         accounts_lt_hash.0.mix_in(&delta_lt_hash);
-        accounts_lt_hash.0.checksum()
     }
 
     /// Calculates the lt hash *of only this slot*
diff --git a/runtime/src/snapshot_bank_utils.rs b/runtime/src/snapshot_bank_utils.rs
index 2547935c52841d..9d6546a9e8cf18 100644
--- a/runtime/src/snapshot_bank_utils.rs
+++ b/runtime/src/snapshot_bank_utils.rs
@@ -930,19 +930,7 @@ fn bank_to_full_snapshot_archive_with(
         .accounts_db
         .set_latest_full_snapshot_slot(bank.slot());
     bank.squash(); // Bank may not be a root
-
-    // Rehashing is not currently supported when the accounts lt hash is enabled.
-    // This is because rehashing will *re-mix-in* all the accounts stored in this bank into the
-    // accounts lt hash!  This is incorrect, as the accounts lt hash would change, even if the bank
-    // was *not* manually modified by the caller.
-    // We can re-allow rehasing if we change the Bank to hold its parent's accounts lt hash plus a
-    // *delta* accounts lt hash, and then Bank::hash_internal_state() will only recalculate the
-    // delta accounts lt hash.
-    // Another option is to consider if manual modification should even be allowed in the first
-    // place. Disallowing it would solve these issues.
-    if !bank.is_accounts_lt_hash_enabled() {
-        bank.rehash(); // Bank accounts may have been manually modified by the caller
-    }
+    bank.rehash(); // Bank may have been manually modified by the caller
     bank.force_flush_accounts_cache();
     bank.clean_accounts();
     let calculated_accounts_hash =
@@ -1005,11 +993,7 @@ pub fn bank_to_incremental_snapshot_archive(
         .accounts_db
         .set_latest_full_snapshot_slot(full_snapshot_slot);
     bank.squash(); // Bank may not be a root
-
-    // See the comment in bank_to_full_snapshot_archive() when calling rehash()
-    if !bank.is_accounts_lt_hash_enabled() {
-        bank.rehash(); // Bank accounts may have been manually modified by the caller
-    }
+    bank.rehash(); // Bank may have been manually modified by the caller
     bank.force_flush_accounts_cache();
     bank.clean_accounts();
     let calculated_incremental_accounts_hash =