feat(archival_state): Add archival block MMR

Add a new field to the archival state, `archival_block_mmr` which is an
archival MMR of all block digests in the canonical chain.

This allows the client to instantly check if a block belongs to the canonical
chain, and thus to faster prune abandoned monitored UTXOs, which is the
motivation for this addition.

In the future, this data structure will allow archival nodes to prove
the canonicity of a block relative to a tip such that light nodes (that
only see the current tip) can verify that any received blocks belong to
the canonical chain.

src/config_models/data_directory.rs

@@ -7,6 +7,7 @@ use directories::ProjectDirs;
 
 use crate::config_models::network::Network;
 use crate::models::database::DATABASE_DIRECTORY_ROOT_NAME;
+use crate::models::state::archival_state::ARCHIVAL_BLOCK_MMR_DIRECTORY_NAME;
 use crate::models::state::archival_state::BLOCK_INDEX_DB_NAME;
 use crate::models::state::archival_state::MUTATOR_SET_DIRECTORY_NAME;
 use crate::models::state::networking_state::BANNED_IPS_DB_NAME;
@@ -137,6 +138,16 @@ impl DataDirectory {
             .join(Path::new(MUTATOR_SET_DIRECTORY_NAME))
     }
 
+    ///////////////////////////////////////////////////////////////////////////
+    ///
+    /// The archival block MMR database director path
+    ///
+    /// This directory lives within `DataDirectory::database_dir_path()`.
+    pub fn archival_block_mmr_dir_path(&self) -> PathBuf {
+        self.database_dir_path()
+            .join(Path::new(ARCHIVAL_BLOCK_MMR_DIRECTORY_NAME))
+    }
+
     ///////////////////////////////////////////////////////////////////////////
     ///
     /// The block body directory.

src/lib.rs

