From 2871b2fad13fe6195c1c44495105d6351fc28bf7 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 15:39:35 -0400
Subject: [PATCH 01/17] Adds zs_iter_opts method

---
 .../src/service/finalized_state/disk_db.rs    | 35 ++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 3772fb7a789..1ca37fc5cc4 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -21,6 +21,7 @@ use std::{
 use itertools::Itertools;
 use rlimit::increase_nofile_limit;
 
+use rocksdb::ReadOptions;
 use zebra_chain::parameters::Network;
 
 use crate::{
@@ -495,11 +496,12 @@ impl DiskDb {
         let range = (start_bound, end_bound);
 
         let mode = Self::zs_iter_mode(&range, reverse);
+        let opts = Self::zs_iter_opts(&range);
 
         // Reading multiple items from iterators has caused database hangs,
         // in previous RocksDB versions
         self.db
-            .iterator_cf(cf, mode)
+            .iterator_cf_opt(cf, opts, mode)
             .map(|result| result.expect("unexpected database failure"))
             .map(|(key, value)| (key.to_vec(), value))
             // Skip excluded "from" bound and empty ranges. The `mode` already skips keys
@@ -514,6 +516,37 @@ impl DiskDb {
             .map(|(key, value)| (K::from_bytes(key), V::from_bytes(value)))
     }
 
+    /// Returns the RocksDB ReadOptions with a lower and upper bound for a range.
+    fn zs_iter_opts<R>(range: &R) -> ReadOptions
+    where
+        R: RangeBounds<Vec<u8>>,
+    {
+        use std::ops::Bound::*;
+
+        let mut opts = ReadOptions::default();
+
+        if let Included(bound) | Excluded(bound) = range.start_bound() {
+            opts.set_iterate_lower_bound(bound.clone());
+        };
+
+        match range.end_bound().cloned() {
+            Included(mut bound) => {
+                // Increment the last byte in the vector that is not u8::MAX, or
+                // skip adding an upper bound if every byte is u8::MAX
+                if let Some(increment_idx) = bound.iter().rposition(|&v| v != u8::MAX) {
+                    bound[increment_idx] += 1;
+                    opts.set_iterate_upper_bound(bound);
+                }
+            }
+            Excluded(bound) => {
+                opts.set_iterate_upper_bound(bound);
+            }
+            Unbounded => {}
+        };
+
+        opts
+    }
+
     /// Returns the RocksDB iterator "from" mode for `range`.
     ///
     /// RocksDB iterators are ordered by increasing key bytes by default.

From d815ae5292328846ffcb852193eeff00f97456be Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 15:50:14 -0400
Subject: [PATCH 02/17] uses checked_add

---
 zebra-state/src/service/finalized_state/disk_db.rs | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 1ca37fc5cc4..fb750aca2db 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -534,7 +534,13 @@ impl DiskDb {
                 // Increment the last byte in the vector that is not u8::MAX, or
                 // skip adding an upper bound if every byte is u8::MAX
                 if let Some(increment_idx) = bound.iter().rposition(|&v| v != u8::MAX) {
-                    bound[increment_idx] += 1;
+                    let increment_byte = bound
+                        .get_mut(increment_idx)
+                        .expect("index should be in bounds");
+                    *increment_byte = increment_byte
+                        .checked_add(1)
+                        .expect("adding 1 should succeed");
+
                     opts.set_iterate_upper_bound(bound);
                 }
             }

From 49123208d0347283e03d97577a0ed19e6c3bc3c0 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 16:05:47 -0400
Subject: [PATCH 03/17] uses zs_range_iter for other read methods

---
 .../src/service/finalized_state/disk_db.rs    | 59 ++++++-------------
 .../finalized_state/zebra_db/shielded.rs      | 12 ++--
 2 files changed, 23 insertions(+), 48 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index fb750aca2db..806cc9536d4 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -195,7 +195,7 @@ pub trait ReadDisk {
     fn zs_first_key_value<C, K, V>(&self, cf: &C) -> Option<(K, V)>
     where
         C: rocksdb::AsColumnFamilyRef,
-        K: FromDisk,
+        K: IntoDisk + FromDisk,
         V: FromDisk;
 
     /// Returns the highest key in `cf`, and the corresponding value.
@@ -204,7 +204,7 @@ pub trait ReadDisk {
     fn zs_last_key_value<C, K, V>(&self, cf: &C) -> Option<(K, V)>
     where
         C: rocksdb::AsColumnFamilyRef,
-        K: FromDisk,
+        K: IntoDisk + FromDisk,
         V: FromDisk;
 
     /// Returns the first key greater than or equal to `lower_bound` in `cf`,
@@ -322,34 +322,22 @@ impl ReadDisk for DiskDb {
     fn zs_first_key_value<C, K, V>(&self, cf: &C) -> Option<(K, V)>
     where
         C: rocksdb::AsColumnFamilyRef,
-        K: FromDisk,
+        K: IntoDisk + FromDisk,
         V: FromDisk,
     {
         // Reading individual values from iterators does not seem to cause database hangs.
-        self.db
-            .iterator_cf(cf, rocksdb::IteratorMode::Start)
-            .next()?
-            .map(|(key_bytes, value_bytes)| {
-                Some((K::from_bytes(key_bytes), V::from_bytes(value_bytes)))
-            })
-            .expect("unexpected database failure")
+        self.zs_range_iter(cf, .., false).next()
     }
 
     #[allow(clippy::unwrap_in_result)]
     fn zs_last_key_value<C, K, V>(&self, cf: &C) -> Option<(K, V)>
     where
         C: rocksdb::AsColumnFamilyRef,
-        K: FromDisk,
+        K: IntoDisk + FromDisk,
         V: FromDisk,
     {
         // Reading individual values from iterators does not seem to cause database hangs.
-        self.db
-            .iterator_cf(cf, rocksdb::IteratorMode::End)
-            .next()?
-            .map(|(key_bytes, value_bytes)| {
-                Some((K::from_bytes(key_bytes), V::from_bytes(value_bytes)))
-            })
-            .expect("unexpected database failure")
+        self.zs_range_iter(cf, .., true).next()
     }
 
     #[allow(clippy::unwrap_in_result)]
@@ -359,17 +347,8 @@ impl ReadDisk for DiskDb {
         K: IntoDisk + FromDisk,
         V: FromDisk,
     {
-        let lower_bound = lower_bound.as_bytes();
-        let from = rocksdb::IteratorMode::From(lower_bound.as_ref(), rocksdb::Direction::Forward);
-
         // Reading individual values from iterators does not seem to cause database hangs.
-        self.db
-            .iterator_cf(cf, from)
-            .next()?
-            .map(|(key_bytes, value_bytes)| {
-                Some((K::from_bytes(key_bytes), V::from_bytes(value_bytes)))
-            })
-            .expect("unexpected database failure")
+        self.zs_range_iter(cf, lower_bound.., false).next()
     }
 
     #[allow(clippy::unwrap_in_result)]
@@ -379,17 +358,8 @@ impl ReadDisk for DiskDb {
         K: IntoDisk + FromDisk,
         V: FromDisk,
     {
-        let upper_bound = upper_bound.as_bytes();
-        let from = rocksdb::IteratorMode::From(upper_bound.as_ref(), rocksdb::Direction::Reverse);
-
         // Reading individual values from iterators does not seem to cause database hangs.
-        self.db
-            .iterator_cf(cf, from)
-            .next()?
-            .map(|(key_bytes, value_bytes)| {
-                Some((K::from_bytes(key_bytes), V::from_bytes(value_bytes)))
-            })
-            .expect("unexpected database failure")
+        self.zs_range_iter(cf, ..=upper_bound, false).next()
     }
 
     fn zs_items_in_range_ordered<C, K, V, R>(&self, cf: &C, range: R) -> BTreeMap<K, V>
@@ -399,7 +369,7 @@ impl ReadDisk for DiskDb {
         V: FromDisk,
         R: RangeBounds<K>,
     {
-        self.zs_range_iter(cf, range).collect()
+        self.zs_range_iter(cf, range, false).collect()
     }
 
     fn zs_items_in_range_unordered<C, K, V, R>(&self, cf: &C, range: R) -> HashMap<K, V>
@@ -409,7 +379,7 @@ impl ReadDisk for DiskDb {
         V: FromDisk,
         R: RangeBounds<K>,
     {
-        self.zs_range_iter(cf, range).collect()
+        self.zs_range_iter(cf, range, false).collect()
     }
 }
 
@@ -432,14 +402,19 @@ impl DiskDb {
     /// Returns an iterator over the items in `cf` in `range`.
     ///
     /// Holding this iterator open might delay block commit transactions.
-    pub fn zs_range_iter<C, K, V, R>(&self, cf: &C, range: R) -> impl Iterator<Item = (K, V)> + '_
+    pub fn zs_range_iter<C, K, V, R>(
+        &self,
+        cf: &C,
+        range: R,
+        reverse: bool,
+    ) -> impl Iterator<Item = (K, V)> + '_
     where
         C: rocksdb::AsColumnFamilyRef,
         K: IntoDisk + FromDisk,
         V: FromDisk,
         R: RangeBounds<K>,
     {
-        self.zs_range_iter_with_direction(cf, range, false)
+        self.zs_range_iter_with_direction(cf, range, reverse)
     }
 
     /// Returns a reverse iterator over the items in `cf` in `range`.
diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
index 75b5db8da64..59788640a33 100644
--- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
+++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs
@@ -209,7 +209,7 @@ impl ZebraDb {
         R: std::ops::RangeBounds<Height>,
     {
         let sapling_trees = self.db.cf_handle("sapling_note_commitment_tree").unwrap();
-        self.db.zs_range_iter(&sapling_trees, range)
+        self.db.zs_range_iter(&sapling_trees, range, false)
     }
 
     /// Returns the Sapling note commitment trees in the reversed range, in decreasing height order.
@@ -282,7 +282,7 @@ impl ZebraDb {
         if let Some(exclusive_end_bound) = exclusive_end_bound {
             list = self
                 .db
-                .zs_range_iter(&sapling_subtrees, start_index..exclusive_end_bound)
+                .zs_range_iter(&sapling_subtrees, start_index..exclusive_end_bound, false)
                 .collect();
         } else {
             // If there is no end bound, just return all the trees.
@@ -291,7 +291,7 @@ impl ZebraDb {
             // the trees run out.)
             list = self
                 .db
-                .zs_range_iter(&sapling_subtrees, start_index..)
+                .zs_range_iter(&sapling_subtrees, start_index.., false)
                 .collect();
         }
 
@@ -376,7 +376,7 @@ impl ZebraDb {
         R: std::ops::RangeBounds<Height>,
     {
         let orchard_trees = self.db.cf_handle("orchard_note_commitment_tree").unwrap();
-        self.db.zs_range_iter(&orchard_trees, range)
+        self.db.zs_range_iter(&orchard_trees, range, false)
     }
 
     /// Returns the Orchard note commitment trees in the reversed range, in decreasing height order.
@@ -449,7 +449,7 @@ impl ZebraDb {
         if let Some(exclusive_end_bound) = exclusive_end_bound {
             list = self
                 .db
-                .zs_range_iter(&orchard_subtrees, start_index..exclusive_end_bound)
+                .zs_range_iter(&orchard_subtrees, start_index..exclusive_end_bound, false)
                 .collect();
         } else {
             // If there is no end bound, just return all the trees.
@@ -458,7 +458,7 @@ impl ZebraDb {
             // the trees run out.)
             list = self
                 .db
-                .zs_range_iter(&orchard_subtrees, start_index..)
+                .zs_range_iter(&orchard_subtrees, start_index.., false)
                 .collect();
         }
 

From 1ce960b2342c37d01db85855285ff6db01c3a36b Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 19:49:47 -0400
Subject: [PATCH 04/17] fixes bug

---
 zebra-state/src/service/finalized_state/disk_db.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 806cc9536d4..68944d79ce1 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -359,7 +359,7 @@ impl ReadDisk for DiskDb {
         V: FromDisk,
     {
         // Reading individual values from iterators does not seem to cause database hangs.
-        self.zs_range_iter(cf, ..=upper_bound, false).next()
+        self.zs_range_iter(cf, ..=upper_bound, true).next()
     }
 
     fn zs_items_in_range_ordered<C, K, V, R>(&self, cf: &C, range: R) -> BTreeMap<K, V>

From bd80289b0fecf4e08d723f9a53a27597634500c9 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 22:05:08 -0400
Subject: [PATCH 05/17] Fixes bug in zs_iter_opts

---
 .../src/service/finalized_state/disk_db.rs    | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 68944d79ce1..b4a791e009e 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -506,16 +506,16 @@ impl DiskDb {
 
         match range.end_bound().cloned() {
             Included(mut bound) => {
-                // Increment the last byte in the vector that is not u8::MAX, or
-                // skip adding an upper bound if every byte is u8::MAX
-                if let Some(increment_idx) = bound.iter().rposition(|&v| v != u8::MAX) {
-                    let increment_byte = bound
-                        .get_mut(increment_idx)
-                        .expect("index should be in bounds");
-                    *increment_byte = increment_byte
-                        .checked_add(1)
-                        .expect("adding 1 should succeed");
-
+                // Skip adding an upper bound if every byte is u8::MAX, or
+                // increment the last byte in the upper bound that is less than u8::MAX,
+                // and clear any bytes after it to increment the big-endian number this
+                // string represents to RocksDB.
+                let is_max_key = bound.iter_mut().rev().all(|v| {
+                    *v = v.wrapping_add(1);
+                    v == &0
+                });
+
+                if !is_max_key {
                     opts.set_iterate_upper_bound(bound);
                 }
             }

From a703d1ce996256b02f13d707d51cde17d8d33c64 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 23:11:51 -0400
Subject: [PATCH 06/17] Adds test & updates method docs

---
 .../src/service/finalized_state/disk_db.rs    | 44 +++++++++++++------
 .../service/finalized_state/disk_db/tests.rs  | 25 +++++++++++
 2 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index b4a791e009e..8016fa3de8b 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -401,6 +401,10 @@ impl DiskWriteBatch {
 impl DiskDb {
     /// Returns an iterator over the items in `cf` in `range`.
     ///
+    /// Accepts a `reverse` argument and creates the iterator with an [`IteratorMode`](rocksdb::IteratorMode)
+    /// of [`End`](rocksdb::IteratorMode::End), or [`From`](rocksdb::IteratorMode::From)
+    /// with [`Direction::Reverse`](rocksdb::Direction::Reverse).
+    ///
     /// Holding this iterator open might delay block commit transactions.
     pub fn zs_range_iter<C, K, V, R>(
         &self,
@@ -496,15 +500,33 @@ impl DiskDb {
     where
         R: RangeBounds<Vec<u8>>,
     {
-        use std::ops::Bound::*;
-
         let mut opts = ReadOptions::default();
+        let (lower_bound, upper_bound) = Self::zs_iter_bounds(range);
+
+        if let Some(bound) = lower_bound {
+            opts.set_iterate_lower_bound(bound);
+        };
+
+        if let Some(bound) = upper_bound {
+            opts.set_iterate_upper_bound(bound);
+        };
+
+        opts
+    }
 
-        if let Included(bound) | Excluded(bound) = range.start_bound() {
-            opts.set_iterate_lower_bound(bound.clone());
+    /// Returns a lower and upper iterate bounds for a range.
+    fn zs_iter_bounds<R>(range: &R) -> (Option<Vec<u8>>, Option<Vec<u8>>)
+    where
+        R: RangeBounds<Vec<u8>>,
+    {
+        use std::ops::Bound::*;
+
+        let lower_bound = match range.start_bound() {
+            Included(bound) | Excluded(bound) => Some(bound.clone()),
+            Unbounded => None,
         };
 
-        match range.end_bound().cloned() {
+        let upper_bound = match range.end_bound().cloned() {
             Included(mut bound) => {
                 // Skip adding an upper bound if every byte is u8::MAX, or
                 // increment the last byte in the upper bound that is less than u8::MAX,
@@ -515,17 +537,13 @@ impl DiskDb {
                     v == &0
                 });
 
-                if !is_max_key {
-                    opts.set_iterate_upper_bound(bound);
-                }
+                (!is_max_key).then_some(bound)
             }
-            Excluded(bound) => {
-                opts.set_iterate_upper_bound(bound);
-            }
-            Unbounded => {}
+            Excluded(bound) => Some(bound),
+            Unbounded => None,
         };
 
-        opts
+        (lower_bound, upper_bound)
     }
 
     /// Returns the RocksDB iterator "from" mode for `range`.
diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index 17613e8b3b5..2b5567217a8 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -24,3 +24,28 @@ impl DiskDb {
         rocksdb::DB::list_cf(&opts, path)
     }
 }
+
+/// Check that the sprout tree database serialization format has not changed.
+#[test]
+fn zs_iter_opts_increments_key_by_one() {
+    let _init_guard = zebra_test::init();
+
+    // let (config, network) = Default::default();
+    // let db = DiskDb::new(&config, network);
+
+    let keys: [u32; 13] = [
+        0, 1, 200, 255, 256, 257, 65535, 65536, 65537, 16777215, 16777216, 16777217, 16777218,
+    ];
+
+    for key in keys {
+        let (_, upper_bound_bytes) = DiskDb::zs_iter_bounds(&..=key.to_be_bytes().to_vec());
+        let upper_bound_bytes = upper_bound_bytes.expect("there should be an upper bound");
+        let upper_bound = u32::from_be_bytes(upper_bound_bytes.try_into().unwrap());
+        let expected_upper_bound = key + 1;
+
+        assert_eq!(
+            expected_upper_bound, upper_bound,
+            "the upper bound should be 1 greater than the original key"
+        );
+    }
+}

From 7814d5fde5bcd7bfaf6fbcef2a159bcfd871e3e2 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 23:20:21 -0400
Subject: [PATCH 07/17] updates docs

---
 zebra-state/src/service/finalized_state/disk_db.rs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 8016fa3de8b..b88fdefe130 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -515,6 +515,11 @@ impl DiskDb {
     }
 
     /// Returns a lower and upper iterate bounds for a range.
+    ///
+    /// Note: Since upper iterate bounds are always exclusive in RocksDB, this method
+    ///       will increment the upper bound by 1 if the end bound of the provided range
+    ///       is inclusive, or will return an upper bound of `None` if the end bound of a
+    ///       provided range is inclusive and already the max key for that column family.
     fn zs_iter_bounds<R>(range: &R) -> (Option<Vec<u8>>, Option<Vec<u8>>)
     where
         R: RangeBounds<Vec<u8>>,
@@ -531,7 +536,7 @@ impl DiskDb {
                 // Skip adding an upper bound if every byte is u8::MAX, or
                 // increment the last byte in the upper bound that is less than u8::MAX,
                 // and clear any bytes after it to increment the big-endian number this
-                // string represents to RocksDB.
+                // Vec represents to RocksDB.
                 let is_max_key = bound.iter_mut().rev().all(|v| {
                     *v = v.wrapping_add(1);
                     v == &0

From b3fc1cd4138fe00bda3652b3a785791811fcf30a Mon Sep 17 00:00:00 2001
From: Arya <aryasolhi@gmail.com>
Date: Wed, 11 Oct 2023 23:28:05 -0400
Subject: [PATCH 08/17] Update
 zebra-state/src/service/finalized_state/disk_db/tests.rs

---
 zebra-state/src/service/finalized_state/disk_db/tests.rs | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index 2b5567217a8..c2d1a04d323 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -30,9 +30,6 @@ impl DiskDb {
 fn zs_iter_opts_increments_key_by_one() {
     let _init_guard = zebra_test::init();
 
-    // let (config, network) = Default::default();
-    // let db = DiskDb::new(&config, network);
-
     let keys: [u32; 13] = [
         0, 1, 200, 255, 256, 257, 65535, 65536, 65537, 16777215, 16777216, 16777217, 16777218,
     ];

From 59dba97c7292bebbbad07407f4ab7ff38253992f Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 12 Oct 2023 13:53:50 -0400
Subject: [PATCH 09/17] Corrects code comment

---
 zebra-state/src/service/finalized_state/disk_db/tests.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index c2d1a04d323..4ba22c39ac7 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -25,7 +25,7 @@ impl DiskDb {
     }
 }
 
-/// Check that the sprout tree database serialization format has not changed.
+/// Check that zs_iter_opts returns an upper bound one greater than provided inclusive end bounds.
 #[test]
 fn zs_iter_opts_increments_key_by_one() {
     let _init_guard = zebra_test::init();

From fcf171f10be906f0c77041401ad1b59d5adadbaf Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Thu, 12 Oct 2023 23:57:29 -0400
Subject: [PATCH 10/17] adds support for variable-sized keys

---
 zebra-state/src/service/finalized_state/disk_db.rs | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index b88fdefe130..7ee2554b226 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -537,12 +537,17 @@ impl DiskDb {
                 // increment the last byte in the upper bound that is less than u8::MAX,
                 // and clear any bytes after it to increment the big-endian number this
                 // Vec represents to RocksDB.
-                let is_max_key = bound.iter_mut().rev().all(|v| {
+                let is_zero = bound.iter_mut().rev().all(|v| {
                     *v = v.wrapping_add(1);
                     v == &0
                 });
 
-                (!is_max_key).then_some(bound)
+                if is_zero {
+                    bound.push(0);
+                    *bound.get_mut(0).expect("should have at least 1 element") += 1;
+                }
+
+                Some(bound)
             }
             Excluded(bound) => Some(bound),
             Unbounded => None,

From 600aa5f6087dc0b6af9d9fcce54b614cedaee5c4 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Fri, 13 Oct 2023 00:14:36 -0400
Subject: [PATCH 11/17] adds test case

---
 .../service/finalized_state/disk_db/tests.rs  | 32 ++++++++++++++++---
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index 4ba22c39ac7..ebfa44ede75 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -30,19 +30,41 @@ impl DiskDb {
 fn zs_iter_opts_increments_key_by_one() {
     let _init_guard = zebra_test::init();
 
-    let keys: [u32; 13] = [
-        0, 1, 200, 255, 256, 257, 65535, 65536, 65537, 16777215, 16777216, 16777217, 16777218,
+    let keys: [u32; 14] = [
+        0,
+        1,
+        200,
+        255,
+        256,
+        257,
+        65535,
+        65536,
+        65537,
+        16777215,
+        16777216,
+        16777217,
+        16777218,
+        u32::MAX,
     ];
 
     for key in keys {
-        let (_, upper_bound_bytes) = DiskDb::zs_iter_bounds(&..=key.to_be_bytes().to_vec());
-        let upper_bound_bytes = upper_bound_bytes.expect("there should be an upper bound");
+        let (_, bytes) = DiskDb::zs_iter_bounds(&..=key.to_be_bytes().to_vec());
+        let mut bytes = bytes.expect("there should be an upper bound");
+        let upper_bound_bytes = bytes.split_off(bytes.len() - 4);
         let upper_bound = u32::from_be_bytes(upper_bound_bytes.try_into().unwrap());
-        let expected_upper_bound = key + 1;
+        let expected_upper_bound = key.wrapping_add(1);
 
         assert_eq!(
             expected_upper_bound, upper_bound,
             "the upper bound should be 1 greater than the original key"
         );
+
+        if expected_upper_bound == 0 {
+            assert_eq!(
+                bytes,
+                vec![1],
+                "there should be an extra byte with a value of 1"
+            );
+        }
     }
 }

From 4d2de6ffdf7ecca93734308076625135de7c2c00 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Tue, 17 Oct 2023 17:57:58 -0400
Subject: [PATCH 12/17] Updates docs

---
 zebra-state/src/service/finalized_state/disk_db.rs | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 7ee2554b226..8f0e1aa6c9f 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -518,8 +518,7 @@ impl DiskDb {
     ///
     /// Note: Since upper iterate bounds are always exclusive in RocksDB, this method
     ///       will increment the upper bound by 1 if the end bound of the provided range
-    ///       is inclusive, or will return an upper bound of `None` if the end bound of a
-    ///       provided range is inclusive and already the max key for that column family.
+    ///       is inclusive.
     fn zs_iter_bounds<R>(range: &R) -> (Option<Vec<u8>>, Option<Vec<u8>>)
     where
         R: RangeBounds<Vec<u8>>,
@@ -533,8 +532,7 @@ impl DiskDb {
 
         let upper_bound = match range.end_bound().cloned() {
             Included(mut bound) => {
-                // Skip adding an upper bound if every byte is u8::MAX, or
-                // increment the last byte in the upper bound that is less than u8::MAX,
+                // Increment the last byte in the upper bound that is less than u8::MAX,
                 // and clear any bytes after it to increment the big-endian number this
                 // Vec represents to RocksDB.
                 let is_zero = bound.iter_mut().rev().all(|v| {

From 9cf67fbb702870f89a95dc98b9e4ccaf6d321332 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Tue, 17 Oct 2023 18:07:00 -0400
Subject: [PATCH 13/17] Applies suggestions from code review

---
 .../src/service/finalized_state/disk_db.rs    | 19 +++++++++----------
 .../service/finalized_state/disk_db/tests.rs  |  2 +-
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db.rs b/zebra-state/src/service/finalized_state/disk_db.rs
index 8f0e1aa6c9f..c042ea6db0d 100644
--- a/zebra-state/src/service/finalized_state/disk_db.rs
+++ b/zebra-state/src/service/finalized_state/disk_db.rs
@@ -401,9 +401,9 @@ impl DiskWriteBatch {
 impl DiskDb {
     /// Returns an iterator over the items in `cf` in `range`.
     ///
-    /// Accepts a `reverse` argument and creates the iterator with an [`IteratorMode`](rocksdb::IteratorMode)
-    /// of [`End`](rocksdb::IteratorMode::End), or [`From`](rocksdb::IteratorMode::From)
-    /// with [`Direction::Reverse`](rocksdb::Direction::Reverse).
+    /// Accepts a `reverse` argument. If it is `true`, creates the iterator with an
+    /// [`IteratorMode`](rocksdb::IteratorMode) of [`End`](rocksdb::IteratorMode::End), or
+    /// [`From`](rocksdb::IteratorMode::From) with [`Direction::Reverse`](rocksdb::Direction::Reverse).
     ///
     /// Holding this iterator open might delay block commit transactions.
     pub fn zs_range_iter<C, K, V, R>(
@@ -532,17 +532,16 @@ impl DiskDb {
 
         let upper_bound = match range.end_bound().cloned() {
             Included(mut bound) => {
-                // Increment the last byte in the upper bound that is less than u8::MAX,
-                // and clear any bytes after it to increment the big-endian number this
-                // Vec represents to RocksDB.
-                let is_zero = bound.iter_mut().rev().all(|v| {
+                // Increment the last byte in the upper bound that is less than u8::MAX, and
+                // clear any bytes after it to increment the next key in lexicographic order
+                // (next big-endian number) this Vec represents to RocksDB.
+                let is_wrapped_overflow = bound.iter_mut().rev().all(|v| {
                     *v = v.wrapping_add(1);
                     v == &0
                 });
 
-                if is_zero {
-                    bound.push(0);
-                    *bound.get_mut(0).expect("should have at least 1 element") += 1;
+                if is_wrapped_overflow {
+                    bound.insert(0, 0x01)
                 }
 
                 Some(bound)
diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index ebfa44ede75..8e71ce64d09 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -51,7 +51,7 @@ fn zs_iter_opts_increments_key_by_one() {
         let (_, bytes) = DiskDb::zs_iter_bounds(&..=key.to_be_bytes().to_vec());
         let mut bytes = bytes.expect("there should be an upper bound");
         let upper_bound_bytes = bytes.split_off(bytes.len() - 4);
-        let upper_bound = u32::from_be_bytes(upper_bound_bytes.try_into().unwrap());
+        let upper_bound = u32::from_be_bytes(upper_bound_bytes.try_into().expect("no added bytes"));
         let expected_upper_bound = key.wrapping_add(1);
 
         assert_eq!(

From 1c595fb1a5a68d7f72ea6226e725a21bc6861401 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Wed, 18 Oct 2023 13:40:29 +1000
Subject: [PATCH 14/17] Add extra checks

---
 .../src/service/finalized_state/disk_db/tests.rs    | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index 8e71ce64d09..e7c9a26919e 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -30,6 +30,7 @@ impl DiskDb {
 fn zs_iter_opts_increments_key_by_one() {
     let _init_guard = zebra_test::init();
 
+    // TODO: add an empty key (`()` type or `[]` when serialized) test case
     let keys: [u32; 14] = [
         0,
         1,
@@ -65,6 +66,18 @@ fn zs_iter_opts_increments_key_by_one() {
                 vec![1],
                 "there should be an extra byte with a value of 1"
             );
+        } else {
+            assert_eq!(
+                key.len(),
+                bytes.len(),
+                "there should be no extra bytes"
+            );
         }
+
+        assert_ne!(
+            bytes[0],
+            0x00,
+            "there must be at least one byte, and the first byte can't be zero"
+        );
     }
 }

From f64456220234c529579744d48e63eecc64e4d3a0 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Wed, 18 Oct 2023 14:11:17 +1000
Subject: [PATCH 15/17] Fix test code and rustfmt

---
 zebra-state/src/service/finalized_state/disk_db/tests.rs | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index e7c9a26919e..c52053ba050 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -68,15 +68,14 @@ fn zs_iter_opts_increments_key_by_one() {
             );
         } else {
             assert_eq!(
-                key.len(),
+                key.to_be_bytes().len(),
                 bytes.len(),
                 "there should be no extra bytes"
             );
         }
 
         assert_ne!(
-            bytes[0],
-            0x00,
+            bytes[0], 0x00,
             "there must be at least one byte, and the first byte can't be zero"
         );
     }

From cbca3343f4c272db2ca49e11d076f6ff70595878 Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 18 Oct 2023 01:00:50 -0400
Subject: [PATCH 16/17] fixes test

---
 .../service/finalized_state/disk_db/tests.rs  | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index c52053ba050..0fd8f857f9c 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -50,9 +50,9 @@ fn zs_iter_opts_increments_key_by_one() {
 
     for key in keys {
         let (_, bytes) = DiskDb::zs_iter_bounds(&..=key.to_be_bytes().to_vec());
-        let mut bytes = bytes.expect("there should be an upper bound");
-        let upper_bound_bytes = bytes.split_off(bytes.len() - 4);
-        let upper_bound = u32::from_be_bytes(upper_bound_bytes.try_into().expect("no added bytes"));
+        let mut extra_bytes = bytes.expect("there should be an upper bound");
+        let bytes = extra_bytes.split_off(extra_bytes.len() - 4);
+        let upper_bound = u32::from_be_bytes(bytes.clone().try_into().expect("should be 4 bytes"));
         let expected_upper_bound = key.wrapping_add(1);
 
         assert_eq!(
@@ -62,21 +62,18 @@ fn zs_iter_opts_increments_key_by_one() {
 
         if expected_upper_bound == 0 {
             assert_eq!(
-                bytes,
+                extra_bytes,
                 vec![1],
                 "there should be an extra byte with a value of 1"
             );
         } else {
-            assert_eq!(
-                key.to_be_bytes().len(),
-                bytes.len(),
-                "there should be no extra bytes"
-            );
+            assert_eq!(extra_bytes.len(), 0, "there should be no extra bytes");
         }
 
         assert_ne!(
-            bytes[0], 0x00,
-            "there must be at least one byte, and the first byte can't be zero"
+            extra_bytes.last().expect("there must be at least one byte"),
+            &0,
+            "the last byte can't be zero"
         );
     }
 }

From da0afea7c44888cd041895e18c02ddadd7e2e5fd Mon Sep 17 00:00:00 2001
From: arya2 <aryasolhi@gmail.com>
Date: Wed, 18 Oct 2023 01:04:10 -0400
Subject: [PATCH 17/17] fixes test

---
 zebra-state/src/service/finalized_state/disk_db/tests.rs | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/zebra-state/src/service/finalized_state/disk_db/tests.rs b/zebra-state/src/service/finalized_state/disk_db/tests.rs
index 0fd8f857f9c..20fecbbf127 100644
--- a/zebra-state/src/service/finalized_state/disk_db/tests.rs
+++ b/zebra-state/src/service/finalized_state/disk_db/tests.rs
@@ -69,11 +69,5 @@ fn zs_iter_opts_increments_key_by_one() {
         } else {
             assert_eq!(extra_bytes.len(), 0, "there should be no extra bytes");
         }
-
-        assert_ne!(
-            extra_bytes.last().expect("there must be at least one byte"),
-            &0,
-            "the last byte can't be zero"
-        );
     }
 }