feat(platform)!: enhanced fetching of compacted address balance changes

packages/rs-drive-proof-verifier/src/proof.rs

@@ -1030,7 +1030,6 @@ impl FromProof<platform::GetRecentCompactedAddressBalanceChangesRequest>
             &proof.grovedb_proof,
             start_block_height,
             limit,
-            false,
             platform_version,
         )
         .map_drive_error(proof, mtd)?;

packages/rs-drive/src/drive/saved_block_transactions/fetch_compacted_address_balances/v0/mod.rs

@@ -16,7 +16,10 @@ pub type CompactedAddressBalanceChanges =
 impl Drive {
     /// Version 0 implementation of fetching compacted address balance changes.
     ///
-    /// Retrieves all compacted address balance change records from `start_block_height` onwards.
+    /// Retrieves all compacted address balance change records where `end_block >= start_block_height`.
+    /// This includes ranges that contain `start_block_height` (e.g., range 400-600 when querying
+    /// from block 505) as well as ranges that start after `start_block_height`.
+    ///
     /// Returns a vector of (start_block, end_block, address_balance_map) tuples.
     pub(super) fn fetch_compacted_address_balance_changes_v0(
         &self,
@@ -27,34 +30,118 @@ impl Drive {
     ) -> Result<CompactedAddressBalanceChanges, Error> {
         let path = Self::saved_compacted_block_transactions_address_balances_path_vec();
 
-        // Create a range query starting from the specified height
-        // Keys are 16 bytes: (start_block, end_block), both big-endian
-        // We query from (start_block_height, 0) onwards
-        let mut start_key = Vec::with_capacity(16);
-        start_key.extend_from_slice(&start_block_height.to_be_bytes());
-        start_key.extend_from_slice(&0u64.to_be_bytes());
+        // Keys are 16 bytes: (start_block, end_block), both big-endian.
+        // We want ranges where end_block >= start_block_height, which includes:
+        // - Ranges that contain start_block_height (e.g., 400-600 contains 505)
+        // - Ranges that start at or after start_block_height
+        //
+        // Strategy:
+        // 1. First query: descending from (start_block_height, u64::MAX) with limit 1
+        //    to find any range where start_block <= start_block_height that might contain it
+        // 2. Second query: ascending from (start_block_height, 0) to get ranges
+        //    that start at or after start_block_height
 
-        let mut query = Query::new();
-        query.insert_range_from(start_key..);
+        let config = bincode::config::standard()
+            .with_big_endian()
+            .with_no_limit();
 
-        let path_query = PathQuery::new(path, SizedQuery::new(query, limit, None));
+        let mut compacted_changes = Vec::new();
+        let limit_usize = limit.map(|l| l as usize);
 
-        let (results, _) = self.grove_get_path_query(
-            &path_query,
+        // Query 1: Find if there's a range containing start_block_height
+        // Query descending from (start_block_height, u64::MAX) with limit 1
+        let mut desc_end_key = Vec::with_capacity(16);
+        desc_end_key.extend_from_slice(&start_block_height.to_be_bytes());
+        desc_end_key.extend_from_slice(&u64::MAX.to_be_bytes());
+
+        let mut desc_query = Query::new_with_direction(false); // descending
+        desc_query.insert_range_to_inclusive(..=desc_end_key);
+
+        let desc_path_query =
+            PathQuery::new(path.clone(), SizedQuery::new(desc_query, Some(1), None));
+
+        let (desc_results, _) = self.grove_get_path_query(
+            &desc_path_query,
             transaction,
             QueryResultType::QueryKeyElementPairResultType,
             &mut vec![],
             &platform_version.drive,
         )?;
 
-        let config = bincode::config::standard()
-            .with_big_endian()
-            .with_no_limit();
+        // Check if we found a range that contains start_block_height
+        if let Some((key, element)) = desc_results.to_key_elements().into_iter().next() {
+            if key.len() != 16 {
+                return Err(Error::Protocol(Box::new(
+                    ProtocolError::CorruptedSerialization(
+                        "invalid compacted block key length, expected 16 bytes".to_string(),
+                    ),
+                )));
+            }
 
-        let mut compacted_changes = Vec::new();
+            let start_block = u64::from_be_bytes(key[0..8].try_into().unwrap());
+            let end_block = u64::from_be_bytes(key[8..16].try_into().unwrap());
+
+            // Only include if end_block >= start_block_height (range contains our block)
+            if end_block >= start_block_height {
+                let Element::Item(serialized_data, _) = element else {
+                    return Err(Error::Protocol(Box::new(
+                        ProtocolError::CorruptedSerialization(
+                            "expected item element for compacted address balances".to_string(),
+                        ),
+                    )));
+                };
+
+                let (address_balances, _): (BTreeMap<PlatformAddress, CreditOperation>, usize) =
+                    bincode::decode_from_slice(&serialized_data, config).map_err(|e| {
+                        Error::Protocol(Box::new(ProtocolError::CorruptedSerialization(format!(
+                            "cannot decode compacted address balances: {}",
+                            e
+                        ))))
+                    })?;
+
+                compacted_changes.push((start_block, end_block, address_balances));
+            }
+        }
+
+        // Check if we've already hit the limit
+        if let Some(l) = limit_usize {
+            if compacted_changes.len() >= l {
+                return Ok(compacted_changes);
+            }
+        }
+
+        // Query 2: Get ranges that start at or after start_block_height (ascending)
+        // Always use (start_block_height, 0) for consistent proof verification
+        // The result may overlap with descending query if descending found a range
+        // starting exactly at start_block_height - we dedupe below
+        let mut asc_start_key = Vec::with_capacity(16);
+        asc_start_key.extend_from_slice(&start_block_height.to_be_bytes());
+        asc_start_key.extend_from_slice(&0u64.to_be_bytes());
+
+        let mut asc_query = Query::new();
+        asc_query.insert_range_from(asc_start_key..);
+
+        let asc_path_query = PathQuery::new(path, SizedQuery::new(asc_query, limit, None));
+
+        let (asc_results, _) = self.grove_get_path_query(
+            &asc_path_query,
+            transaction,
+            QueryResultType::QueryKeyElementPairResultType,
+            &mut vec![],
+            &platform_version.drive,
+        )?;
+
+        // Track the start_block from descending query to avoid duplicates
+        let desc_start_block = compacted_changes.first().map(|(start, _, _)| *start);
+
+        for (key, element) in asc_results.to_key_elements() {
+            // Check if we've reached the limit
+            if let Some(l) = limit_usize {
+                if compacted_changes.len() >= l {
+                    break;
+                }
+            }
 
-        for (key, element) in results.to_key_elements() {
-            // Parse start_block and end_block from key (16 bytes total, both big-endian)
             if key.len() != 16 {
                 return Err(Error::Protocol(Box::new(
                     ProtocolError::CorruptedSerialization(
@@ -66,7 +153,11 @@ impl Drive {
             let start_block = u64::from_be_bytes(key[0..8].try_into().unwrap());
             let end_block = u64::from_be_bytes(key[8..16].try_into().unwrap());
 
-            // Get the serialized data from the Item element
+            // Skip if this is the same range we got from descending query
+            if Some(start_block) == desc_start_block {
+                continue;
+            }
+
             let Element::Item(serialized_data, _) = element else {
                 return Err(Error::Protocol(Box::new(
                     ProtocolError::CorruptedSerialization(
@@ -75,7 +166,6 @@ impl Drive {
                 )));
             };
 
-            // Deserialize the address balance map
             let (address_balances, _): (BTreeMap<PlatformAddress, CreditOperation>, usize) =
                 bincode::decode_from_slice(&serialized_data, config).map_err(|e| {
                     Error::Protocol(Box::new(ProtocolError::CorruptedSerialization(format!(
@@ -91,6 +181,12 @@ impl Drive {
     }
 
     /// Version 0 implementation for proving compacted address balance changes.
+    ///
+    /// Uses a two-step approach:
+    /// 1. First query (non-proving): descending to find any range containing start_block_height
+    /// 2. Second query (proving): ascending from the found start_block or start_block_height
+    ///
+    /// This ensures the proof covers all relevant ranges efficiently.
     pub(super) fn prove_compacted_address_balance_changes_v0(
         &self,
         start_block_height: u64,
@@ -100,10 +196,55 @@ impl Drive {
     ) -> Result<Vec<u8>, Error> {
         let path = Self::saved_compacted_block_transactions_address_balances_path_vec();
 
-        // Create a range query starting from the specified height
-        let mut start_key = Vec::with_capacity(16);
-        start_key.extend_from_slice(&start_block_height.to_be_bytes());
-        start_key.extend_from_slice(&0u64.to_be_bytes());
+        // Step 1: Non-proving descending query to find any range containing start_block_height
+        let mut desc_end_key = Vec::with_capacity(16);
+        desc_end_key.extend_from_slice(&start_block_height.to_be_bytes());
+        desc_end_key.extend_from_slice(&u64::MAX.to_be_bytes());
+
+        let mut desc_query = Query::new_with_direction(false); // descending
+        desc_query.insert_range_to_inclusive(..=desc_end_key);
+
+        let desc_path_query =
+            PathQuery::new(path.clone(), SizedQuery::new(desc_query, Some(1), None));
+
+        let (desc_results, _) = self.grove_get_path_query(
+            &desc_path_query,
+            transaction,
+            QueryResultType::QueryKeyElementPairResultType,
+            &mut vec![],
+            &platform_version.drive,
+        )?;
+
+        // Determine the actual start key for the proved query
+        // If we found a containing range, use its exact key
+        // Otherwise use (start_block_height, start_block_height) since end_block >= start_block always
+        let start_key = if let Some((key, _)) = desc_results.to_key_elements().into_iter().next() {
+            if key.len() == 16 {
+                let end_block = u64::from_be_bytes(key[8..16].try_into().unwrap());
+                // If this range contains start_block_height, use its exact key
+                if end_block >= start_block_height {
+                    key
+                } else {
+                    // No containing range, use (start_block_height, start_block_height)
+                    let mut key = Vec::with_capacity(16);
+                    key.extend_from_slice(&start_block_height.to_be_bytes());
+                    key.extend_from_slice(&start_block_height.to_be_bytes());
+                    key
+                }
+            } else {
+                let mut key = Vec::with_capacity(16);
+                key.extend_from_slice(&start_block_height.to_be_bytes());
+                key.extend_from_slice(&start_block_height.to_be_bytes());
+                key
+            }
+        } else {
+            let mut key = Vec::with_capacity(16);
+            key.extend_from_slice(&start_block_height.to_be_bytes());
+            key.extend_from_slice(&start_block_height.to_be_bytes());
+            key
+        };
+
+        // Step 2: Proved ascending query from start_key
 
         let mut query = Query::new();
         query.insert_range_from(start_key..);

packages/rs-drive/src/drive/saved_block_transactions/store_address_balances/v0/mod.rs

@@ -1090,4 +1090,132 @@ mod tests {
             panic!("expected Item element for block ranges");
         }
     }
+
+    #[test]
+    fn should_fetch_compacted_range_when_start_height_falls_within_range() {
+        let drive = setup_drive_with_initial_state_structure(None);
+        // Set compaction threshold to 4 blocks so we can trigger compaction easily
+        let platform_version = create_test_platform_version_for_compaction(4, 1000);
+
+        // Store 4 blocks (100-103) to trigger compaction into range (100, 103)
+        for block_height in 100u64..=103 {
+            let mut balances = BTreeMap::new();
+            balances.insert(ADDR_1, CreditOperation::AddToCredits(1000));
+
+            drive
+                .store_address_balances_for_block_v0(
+                    &balances,
+                    block_height,
+                    1700000000000,
+                    None,
+                    &platform_version,
+                )
+                .expect("should store block");
+        }
+
+        // Verify compaction happened - fetch from block 0 should return the compacted range
+        let all_results = drive
+            .fetch_compacted_address_balance_changes(0, None, None, &platform_version)
+            .expect("should fetch compacted changes");
+
+        assert_eq!(all_results.len(), 1, "should have 1 compacted entry");
+        let (start, end, _) = &all_results[0];
+        assert_eq!(*start, 100, "compacted range should start at 100");
+        assert_eq!(*end, 103, "compacted range should end at 103");
+
+        // Now query with start_height=101, which falls WITHIN the compacted range (100-103)
+        // We should still get this range since block 101 is within it
+        let results_from_101 = drive
+            .fetch_compacted_address_balance_changes(101, None, None, &platform_version)
+            .expect("should fetch compacted changes from 101");
+
+        assert_eq!(
+            results_from_101.len(),
+            1,
+            "should include compacted range (100, 103) when querying from block 101 \
+             since 101 falls within that range"
+        );
+    }
+
+    #[test]
+    fn should_prove_and_verify_compacted_address_balances() {
+        use crate::drive::Drive;
+
+        let drive = setup_drive_with_initial_state_structure(None);
+        // Set compaction threshold to 4 blocks
+        let platform_version = create_test_platform_version_for_compaction(4, 1000);
+
+        // Store 4 blocks (100-103) to trigger compaction into range (100, 103)
+        for block_height in 100u64..=103 {
+            let mut balances = BTreeMap::new();
+            balances.insert(ADDR_1, CreditOperation::AddToCredits(1000));
+
+            drive
+                .store_address_balances_for_block_v0(
+                    &balances,
+                    block_height,
+                    1700000000000,
+                    None,
+                    &platform_version,
+                )
+                .expect("should store block");
+        }
+
+        // Generate proof with limit (required for verify_query_with_absence_proof)
+        let limit = Some(100u16);
+        let proof = drive
+            .prove_compacted_address_balance_changes(0, limit, None, &platform_version)
+            .expect("should generate proof");
+
+        // Verify proof
+        let (root_hash, verified_changes) =
+            Drive::verify_compacted_address_balance_changes(&proof, 0, limit, &platform_version)
+                .expect("should verify proof");
+
+        assert!(!root_hash.is_empty(), "root hash should not be empty");
+        assert_eq!(verified_changes.len(), 1, "should have 1 compacted entry");
+        assert_eq!(verified_changes[0].0, 100, "start block should be 100");
+        assert_eq!(verified_changes[0].1, 103, "end block should be 103");
+    }
+
+    #[test]
+    fn should_prove_and_verify_with_containing_range() {
+        use crate::drive::Drive;
+
+        let drive = setup_drive_with_initial_state_structure(None);
+        // Set compaction threshold to 4 blocks
+        let platform_version = create_test_platform_version_for_compaction(4, 1000);
+
+        // Store 4 blocks (100-103) to trigger compaction
+        for block_height in 100u64..=103 {
+            let mut balances = BTreeMap::new();
+            balances.insert(ADDR_1, CreditOperation::AddToCredits(1000));
+
+            drive
+                .store_address_balances_for_block_v0(
+                    &balances,
+                    block_height,
+                    1700000000000,
+                    None,
+                    &platform_version,
+                )
+                .expect("should store block");
+        }
+
+        // Generate proof starting from block 101 (which falls within range 100-103)
+        let limit = Some(100u16);
+        let proof = drive
+            .prove_compacted_address_balance_changes(101, limit, None, &platform_version)
+            .expect("should generate proof");
+
+        // Verify proof - should include range (100, 103) since 101 is within it
+        let (root_hash, verified_changes) =
+            Drive::verify_compacted_address_balance_changes(&proof, 101, limit, &platform_version)
+                .expect("should verify proof");
+
+        assert!(!root_hash.is_empty(), "root hash should not be empty");
+        assert_eq!(verified_changes.len(), 1, "should have 1 compacted entry");
+        assert_eq!(verified_changes[0].0, 100, "start block should be 100");
+        assert_eq!(verified_changes[0].1, 103, "end block should be 103");
+    }
 }