@@ -116,10 +116,14 @@ pub async fn initialize(cli_args: cli_args::Args) -> Result<i32> {
     let archival_mutator_set = ArchivalState::initialize_mutator_set(&data_dir).await?;
     info!("Got archival mutator set");
 
+    let archival_block_mmr = ArchivalState::initialize_archival_block_mmr(&data_dir).await?;
+    info!("Got archival block MMR");
+
     let archival_state = ArchivalState::new(
         data_dir,
         block_index_db,
         archival_mutator_set,
+        archival_block_mmr,
         cli_args.network,
     )
     .await;

src/models/state/archival_state.rs

@@ -18,6 +18,8 @@ use crate::config_models::data_directory::DataDirectory;
 use crate::config_models::network::Network;
 use crate::database::create_db_if_missing;
 use crate::database::storage::storage_schema::traits::*;
+use crate::database::storage::storage_schema::DbtVec;
+use crate::database::storage::storage_schema::SimpleRustyStorage;
 use crate::database::NeptuneLevelDb;
 use crate::database::WriteBatchAsync;
 use crate::models::blockchain::block::block_header::BlockHeader;
@@ -33,13 +35,17 @@ use crate::models::database::FileRecord;
 use crate::models::database::LastFileRecord;
 use crate::prelude::twenty_first;
 use crate::util_types::mutator_set::addition_record::AdditionRecord;
+use crate::util_types::mutator_set::archival_mmr::ArchivalMmr;
 use crate::util_types::mutator_set::mutator_set_accumulator::MutatorSetAccumulator;
 use crate::util_types::mutator_set::removal_record::AbsoluteIndexSet;
 use crate::util_types::mutator_set::removal_record::RemovalRecord;
 use crate::util_types::mutator_set::rusty_archival_mutator_set::RustyArchivalMutatorSet;
 
+type ArchivalBlockMmr = ArchivalMmr<DbtVec<Digest>>;
+
 pub const BLOCK_INDEX_DB_NAME: &str = "block_index";
 pub const MUTATOR_SET_DIRECTORY_NAME: &str = "mutator_set";
+pub(crate) const ARCHIVAL_BLOCK_MMR_DIRECTORY_NAME: &str = "archival_block_mmr";
 
 /// Provides interface to historic blockchain data which consists of
 ///  * block-data stored in individual files (append-only)
@@ -70,6 +76,9 @@ pub struct ArchivalState {
     // The archival mutator set is persisted to one database that also records a sync label,
     // which corresponds to the hash of the block to which the mutator set is synced.
     pub archival_mutator_set: RustyArchivalMutatorSet,
+
+    /// Archival-MMR of the block digests belonging to the canonical chain.
+    pub(crate) archival_block_mmr: ArchivalBlockMmr,
 }
 
 // The only reason we have this `Debug` implementation is that it's required
@@ -102,7 +111,7 @@ impl ArchivalState {
     }
 
     /// Initialize an `ArchivalMutatorSet` by opening or creating its databases.
-    pub async fn initialize_mutator_set(
+    pub(crate) async fn initialize_mutator_set(
         data_dir: &DataDirectory,
     ) -> Result<RustyArchivalMutatorSet> {
         let ms_db_dir_path = data_dir.mutator_set_database_dir_path();
@@ -133,6 +142,46 @@ impl ArchivalState {
         Ok(archival_set)
     }
 
+    pub(crate) async fn initialize_archival_block_mmr(
+        data_dir: &DataDirectory,
+    ) -> Result<ArchivalBlockMmr> {
+        let abmmr_dir_path = data_dir.archival_block_mmr_dir_path();
+        DataDirectory::create_dir_if_not_exists(&abmmr_dir_path).await?;
+
+        let path = abmmr_dir_path.clone();
+        let result = NeptuneLevelDb::new(&path, &create_db_if_missing()).await;
+
+        let db = match result {
+            Ok(db) => db,
+            Err(e) => {
+                tracing::error!(
+                    "Could not open archival MMR database at {}: {e}",
+                    abmmr_dir_path.display()
+                );
+                panic!(
+                    "Could not open database; do not know how to proceed. Panicking.\n\
+                    If you suspect the database may be corrupted, consider renaming the directory {}\
+                     or removing it altogether. Or perhaps a node is already running?",
+                     abmmr_dir_path.display()
+                );
+            }
+        };
+
+        let mut storage = SimpleRustyStorage::new_with_callback(
+            db,
+            "archival-block-mmr-Schema",
+            crate::LOG_TOKIO_LOCK_EVENT_CB,
+        );
+
+        // We do not need a sync-label since the last leaf of the MMR will
+        // be the sync-label, i.e., the block digest of the latest block added.
+        let abmmr = storage.schema.new_vec::<Digest>("archival_block_mmr").await;
+
+        let archival_bmmr = ArchivalBlockMmr::new(abmmr).await;
+
+        Ok(archival_bmmr)
+    }
+
     /// Find the path connecting two blocks. Every path involves
     /// going down some number of steps and then going up some number
     /// of steps. So this function returns two lists: the list of
@@ -196,10 +245,11 @@ impl ArchivalState {
         (leaving, luca, arriving)
     }
 
-    pub async fn new(
+    pub(crate) async fn new(
         data_dir: DataDirectory,
         block_index_db: NeptuneLevelDb<BlockIndexKey, BlockIndexValue>,
         mut archival_mutator_set: RustyArchivalMutatorSet,
+        mut archival_block_mmr: ArchivalBlockMmr,
         network: Network,
     ) -> Self {
         let genesis_block = Box::new(Block::genesis_block(network));
@@ -218,11 +268,17 @@ impl ArchivalState {
             archival_mutator_set.persist().await;
         }
 
+        // Add genesis block digest to archival MMR, if empty.
+        if archival_block_mmr.is_empty().await {
+            archival_block_mmr.append(genesis_block.hash()).await;
+        }
+
         Self {
             data_dir,
             block_index_db,
             genesis_block,
             archival_mutator_set,
+            archival_block_mmr,
         }
     }
 
@@ -393,6 +449,46 @@ impl ArchivalState {
         Ok(())
     }
 
+    /// Add a new block as tip for the archival block MMR.
+    ///
+    /// All predecessors of this block must be known and stored in the block
+    /// index database for this update to work.
+    pub(crate) async fn add_to_archival_block_mmr(&mut self, new_block: &Block) {
+        // Roll back to length of parent (accounting for genesis block),
+        // then add new digest.
+        let num_leafs_prior_to_this_block = new_block.header().height.into();
+        self.archival_block_mmr
+            .prune_to_num_leafs(num_leafs_prior_to_this_block)
+            .await;
+
+        let latest_leaf = self
+            .archival_block_mmr
+            .get_latest_leaf()
+            .await
+            .expect("block MMR must always have at least one leaf");
+        if new_block.header().prev_block_digest != latest_leaf {
+            let (backwards, _, forwards) = self
+                .find_path(latest_leaf, new_block.header().prev_block_digest)
+                .await;
+            for _ in backwards {
+                self.archival_block_mmr.remove_last_leaf_async().await;
+            }
+            for digest in forwards {
+                self.archival_block_mmr.append(digest).await;
+            }
+        }
+
+        assert_eq!(
+            new_block.header().prev_block_digest,
+            self.archival_block_mmr
+                .get_latest_leaf()
+                .await
+                .expect("block MMR must always have at least one leaf"),
+            "Archival block-MMR must be in a consistent state. Try deleting this database to have it rebuilt."
+        );
+        self.archival_block_mmr.append(new_block.hash()).await;
+    }
+
     async fn get_block_from_block_record(&self, block_record: BlockRecord) -> Result<Block> {
         // Get path of file for block
         let block_file_path: PathBuf = self
@@ -1197,7 +1293,11 @@ mod archival_state_tests {
             .await
             .unwrap();
 
-        ArchivalState::new(data_dir, block_index_db, ams, network).await
+        let archival_block_mmr = ArchivalState::initialize_archival_block_mmr(&data_dir)
+            .await
+            .unwrap();
+
+        ArchivalState::new(data_dir, block_index_db, ams, archival_block_mmr, network).await
     }
 
     #[traced_test]

src/models/state/mod.rs

@@ -1338,6 +1338,11 @@ impl GlobalState {
             .write_block_as_tip(&new_block)
             .await?;
 
+        self.chain
+            .archival_state_mut()
+            .add_to_archival_block_mmr(&new_block)
+            .await;
+
         // update the mutator set with the UTXOs from this block
         self.chain
             .archival_state_mut()
@@ -2615,6 +2620,8 @@ mod global_state_tests {
     }
 
     mod state_update_on_reorganizations {
+        use twenty_first::prelude::Mmr;
+
         use super::*;
 
         async fn assert_correct_global_state(
@@ -2654,6 +2661,35 @@ mod global_state_tests {
                 "Expected block must be returned"
             );
 
+            assert_eq!(
+                expected_tip_digest,
+                global_state
+                    .chain
+                    .archival_state()
+                    .archival_block_mmr
+                    .get_latest_leaf()
+                    .await
+                    .unwrap(),
+                "Latest leaf in archival block MMR must match expected block"
+            );
+
+            // Verify that archival-block MMR matches that of block
+            {
+                let mut expected_archival_block_mmr_value =
+                    expected_tip.body().block_mmr_accumulator.clone();
+                expected_archival_block_mmr_value.append(expected_tip_digest);
+                assert_eq!(
+                    expected_archival_block_mmr_value,
+                    global_state
+                        .chain
+                        .archival_state()
+                        .archival_block_mmr
+                        .to_accumulator_async()
+                        .await,
+                    "archival block-MMR must match that in tip after adding tip digest"
+                );
+            }
+
             let tip_height = expected_tip.header().height;
             assert_eq!(
                 expected_num_blocks_at_tip_height,

src/tests/shared.rs

@@ -270,6 +270,7 @@ pub(crate) async fn get_test_genesis_setup(
     ))
 }
 
+/// Set a new block as tip
 pub(crate) async fn add_block_to_archival_state(
     archival_state: &mut ArchivalState,
     new_block: Block,
@@ -278,6 +279,8 @@ pub(crate) async fn add_block_to_archival_state(
 
     archival_state.update_mutator_set(&new_block).await.unwrap();
 
+    archival_state.add_to_archival_block_mmr(&new_block).await;
+
     Ok(())
 }
 
@@ -765,7 +768,18 @@ pub async fn mock_genesis_archival_state(
         .await
         .unwrap();
 
-    let archival_state = ArchivalState::new(data_dir.clone(), block_index_db, ams, network).await;
+    let archival_block_mmr = ArchivalState::initialize_archival_block_mmr(&data_dir)
+        .await
+        .unwrap();
+
+    let archival_state = ArchivalState::new(
+        data_dir.clone(),
+        block_index_db,
+        ams,
+        archival_block_mmr,
+        network,
+    )
+    .await;
 
     (archival_state, peer_db, data_dir)
 }