Dk prep refactoring

consensus/Cargo.toml

@@ -40,6 +40,7 @@ rayon.workspace = true
 rocksdb.workspace = true
 secp256k1.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 smallvec.workspace = true
 thiserror.workspace = true
 tokio.workspace = true

consensus/src/model/services/reachability.rs

@@ -11,24 +11,39 @@ use kaspa_hashes::Hash;
 pub trait ReachabilityService {
     /// Checks if `this` block is a chain ancestor of `queried` block (i.e., `this ∈ chain(queried) ∪ {queried}`).
     /// Note that we use the graph theory convention here which defines that a block is also an ancestor of itself.
-    fn is_chain_ancestor_of(&self, this: Hash, queried: Hash) -> bool;
+    fn is_chain_ancestor_of(&self, this: Hash, queried: Hash) -> bool {
+        self.try_is_chain_ancestor_of(this, queried).unwrap()
+    }
+
+    /// Checks if `this` block is a chain ancestor of all blocks in `queried`.
+    /// See [`Self::is_chain_ancestor_of`] as well.
+    fn is_chain_ancestor_of_all(&self, this: Hash, queried: &[Hash]) -> bool {
+        queried.iter().all(|&hash| self.try_is_chain_ancestor_of(this, hash).unwrap())
+    }
+
+    /// Result version of [`Self::is_chain_ancestor_of`] (avoids unwrapping internally)
+    fn try_is_chain_ancestor_of(&self, this: Hash, queried: Hash) -> Result<bool>;
 
     /// Result version of [`Self::is_dag_ancestor_of`] (avoids unwrapping internally)
-    fn is_dag_ancestor_of_result(&self, this: Hash, queried: Hash) -> Result<bool>;
+    fn try_is_dag_ancestor_of(&self, this: Hash, queried: Hash) -> Result<bool>;
 
     /// Returns true if `this` is a DAG ancestor of `queried` (i.e., `queried ∈ future(this) ∪ {this}`).
     /// Note: this method will return true if `this == queried`.
     /// The complexity of this method is `O(log(|future_covering_set(this)|))`
-    fn is_dag_ancestor_of(&self, this: Hash, queried: Hash) -> bool;
+    fn is_dag_ancestor_of(&self, this: Hash, queried: Hash) -> bool {
+        self.try_is_dag_ancestor_of(this, queried).unwrap()
+    }
 
     /// Checks if `this` is DAG ancestor of any of the blocks in `queried`. See [`Self::is_dag_ancestor_of`] as well.
     fn is_dag_ancestor_of_any(&self, this: Hash, queried: &mut impl Iterator<Item = Hash>) -> bool;
 
     /// Checks if any of the blocks in `list` is DAG ancestor of `queried`. See [`Self::is_dag_ancestor_of`] as well.
-    fn is_any_dag_ancestor(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> bool;
+    fn is_any_dag_ancestor(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> bool {
+        self.try_is_any_dag_ancestor(list, queried).unwrap()
+    }
 
     /// Result version of [`Self::is_any_dag_ancestor`] (avoids unwrapping internally)
-    fn is_any_dag_ancestor_result(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> Result<bool>;
+    fn try_is_any_dag_ancestor(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> Result<bool>;
 
     /// Finds the tree child of `ancestor` which is also a chain ancestor of `descendant`.
     /// (A "tree child of X" is a block which X is its chain parent)
@@ -42,27 +57,19 @@ pub trait ReachabilityService {
 }
 
 impl<T: ReachabilityStoreReader + ?Sized> ReachabilityService for T {
-    fn is_chain_ancestor_of(&self, this: Hash, queried: Hash) -> bool {
-        inquirer::is_chain_ancestor_of(self, this, queried).unwrap()
+    fn try_is_chain_ancestor_of(&self, this: Hash, queried: Hash) -> Result<bool> {
+        inquirer::is_chain_ancestor_of(self, this, queried)
     }
 
-    fn is_dag_ancestor_of_result(&self, this: Hash, queried: Hash) -> Result<bool> {
+    fn try_is_dag_ancestor_of(&self, this: Hash, queried: Hash) -> Result<bool> {
         inquirer::is_dag_ancestor_of(self, this, queried)
     }
 
-    fn is_dag_ancestor_of(&self, this: Hash, queried: Hash) -> bool {
-        inquirer::is_dag_ancestor_of(self, this, queried).unwrap()
-    }
-
     fn is_dag_ancestor_of_any(&self, this: Hash, queried: &mut impl Iterator<Item = Hash>) -> bool {
         queried.any(|hash| inquirer::is_dag_ancestor_of(self, this, hash).unwrap())
     }
 
-    fn is_any_dag_ancestor(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> bool {
-        list.any(|hash| inquirer::is_dag_ancestor_of(self, hash, queried).unwrap())
-    }
-
-    fn is_any_dag_ancestor_result(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> Result<bool> {
+    fn try_is_any_dag_ancestor(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> Result<bool> {
         for hash in list {
             if inquirer::is_dag_ancestor_of(self, hash, queried)? {
                 return Ok(true);
@@ -97,35 +104,25 @@ impl<T: ReachabilityStoreReader + ?Sized> MTReachabilityService<T> {
 }
 
 impl<T: ReachabilityStoreReader + ?Sized> ReachabilityService for MTReachabilityService<T> {
-    fn is_chain_ancestor_of(&self, this: Hash, queried: Hash) -> bool {
+    fn try_is_chain_ancestor_of(&self, this: Hash, queried: Hash) -> Result<bool> {
         let read_guard = self.store.read();
-        inquirer::is_chain_ancestor_of(read_guard.deref(), this, queried).unwrap()
+        inquirer::is_chain_ancestor_of(read_guard.deref(), this, queried)
     }
 
-    fn is_dag_ancestor_of_result(&self, this: Hash, queried: Hash) -> Result<bool> {
+    fn try_is_dag_ancestor_of(&self, this: Hash, queried: Hash) -> Result<bool> {
         let read_guard = self.store.read();
         inquirer::is_dag_ancestor_of(read_guard.deref(), this, queried)
     }
 
-    fn is_dag_ancestor_of(&self, this: Hash, queried: Hash) -> bool {
-        let read_guard = self.store.read();
-        inquirer::is_dag_ancestor_of(read_guard.deref(), this, queried).unwrap()
-    }
-
-    fn is_any_dag_ancestor(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> bool {
-        let read_guard = self.store.read();
-        list.any(|hash| inquirer::is_dag_ancestor_of(read_guard.deref(), hash, queried).unwrap())
-    }
-
-    fn is_any_dag_ancestor_result(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> Result<bool> {
-        self.store.read().is_any_dag_ancestor_result(list, queried)
-    }
-
     fn is_dag_ancestor_of_any(&self, this: Hash, queried: &mut impl Iterator<Item = Hash>) -> bool {
         let read_guard = self.store.read();
         queried.any(|hash| inquirer::is_dag_ancestor_of(read_guard.deref(), this, hash).unwrap())
     }
 
+    fn try_is_any_dag_ancestor(&self, list: &mut impl Iterator<Item = Hash>, queried: Hash) -> Result<bool> {
+        self.store.read().try_is_any_dag_ancestor(list, queried)
+    }
+
     fn get_next_chain_ancestor(&self, descendant: Hash, ancestor: Hash) -> Hash {
         let read_guard = self.store.read();
         inquirer::get_next_chain_ancestor(read_guard.deref(), descendant, ancestor).unwrap()

consensus/src/model/stores/headers.rs

@@ -1,7 +1,9 @@
+use std::cell::RefCell;
 use std::sync::Arc;
 
+use kaspa_consensus_core::{BlockHashMap, HashMapCustomHasher};
 use kaspa_consensus_core::{BlockHasher, BlockLevel, header::Header};
-use kaspa_database::prelude::{BatchDbWriter, CachedDbAccess};
+use kaspa_database::prelude::{BatchDbWriter, CachedDbAccess, DbKey};
 use kaspa_database::prelude::{CachePolicy, DB};
 use kaspa_database::prelude::{StoreError, StoreResult};
 use kaspa_database::registry::DatabaseStorePrefixes;
@@ -222,3 +224,93 @@ impl HeaderStore for DbHeadersStore {
         Ok(())
     }
 }
+
+#[derive(Clone)]
+pub struct MemoryHeaderStore {
+    header_map: RefCell<BlockHashMap<Arc<kaspa_consensus_core::header::Header>>>,
+}
+
+impl Default for MemoryHeaderStore {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl MemoryHeaderStore {
+    pub fn new() -> Self {
+        Self { header_map: RefCell::new(BlockHashMap::new()) }
+    }
+}
+
+impl HeaderStoreReader for MemoryHeaderStore {
+    fn get_bits(&self, hash: Hash) -> std::result::Result<u32, StoreError> {
+        // Prefer stored header if available
+        let map = self.header_map.borrow();
+        if let Some(h) = map.get(&hash) {
+            return Ok(h.bits);
+        }
+        Err(StoreError::KeyNotFound(DbKey::new(DatabaseStorePrefixes::Headers.as_ref(), hash)))
+    }
+
+    fn get_blue_score(&self, hash: Hash) -> std::result::Result<u64, StoreError> {
+        let map = self.header_map.borrow();
+        if let Some(h) = map.get(&hash) {
+            return Ok(h.blue_score);
+        }
+        Err(StoreError::KeyNotFound(DbKey::new(DatabaseStorePrefixes::Headers.as_ref(), hash)))
+    }
+
+    fn get_compact_header_data(
+        &self,
+        hash: Hash,
+    ) -> std::result::Result<crate::model::stores::headers::CompactHeaderData, StoreError> {
+        let map = self.header_map.borrow();
+        if let Some(h) = map.get(&hash) {
+            let compact: crate::model::stores::headers::CompactHeaderData =
+                crate::model::stores::headers::CompactHeaderData::from(h.as_ref());
+            return Ok(compact);
+        }
+        Err(StoreError::KeyNotFound(DbKey::new(DatabaseStorePrefixes::HeadersCompact.as_ref(), hash)))
+    }
+
+    fn get_daa_score(&self, hash: Hash) -> std::result::Result<u64, StoreError> {
+        let map = self.header_map.borrow();
+        if let Some(h) = map.get(&hash) {
+            return Ok(h.daa_score);
+        }
+        Err(StoreError::KeyNotFound(DbKey::new(DatabaseStorePrefixes::Headers.as_ref(), hash)))
+    }
+
+    fn get_header(&self, hash: Hash) -> std::result::Result<Arc<kaspa_consensus_core::header::Header>, StoreError> {
+        let map = self.header_map.borrow();
+        if let Some(h) = map.get(&hash) {
+            return Ok(Arc::clone(h));
+        }
+        Err(StoreError::KeyNotFound(DbKey::new(DatabaseStorePrefixes::Headers.as_ref(), hash)))
+    }
+
+    fn get_header_with_block_level(
+        &self,
+        hash: Hash,
+    ) -> std::result::Result<crate::model::stores::headers::HeaderWithBlockLevel, StoreError> {
+        let map = self.header_map.borrow();
+        if let Some(h) = map.get(&hash) {
+            return Ok(crate::model::stores::headers::HeaderWithBlockLevel { header: Arc::clone(h), block_level: 0 });
+        }
+        Err(StoreError::KeyNotFound(DbKey::new(DatabaseStorePrefixes::Headers.as_ref(), hash)))
+    }
+
+    fn get_timestamp(&self, hash: Hash) -> std::result::Result<u64, StoreError> {
+        let map = self.header_map.borrow();
+        if let Some(h) = map.get(&hash) {
+            return Ok(h.timestamp);
+        }
+        Err(StoreError::KeyNotFound(DbKey::new(DatabaseStorePrefixes::Headers.as_ref(), hash)))
+    }
+}
+
+impl MemoryHeaderStore {
+    pub fn insert(&self, header: Arc<Header>) {
+        self.header_map.borrow_mut().insert(header.hash, header);
+    }
+}