feat(drive): add shielded pool storage, actions, and verification (Medusa part 2)

packages/rs-drive-abci/tests/strategy_tests/verify_state_transitions.rs

@@ -1445,6 +1445,12 @@ pub(crate) fn verify_state_transitions_were_or_were_not_executed(
                     }
                 }
                 StateTransitionAction::BumpAddressInputNoncesAction(_) => {}
+                StateTransitionAction::ShieldAction(_) => {}
+                StateTransitionAction::ShieldedTransferAction(_) => {}
+                StateTransitionAction::UnshieldAction(_) => {}
+                StateTransitionAction::ShieldFromAssetLockAction(_) => {}
+                StateTransitionAction::ShieldedWithdrawalAction(_) => {}
+                StateTransitionAction::PenalizeShieldedPoolAction(_) => {}
             }
         } else {
             // if we don't have an action this means there was a problem in the validation of the state transition

packages/rs-drive/src/drive/credit_pools/mod.rs

@@ -261,7 +261,7 @@ mod tests {
 
             for (i, operation) in batch.into_iter().enumerate() {
                 assert_eq!(
-                    operation.key.expect("key should be present").get_key(),
+                    operation.key.expect("expected key").get_key(),
                     KEY_POOL_STORAGE_FEES
                 );
 

packages/rs-drive/src/drive/credit_pools/pending_epoch_refunds/methods/add_delete_pending_epoch_refunds_except_specified/v0/mod.rs

@@ -128,7 +128,7 @@ mod tests {
 
             assert_eq!(operation.path.to_path(), pending_epoch_refunds_path_vec());
 
-            let epoch_index_key = operation.key.expect("key should be present").get_key();
+            let epoch_index_key = operation.key.expect("expected key").get_key();
             let epoch_index = u16::from_be_bytes(
                 epoch_index_key
                     .try_into()

packages/rs-drive/src/drive/initialization/v3/mod.rs

@@ -1,23 +1,140 @@
-//! Drive Initialization v3 — adds shielded pool trees.
+//! Drive Initialization
 
-use crate::drive::Drive;
+use crate::drive::saved_block_transactions::{
+    COMPACTED_NULLIFIERS_KEY_U8, NULLIFIERS_EXPIRATION_TIME_KEY_U8, NULLIFIERS_KEY_U8,
+};
+use crate::drive::shielded::paths::*;
+use crate::drive::{Drive, RootTree};
 use crate::error::Error;
+use crate::util::batch::grovedb_op_batch::GroveDbOpBatchV0Methods;
+use crate::util::batch::GroveDbOpBatch;
 use dpp::version::PlatformVersion;
-use grovedb::TransactionArg;
+use grovedb::{Element, TransactionArg, TreeType};
+use grovedb_path::SubtreePath;
 
 impl Drive {
-    /// Creates the initial state structure (v3).
-    ///
-    /// Extends v2 by initializing shielded pool trees (commitment tree,
-    /// nullifier set, anchor history).
+    /// Creates the initial state structure.
     pub(super) fn create_initial_state_structure_v3(
         &self,
         transaction: TransactionArg,
         platform_version: &PlatformVersion,
     ) -> Result<(), Error> {
-        // v3 currently delegates to v2.
-        // Shielded pool tree initialization will be added here once the
-        // RootTree variants for the shielded pool are defined.
-        self.create_initial_state_structure_v2(transaction, platform_version)
+        let drive_version = &platform_version.drive;
+        self.create_initial_state_structure_top_level_0(transaction, platform_version)?;
+
+        self.grove_insert_empty_tree(
+            SubtreePath::empty(),
+            &[RootTree::GroupActions as u8],
+            TreeType::NormalTree,
+            transaction,
+            None,
+            &mut vec![],
+            drive_version,
+        )?;
+
+        // AddressBalances top-level sum tree (introduced in v2)
+        self.grove_insert_empty_tree(
+            SubtreePath::empty(),
+            &[RootTree::AddressBalances as u8],
+            TreeType::SumTree,
+            transaction,
+            None,
+            &mut vec![],
+            drive_version,
+        )?;
+
+        // SavedBlockTransactions for address-based transaction sync
+        self.grove_insert_empty_tree(
+            SubtreePath::empty(),
+            &[RootTree::SavedBlockTransactions as u8],
+            TreeType::NormalTree,
+            transaction,
+            None,
+            &mut vec![],
+            drive_version,
+        )?;
+
+        // On lower layers we can use batching
+
+        let mut batch =
+            self.create_initial_state_structure_lower_layers_operations_0(platform_version)?;
+
+        self.initial_state_structure_lower_layers_add_operations_2(&mut batch, platform_version)?;
+
+        // Add shielded pool structures
+        self.initial_state_structure_shielded_pool_operations(&mut batch)?;
+
+        self.grove_apply_batch(batch, false, transaction, drive_version)?;
+
+        Ok(())
+    }
+
+    /// Adds shielded pool batch operations for initialization.
+    pub(in crate::drive::initialization) fn initial_state_structure_shielded_pool_operations(
+        &self,
+        batch: &mut GroveDbOpBatch,
+    ) -> Result<(), Error> {
+        // 1. Shielded credit pool SumTree under AddressBalances
+        batch.add_insert(
+            vec![vec![RootTree::AddressBalances as u8]],
+            vec![SHIELDED_CREDIT_POOL_KEY_U8],
+            Element::empty_sum_tree(),
+        );
+
+        // 2. Notes tree (CommitmentTree = CountTree items + Sinsemilla Frontier)
+        batch.add_insert(
+            shielded_credit_pool_path_vec(),
+            vec![SHIELDED_NOTES_KEY],
+            Element::empty_commitment_tree(SHIELDED_NOTES_CHUNK_POWER),
+        );
+
+        // 3. Nullifiers tree (ProvableCountTree)
+        batch.add_insert(
+            shielded_credit_pool_path_vec(),
+            vec![SHIELDED_NULLIFIERS_KEY],
+            Element::empty_provable_count_tree(),
+        );
+
+        // 4. Total balance SumItem(0)
+        batch.add_insert(
+            shielded_credit_pool_path_vec(),
+            vec![SHIELDED_TOTAL_BALANCE_KEY],
+            Element::new_sum_item(0),
+        );
+
+        // 5. Anchors tree (NormalTree) inside pool: block_height_be → anchor_bytes
+        batch.add_insert(
+            shielded_credit_pool_path_vec(),
+            vec![SHIELDED_ANCHORS_IN_POOL_KEY],
+            Element::empty_tree(),
+        );
+
+        // 6. Nullifiers CountSumTree under SavedBlockTransactions for storing
+        // per-block nullifier lists. Each item is an ItemWithSumItem
+        // (serialized Vec<[u8;32]> + nullifier count as sum).
+        batch.add_insert(
+            Self::saved_block_transactions_path(),
+            vec![NULLIFIERS_KEY_U8],
+            Element::empty_count_sum_tree(),
+        );
+
+        // 7. Compacted nullifiers NormalTree under SavedBlockTransactions for storing
+        // compacted/aggregated nullifier lists spanning multiple blocks.
+        // Key: (start_block, end_block) as 16 bytes, Value: serialized Vec<[u8;32]>
+        batch.add_insert(
+            Self::saved_block_transactions_path(),
+            vec![COMPACTED_NULLIFIERS_KEY_U8],
+            Element::empty_tree(),
+        );
+
+        // 8. Nullifiers expiration time NormalTree under SavedBlockTransactions for storing
+        // expiration timestamps for compacted nullifier ranges.
+        batch.add_insert(
+            Self::saved_block_transactions_path(),
+            vec![NULLIFIERS_EXPIRATION_TIME_KEY_U8],
+            Element::empty_tree(),
+        );
+
+        Ok(())
     }
 }

packages/rs-drive/src/drive/mod.rs

@@ -71,6 +71,9 @@ pub mod address_funds;
 /// Saved block transactions module
 #[cfg(feature = "server")]
 pub mod saved_block_transactions;
+/// Shielded pools module
+#[cfg(any(feature = "server", feature = "verify"))]
+pub mod shielded;
 /// Token module
 #[cfg(any(feature = "server", feature = "verify"))]
 pub mod tokens;

packages/rs-drive/src/drive/saved_block_transactions/cleanup_expired_nullifiers/mod.rs

@@ -0,0 +1,49 @@
+mod v0;
+
+use crate::drive::Drive;
+use crate::error::drive::DriveError;
+use crate::error::Error;
+use grovedb::TransactionArg;
+use platform_version::version::PlatformVersion;
+
+impl Drive {
+    /// Cleans up expired compacted nullifier entries.
+    ///
+    /// This function queries the nullifier expiration time tree for entries with
+    /// expiration time <= current_block_time_ms, then deletes:
+    /// 1. The corresponding compacted nullifier entries
+    /// 2. The expiration time entries themselves
+    ///
+    /// # Arguments
+    /// * `current_block_time_ms` - The current block time in milliseconds
+    /// * `transaction` - Optional database transaction
+    /// * `platform_version` - The platform version
+    ///
+    /// # Returns
+    /// * `Ok(usize)` - The number of compacted entries that were cleaned up
+    /// * `Err` - An error occurred
+    pub fn cleanup_expired_nullifiers(
+        &self,
+        current_block_time_ms: u64,
+        transaction: TransactionArg,
+        platform_version: &PlatformVersion,
+    ) -> Result<usize, Error> {
+        match platform_version
+            .drive
+            .methods
+            .saved_block_transactions
+            .cleanup_expired_nullifiers
+        {
+            0 => self.cleanup_expired_nullifiers_v0(
+                current_block_time_ms,
+                transaction,
+                platform_version,
+            ),
+            version => Err(Error::Drive(DriveError::UnknownVersionMismatch {
+                method: "cleanup_expired_nullifiers".to_string(),
+                known_versions: vec![0],
+                received: version,
+            })),
+        }
+    }
+}

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

@@ -0,0 +1,95 @@
+use crate::drive::Drive;
+use crate::error::Error;
+use crate::util::batch::grovedb_op_batch::GroveDbOpBatchV0Methods;
+use crate::util::batch::GroveDbOpBatch;
+use dpp::ProtocolError;
+use grovedb::query_result_type::QueryResultType;
+use grovedb::{PathQuery, Query, QueryItem, SizedQuery, TransactionArg};
+use platform_version::version::PlatformVersion;
+
+impl Drive {
+    /// Version 0 implementation of cleaning up expired compacted nullifier entries.
+    ///
+    /// Queries for all expiration entries with time <= current_block_time_ms,
+    /// then deletes the corresponding compacted entries and the expiration entries.
+    pub(super) fn cleanup_expired_nullifiers_v0(
+        &self,
+        current_block_time_ms: u64,
+        transaction: TransactionArg,
+        platform_version: &PlatformVersion,
+    ) -> Result<usize, Error> {
+        let expiration_path = Self::saved_nullifiers_expiration_time_path_vec();
+
+        // Query all entries with expiration time <= current_block_time_ms
+        let mut query = Query::new();
+        // Range from 0 to current_block_time_ms (inclusive)
+        query.insert_item(QueryItem::RangeToInclusive(
+            ..=current_block_time_ms.to_be_bytes().to_vec(),
+        ));
+
+        let path_query =
+            PathQuery::new(expiration_path.clone(), SizedQuery::new(query, None, None));
+
+        let (results, _) = self.grove_get_path_query(
+            &path_query,
+            transaction,
+            QueryResultType::QueryKeyElementPairResultType,
+            &mut vec![],
+            &platform_version.drive,
+        )?;
+
+        let key_elements = results.to_key_elements();
+
+        if key_elements.is_empty() {
+            return Ok(0);
+        }
+
+        let config = bincode::config::standard()
+            .with_big_endian()
+            .with_no_limit();
+
+        let mut batch = GroveDbOpBatch::new();
+        let mut total_cleaned = 0usize;
+
+        let compacted_path = Self::saved_compacted_block_transactions_nullifiers_path_vec();
+
+        for (expiration_key, element) in key_elements {
+            // Get the vec of block ranges from the element
+            let grovedb::Element::Item(serialized_ranges, _) = element else {
+                return Err(Error::Protocol(Box::new(
+                    ProtocolError::CorruptedSerialization(
+                        "expected item element for expiration block ranges".to_string(),
+                    ),
+                )));
+            };
+
+            // Deserialize the vec of block ranges
+            let (ranges, _): (Vec<(u64, u64)>, usize) =
+                bincode::decode_from_slice(&serialized_ranges, config).map_err(|e| {
+                    Error::Protocol(Box::new(ProtocolError::CorruptedSerialization(format!(
+                        "cannot decode expiration block ranges: {}",
+                        e
+                    ))))
+                })?;
+
+            // Delete each compacted nullifier entry
+            for (start_block, end_block) in &ranges {
+                let mut compacted_key = Vec::with_capacity(16);
+                compacted_key.extend_from_slice(&start_block.to_be_bytes());
+                compacted_key.extend_from_slice(&end_block.to_be_bytes());
+
+                batch.add_delete(compacted_path.clone(), compacted_key);
+                total_cleaned += 1;
+            }
+
+            // Delete the expiration entry itself
+            batch.add_delete(expiration_path.clone(), expiration_key);
+        }
+
+        if !batch.is_empty() {
+            self.grove_apply_batch(batch, false, transaction, &platform_version.drive)?;
+        }
+
+        Ok(total_cleaned)
+    }
+}