From f73e5465afc9890cdb0209fd56e05a5ac730014b Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 17 Jul 2024 12:06:54 +0400
Subject: [PATCH] Fix MMR position naming.

---
 .../src/mmr/request_handler.rs                | 25 +++----
 crates/subspace-service/src/mmr/sync.rs       | 65 +++++++++++--------
 2 files changed, 50 insertions(+), 40 deletions(-)

diff --git a/crates/subspace-service/src/mmr/request_handler.rs b/crates/subspace-service/src/mmr/request_handler.rs
index fc7ddb6a51..470382edb4 100644
--- a/crates/subspace-service/src/mmr/request_handler.rs
+++ b/crates/subspace-service/src/mmr/request_handler.rs
@@ -32,7 +32,6 @@ use std::hash::{Hash, Hasher};
 use std::marker::PhantomData;
 use std::sync::Arc;
 use std::time::Duration;
-use subspace_core_primitives::BlockNumber;
 use tracing::{debug, error, trace};
 
 const MAX_NUMBER_OF_SAME_REQUESTS_PER_PEER: usize = 2;
@@ -90,14 +89,14 @@ pub fn generate_protocol_name<Hash: AsRef<[u8]>>(
 #[derive(Eq, PartialEq, Clone)]
 struct SeenRequestsKey {
     peer: PeerId,
-    block_number: BlockNumber,
+    starting_position: u32,
 }
 
 #[allow(clippy::derived_hash_with_manual_eq)]
 impl Hash for SeenRequestsKey {
     fn hash<H: Hasher>(&self, state: &mut H) {
         self.peer.hash(state);
-        self.block_number.hash(state);
+        self.starting_position.hash(state);
     }
 }
 
@@ -105,17 +104,17 @@ impl Hash for SeenRequestsKey {
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, Encode, Decode, Debug)]
 pub struct MmrRequest {
-    /// Starting block for MMR
-    pub starting_block: BlockNumber,
-    /// Max returned data items
-    pub limit: BlockNumber,
+    /// Starting position for MMR node.
+    pub starting_position: u32,
+    /// Max returned nodes.
+    pub limit: u32,
 }
 
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, Encode, Decode, Debug)]
 pub struct MmrResponse {
-    /// MMR-leafs related to block number
-    pub mmr_data: BTreeMap<BlockNumber, Vec<u8>>,
+    /// MMR-nodes related to node position
+    pub mmr_data: BTreeMap<u32, Vec<u8>>,
 }
 
 /// The value of [`StateRequestHandler::seen_requests`].
@@ -213,7 +212,7 @@ where
 
         let key = SeenRequestsKey {
             peer: *peer,
-            block_number: request.starting_block,
+            starting_position: request.starting_position,
         };
 
         let mut reputation_changes = Vec::new();
