fix: early limit reduction

grovedb/src/batch/key_info.rs

@@ -100,20 +100,7 @@ impl PartialEq<&[u8]> for KeyInfo {
 #[cfg(feature = "full")]
 impl PartialOrd<Self> for KeyInfo {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        match self.as_slice().partial_cmp(other.as_slice()) {
-            None => None,
-            Some(ord) => match ord {
-                Ordering::Less => Some(Ordering::Less),
-                Ordering::Equal => {
-                    let other_len = other.max_length();
-                    match self.max_length().partial_cmp(&other_len) {
-                        None => Some(Ordering::Equal),
-                        Some(ord) => Some(ord),
-                    }
-                }
-                Ordering::Greater => Some(Ordering::Greater),
-            },
-        }
+        Some(self.cmp(other))
     }
 }
 

grovedb/src/batch/mod.rs

@@ -168,19 +168,19 @@ pub enum Op {
 
 impl PartialOrd for Op {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        match (self, other) {
-            (Op::Delete, Op::Insert { .. }) => Some(Ordering::Less),
-            (Op::Delete, Op::Replace { .. }) => Some(Ordering::Less),
-            (Op::Insert { .. }, Op::Delete) => Some(Ordering::Greater),
-            (Op::Replace { .. }, Op::Delete) => Some(Ordering::Greater),
-            _ => Some(Ordering::Equal),
-        }
+        Some(self.cmp(other))
     }
 }
 
 impl Ord for Op {
     fn cmp(&self, other: &Self) -> Ordering {
-        self.partial_cmp(other).expect("all ops have order")
+        match (self, other) {
+            (Op::Delete, Op::Insert { .. }) => Ordering::Less,
+            (Op::Delete, Op::Replace { .. }) => Ordering::Less,
+            (Op::Insert { .. }, Op::Delete) => Ordering::Greater,
+            (Op::Replace { .. }, Op::Delete) => Ordering::Greater,
+            _ => Ordering::Equal,
+        }
     }
 }
 

grovedb/src/operations/proof/generate.rs

@@ -47,8 +47,8 @@
 use crate::{
     element::helpers::raw_decode,
     operations::proof::util::{
-        reduce_limit_and_offset_by, write_slice_of_slice_to_slice, write_slice_to_vec,
-        write_to_vec, ProofTokenType, EMPTY_TREE_HASH,
+        increase_limit_and_offset_by, reduce_limit_and_offset_by, write_slice_of_slice_to_slice,
+        write_slice_to_vec, write_to_vec, ProofTokenType, EMPTY_TREE_HASH,
     },
     reference_path::path_from_reference_path_type,
     versioning::{prepend_version_to_bytes, PROOF_VERSION},
@@ -169,16 +169,16 @@
 
     /// Perform a pre-order traversal of the tree based on the provided
     /// subqueries
     fn prove_subqueries(
         &self,
         proofs: &mut Vec<u8>,
         path: Vec<&[u8]>,
         query: &PathQuery,
         current_limit: &mut Option<u16>,
         current_offset: &mut Option<u16>,
         is_first_call: bool,
         is_verbose: bool,
     ) -> CostResult<(), Error> {
         let mut cost = OperationCost::default();
         let mut to_add_to_result_set: u16 = 0;
 
@@ -216,6 +216,9 @@
 
         let mut is_leaf_tree = true;
 
+        let mut offset_inc = 0;
+        let mut limit_inc = 0;
+
         let mut kv_iterator = KVIterator::new(subtree.storage.raw_iter(), &query.query.query)
             .unwrap_add_cost(&mut cost);
 
@@ -231,7 +234,13 @@
                     if subquery_value.is_none() && subquery_path.is_none() {
                         // this element should be added to the result set
                         // hence we have to update the limit and offset value
-                        reduce_limit_and_offset_by(current_limit, current_offset, 1);
+                        let reduced_offset =
+                            reduce_limit_and_offset_by(current_limit, current_offset, 1);
+                        if reduced_offset {
+                            offset_inc += 1;
+                        } else {
+                            limit_inc += 1;
+                        }
                         continue;
                     }
 
@@ -411,6 +420,7 @@
         if is_leaf_tree {
             // if no useful subtree, then we care about the result set of this subtree.
             // apply the sized query
+            increase_limit_and_offset_by(current_limit, current_offset, limit_inc, offset_inc);
             let limit_offset = cost_return_on_error!(
                 &mut cost,
                 self.generate_and_store_merk_proof(
@@ -475,17 +485,17 @@
 
     /// Generates query proof given a subtree and appends the result to a proof
     /// list
     fn generate_and_store_merk_proof<'a, S, B>(
         &self,
         path: &SubtreePath<B>,
         subtree: &'a Merk<S>,
         query: &Query,
         limit_offset: LimitOffset,
         proof_token_type: ProofTokenType,
         proofs: &mut Vec<u8>,
         is_verbose: bool,
         key: &[u8],
     ) -> CostResult<(Option<u16>, Option<u16>), Error>
     where
         S: StorageContext<'a> + 'a,
         B: AsRef<[u8]>,

grovedb/src/operations/proof/util.rs

@@ -333,6 +333,20 @@ pub fn reduce_limit_and_offset_by(
     skip_limit
 }
 
+pub fn increase_limit_and_offset_by(
+    limit: &mut Option<u16>,
+    offset: &mut Option<u16>,
+    limit_inc: u16,
+    offset_inc: u16,
+) {
+    if let Some(offset_value) = *offset {
+        *offset = Some(offset_value + offset_inc);
+    }
+    if let Some(limit_value) = *limit {
+        *limit = Some(limit_value + limit_inc);
+    }
+}
+
 /// Proved path-key-values
 pub type ProvedPathKeyValues = Vec<ProvedPathKeyValue>;