change(state): Deduplicate note commitment trees in non-finalized state (#7239)

* Remove duplicate asserts

There are the same two asserts above the two removed ones.

* Remove workarounds for inserting trees into NFS

NFS = non finalized state

* Use correct height for constructing new chain

We were using the height of the last block instead of the initial block
to construct a new chain.

* Don't push the 0th block into a chain

* Don't commit two blocks at the same height

* Add helpers for heights

* Support the retrieval of deduped Sprout trees

* Dedup Sprout trees

* Refactor docs for adding & removing Sprout trees

* Support the retrieval of deduped Sapling trees

* Dedup Sapling trees

* Refactor docs for adding & removing Sapling trees

* Support the retrieval of deduped Orchard trees

* Dedup Orchard trees

* Refactor docs for adding & removing Orchard trees

* Make the docs for storing trees clearer

zebra-chain/src/block/height.rs

@@ -65,6 +65,29 @@ impl Height {
     /// previous to Nu5 and in non-coinbase transactions from Nu5 activation
     /// height and above.
     pub const MAX_EXPIRY_HEIGHT: Height = Height(499_999_999);
+
+    /// Returns the next [`Height`].
+    ///
+    /// # Panics
+    ///
+    /// - If the current height is at its maximum.
+    pub fn next(self) -> Self {
+        (self + 1).expect("Height should not be at its maximum.")
+    }
+
+    /// Returns the previous [`Height`].
+    ///
+    /// # Panics
+    ///
+    /// - If the current height is at its minimum.
+    pub fn previous(self) -> Self {
+        (self - 1).expect("Height should not be at its minimum.")
+    }
+
+    /// Returns `true` if the [`Height`] is at its minimum.
+    pub fn is_min(self) -> bool {
+        self == Self::MIN
+    }
 }
 
 /// A difference between two [`Height`]s, possibly negative.

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

@@ -500,14 +500,25 @@ impl Chain {
         let height =
             hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?;
 
-        self.sprout_trees_by_height.get(&height).cloned()
+        self.sprout_trees_by_height
+            .range(..=height)
+            .next_back()
+            .map(|(_height, tree)| tree.clone())
     }
 
-    /// Add the Sprout `tree` to the tree and anchor indexes at `height`.
+    /// Adds the Sprout `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
+    ///
     /// - the height of a new block that has just been added to the chain tip, or
-    /// - the finalized tip height: the height of the parent of the first block of a new chain.
+    /// - the finalized tip height—the height of the parent of the first block of a new chain.
+    ///
+    /// Stores only the first tree in each series of identical trees.
+    ///
+    /// # Panics
+    ///
+    /// - If there's a tree already stored at `height`.
+    /// - If there's an anchor already stored at `height`.
     fn add_sprout_tree_and_anchor(
         &mut self,
         height: Height,
@@ -521,11 +532,20 @@ impl Chain {
         let anchor = tree.root();
         trace!(?height, ?anchor, "adding sprout tree");
 
-        assert_eq!(
+        // Don't add a new tree unless it differs from the previous one or there's no previous tree.
-            self.sprout_trees_by_height.insert(height, tree.clone()),
+        if height.is_min()
-            None,
+            || self
-            "incorrect overwrite of sprout tree: trees must be reverted then inserted",
+                .sprout_tree(height.previous().into())
-        );
+                .map_or(true, |prev_tree| prev_tree != tree)
+        {
+            assert_eq!(
+                self.sprout_trees_by_height.insert(height, tree.clone()),
+                None,
+                "incorrect overwrite of sprout tree: trees must be reverted then inserted",
+            );
+        }
+
+        // Store the root.
         assert_eq!(
             self.sprout_anchors_by_height.insert(height, anchor),
             None,
@@ -538,24 +558,35 @@ impl Chain {
         self.sprout_trees_by_anchor.insert(anchor, tree);
     }
 
-    /// Remove the Sprout tree and anchor indexes at `height`.
+    /// Removes the Sprout tree and anchor indexes at `height`.
     ///
     /// `height` can be at two different [`RevertPosition`]s in the chain:
-    /// - a tip block above a chain fork: only that height is removed, or
+    ///
-    /// - a root block: all trees and anchors below that height are removed,
+    /// - a tip block above a chain fork—only the tree and anchor at that height are removed, or
-    ///   including temporary finalized tip trees.
+    /// - a root block—all trees and anchors at and below that height are removed, including
+    ///   temporary finalized tip trees.
+    ///
+    ///   # Panics
+    ///
+    ///  - If the anchor being removed is not present.
+    ///  - If there is no tree at `height`.
     fn remove_sprout_tree_and_anchor(&mut self, position: RevertPosition, height: Height) {
-        let removed_heights: Vec<Height> = if position == RevertPosition::Root {
+        let (removed_heights, highest_removed_tree) = if position == RevertPosition::Root {
-            // Remove all trees and anchors at or below the removed block.
+            (
-            // This makes sure the temporary trees from finalized tip forks are removed.
+                // Remove all trees and anchors at or below the removed block.
-            self.sprout_anchors_by_height
+                // This makes sure the temporary trees from finalized tip forks are removed.
-                .keys()
+                self.sprout_anchors_by_height
-                .cloned()
+                    .keys()
-                .filter(|index_height| *index_height <= height)
+                    .cloned()
-                .collect()
+                    .filter(|index_height| *index_height <= height)
+                    .collect(),
+                // Cache the highest (rightmost) tree before its removal.
+                self.sprout_tree(height.into()),
+            )
         } else {
             // Just remove the reverted tip trees and anchors.
-            vec![height]
+            // We don't need to cache the highest (rightmost) tree.
+            (vec![height], None)
         };
 
         for height in &removed_heights {
@@ -563,9 +594,8 @@ impl Chain {
                 .sprout_anchors_by_height
                 .remove(height)
                 .expect("Sprout anchor must be present if block was added to chain");
-            self.sprout_trees_by_height
+
-                .remove(height)
+            self.sprout_trees_by_height.remove(height);
-                .expect("Sprout note commitment tree must be present if block was added to chain");
 
             trace!(?height, ?position, ?anchor, "removing sprout tree");
 
@@ -579,6 +609,26 @@ impl Chain {
                 self.sprout_trees_by_anchor.remove(&anchor);
             }
         }
+
+        // # Invariant
+        //
+        // The height following after the removed heights in a non-empty non-finalized state must
+        // always have its tree.
+        //
+        // The loop above can violate the invariant, and if `position` is [`RevertPosition::Root`],
+        // it will always violate the invariant. We restore the invariant by storing the highest
+        // (rightmost) removed tree just above `height` if there is no tree at that height.
+        if !self.is_empty() && height < self.non_finalized_tip_height() {
+            let next_height = height.next();
+
+            if self.sprout_trees_by_height.get(&next_height).is_none() {
+                // TODO: Use `try_insert` once it stabilises.
+                self.sprout_trees_by_height.insert(
+                    next_height,
+                    highest_removed_tree.expect("There should be a cached removed tree."),
+                );
+            }
+        }
     }
 
     /// Returns the Sapling note commitment tree of the tip of this [`Chain`],
@@ -607,14 +657,25 @@ impl Chain {
         let height =
             hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?;
 
-        self.sapling_trees_by_height.get(&height).cloned()
+        self.sapling_trees_by_height
+            .range(..=height)
+            .next_back()
+            .map(|(_height, tree)| tree.clone())
     }
 
-    /// Add the Sapling `tree` to the tree and anchor indexes at `height`.
+    /// Adds the Sapling `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
+    ///
     /// - the height of a new block that has just been added to the chain tip, or
-    /// - the finalized tip height: the height of the parent of the first block of a new chain.
+    /// - the finalized tip height—the height of the parent of the first block of a new chain.
+    ///
+    /// Stores only the first tree in each series of identical trees.
+    ///
+    /// # Panics
+    ///
+    /// - If there's a tree already stored at `height`.
+    /// - If there's an anchor already stored at `height`.
     fn add_sapling_tree_and_anchor(
         &mut self,
         height: Height,
@@ -623,11 +684,20 @@ impl Chain {
         let anchor = tree.root();
         trace!(?height, ?anchor, "adding sapling tree");
 
-        assert_eq!(
+        // Don't add a new tree unless it differs from the previous one or there's no previous tree.
-            self.sapling_trees_by_height.insert(height, tree),
+        if height.is_min()
-            None,
+            || self
-            "incorrect overwrite of sapling tree: trees must be reverted then inserted",
+                .sapling_tree(height.previous().into())
-        );
+                .map_or(true, |prev_tree| prev_tree != tree)
+        {
+            assert_eq!(
+                self.sapling_trees_by_height.insert(height, tree),
+                None,
+                "incorrect overwrite of sapling tree: trees must be reverted then inserted",
+            );
+        }
+
+        // Store the root.
         assert_eq!(
             self.sapling_anchors_by_height.insert(height, anchor),
             None,
@@ -639,24 +709,35 @@ impl Chain {
         self.sapling_anchors.insert(anchor);
     }
 
-    /// Remove the Sapling tree and anchor indexes at `height`.
+    /// Removes the Sapling tree and anchor indexes at `height`.
     ///
     /// `height` can be at two different [`RevertPosition`]s in the chain:
-    /// - a tip block above a chain fork: only that height is removed, or
+    ///
-    /// - a root block: all trees and anchors below that height are removed,
+    /// - a tip block above a chain fork—only the tree and anchor at that height are removed, or
-    ///   including temporary finalized tip trees.
+    /// - a root block—all trees and anchors at and below that height are removed, including
+    ///   temporary finalized tip trees.
+    ///
+    ///   # Panics
+    ///
+    ///  - If the anchor being removed is not present.
+    ///  - If there is no tree at `height`.
     fn remove_sapling_tree_and_anchor(&mut self, position: RevertPosition, height: Height) {
-        let removed_heights: Vec<Height> = if position == RevertPosition::Root {
+        let (removed_heights, highest_removed_tree) = if position == RevertPosition::Root {
-            // Remove all trees and anchors at or below the removed block.
+            (
-            // This makes sure the temporary trees from finalized tip forks are removed.
+                // Remove all trees and anchors at or below the removed block.
-            self.sapling_anchors_by_height
+                // This makes sure the temporary trees from finalized tip forks are removed.
-                .keys()
+                self.sapling_anchors_by_height
-                .cloned()
+                    .keys()
-                .filter(|index_height| *index_height <= height)
+                    .cloned()
-                .collect()
+                    .filter(|index_height| *index_height <= height)
+                    .collect(),
+                // Cache the highest (rightmost) tree before its removal.
+                self.sapling_tree(height.into()),
+            )
         } else {
             // Just remove the reverted tip trees and anchors.
-            vec![height]
+            // We don't need to cache the highest (rightmost) tree.
+            (vec![height], None)
         };
 
         for height in &removed_heights {
@@ -664,9 +745,8 @@ impl Chain {
                 .sapling_anchors_by_height
                 .remove(height)
                 .expect("Sapling anchor must be present if block was added to chain");
-            self.sapling_trees_by_height
+
-                .remove(height)
+            self.sapling_trees_by_height.remove(height);
-                .expect("Sapling note commitment tree must be present if block was added to chain");
 
             trace!(?height, ?position, ?anchor, "removing sapling tree");
 
@@ -677,6 +757,26 @@ impl Chain {
                 "Sapling anchor must be present if block was added to chain"
             );
         }
+
+        // # Invariant
+        //
+        // The height following after the removed heights in a non-empty non-finalized state must
+        // always have its tree.
+        //
+        // The loop above can violate the invariant, and if `position` is [`RevertPosition::Root`],
+        // it will always violate the invariant. We restore the invariant by storing the highest
+        // (rightmost) removed tree just above `height` if there is no tree at that height.
+        if !self.is_empty() && height < self.non_finalized_tip_height() {
+            let next_height = height.next();
+
+            if self.sapling_trees_by_height.get(&next_height).is_none() {
+                // TODO: Use `try_insert` once it stabilises.
+                self.sapling_trees_by_height.insert(
+                    next_height,
+                    highest_removed_tree.expect("There should be a cached removed tree."),
+                );
+            }
+        }
     }
 
     /// Returns the Orchard note commitment tree of the tip of this [`Chain`],
@@ -706,14 +806,25 @@ impl Chain {
         let height =
             hash_or_height.height_or_else(|hash| self.height_by_hash.get(&hash).cloned())?;
 
-        self.orchard_trees_by_height.get(&height).cloned()
+        self.orchard_trees_by_height
+            .range(..=height)
+            .next_back()
+            .map(|(_height, tree)| tree.clone())
     }
 
-    /// Add the Orchard `tree` to the tree and anchor indexes at `height`.
+    /// Adds the Orchard `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
+    ///
     /// - the height of a new block that has just been added to the chain tip, or
-    /// - the finalized tip height: the height of the parent of the first block of a new chain.
+    /// - the finalized tip height—the height of the parent of the first block of a new chain.
+    ///
+    /// Stores only the first tree in each series of identical trees.
+    ///
+    /// # Panics
+    ///
+    /// - If there's a tree already stored at `height`.
+    /// - If there's an anchor already stored at `height`.
     fn add_orchard_tree_and_anchor(
         &mut self,
         height: Height,
@@ -727,11 +838,20 @@ impl Chain {
         let anchor = tree.root();
         trace!(?height, ?anchor, "adding orchard tree");
 
-        assert_eq!(
+        // Don't add a new tree unless it differs from the previous one or there's no previous tree.
-            self.orchard_trees_by_height.insert(height, tree),
+        if height.is_min()
-            None,
+            || self
-            "incorrect overwrite of orchard tree: trees must be reverted then inserted",
+                .orchard_tree(height.previous().into())
-        );
+                .map_or(true, |prev_tree| prev_tree != tree)
+        {
+            assert_eq!(
+                self.orchard_trees_by_height.insert(height, tree),
+                None,
+                "incorrect overwrite of orchard tree: trees must be reverted then inserted",
+            );
+        }
+
+        // Store the root.
         assert_eq!(
             self.orchard_anchors_by_height.insert(height, anchor),
             None,
@@ -743,24 +863,35 @@ impl Chain {
         self.orchard_anchors.insert(anchor);
     }
 
-    /// Remove the Orchard tree and anchor indexes at `height`.
+    /// Removes the Orchard tree and anchor indexes at `height`.
     ///
     /// `height` can be at two different [`RevertPosition`]s in the chain:
-    /// - a tip block above a chain fork: only that height is removed, or
+    ///
-    /// - a root block: all trees and anchors below that height are removed,
+    /// - a tip block above a chain fork—only the tree and anchor at that height are removed, or
-    ///   including temporary finalized tip trees.
+    /// - a root block—all trees and anchors at and below that height are removed, including
+    ///   temporary finalized tip trees.
+    ///
+    ///   # Panics
+    ///
+    ///  - If the anchor being removed is not present.
+    ///  - If there is no tree at `height`.
     fn remove_orchard_tree_and_anchor(&mut self, position: RevertPosition, height: Height) {
-        let removed_heights: Vec<Height> = if position == RevertPosition::Root {
+        let (removed_heights, highest_removed_tree) = if position == RevertPosition::Root {
-            // Remove all trees and anchors at or below the removed block.
+            (
-            // This makes sure the temporary trees from finalized tip forks are removed.
+                // Remove all trees and anchors at or below the removed block.
-            self.orchard_anchors_by_height
+                // This makes sure the temporary trees from finalized tip forks are removed.
-                .keys()
+                self.orchard_anchors_by_height
-                .cloned()
+                    .keys()
-                .filter(|index_height| *index_height <= height)
+                    .cloned()
-                .collect()
+                    .filter(|index_height| *index_height <= height)
+                    .collect(),
+                // Cache the highest (rightmost) tree before its removal.
+                self.orchard_tree(height.into()),
+            )
         } else {
             // Just remove the reverted tip trees and anchors.
-            vec![height]
+            // We don't need to cache the highest (rightmost) tree.
+            (vec![height], None)
         };
 
         for height in &removed_heights {
@@ -768,9 +899,8 @@ impl Chain {
                 .orchard_anchors_by_height
                 .remove(height)
                 .expect("Orchard anchor must be present if block was added to chain");
-            self.orchard_trees_by_height
+
-                .remove(height)
+            self.orchard_trees_by_height.remove(height);
-                .expect("Orchard note commitment tree must be present if block was added to chain");
 
             trace!(?height, ?position, ?anchor, "removing orchard tree");
 
@@ -781,6 +911,26 @@ impl Chain {
                 "Orchard anchor must be present if block was added to chain"
             );
         }
+
+        // # Invariant
+        //
+        // The height following after the removed heights in a non-empty non-finalized state must
+        // always have its tree.
+        //
+        // The loop above can violate the invariant, and if `position` is [`RevertPosition::Root`],
+        // it will always violate the invariant. We restore the invariant by storing the highest
+        // (rightmost) removed tree just above `height` if there is no tree at that height.
+        if !self.is_empty() && height < self.non_finalized_tip_height() {
+            let next_height = height.next();
+
+            if self.orchard_trees_by_height.get(&next_height).is_none() {
+                // TODO: Use `try_insert` once it stabilises.
+                self.orchard_trees_by_height.insert(
+                    next_height,
+                    highest_removed_tree.expect("There should be a cached removed tree."),
+                );
+            }
+        }
     }
 
     /// Returns the History tree of the tip of this [`Chain`],