fix(state): Remove workarounds for storing trees

zebra-chain/src/block/height.rs

@@ -205,9 +205,6 @@ fn operator_tests() {
     assert_eq!(None, Height(i32::MAX as u32) + 1);
     assert_eq!(None, Height(u32::MAX) + 0);
 
-    assert_eq!(Some(Height(2)), Height(1) + 1);
-    assert_eq!(None, Height::MAX + 1);
-
     // Adding negative numbers
     assert_eq!(Some(Height(1)), Height(2) + -1);
     assert_eq!(Some(Height(0)), Height(1) + -1);

zebra-state/src/service/arbitrary.rs

@@ -173,7 +173,10 @@ impl Strategy for PreparedChain {
         }
 
         let chain = chain.clone().expect("should be generated");
-        let count = (1..chain.1.len()).new_tree(runner)?;
+        // The generated chain should contain at least two blocks:
+        // 1. the zeroth genesis block, and
+        // 2. a first block.
+        let count = (2..chain.1.len()).new_tree(runner)?;
         Ok(PreparedChainTree {
             chain: chain.1,
             count,

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

@@ -521,26 +521,16 @@ impl Chain {
         let anchor = tree.root();
         trace!(?height, ?anchor, "adding sprout tree");
 
-        // TODO: fix test code that incorrectly overwrites trees
-        #[cfg(not(test))]
-        {
-            assert_eq!(
-                self.sprout_trees_by_height.insert(height, tree.clone()),
-                None,
-                "incorrect overwrite of sprout tree: trees must be reverted then inserted",
-            );
-            assert_eq!(
-                self.sprout_anchors_by_height.insert(height, anchor),
-                None,
-                "incorrect overwrite of sprout anchor: anchors must be reverted then inserted",
-            );
-        }
-
-        #[cfg(test)]
-        {
-            self.sprout_trees_by_height.insert(height, tree.clone());
-            self.sprout_anchors_by_height.insert(height, anchor);
-        }
+        assert_eq!(
+            self.sprout_trees_by_height.insert(height, tree.clone()),
+            None,
+            "incorrect overwrite of sprout tree: trees must be reverted then inserted",
+        );
+        assert_eq!(
+            self.sprout_anchors_by_height.insert(height, anchor),
+            None,
+            "incorrect overwrite of sprout anchor: anchors must be reverted then inserted",
+        );
 
         // Multiple inserts are expected here,
         // because the anchors only change if a block has shielded transactions.
@@ -633,26 +623,16 @@ impl Chain {
         let anchor = tree.root();
         trace!(?height, ?anchor, "adding sapling tree");
 
-        // TODO: fix test code that incorrectly overwrites trees
-        #[cfg(not(test))]
-        {
-            assert_eq!(
-                self.sapling_trees_by_height.insert(height, tree),
-                None,
-                "incorrect overwrite of sapling tree: trees must be reverted then inserted",
-            );
-            assert_eq!(
-                self.sapling_anchors_by_height.insert(height, anchor),
-                None,
-                "incorrect overwrite of sapling anchor: anchors must be reverted then inserted",
-            );
-        }
-
-        #[cfg(test)]
-        {
-            self.sapling_trees_by_height.insert(height, tree);
-            self.sapling_anchors_by_height.insert(height, anchor);
-        }
+        assert_eq!(
+            self.sapling_trees_by_height.insert(height, tree),
+            None,
+            "incorrect overwrite of sapling tree: trees must be reverted then inserted",
+        );
+        assert_eq!(
+            self.sapling_anchors_by_height.insert(height, anchor),
+            None,
+            "incorrect overwrite of sapling anchor: anchors must be reverted then inserted",
+        );
 
         // Multiple inserts are expected here,
         // because the anchors only change if a block has shielded transactions.
@@ -747,26 +727,16 @@ impl Chain {
         let anchor = tree.root();
         trace!(?height, ?anchor, "adding orchard tree");
 
-        // TODO: fix test code that incorrectly overwrites trees
-        #[cfg(not(test))]
-        {
-            assert_eq!(
-                self.orchard_trees_by_height.insert(height, tree),
-                None,
-                "incorrect overwrite of orchard tree: trees must be reverted then inserted",
-            );
-            assert_eq!(
-                self.orchard_anchors_by_height.insert(height, anchor),
-                None,
-                "incorrect overwrite of orchard anchor: anchors must be reverted then inserted",
-            );
-        }
-
-        #[cfg(test)]
-        {
-            self.orchard_trees_by_height.insert(height, tree);
-            self.orchard_anchors_by_height.insert(height, anchor);
-        }
+        assert_eq!(
+            self.orchard_trees_by_height.insert(height, tree),
+            None,
+            "incorrect overwrite of orchard tree: trees must be reverted then inserted",
+        );
+        assert_eq!(
+            self.orchard_anchors_by_height.insert(height, anchor),
+            None,
+            "incorrect overwrite of orchard anchor: anchors must be reverted then inserted",
+        );
 
         // Multiple inserts are expected here,
         // because the anchors only change if a block has shielded transactions.
@@ -850,16 +820,11 @@ impl Chain {
         // Use the previously cached root which was calculated in parallel.
         trace!(?height, "adding history tree");
 
-        // TODO: fix test code that incorrectly overwrites trees
-        #[cfg(not(test))]
         assert_eq!(
             self.history_trees_by_height.insert(height, tree),
             None,
             "incorrect overwrite of history tree: trees must be reverted then inserted",
         );
-
-        #[cfg(test)]
-        self.history_trees_by_height.insert(height, tree);
     }
 
     /// Remove the History tree index at `height`.

zebra-state/src/service/non_finalized_state/tests/prop.rs

@@ -53,7 +53,7 @@ fn push_genesis_chain() -> Result<()> {
 
             chain_values.insert(None, (None, only_chain.chain_value_pools.into()));
 
-            for block in chain.iter().take(count).cloned() {
+            for block in chain.iter().take(count).skip(1).cloned() {
                 let block =
                 ContextuallyVerifiedBlock::with_block_and_spent_utxos(
                         block,
@@ -72,7 +72,7 @@ fn push_genesis_chain() -> Result<()> {
                 chain_values.insert(block.height.into(), (block.chain_value_pool_change.into(), only_chain.chain_value_pools.into()));
             }
 
-            prop_assert_eq!(only_chain.blocks.len(), count);
+            prop_assert_eq!(only_chain.blocks.len(), count - 1);
         });
 
     Ok(())
@@ -150,7 +150,7 @@ fn forked_equals_pushed_genesis() -> Result<()> {
             empty_tree.clone(),
             ValueBalance::zero(),
         );
-        for block in chain.iter().take(fork_at_count).cloned() {
+        for block in chain.iter().take(fork_at_count).skip(1).cloned() {
             let block = ContextuallyVerifiedBlock::with_block_and_spent_utxos(
                 block,
                 partial_chain.unspent_utxos(),
@@ -170,7 +170,7 @@ fn forked_equals_pushed_genesis() -> Result<()> {
             empty_tree,
             ValueBalance::zero(),
         );
-        for block in chain.iter().cloned() {
+        for block in chain.iter().skip(1).cloned() {
             let block =
             ContextuallyVerifiedBlock::with_block_and_spent_utxos(block, full_chain.unspent_utxos())?;
             full_chain = full_chain
@@ -460,7 +460,7 @@ fn rejection_restores_internal_state_genesis() -> Result<()> {
                                .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
     |((chain, valid_count, network, mut bad_block) in (PreparedChain::default(), any::<bool>(), any::<bool>())
       .prop_flat_map(|((chain, valid_count, network, _history_tree), is_nu5, is_v5)| {
-          let next_height = chain[valid_count - 1].height;
+          let next_height = chain[valid_count].height;
           (
               Just(chain),
               Just(valid_count),
@@ -486,7 +486,7 @@ fn rejection_restores_internal_state_genesis() -> Result<()> {
         // use `valid_count` as the number of valid blocks before an invalid block
         let valid_tip_height = chain[valid_count - 1].height;
         let valid_tip_hash = chain[valid_count - 1].hash;
-        let mut chain = chain.iter().take(valid_count).cloned();
+        let mut chain = chain.iter().take(valid_count).skip(1).cloned();
 
         prop_assert!(state.eq_internal_state(&state));
 

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -65,6 +65,9 @@ fn construct_many() -> Result<()> {
 
     let mut block: Arc<Block> =
         zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
+    let initial_height = block
+        .coinbase_height()
+        .expect("Block 434873 should have its height in its coinbase tx.");
     let mut blocks = vec![];
 
     while blocks.len() < 100 {
@@ -75,7 +78,7 @@ fn construct_many() -> Result<()> {
 
     let mut chain = Chain::new(
         Network::Mainnet,
-        Height(block.coinbase_height().unwrap().0 - 1),
+        (initial_height - 1).expect("Initial height should be at least 1."),
         Default::default(),
         Default::default(),
         Default::default(),