Fix retrieval of block ranges when computing merkle root

Alenar · Alenar · commit e2c42cfcc3cd · 2024-07-04T16:14:40.000+02:00
It was missing one block range when retrieving data.

With the given block ranges in DB (end is exclusive):
[ (15..30),
  (30..45),
  (45..60) ]

Before the fix it would return when asked block number:
 - `44`: [ (15..30) ] -&gt; WRONG, should also include (30..45)
 - `45`: [ (15..30) ]  -&gt; WRONG, should also include (30..45)
 - `46`: [ (15..30), (30..45) ] -&gt; WRONG, should also include (45..60)
diff --git a/internal/mithril-persistence/src/database/query/block_range_root/delete_block_range_root.rs b/internal/mithril-persistence/src/database/query/block_range_root/delete_block_range_root.rs
@@ -49,18 +49,13 @@ mod tests {
     use mithril_common::crypto_helper::MKTreeNode;
     use mithril_common::entities::BlockRange;
 
-    use crate::database::query::{GetBlockRangeRootQuery, InsertBlockRangeRootQuery};
+    use crate::database::query::block_range_root::test_helper::insert_block_range_roots;
+    use crate::database::query::GetBlockRangeRootQuery;
     use crate::database::test_helper::cardano_tx_db_connection;
-    use crate::sqlite::{ConnectionExtensions, SqliteConnection};
+    use crate::sqlite::ConnectionExtensions;
 
     use super::*;
 
-    fn insert_block_range_roots(connection: &SqliteConnection, records: Vec<BlockRangeRootRecord>) {
-        connection
-            .fetch_first(InsertBlockRangeRootQuery::insert_many(records).unwrap())
-            .unwrap();
-    }
-
     fn block_range_root_dataset() -> Vec<BlockRangeRootRecord> {
         [
             (
diff --git a/internal/mithril-persistence/src/database/query/block_range_root/get_block_range_root.rs b/internal/mithril-persistence/src/database/query/block_range_root/get_block_range_root.rs
@@ -1,6 +1,7 @@
-use mithril_common::entities::BlockNumber;
 use sqlite::Value;
 
+use mithril_common::entities::BlockNumber;
+
 use crate::database::record::BlockRangeRootRecord;
 use crate::sqlite::{Query, SourceAlias, SqLiteEntity, WhereCondition};
 
@@ -16,12 +17,9 @@ impl GetBlockRangeRootQuery {
         }
     }
 
-    pub fn up_to_block_number(up_to_or_equal_end_block_number: BlockNumber) -> Self {
+    pub fn contains_or_below_block_number(block_number: BlockNumber) -> Self {
         Self {
-            condition: WhereCondition::new(
-                "end <= ?*",
-                vec![Value::Integer(up_to_or_equal_end_block_number as i64)],
-            ),
+            condition: WhereCondition::new("start < ?*", vec![Value::Integer(block_number as i64)]),
         }
     }
 }
@@ -40,3 +38,100 @@ impl Query for GetBlockRangeRootQuery {
         format!("select {projection} from block_range_root where {condition} order by start, end")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use mithril_common::crypto_helper::MKTreeNode;
+    use mithril_common::entities::BlockRange;
+
+    use crate::database::query::block_range_root::test_helper::insert_block_range_roots;
+    use crate::database::query::GetBlockRangeRootQuery;
+    use crate::database::test_helper::cardano_tx_db_connection;
+    use crate::sqlite::ConnectionExtensions;
+
+    use super::*;
+
+    fn block_range_root_dataset() -> Vec<BlockRangeRootRecord> {
+        [
+            (
+                BlockRange::from_block_number(15),
+                MKTreeNode::from_hex("AAAA").unwrap(),
+            ),
+            (
+                BlockRange::from_block_number(30),
+                MKTreeNode::from_hex("BBBB").unwrap(),
+            ),
+            (
+                BlockRange::from_block_number(45),
+                MKTreeNode::from_hex("CCCC").unwrap(),
+            ),
+        ]
+        .into_iter()
+        .map(BlockRangeRootRecord::from)
+        .collect()
+    }
+
+    #[test]
+    fn test_get_contains_or_below_block_number_with_empty_db() {
+        let connection = cardano_tx_db_connection().unwrap();
+
+        let cursor: Vec<BlockRangeRootRecord> = connection
+            .fetch_collect(GetBlockRangeRootQuery::contains_or_below_block_number(100))
+            .unwrap();
+        assert_eq!(Vec::<BlockRangeRootRecord>::new(), cursor);
+    }
+
+    #[test]
+    fn test_get_contains_or_below_block_number_higher_than_the_highest_stored_block_range() {
+        let connection = cardano_tx_db_connection().unwrap();
+        let dataset = block_range_root_dataset();
+        insert_block_range_roots(&connection, dataset.clone());
+
+        let cursor: Vec<BlockRangeRootRecord> = connection
+            .fetch_collect(GetBlockRangeRootQuery::contains_or_below_block_number(
+                10_000,
+            ))
+            .unwrap();
+
+        assert_eq!(dataset, cursor);
+    }
+
+    #[test]
+    fn test_get_contains_or_below_block_number_below_end_of_the_third_block_range() {
+        let connection = cardano_tx_db_connection().unwrap();
+        let dataset = block_range_root_dataset();
+        insert_block_range_roots(&connection, dataset.clone());
+
+        let cursor: Vec<BlockRangeRootRecord> = connection
+            .fetch_collect(GetBlockRangeRootQuery::contains_or_below_block_number(44))
+            .unwrap();
+
+        assert_eq!(&dataset[0..2], &cursor);
+    }
+
+    #[test]
+    fn test_get_contains_or_below_block_number_equal_to_end_of_the_third_block_range() {
+        let connection = cardano_tx_db_connection().unwrap();
+        let dataset = block_range_root_dataset();
+        insert_block_range_roots(&connection, dataset.clone());
+
+        let cursor: Vec<BlockRangeRootRecord> = connection
+            .fetch_collect(GetBlockRangeRootQuery::contains_or_below_block_number(45))
+            .unwrap();
+
+        assert_eq!(&dataset[0..2], &cursor);
+    }
+
+    #[test]
+    fn test_get_contains_or_below_block_number_after_end_of_the_third_block_range() {
+        let connection = cardano_tx_db_connection().unwrap();
+        let dataset = block_range_root_dataset();
+        insert_block_range_roots(&connection, dataset.clone());
+
+        let cursor: Vec<BlockRangeRootRecord> = connection
+            .fetch_collect(GetBlockRangeRootQuery::contains_or_below_block_number(46))
+            .unwrap();
+
+        assert_eq!(dataset, cursor);
+    }
+}
diff --git a/internal/mithril-persistence/src/database/query/block_range_root/mod.rs b/internal/mithril-persistence/src/database/query/block_range_root/mod.rs
@@ -7,3 +7,20 @@ pub use delete_block_range_root::*;
 pub use get_block_range_root::*;
 pub use get_interval_without_block_range::*;
 pub use insert_block_range::*;
+
+#[cfg(test)]
+mod test_helper {
+    use crate::database::record::BlockRangeRootRecord;
+    use crate::sqlite::{ConnectionExtensions, SqliteConnection};
+
+    use super::*;
+
+    pub fn insert_block_range_roots(
+        connection: &SqliteConnection,
+        records: Vec<BlockRangeRootRecord>,
+    ) {
+        connection
+            .fetch_first(InsertBlockRangeRootQuery::insert_many(records).unwrap())
+            .unwrap();
+    }
+}
diff --git a/internal/mithril-persistence/src/database/repository/cardano_transaction_repository.rs b/internal/mithril-persistence/src/database/repository/cardano_transaction_repository.rs
@@ -146,16 +146,17 @@ impl CardanoTransactionRepository {
             )
     }
 
-    /// Retrieve all the Block Range Roots in database up to the given end block number included.
+    /// Retrieve all the Block Range Roots in database up to the block range that contains the given
+    /// block number.
     pub async fn retrieve_block_range_roots_up_to(
         &self,
-        up_to_or_equal_end_block_number: BlockNumber,
+        block_number: BlockNumber,
     ) -> StdResult<Box<dyn Iterator<Item = (BlockRange, MKTreeNode)> + '_>> {
         let block_range_roots = self
             .connection_pool
             .connection()?
-            .fetch(GetBlockRangeRootQuery::up_to_block_number(
-                up_to_or_equal_end_block_number,
+            .fetch(GetBlockRangeRootQuery::contains_or_below_block_number(
+                block_number,
             ))?
             .map(|record| -> (BlockRange, MKTreeNode) { record.into() })
             .collect::<Vec<_>>(); // TODO: remove this collect to return the iterator directly
@@ -325,13 +326,9 @@ impl CardanoTransactionRepository {
 impl BlockRangeRootRetriever for CardanoTransactionRepository {
     async fn retrieve_block_range_roots<'a>(
         &'a self,
-        up_to_or_equal_beacon: BlockNumber,
+        up_to_beacon: BlockNumber,
     ) -> StdResult<Box<dyn Iterator<Item = (BlockRange, MKTreeNode)> + 'a>> {
-        let iterator = self
-            .retrieve_block_range_roots_up_to(up_to_or_equal_beacon)
-            .await?;
-
-        Ok(Box::new(iterator))
+        self.retrieve_block_range_roots_up_to(up_to_beacon).await
     }
 }
 
@@ -973,50 +970,14 @@ mod tests {
             .await
             .unwrap();
 
-        // Retrieve with a block far higher than the highest block range - should return all
-        {
-            let retrieved_block_ranges = repository
-                .retrieve_block_range_roots_up_to(1000)
-                .await
-                .unwrap();
-            assert_eq!(
-                block_range_roots,
-                retrieved_block_ranges.collect::<Vec<_>>()
-            );
-        }
-        // Retrieve with a block bellow than the smallest block range - should return none
-        {
-            let retrieved_block_ranges = repository
-                .retrieve_block_range_roots_up_to(2)
-                .await
-                .unwrap();
-            assert_eq!(
-                Vec::<(BlockRange, MKTreeNode)>::new(),
-                retrieved_block_ranges.collect::<Vec<_>>()
-            );
-        }
-        // Right below the end of the second block range - should return first of the three
-        {
-            let retrieved_block_ranges = repository
-                .retrieve_block_range_roots_up_to(44)
-                .await
-                .unwrap();
-            assert_eq!(
-                vec![block_range_roots[0].clone()],
-                retrieved_block_ranges.collect::<Vec<_>>()
-            );
-        }
-        // The given block is matched to the end (included) - should return the two of the three
-        {
-            let retrieved_block_ranges = repository
-                .retrieve_block_range_roots_up_to(45)
-                .await
-                .unwrap();
-            assert_eq!(
-                block_range_roots[0..=1].to_vec(),
-                retrieved_block_ranges.collect::<Vec<_>>()
-            );
-        }
+        let retrieved_block_ranges = repository
+            .retrieve_block_range_roots_up_to(45)
+            .await
+            .unwrap();
+        assert_eq!(
+            block_range_roots[0..2].to_vec(),
+            retrieved_block_ranges.collect::<Vec<_>>()
+        );
     }
 
     #[tokio::test]
diff --git a/mithril-common/src/signable_builder/cardano_transactions.rs b/mithril-common/src/signable_builder/cardano_transactions.rs
@@ -29,7 +29,7 @@ pub trait BlockRangeRootRetriever: Send + Sync {
     /// Returns a Merkle map of the block ranges roots up to a given beacon
     async fn retrieve_block_range_roots<'a>(
         &'a self,
-        up_to_or_equal_beacon: BlockNumber,
+        up_to_beacon: BlockNumber,
     ) -> StdResult<Box<dyn Iterator<Item = (BlockRange, MKTreeNode)> + 'a>>;
 
     /// Returns a Merkle map of the block ranges roots up to a given beacon
