change(state): Put chain data into an inner struct (#7608)

* Add `ChainInner`

* Impl `Deref` & `DerefMut` for `Chain`

* Remove an obsolete note

* Derive `Eq` for `ChainInner`

Co-authored-by: teor <teor@riseup.net>

* Fix docs

---------

Co-authored-by: teor <teor@riseup.net>

zebra-state/src/service/non_finalized_state/chain.rs

@@ -4,7 +4,7 @@
 use std::{
     cmp::Ordering,
     collections::{BTreeMap, BTreeSet, HashMap, HashSet},
-    ops::{Deref, RangeInclusive},
+    ops::{Deref, DerefMut, RangeInclusive},
     sync::Arc,
 };
 
@@ -41,13 +41,34 @@ pub mod index;
 /// to a non-finalized chain tip.
 #[derive(Clone, Debug)]
 pub struct Chain {
-    // Note: `eq_internal_state()` must be updated every time a field is added to [`Chain`].
-
     // Config
     //
     /// The configured network for this chain.
     network: Network,
 
+    /// The internal state of this chain.
+    inner: ChainInner,
+
+    // Diagnostics
+    //
+    /// The last height this chain forked at. Diagnostics only.
+    ///
+    /// This field is only used for metrics, it is not consensus-critical, and it is not checked
+    /// for equality.
+    ///
+    /// We keep the same last fork height in both sides of a clone, because every new block clones
+    /// a chain, even if it's just growing that chain.
+    pub(super) last_fork_height: Option<Height>,
+    // # Note
+    //
+    // Most diagnostics are implemented on the NonFinalizedState, rather than each chain.
+    // Some diagnostics only use the best chain, and others need to modify the Chain state,
+    // but that's difficult with `Arc<Chain>`s.
+}
+
+/// The internal state of [`Chain`].
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ChainInner {
     // Blocks, heights, hashes, and transaction locations
     //
     /// The contextually valid blocks which form this non-finalized partial chain, in height order.
@@ -185,22 +206,6 @@ pub struct Chain {
     /// When a new chain is created from the finalized tip,
     /// it is initialized with the finalized tip chain value pool balances.
     pub(crate) chain_value_pools: ValueBalance<NonNegative>,
-
-    // Diagnostics
-    //
-    /// The last height this chain forked at. Diagnostics only.
-    ///
-    /// This field is only used for metrics, it is not consensus-critical, and it is not checked
-    /// for equality.
-    ///
-    /// We keep the same last fork height in both sides of a clone, because every new block clones
-    /// a chain, even if it's just growing that chain.
-    pub(super) last_fork_height: Option<Height>,
-    // # Note
-    //
-    // Most diagnostics are implemented on the NonFinalizedState, rather than each chain.
-    // Some diagnostics only use the best chain, and others need to modify the Chain state,
-    // but that's difficult with `Arc<Chain>`s.
 }
 
 impl Chain {
@@ -214,8 +219,7 @@ impl Chain {
         history_tree: Arc<HistoryTree>,
         finalized_tip_chain_value_pools: ValueBalance<NonNegative>,
     ) -> Self {
-        let mut chain = Self {
+        let inner = ChainInner {
-            network,
             blocks: Default::default(),
             height_by_hash: Default::default(),
             tx_loc_by_hash: Default::default(),
@@ -240,6 +244,11 @@ impl Chain {
             partial_cumulative_work: Default::default(),
             history_trees_by_height: Default::default(),
             chain_value_pools: finalized_tip_chain_value_pools,
+        };
+
+        let mut chain = Self {
+            network,
+            inner,
             last_fork_height: None,
         };
 
@@ -263,49 +272,7 @@ impl Chain {
     /// even if the blocks in the two chains are equal.
     #[cfg(any(test, feature = "proptest-impl"))]
     pub fn eq_internal_state(&self, other: &Chain) -> bool {
-        // blocks, heights, hashes
+        self.inner == other.inner
-        self.blocks == other.blocks &&
-            self.height_by_hash == other.height_by_hash &&
-            self.tx_loc_by_hash == other.tx_loc_by_hash &&
-
-            // transparent UTXOs
-            self.created_utxos == other.created_utxos &&
-            self.spent_utxos == other.spent_utxos &&
-
-            // note commitment trees
-            self.sprout_trees_by_anchor == other.sprout_trees_by_anchor &&
-            self.sprout_trees_by_height == other.sprout_trees_by_height &&
-            self.sapling_trees_by_height == other.sapling_trees_by_height &&
-            self.orchard_trees_by_height == other.orchard_trees_by_height &&
-
-            // history trees
-            self.history_trees_by_height == other.history_trees_by_height &&
-
-            // note commitment subtrees
-            self.sapling_subtrees == other.sapling_subtrees &&
-            self.orchard_subtrees == other.orchard_subtrees &&
-
-            // anchors
-            self.sprout_anchors == other.sprout_anchors &&
-            self.sprout_anchors_by_height == other.sprout_anchors_by_height &&
-            self.sapling_anchors == other.sapling_anchors &&
-            self.sapling_anchors_by_height == other.sapling_anchors_by_height &&
-            self.orchard_anchors == other.orchard_anchors &&
-            self.orchard_anchors_by_height == other.orchard_anchors_by_height &&
-
-            // nullifiers
-            self.sprout_nullifiers == other.sprout_nullifiers &&
-            self.sapling_nullifiers == other.sapling_nullifiers &&
-            self.orchard_nullifiers == other.orchard_nullifiers &&
-
-            // transparent address indexes
-            self.partial_transparent_transfers == other.partial_transparent_transfers &&
-
-            // proof of work
-            self.partial_cumulative_work == other.partial_cumulative_work &&
-
-            // chain value pool balances
-            self.chain_value_pools == other.chain_value_pools
     }
 
     /// Returns the last fork height if that height is still in the non-finalized state.
@@ -1580,6 +1547,20 @@ impl Chain {
     }
 }
 
+impl Deref for Chain {
+    type Target = ChainInner;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl DerefMut for Chain {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
 /// The revert position being performed on a chain.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 enum RevertPosition {
@@ -2096,14 +2077,14 @@ impl UpdateWith<ValueBalance<NegativeAllowed>> for Chain {
     /// When forking from the tip, subtract the block's chain value pool change.
     ///
     /// When finalizing the root, leave the chain value pool balances unchanged.
-    /// [`Self::chain_value_pools`] tracks the chain value pools for all
+    /// [`ChainInner::chain_value_pools`] tracks the chain value pools for all finalized blocks, and
-    /// finalized blocks, and the non-finalized blocks in this chain. So
+    /// the non-finalized blocks in this chain. So finalizing the root doesn't change the set of
-    /// finalizing the root doesn't change the set of blocks it tracks.
+    /// blocks it tracks.
     ///
     /// # Panics
     ///
-    /// Panics if the chain pool value balance is invalid
+    /// Panics if the chain pool value balance is invalid after we subtract the block value pool
-    /// after we subtract the block value pool change.
+    /// change.
     fn revert_chain_with(
         &mut self,
         block_value_pool_change: &ValueBalance<NegativeAllowed>,