Check for sprout tree roots, key format, and de-duplication on startup

zebra-state/src/service/finalized_state/disk_format.rs

@@ -110,6 +110,44 @@ impl FromDisk for () {
     }
 }
 
+/// Access database keys or values as raw bytes.
+/// Mainly for use in tests or runtime checks.
+#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
+pub struct RawBytes(Vec<u8>);
+
+// Note: don't implement From or Into for RawBytes, because it makes it harder to spot in reviews.
+// Instead, implement IntoDisk and FromDisk on the original type, or a specific wrapper type.
+
+impl RawBytes {
+    /// Create a new raw byte key or data.
+    ///
+    /// Mainly for use in tests or runtime checks.
+    /// These methods
+    pub fn new_raw_bytes(bytes: Vec<u8>) -> Self {
+        Self(bytes)
+    }
+
+    /// Create a new raw byte key or data.
+    /// Mainly for use in tests.
+    pub fn raw_bytes(&self) -> &Vec<u8> {
+        &self.0
+    }
+}
+
+impl IntoDisk for RawBytes {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.raw_bytes().clone()
+    }
+}
+
+impl FromDisk for RawBytes {
+    fn from_bytes(bytes: impl AsRef<[u8]>) -> Self {
+        Self::new_raw_bytes(bytes.as_ref().to_vec())
+    }
+}
+
 // Serialization Modification Functions
 
 /// Truncates `mem_bytes` to `disk_len`, by removing zero bytes from the start of the slice.

zebra-state/src/service/finalized_state/disk_format/upgrade.rs

@@ -405,6 +405,67 @@ impl DbFormatChange {
         // We always run this test, even if the state has supposedly been upgraded.
         let mut problem_found = false;
 
+        for (root, tree) in upgrade_db.sprout_trees_full_map() {
+            if tree.cached_root().is_none() {
+                // TODO: replace this with a panic because it indicates an unrecoverable
+                //       bug, which should fail the tests immediately
+                error!(
+                    key_root = ?root,
+                    tree_root = ?tree.root(),
+                    "found un-cached sprout tree root after running genesis tree root fix"
+                );
+
+                problem_found = true;
+            }
+        }
+
+        let mut prev_key = None;
+        let mut prev_tree: Option<Arc<sprout::tree::NoteCommitmentTree>> = None;
+        for (key, tree) in upgrade_db.sprout_trees_full_tip() {
+            if tree.cached_root().is_none() {
+                // TODO: replace this with a panic because it indicates an unrecoverable
+                //       bug, which should fail the tests immediately
+                error!(
+                    key = ?key,
+                    tree_root = ?tree.root(),
+                    "found un-cached sprout tree root after running genesis tree root fix"
+                );
+
+                problem_found = true;
+            }
+
+            // The tip tree should be indexed by `()` (which serializes to an empty array).
+            if !key.raw_bytes().is_empty() {
+                // TODO: replace this with a panic because it indicates an unrecoverable
+                //       bug, which should fail the tests immediately
+                error!(
+                    key = ?key,
+                    tree_root = ?tree.root(),
+                    "found incorrect sprout tree key format after running genesis tree root fix"
+                );
+
+                problem_found = true;
+            }
+
+            // There should only be one tip tree in this column family.
+            if prev_tree.is_some() {
+                // TODO: replace this with a panic because it indicates an unrecoverable
+                //       bug, which should fail the tests immediately
+                error!(
+                    key = ?key,
+                    tree_root = ?tree.root(),
+                    prev_key = ?prev_key.unwrap(),
+                    prev_tree_root = ?prev_tree.unwrap().root(),
+                    "found duplicate sprout trees after running de-duplicate tree upgrade"
+                );
+
+                problem_found = true;
+            }
+
+            prev_key = Some(key);
+            prev_tree = Some(tree);
+        }
+
         let mut prev_height = None;
         let mut prev_tree = None;
         for (height, tree) in upgrade_db.sapling_tree_by_height_range(..) {

zebra-state/src/service/finalized_state/zebra_db/shielded.rs

@@ -23,6 +23,7 @@ use crate::{
     request::SemanticallyVerifiedBlockWithTrees,
     service::finalized_state::{
         disk_db::{DiskDb, DiskWriteBatch, ReadDisk, WriteDisk},
+        disk_format::RawBytes,
         zebra_db::ZebraDb,
     },
     BoxError, SemanticallyVerifiedBlock,
@@ -117,6 +118,16 @@ impl ZebraDb {
             .zs_items_in_range_unordered(&sprout_anchors_handle, ..)
     }
 
+    /// Returns all the Sprout note commitment tip trees.
+    /// We only store one sprout tree for the tip, so this method is mainly used in tests.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn sprout_trees_full_tip(
+        &self,
+    ) -> impl Iterator<Item = (RawBytes, Arc<sprout::tree::NoteCommitmentTree>)> + '_ {
+        let sprout_trees = self.db.cf_handle("sprout_note_commitment_tree").unwrap();
+        self.db.zs_range_iter(&sprout_trees, ..)
+    }
+
     // # Sapling trees
 
     /// Returns the Sapling note commitment tree of the finalized tip or the empty tree if the state