@@ -244,7 +243,9 @@ where
             Err(())
         } else {
             let mut mmr_data = BTreeMap::new();
-            for block_number in request.starting_block..(request.starting_block + request.limit) {
+            for block_number in
+                request.starting_position..(request.starting_position + request.limit)
+            {
                 let canon_key = get_offchain_key(block_number.into());
                 let storage_value = self
                     .offchain_db
@@ -282,7 +283,7 @@ where
 
 #[derive(Debug, thiserror::Error)]
 enum HandleRequestError {
-    #[error("Invalid request: max MMR items limit exceeded.")]
+    #[error("Invalid request: max MMR nodes limit exceeded.")]
     MaxItemsLimitExceeded,
 
     #[error(transparent)]
diff --git a/crates/subspace-service/src/mmr/sync.rs b/crates/subspace-service/src/mmr/sync.rs
index 7f15527bca..ac40d661cd 100644
--- a/crates/subspace-service/src/mmr/sync.rs
+++ b/crates/subspace-service/src/mmr/sync.rs
@@ -9,10 +9,10 @@ use sc_network_sync::SyncingService;
 use sp_blockchain::HeaderBackend;
 use sp_core::offchain::storage::OffchainDb;
 use sp_core::offchain::{DbExternalities, OffchainStorage, StorageKind};
+use sp_mmr_primitives::utils::NodesUtils;
 use sp_runtime::traits::Block as BlockT;
 use std::sync::Arc;
 use std::time::Duration;
-use subspace_core_primitives::BlockNumber;
 use tokio::time::sleep;
 use tracing::{debug, error, trace};
 
@@ -37,29 +37,29 @@ pub async fn mmr_sync<Block, Client, NR, OS>(
 
     let mut offchain_db = OffchainDb::new(offchain_storage);
 
-    // Look for existing local MMR-entries
-    let mut starting_block = {
-        let mut starting_block: Option<BlockNumber> = None;
-        for block_number in 0..=BlockNumber::MAX {
-            let canon_key = get_offchain_key(block_number.into());
+    // Look for existing local MMR-nodes
+    let mut starting_position = {
+        let mut starting_position: Option<u32> = None;
+        for position in 0..=u32::MAX {
+            let canon_key = get_offchain_key(position.into());
             if offchain_db
                 .local_storage_get(StorageKind::PERSISTENT, &canon_key)
                 .is_none()
             {
-                starting_block = Some(block_number);
+                starting_position = Some(position);
                 break;
             }
         }
 
-        match starting_block {
+        match starting_position {
             None => {
-                error!("Can't get starting MMR block - MMR storage is corrupted.");
+                error!("Can't get starting MMR position - MMR storage is corrupted.");
                 return;
             }
-            Some(last_processed_block) => {
-                debug!("MMR-sync last processed block: {last_processed_block}");
+            Some(last_processed_position) => {
+                debug!("MMR-sync last processed position: {last_processed_position}");
 
-                last_processed_block
+                last_processed_position
             }
         }
     };
@@ -87,14 +87,25 @@ pub async fn mmr_sync<Block, Client, NR, OS>(
 
             // Request MMR until target block reached.
             loop {
-                let target_block_number = {
+                let target_position = {
                     let best_block = sync_service.best_seen_block().await;
 
                     match best_block {
                         Ok(Some(block)) => {
-                            debug!("MMR-sync. Best seen block={block}");
+                            let block_number: u32 = block
+                                .try_into()
+                                .map_err(|_| "Can't convert block number to u32")
+                                .expect("We convert BlockNumber which is defined as u32.");
+                            let nodes = NodesUtils::new(block_number.into());
 
-                            block
+                            let target_position = nodes.size().saturating_sub(1);
+
+                            debug!(
+                                "MMR-sync. Best seen block={}, Node target position={}",
+                                block_number, target_position
+                            );
+
+                            target_position
                         }
                         Ok(None) => {
                             debug!("Can't obtain best sync block for MMR-sync.");
@@ -108,7 +119,7 @@ pub async fn mmr_sync<Block, Client, NR, OS>(
                 };
 
                 let request = MmrRequest {
-                    starting_block,
+                    starting_position,
                     limit: MAX_MMR_ITEMS,
                 };
                 let response =
@@ -124,22 +135,20 @@ pub async fn mmr_sync<Block, Client, NR, OS>(
                             break;
                         }
 
-                        // Save the MMR-items from response to the local storage
-                        'data: for (block_number, data) in response.mmr_data.iter() {
+                        // Save the MMR-nodes from response to the local storage
+                        'data: for (position, data) in response.mmr_data.iter() {
                             // Ensure continuous sync
-                            if *block_number == starting_block {
-                                let canon_key = get_offchain_key((*block_number).into());
+                            if *position == starting_position {
+                                let canon_key = get_offchain_key((*position).into());
                                 offchain_db.local_storage_set(
                                     StorageKind::PERSISTENT,
                                     &canon_key,
                                     data,
                                 );
 
-                                starting_block += 1;
+                                starting_position += 1;
                             } else {
-                                debug!(
-                                    "MMR-sync gap detected={peer_id}, block_number={block_number}",
-                                );
+                                debug!("MMR-sync gap detected={peer_id}, position={position}",);
                                 break 'data; // We don't support gaps in MMR data
                             }
                         }
@@ -151,15 +160,15 @@ pub async fn mmr_sync<Block, Client, NR, OS>(
                     }
                 }
 
-                // Actual MMR-items may exceed this number, however, we will catch up with the rest
+                // Actual MMR-nodes may exceed this number, however, we will catch up with the rest
                 // when we sync the remaining data (consensus and domain chains).
-                if target_block_number <= starting_block.into() {
-                    debug!("Target block number reached: {target_block_number}");
+                if target_position <= starting_position.into() {
+                    debug!("Target position reached: {target_position}");
                     break 'outer;
                 }
             }
         }
-        debug!("No synced peers to handle the MMR-sync. Pausing...",);
+        debug!(%starting_position, "No synced peers to handle the MMR-sync. Pausing...",);
         sleep(SYNC_PAUSE).await;
     }