diff --git a/grovedb/src/operations/proof/generate.rs b/grovedb/src/operations/proof/generate.rs index ad589b5f..82852359 100644 --- a/grovedb/src/operations/proof/generate.rs +++ b/grovedb/src/operations/proof/generate.rs @@ -47,7 +47,7 @@ use grovedb_storage::StorageContext; use crate::{ element::helpers::raw_decode, operations::proof::util::{ - increase_limit_by, reduce_limit_and_offset_by, write_slice_of_slice_to_slice, + 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, @@ -215,7 +215,9 @@ impl GroveDb { } let mut is_leaf_tree = true; - let mut leaf_element_count = 0; + + 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); @@ -232,8 +234,13 @@ impl GroveDb { 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); - leaf_element_count += 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; } @@ -413,7 +420,7 @@ impl GroveDb { if is_leaf_tree { // if no useful subtree, then we care about the result set of this subtree. // apply the sized query - increase_limit_by(current_limit, leaf_element_count); + 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( diff --git a/grovedb/src/operations/proof/util.rs b/grovedb/src/operations/proof/util.rs index 6e2640dd..0ea9ad26 100644 --- a/grovedb/src/operations/proof/util.rs +++ b/grovedb/src/operations/proof/util.rs @@ -333,9 +333,17 @@ pub fn reduce_limit_and_offset_by( skip_limit } -pub fn increase_limit_by(limit: &mut Option, n: u16) { +pub fn increase_limit_and_offset_by( + limit: &mut Option, + offset: &mut Option, + 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 + n); + *limit = Some(limit_value + limit_inc); } }