From f58faa4c5bbfe76121539d849ca280254575b8ef Mon Sep 17 00:00:00 2001
From: Mark Logan <103447440+mystenmark@users.noreply.github.com>
Date: Thu, 27 Jun 2024 20:41:08 -0700
Subject: [PATCH] [Cherry-pick] Fix epoch flag logic for state accum v2
 (#18448) (#18449)

Create new epoch flags for state accum v2, since the old one was already
written in some places. Also fix chain-specific logic

Cherry pick of #18448
---
 .../authority/authority_per_epoch_store.rs    |  9 +++++---
 .../authority/epoch_start_configuration.rs    | 22 ++++++++++++++++---
 crates/sui-core/src/state_accumulator.rs      |  6 ++---
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/crates/sui-core/src/authority/authority_per_epoch_store.rs b/crates/sui-core/src/authority/authority_per_epoch_store.rs
index e6e7315ad7224..f54c3c0da6e83 100644
--- a/crates/sui-core/src/authority/authority_per_epoch_store.rs
+++ b/crates/sui-core/src/authority/authority_per_epoch_store.rs
@@ -916,9 +916,12 @@ impl AuthorityPerEpochStore {
     }
 
     pub fn state_accumulator_v2_enabled(&self) -> bool {
-        self.epoch_start_configuration
-            .flags()
-            .contains(&EpochFlag::StateAccumulatorV2Enabled)
+        let flag = match self.get_chain_identifier().chain() {
+            Chain::Unknown | Chain::Testnet => EpochFlag::StateAccumulatorV2EnabledTestnet,
+            Chain::Mainnet => EpochFlag::StateAccumulatorV2EnabledMainnet,
+        };
+
+        self.epoch_start_configuration.flags().contains(&flag)
     }
 
     /// Returns `&Arc<EpochStartConfiguration>`
diff --git a/crates/sui-core/src/authority/epoch_start_configuration.rs b/crates/sui-core/src/authority/epoch_start_configuration.rs
index 4a35b467a5fb5..b7ca15d2d6d47 100644
--- a/crates/sui-core/src/authority/epoch_start_configuration.rs
+++ b/crates/sui-core/src/authority/epoch_start_configuration.rs
@@ -51,7 +51,13 @@ pub enum EpochFlag {
     _ObjectLockSplitTablesDeprecated,
 
     WritebackCacheEnabled,
-    StateAccumulatorV2Enabled,
+
+    // This flag was "burned" because it was deployed with a broken version of the code. The
+    // new flags below are required to enable state accumulator v2
+    _StateAccumulatorV2EnabledDeprecated,
+
+    StateAccumulatorV2EnabledTestnet,
+    StateAccumulatorV2EnabledMainnet,
 }
 
 impl EpochFlag {
@@ -78,7 +84,9 @@ impl EpochFlag {
         }
 
         if enable_state_accumulator_v2 {
-            new_flags.push(EpochFlag::StateAccumulatorV2Enabled);
+            new_flags.push(EpochFlag::StateAccumulatorV2EnabledTestnet);
+            // TODO: enable on mainnet
+            // new_flags.push(EpochFlag::StateAccumulatorV2EnabledMainnet);
         }
 
         new_flags
@@ -99,7 +107,15 @@ impl fmt::Display for EpochFlag {
                 write!(f, "ObjectLockSplitTables (DEPRECATED)")
             }
             EpochFlag::WritebackCacheEnabled => write!(f, "WritebackCacheEnabled"),
-            EpochFlag::StateAccumulatorV2Enabled => write!(f, "StateAccumulatorV2Enabled"),
+            EpochFlag::_StateAccumulatorV2EnabledDeprecated => {
+                write!(f, "StateAccumulatorV2EnabledDeprecated (DEPRECATED)")
+            }
+            EpochFlag::StateAccumulatorV2EnabledTestnet => {
+                write!(f, "StateAccumulatorV2EnabledTestnet")
+            }
+            EpochFlag::StateAccumulatorV2EnabledMainnet => {
+                write!(f, "StateAccumulatorV2EnabledMainnet")
+            }
         }
     }
 }
diff --git a/crates/sui-core/src/state_accumulator.rs b/crates/sui-core/src/state_accumulator.rs
index 8ec3b5f184b6d..0ba94875a0d9b 100644
--- a/crates/sui-core/src/state_accumulator.rs
+++ b/crates/sui-core/src/state_accumulator.rs
@@ -5,7 +5,7 @@ use itertools::Itertools;
 use mysten_metrics::monitored_scope;
 use prometheus::{register_int_gauge_with_registry, IntGauge, Registry};
 use serde::Serialize;
-use sui_protocol_config::{Chain, ProtocolConfig};
+use sui_protocol_config::ProtocolConfig;
 use sui_types::base_types::{ObjectID, ObjectRef, SequenceNumber, VersionNumber};
 use sui_types::committee::EpochId;
 use sui_types::digests::{ObjectDigest, TransactionDigest};
@@ -389,9 +389,7 @@ impl StateAccumulator {
         epoch_store: &Arc<AuthorityPerEpochStore>,
         metrics: Arc<StateAccumulatorMetrics>,
     ) -> Self {
-        if epoch_store.state_accumulator_v2_enabled()
-            && epoch_store.get_chain_identifier().chain() != Chain::Mainnet
-        {
+        if epoch_store.state_accumulator_v2_enabled() {
             StateAccumulator::V2(StateAccumulatorV2::new(store, metrics))
         } else {
             StateAccumulator::V1(StateAccumulatorV1::new(store, metrics))