fix(rpc): Fix bugs and performance of `getnetworksolps` & `getnetworkhashps` RPCs (#7647)

* Handle negative and zero getnetworksolsps arguments correctly

* Simplify the solsps loop structure

* Simplify loop by avoiding manual iteration and peeking

* Avoid division by zero

* Use min and max times rather than first and last times

* Refactor block iterators so they are more efficient

* Make finding chains easier

* Simplify block iteration code

* Remove implemented TODO comments

* Simplify internal iterator state

* Implement iteration by any chain item

* Iterate block headers rather than full blocks

* Ignore code that is (sometimes) dead

* Fix a dead code warning

* Add a no blocks in state error constant

* Check result values in the RPC test

* Fix invalid calculation handling

zebra-chain/src/block/height.rs

@@ -80,7 +80,6 @@ impl Height {
     /// # Panics
     ///
     /// - If the current height is at its maximum.
-    // TODO Return an error instead of panicking #7263.
     pub fn next(self) -> Result<Self, HeightError> {
         (self + 1).ok_or(HeightError::Overflow)
     }
@@ -90,7 +89,6 @@ impl Height {
     /// # Panics
     ///
     /// - If the current height is at its minimum.
-    // TODO Return an error instead of panicking #7263.
     pub fn previous(self) -> Result<Self, HeightError> {
         (self - 1).ok_or(HeightError::Underflow)
     }
@@ -99,6 +97,11 @@ impl Height {
     pub fn is_min(self) -> bool {
         self == Self::MIN
     }
+
+    /// Returns the value as a `usize`.
+    pub fn as_usize(self) -> usize {
+        self.0.try_into().expect("fits in usize")
+    }
 }
 
 /// A difference between two [`Height`]s, possibly negative.
@@ -169,6 +172,14 @@ impl TryIntoHeight for String {
     }
 }
 
+impl TryIntoHeight for i32 {
+    type Error = BoxError;
+
+    fn try_into_height(&self) -> Result<Height, Self::Error> {
+        u32::try_from(*self)?.try_into().map_err(Into::into)
+    }
+}
+
 // We don't implement Add<u32> or Sub<u32>, because they cause type inference issues for integer constants.
 
 impl Sub<Height> for Height {

zebra-chain/src/work/difficulty.rs

@@ -129,6 +129,11 @@ pub struct ExpandedDifficulty(U256);
 pub struct Work(u128);
 
 impl Work {
+    /// Returns a value representing no work.
+    pub fn zero() -> Self {
+        Self(0)
+    }
+
     /// Return the inner `u128` value.
     pub fn as_u128(self) -> u128 {
         self.0
@@ -660,6 +665,11 @@ impl std::ops::Add for Work {
 pub struct PartialCumulativeWork(u128);
 
 impl PartialCumulativeWork {
+    /// Returns a value representing no work.
+    pub fn zero() -> Self {
+        Self(0)
+    }
+
     /// Return the inner `u128` value.
     pub fn as_u128(self) -> u128 {
         self.0

zebra-rpc/src/constants.rs

@@ -1,6 +1,6 @@
 //! Constants for RPC methods and server responses.
 
-use jsonrpc_core::ErrorCode;
+use jsonrpc_core::{Error, ErrorCode};
 
 /// The RPC error code used by `zcashd` for incorrect RPC parameters.
 ///
@@ -17,5 +17,24 @@ pub const INVALID_PARAMETERS_ERROR_CODE: ErrorCode = ErrorCode::ServerError(-1);
 /// <https://github.com/zcash/lightwalletd/blob/v0.4.16/common/common.go#L287-L290>
 pub const MISSING_BLOCK_ERROR_CODE: ErrorCode = ErrorCode::ServerError(-8);
 
+/// The RPC error code used by `zcashd` when there are no blocks in the state.
+///
+/// `lightwalletd` expects error code `0` when there are no blocks in the state.
+//
+// TODO: find the source code that expects or generates this error
+pub const NO_BLOCKS_IN_STATE_ERROR_CODE: ErrorCode = ErrorCode::ServerError(0);
+
+/// The RPC error used by `zcashd` when there are no blocks in the state.
+//
+// TODO: find the source code that expects or generates this error text, if there is any
+//       replace literal Error { ... } with this error
+pub fn no_blocks_in_state_error() -> Error {
+    Error {
+        code: NO_BLOCKS_IN_STATE_ERROR_CODE,
+        message: "No blocks in state".to_string(),
+        data: None,
+    }
+}
+
 /// When logging parameter data, only log this much data.
 pub const MAX_PARAMS_LOG_LENGTH: usize = 100;

zebra-rpc/src/methods/get_block_template_rpcs.rs

@@ -11,10 +11,10 @@ use zcash_address::{self, unified::Encoding, TryFromAddress};
 
 use zebra_chain::{
     amount::Amount,
-    block::{self, Block, Height},
+    block::{self, Block, Height, TryIntoHeight},
     chain_sync_status::ChainSyncStatus,
     chain_tip::ChainTip,
-    parameters::Network,
+    parameters::{Network, POW_AVERAGING_WINDOW},
     primitives,
     serialization::ZcashDeserializeInto,
     transparent::{
@@ -142,27 +142,32 @@ pub trait GetBlockTemplateRpc {
     #[rpc(name = "getmininginfo")]
     fn get_mining_info(&self) -> BoxFuture<Result<get_mining_info::Response>>;
 
-    /// Returns the estimated network solutions per second based on the last `num_blocks` before `height`.
+    /// Returns the estimated network solutions per second based on the last `num_blocks` before
-    /// If `num_blocks` is not supplied, uses 120 blocks.
+    /// `height`.
-    /// If `height` is not supplied or is 0, uses the tip height.
+    ///
+    /// If `num_blocks` is not supplied, uses 120 blocks. If it is 0 or -1, uses the difficulty
+    /// averaging window.
+    /// If `height` is not supplied or is -1, uses the tip height.
     ///
     /// zcashd reference: [`getnetworksolps`](https://zcash.github.io/rpc/getnetworksolps.html)
     #[rpc(name = "getnetworksolps")]
     fn get_network_sol_ps(
         &self,
-        num_blocks: Option<usize>,
+        num_blocks: Option<i32>,
         height: Option<i32>,
     ) -> BoxFuture<Result<u64>>;
 
-    /// Returns the estimated network solutions per second based on the last `num_blocks` before `height`.
+    /// Returns the estimated network solutions per second based on the last `num_blocks` before
-    /// If `num_blocks` is not supplied, uses 120 blocks.
+    /// `height`.
-    /// If `height` is not supplied or is 0, uses the tip height.
+    ///
+    /// This method name is deprecated, use [`getnetworksolps`](Self::get_network_sol_ps) instead.
+    /// See that method for details.
     ///
     /// zcashd reference: [`getnetworkhashps`](https://zcash.github.io/rpc/getnetworkhashps.html)
     #[rpc(name = "getnetworkhashps")]
     fn get_network_hash_ps(
         &self,
-        num_blocks: Option<usize>,
+        num_blocks: Option<i32>,
         height: Option<i32>,
     ) -> BoxFuture<Result<u64>> {
         self.get_network_sol_ps(num_blocks, height)
@@ -838,13 +843,22 @@ where
 
     fn get_network_sol_ps(
         &self,
-        num_blocks: Option<usize>,
+        num_blocks: Option<i32>,
         height: Option<i32>,
     ) -> BoxFuture<Result<u64>> {
-        let num_blocks = num_blocks
+        // Default number of blocks is 120 if not supplied.
-            .map(|num_blocks| num_blocks.max(1))
+        let mut num_blocks = num_blocks.unwrap_or(DEFAULT_SOLUTION_RATE_WINDOW_SIZE);
-            .unwrap_or(DEFAULT_SOLUTION_RATE_WINDOW_SIZE);
+        // But if it is 0 or negative, it uses the proof of work averaging window.
-        let height = height.and_then(|height| (height > 1).then_some(Height(height as u32)));
+        if num_blocks < 1 {
+            num_blocks = i32::try_from(POW_AVERAGING_WINDOW).expect("fits in i32");
+        }
+        let num_blocks =
+            usize::try_from(num_blocks).expect("just checked for negatives, i32 fits in usize");
+
+        // Default height is the tip height if not supplied. Negative values also mean the tip
+        // height. Since negative values aren't valid heights, we can just use the conversion.
+        let height = height.and_then(|height| height.try_into_height().ok());
+
         let mut state = self.state.clone();
 
         async move {
@@ -861,11 +875,9 @@ where
                 })?;
 
             let solution_rate = match response {
-                ReadResponse::SolutionRate(solution_rate) => solution_rate.ok_or(Error {
+                // zcashd returns a 0 rate when the calculation is invalid
-                    code: ErrorCode::ServerError(0),
+                ReadResponse::SolutionRate(solution_rate) => solution_rate.unwrap_or(0),
-                    message: "No blocks in state".to_string(),
+
-                    data: None,
-                })?,
                 _ => unreachable!("unmatched response to a solution rate request"),
             };
 

zebra-rpc/src/methods/get_block_template_rpcs/constants.rs

@@ -54,7 +54,7 @@ pub const NOT_SYNCED_ERROR_CODE: ErrorCode = ErrorCode::ServerError(-10);
 /// The default window size specifying how many blocks to check when estimating the chain's solution rate.
 ///
 /// Based on default value in zcashd.
-pub const DEFAULT_SOLUTION_RATE_WINDOW_SIZE: usize = 120;
+pub const DEFAULT_SOLUTION_RATE_WINDOW_SIZE: i32 = 120;
 
 /// The funding stream order in `zcashd` RPC responses.
 ///

zebra-rpc/src/methods/tests/vectors.rs

@@ -1148,32 +1148,39 @@ async fn rpc_getnetworksolps() {
     );
 
     let get_network_sol_ps_inputs = [
-        (None, None),
+        // num_blocks, height, return value
-        (Some(0), None),
+        (None, None, Ok(2)),
-        (Some(0), Some(0)),
+        (Some(-4), None, Ok(2)),
-        (Some(0), Some(-1)),
+        (Some(-3), Some(0), Ok(0)),
-        (Some(0), Some(10)),
+        (Some(-2), Some(-4), Ok(2)),
-        (Some(0), Some(i32::MAX)),
+        (Some(-1), Some(10), Ok(2)),
-        (Some(1), None),
+        (Some(-1), Some(i32::MAX), Ok(2)),
-        (Some(1), Some(0)),
+        (Some(0), None, Ok(2)),
-        (Some(1), Some(-1)),
+        (Some(0), Some(0), Ok(0)),
-        (Some(1), Some(10)),
+        (Some(0), Some(-3), Ok(2)),
-        (Some(1), Some(i32::MAX)),
+        (Some(0), Some(10), Ok(2)),
-        (Some(usize::MAX), None),
+        (Some(0), Some(i32::MAX), Ok(2)),
-        (Some(usize::MAX), Some(0)),
+        (Some(1), None, Ok(4096)),
-        (Some(usize::MAX), Some(-1)),
+        (Some(1), Some(0), Ok(0)),
-        (Some(usize::MAX), Some(10)),
+        (Some(1), Some(-2), Ok(4096)),
-        (Some(usize::MAX), Some(i32::MAX)),
+        (Some(1), Some(10), Ok(4096)),
+        (Some(1), Some(i32::MAX), Ok(4096)),
+        (Some(i32::MAX), None, Ok(2)),
+        (Some(i32::MAX), Some(0), Ok(0)),
+        (Some(i32::MAX), Some(-1), Ok(2)),
+        (Some(i32::MAX), Some(10), Ok(2)),
+        (Some(i32::MAX), Some(i32::MAX), Ok(2)),
     ];
 
-    for (num_blocks_input, height_input) in get_network_sol_ps_inputs {
+    for (num_blocks_input, height_input, return_value) in get_network_sol_ps_inputs {
         let get_network_sol_ps_result = get_block_template_rpc
             .get_network_sol_ps(num_blocks_input, height_input)
             .await;
-        assert!(
+        assert_eq!(
-            get_network_sol_ps_result
+            get_network_sol_ps_result, return_value,
-                .is_ok(),
+            "get_network_sol_ps({num_blocks_input:?}, {height_input:?}) result\n\
-            "get_network_sol_ps({num_blocks_input:?}, {height_input:?}) call with should be ok, got: {get_network_sol_ps_result:?}"
+             should be {return_value:?},\n\
+             got: {get_network_sol_ps_result:?}"
         );
     }
 }

zebra-state/src/request.rs

@@ -938,9 +938,10 @@ pub enum ReadRequest {
     ///
     /// Returns [`ReadResponse::SolutionRate`]
     SolutionRate {
-        /// Specifies over difficulty averaging window.
+        /// The number of blocks to calculate the average difficulty for.
         num_blocks: usize,
-        /// Optionally estimate the network speed at the time when a certain block was found
+        /// Optionally estimate the network solution rate at the time when this height was mined.
+        /// Otherwise, estimate at the current tip height.
         height: Option<block::Height>,
     },
 

zebra-state/src/service/block_iter.rs

@@ -1,102 +1,86 @@
 //! Iterators for blocks in the non-finalized and finalized state.
 
-use std::sync::Arc;
+use std::{marker::PhantomData, sync::Arc};
 
-use zebra_chain::block::{self, Block};
+use zebra_chain::block::{self, Block, Height};
 
-use crate::{service::non_finalized_state::NonFinalizedState, HashOrHeight};
+use crate::{
+    service::{
+        finalized_state::ZebraDb,
+        non_finalized_state::{Chain, NonFinalizedState},
+        read,
+    },
+    HashOrHeight,
+};
 
-use super::finalized_state::ZebraDb;
+/// Generic state chain iterator, which iterates by block height or hash.
-
+/// Can be used for blocks, block headers, or any type indexed by [`HashOrHeight`].
-/// Iterator for state blocks.
 ///
-/// Starts at any block in any non-finalized or finalized chain,
+/// Starts at any hash or height in any non-finalized or finalized chain,
 /// and iterates in reverse height order. (Towards the genesis block.)
-pub(crate) struct Iter<'a> {
+#[derive(Clone, Debug)]
-    pub(super) non_finalized_state: &'a NonFinalizedState,
+pub(crate) struct Iter<Item: ChainItem> {
-    pub(super) db: &'a ZebraDb,
+    /// The non-finalized chain fork we're iterating, if the iterator is in the non-finalized state.
-    pub(super) state: IterState,
+    ///
+    /// This is a cloned copy of a potentially out-of-date chain fork.
+    pub(super) chain: Option<Arc<Chain>>,
+
+    /// The finalized database we're iterating.
+    ///
+    /// This is the shared live database instance, which can concurrently write blocks.
+    pub(super) db: ZebraDb,
+
+    /// The height of the item which will be yielded by `Iterator::next()`.
+    pub(super) height: Option<Height>,
+
+    /// An internal marker type that tells the Rust type system what we're iterating.
+    iterable: PhantomData<Item::Type>,
 }
 
-pub(super) enum IterState {
+impl<Item> Iter<Item>
-    NonFinalized(block::Hash),
+where
-    Finalized(block::Height),
+    Item: ChainItem,
-    Finished,
+{
-}
+    /// Returns an item by height, and updates the iterator's internal state to point to the
+    /// previous height.
+    fn yield_by_height(&mut self) -> Option<Item::Type> {
+        let current_height = self.height?;
 
-impl Iter<'_> {
+        // TODO:
-    fn next_non_finalized_block(&mut self) -> Option<Arc<Block>> {
+        // Check if the root of the chain connects to the finalized state. Cloned chains can become
-        let Iter {
+        // disconnected if they are concurrently pruned by a finalized block from another chain
-            non_finalized_state,
+        // fork. If that happens, the iterator is invalid and should stop returning items.
-            db: _,
+        //
-            state,
+        // Currently, we skip from the disconnected chain root to the previous height in the
-        } = self;
+        // finalized state, which is usually ok, but could cause consensus or light wallet bugs.
+        let item = Item::read(self.chain.as_ref(), &self.db, current_height);
 
-        let hash = match state {
+        // The iterator is finished if the current height is genesis.
-            IterState::NonFinalized(hash) => *hash,
+        self.height = current_height.previous().ok();
-            IterState::Finalized(_) | IterState::Finished => unreachable!(),
-        };
 
-        if let Some(block) = non_finalized_state.any_block_by_hash(hash) {
+        // Drop the chain if we've finished using it.
-            let hash = block.header.previous_block_hash;
+        if let Some(chain) = self.chain.as_ref() {
-            self.state = IterState::NonFinalized(hash);
+            if let Some(height) = self.height {
-            Some(block)
+                if !chain.contains_block_height(height) {
-        } else {
+                    std::mem::drop(self.chain.take());
-            None
+                }
-        }
-    }
-
-    #[allow(clippy::unwrap_in_result)]
-    fn next_finalized_block(&mut self) -> Option<Arc<Block>> {
-        let Iter {
-            non_finalized_state: _,
-            db,
-            state,
-        } = self;
-
-        let hash_or_height: HashOrHeight = match *state {
-            IterState::Finalized(height) => height.into(),
-            IterState::NonFinalized(hash) => hash.into(),
-            IterState::Finished => unreachable!(),
-        };
-
-        if let Some(block) = db.block(hash_or_height) {
-            let height = block
-                .coinbase_height()
-                .expect("valid blocks have a coinbase height");
-
-            if let Some(next_height) = height - 1 {
-                self.state = IterState::Finalized(next_height);
             } else {
-                self.state = IterState::Finished;
+                std::mem::drop(self.chain.take());
             }
-
-            Some(block)
-        } else {
-            self.state = IterState::Finished;
-            None
         }
-    }
 
-    /// Return the height for the block at `hash` in any chain.
+        item
-    fn any_height_by_hash(&self, hash: block::Hash) -> Option<block::Height> {
-        self.non_finalized_state
-            .any_height_by_hash(hash)
-            .or_else(|| self.db.height(hash))
     }
 }
 
-impl Iterator for Iter<'_> {
+impl<Item> Iterator for Iter<Item>
-    type Item = Arc<Block>;
+where
+    Item: ChainItem,
+{
+    type Item = Item::Type;
 
     fn next(&mut self) -> Option<Self::Item> {
-        match self.state {
+        self.yield_by_height()
-            IterState::NonFinalized(_) => self
-                .next_non_finalized_block()
-                .or_else(|| self.next_finalized_block()),
-            IterState::Finalized(_) => self.next_finalized_block(),
-            IterState::Finished => None,
-        }
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -105,34 +89,121 @@ impl Iterator for Iter<'_> {
     }
 }
 
-impl std::iter::FusedIterator for Iter<'_> {}
+impl<Item> ExactSizeIterator for Iter<Item>
-
+where
-impl ExactSizeIterator for Iter<'_> {
+    Item: ChainItem,
+{
     fn len(&self) -> usize {
-        match self.state {
+        // Add one to the height for the genesis block.
-            IterState::NonFinalized(hash) => self
+        //
-                .any_height_by_hash(hash)
+        // TODO:
-                .map(|height| (height.0 + 1) as _)
+        // If the Item can skip heights, or return multiple items per block, we can't calculate
-                .unwrap_or(0),
+        // its length using the block height. For example, subtree end height iterators, or
-            IterState::Finalized(height) => (height.0 + 1) as _,
+        // transaction iterators.
-            IterState::Finished => 0,
+        //
-        }
+        // TODO:
+        // Check if the root of the chain connects to the finalized state. If that happens, the
+        // iterator is invalid and the length should be zero. See the comment in yield_by_height()
+        // for details.
+        self.height.map_or(0, |height| height.as_usize() + 1)
     }
 }
 
-/// Return an iterator over the relevant chain of the block identified by
+// TODO:
-/// `hash`, in order from the largest height to the genesis block.
+// If the Item can return None before it gets to genesis, it is not fused. For example, subtree
-///
+// end height iterators.
-/// The block identified by `hash` is included in the chain of blocks yielded
+impl<Item> std::iter::FusedIterator for Iter<Item> where Item: ChainItem {}
-/// by the iterator. `hash` can come from any chain.
+
-pub(crate) fn any_ancestor_blocks<'a>(
+/// A trait that implements iteration for a specific chain type.
-    non_finalized_state: &'a NonFinalizedState,
+pub(crate) trait ChainItem {
-    db: &'a ZebraDb,
+    type Type;
-    hash: block::Hash,
+
-) -> Iter<'a> {
+    /// Read the `Type` at `height` from the non-finalized `chain` or finalized `db`.
-    Iter {
+    fn read(chain: Option<&Arc<Chain>>, db: &ZebraDb, height: Height) -> Option<Self::Type>;
-        non_finalized_state,
+}
-        db,
+
-        state: IterState::NonFinalized(hash),
+// Block iteration
+
+impl ChainItem for Block {
+    type Type = Arc<Block>;
+
+    fn read(chain: Option<&Arc<Chain>>, db: &ZebraDb, height: Height) -> Option<Self::Type> {
+        read::block(chain, db, height.into())
+    }
+}
+
+// Block header iteration
+
+impl ChainItem for block::Header {
+    type Type = Arc<block::Header>;
+
+    fn read(chain: Option<&Arc<Chain>>, db: &ZebraDb, height: Height) -> Option<Self::Type> {
+        read::block_header(chain, db, height.into())
+    }
+}
+
+/// Returns a block iterator over the relevant chain containing `hash`,
+/// in order from the largest height to genesis.
+///
+/// The block with `hash` is included in the iterator.
+/// `hash` can come from any chain or `db`.
+///
+/// Use [`any_chain_ancestor_iter()`] in new code.
+pub(crate) fn any_ancestor_blocks(
+    non_finalized_state: &NonFinalizedState,
+    db: &ZebraDb,
+    hash: block::Hash,
+) -> Iter<Block> {
+    any_chain_ancestor_iter(non_finalized_state, db, hash)
+}
+
+/// Returns a generic chain item iterator over the relevant chain containing `hash`,
+/// in order from the largest height to genesis.
+///
+/// The item with `hash` is included in the iterator.
+/// `hash` can come from any chain or `db`.
+pub(crate) fn any_chain_ancestor_iter<Item>(
+    non_finalized_state: &NonFinalizedState,
+    db: &ZebraDb,
+    hash: block::Hash,
+) -> Iter<Item>
+where
+    Item: ChainItem,
+{
+    // We need to look up the relevant chain, and the height for the hash.
+    let chain = non_finalized_state.find_chain(|chain| chain.contains_block_hash(hash));
+    let height = read::height_by_hash(chain.as_ref(), db, hash);
+
+    Iter {
+        chain,
+        db: db.clone(),
+        height,
+        iterable: PhantomData,
+    }
+}
+
+/// Returns a generic chain item iterator over a `chain` containing `hash_or_height`,
+/// in order from the largest height to genesis.
+///
+/// The item with `hash_or_height` is included in the iterator.
+/// `hash_or_height` must be in `chain` or `db`.
+#[allow(dead_code)]
+pub(crate) fn known_chain_ancestor_iter<Item>(
+    chain: Option<Arc<Chain>>,
+    db: &ZebraDb,
+    hash_or_height: HashOrHeight,
+) -> Iter<Item>
+where
+    Item: ChainItem,
+{
+    // We need to look up the height for the hash.
+    let height =
+        hash_or_height.height_or_else(|hash| read::height_by_hash(chain.as_ref(), db, hash));
+
+    Iter {
+        chain,
+        db: db.clone(),
+        height,
+        iterable: PhantomData,
     }
 }

zebra-state/src/service/finalized_state/zebra_db/block.rs

@@ -97,6 +97,23 @@ impl ZebraDb {
         self.db.zs_get(&height_by_hash, &hash)
     }
 
+    /// Returns the previous block hash for the given block hash in the finalized state.
+    #[allow(dead_code)]
+    pub fn prev_block_hash_for_hash(&self, hash: block::Hash) -> Option<block::Hash> {
+        let height = self.height(hash)?;
+        let prev_height = height.previous().ok()?;
+
+        self.hash(prev_height)
+    }
+
+    /// Returns the previous block height for the given block hash in the finalized state.
+    #[allow(dead_code)]
+    pub fn prev_block_height_for_hash(&self, hash: block::Hash) -> Option<block::Height> {
+        let height = self.height(hash)?;
+
+        height.previous().ok()
+    }
+
     /// Returns the [`block::Header`] with [`block::Hash`] or
     /// [`Height`], if it exists in the finalized chain.
     //

zebra-state/src/service/non_finalized_state.rs

@@ -436,16 +436,20 @@ impl NonFinalizedState {
             .any(|chain| chain.height_by_hash.contains_key(hash))
     }
 
-    /// Removes and returns the first chain satisfying the given predicate.
+    /// Returns the first chain satisfying the given predicate.
     ///
     /// If multiple chains satisfy the predicate, returns the chain with the highest difficulty.
     /// (Using the tip block hash tie-breaker.)
-    fn find_chain<P>(&mut self, mut predicate: P) -> Option<&Arc<Chain>>
+    pub fn find_chain<P>(&self, mut predicate: P) -> Option<Arc<Chain>>
     where
         P: FnMut(&Chain) -> bool,
     {
         // Reverse the iteration order, to find highest difficulty chains first.
-        self.chain_set.iter().rev().find(|chain| predicate(chain))
+        self.chain_set
+            .iter()
+            .rev()
+            .find(|chain| predicate(chain))
+            .cloned()
     }
 
     /// Returns the [`transparent::Utxo`] pointed to by the given
@@ -460,7 +464,9 @@ impl NonFinalizedState {
     }
 
     /// Returns the `block` with the given hash in any chain.
+    #[allow(dead_code)]
     pub fn any_block_by_hash(&self, hash: block::Hash) -> Option<Arc<Block>> {
+        // This performs efficiently because the number of chains is limited to 10.
         for chain in self.chain_set.iter().rev() {
             if let Some(prepared) = chain
                 .height_by_hash
@@ -474,6 +480,14 @@ impl NonFinalizedState {
         None
     }
 
+    /// Returns the previous block hash for the given block hash in any chain.
+    #[allow(dead_code)]
+    pub fn any_prev_block_hash_for_hash(&self, hash: block::Hash) -> Option<block::Hash> {
+        // This performs efficiently because the blocks are in memory.
+        self.any_block_by_hash(hash)
+            .map(|block| block.header.previous_block_hash)
+    }
+
     /// Returns the hash for a given `block::Height` if it is present in the best chain.
     #[allow(dead_code)]
     pub fn best_hash(&self, height: block::Height) -> Option<block::Hash> {
@@ -510,6 +524,7 @@ impl NonFinalizedState {
     }
 
     /// Returns the height of `hash` in any chain.
+    #[allow(dead_code)]
     pub fn any_height_by_hash(&self, hash: block::Hash) -> Option<block::Height> {
         for chain in self.chain_set.iter().rev() {
             if let Some(height) = chain.height_by_hash.get(&hash) {
@@ -568,10 +583,7 @@ impl NonFinalizedState {
     /// The chain can be an existing chain in the non-finalized state, or a freshly
     /// created fork.
     #[allow(clippy::unwrap_in_result)]
-    fn parent_chain(
+    fn parent_chain(&self, parent_hash: block::Hash) -> Result<Arc<Chain>, ValidateContextError> {
-        &mut self,
-        parent_hash: block::Hash,
-    ) -> Result<Arc<Chain>, ValidateContextError> {
         match self.find_chain(|chain| chain.non_finalized_tip_hash() == parent_hash) {
             // Clone the existing Arc<Chain> in the non-finalized state
             Some(chain) => Ok(chain.clone()),

zebra-state/src/service/non_finalized_state/chain.rs

@@ -450,8 +450,28 @@ impl Chain {
 
     /// Returns true is the chain contains the given block hash.
     /// Returns false otherwise.
-    pub fn contains_block_hash(&self, hash: &block::Hash) -> bool {
+    pub fn contains_block_hash(&self, hash: block::Hash) -> bool {
-        self.height_by_hash.contains_key(hash)
+        self.height_by_hash.contains_key(&hash)
+    }
+
+    /// Returns true is the chain contains the given block height.
+    /// Returns false otherwise.
+    pub fn contains_block_height(&self, height: Height) -> bool {
+        self.blocks.contains_key(&height)
+    }
+
+    /// Returns true is the chain contains the given block hash or height.
+    /// Returns false otherwise.
+    #[allow(dead_code)]
+    pub fn contains_hash_or_height(&self, hash_or_height: impl Into<HashOrHeight>) -> bool {
+        use HashOrHeight::*;
+
+        let hash_or_height = hash_or_height.into();
+
+        match hash_or_height {
+            Hash(hash) => self.contains_block_hash(hash),
+            Height(height) => self.contains_block_height(height),
+        }
     }
 
     /// Returns the non-finalized tip block height and hash.

zebra-state/src/service/read/difficulty.rs

@@ -9,12 +9,13 @@ use zebra_chain::{
     history_tree::HistoryTree,
     parameters::{Network, NetworkUpgrade, POST_BLOSSOM_POW_TARGET_SPACING},
     serialization::{DateTime32, Duration32},
-    work::difficulty::{CompactDifficulty, PartialCumulativeWork},
+    work::difficulty::{CompactDifficulty, PartialCumulativeWork, Work},
 };
 
 use crate::{
     service::{
         any_ancestor_blocks,
+        block_iter::any_chain_ancestor_iter,
         check::{
             difficulty::{
                 BLOCK_MAX_TIME_SINCE_MEDIAN, POW_ADJUSTMENT_BLOCK_SPAN, POW_MEDIAN_BLOCK_SPAN,
@@ -87,42 +88,55 @@ pub fn solution_rate(
     num_blocks: usize,
     start_hash: Hash,
 ) -> Option<u128> {
-    // Take 1 extra block for calculating the number of seconds between when mining on the first block likely started.
+    // Take 1 extra header for calculating the number of seconds between when mining on the first
-    // The work for the last block in this iterator is not added to `total_work`.
+    // block likely started. The work for the extra header is not added to `total_work`.
-    let mut block_iter = any_ancestor_blocks(non_finalized_state, db, start_hash)
+    //
-        .take(num_blocks.checked_add(1).unwrap_or(num_blocks))
+    // Since we can't take more headers than are actually in the chain, this automatically limits
-        .peekable();
+    // `num_blocks` to the chain length, like `zcashd` does.
+    let mut header_iter =
+        any_chain_ancestor_iter::<block::Header>(non_finalized_state, db, start_hash)
+            .take(num_blocks.checked_add(1).unwrap_or(num_blocks))
+            .peekable();
 
-    let get_work = |block: Arc<Block>| {
+    let get_work = |header: &block::Header| {
-        block
+        header
-            .header
             .difficulty_threshold
             .to_work()
             .expect("work has already been validated")
     };
 
-    let block = block_iter.next()?;
+    // If there are no blocks in the range, we can't return a useful result.
-    let last_block_time = block.header.time;
+    let last_header = header_iter.peek()?;
 
-    let mut total_work: PartialCumulativeWork = get_work(block).into();
+    // Initialize the cumulative variables.
+    let mut min_time = last_header.time;
+    let mut max_time = last_header.time;
 
-    loop {
+    let mut last_work = Work::zero();
-        // Return `None` if the iterator doesn't yield a second item.
+    let mut total_work = PartialCumulativeWork::zero();
-        let block = block_iter.next()?;
 
-        if block_iter.peek().is_some() {
+    for header in header_iter {
-            // Add the block's work to `total_work` if it's not the last item in the iterator.
+        min_time = min_time.min(header.time);
-            // The last item in the iterator is only used to estimate when mining on the first block
+        max_time = max_time.max(header.time);
-            // in the window of `num_blocks` likely started.
+
-            total_work += get_work(block);
+        last_work = get_work(&header);
-        } else {
+        total_work += last_work;
-            let first_block_time = block.header.time;
-            let duration_between_first_and_last_block = last_block_time - first_block_time;
-            return Some(
-                total_work.as_u128() / duration_between_first_and_last_block.num_seconds() as u128,
-            );
-        }
     }
+
+    // We added an extra header so we could estimate when mining on the first block
+    // in the window of `num_blocks` likely started. But we don't want to add the work
+    // for that header.
+    total_work -= last_work;
+
+    let work_duration = (max_time - min_time).num_seconds();
+
+    // Avoid division by zero errors and negative average work.
+    // This also handles the case where there's only one block in the range.
+    if work_duration <= 0 {
+        return None;
+    }
+
+    Some(total_work.as_u128() / work_duration as u128)
 }
 
 /// Do a consistency check by checking the finalized tip before and after all other database

zebra-state/src/service/read/find.rs

@@ -108,7 +108,7 @@ pub fn non_finalized_state_contains_block_hash(
     hash: block::Hash,
 ) -> Option<KnownBlock> {
     let mut chains_iter = non_finalized_state.chain_iter();
-    let is_hash_in_chain = |chain: &Arc<Chain>| chain.contains_block_hash(&hash);
+    let is_hash_in_chain = |chain: &Arc<Chain>| chain.contains_block_hash(hash);
 
     // Equivalent to `chain_set.iter().next_back()` in `NonFinalizedState.best_chain()` method.
     let best_chain = chains_iter.next();