Add internal iterator API for accessing relevant chain blocks (#1271)

* Add internal iterator API for accessing relevant chain blocks
* get blocks from all chains in non_finalized state
* Impl FusedIterator for service::Iter
* impl ExactSizedIterator for service::Iter
* let size_hint find heights in side chains

Co-authored-by: teor <teor@riseup.net>

zebra-state/src/service.rs

@@ -252,10 +252,122 @@ impl StateService {
         self.mem.hash(height).or_else(|| self.sled.hash(height))
     }
 
+    /// Return the height for the block at `hash` in any chain.
+    pub fn height_by_hash(&self, hash: block::Hash) -> Option<block::Height> {
+        self.mem
+            .height_by_hash(hash)
+            .or_else(|| self.sled.height(hash))
+    }
+
     /// Return the utxo pointed to by `outpoint` if it exists in any chain.
     pub fn utxo(&self, outpoint: &transparent::OutPoint) -> Option<transparent::Output> {
         self.mem.utxo(outpoint).or_else(|| self.sled.utxo(outpoint))
     }
+
+    /// Return an iterator over the relevant chain of the block identified by
+    /// `hash`.
+    ///
+    /// The block identified by `hash` is included in the chain of blocks yielded
+    /// by the iterator.
+    #[allow(dead_code)]
+    pub fn chain(&self, hash: block::Hash) -> Iter<'_> {
+        Iter {
+            service: self,
+            state: IterState::NonFinalized(hash),
+        }
+    }
+}
+
+struct Iter<'a> {
+    service: &'a StateService,
+    state: IterState,
+}
+
+enum IterState {
+    NonFinalized(block::Hash),
+    Finalized(block::Height),
+    Finished,
+}
+
+impl Iter<'_> {
+    fn next_non_finalized_block(&mut self) -> Option<Arc<Block>> {
+        let Iter { service, state } = self;
+
+        let hash = match state {
+            IterState::NonFinalized(hash) => *hash,
+            IterState::Finalized(_) | IterState::Finished => unreachable!(),
+        };
+
+        if let Some(block) = service.mem.block_by_hash(hash) {
+            let hash = block.header.previous_block_hash;
+            self.state = IterState::NonFinalized(hash);
+            Some(block)
+        } else {
+            None
+        }
+    }
+
+    fn next_finalized_block(&mut self) -> Option<Arc<Block>> {
+        let Iter { service, 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) = service.sled.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;
+            }
+
+            Some(block)
+        } else {
+            self.state = IterState::Finished;
+            None
+        }
+    }
+}
+
+impl Iterator for Iter<'_> {
+    type Item = Arc<Block>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.state {
+            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>) {
+        let len = self.len();
+        (len, Some(len))
+    }
+}
+
+impl std::iter::FusedIterator for Iter<'_> {}
+
+impl ExactSizeIterator for Iter<'_> {
+    fn len(&self) -> usize {
+        match self.state {
+            IterState::NonFinalized(hash) => self
+                .service
+                .height_by_hash(hash)
+                .map(|height| (height.0 + 1) as _)
+                .unwrap_or(0),
+            IterState::Finalized(height) => (height.0 + 1) as _,
+            IterState::Finished => 0,
+        }
+    }
 }
 
 impl Service<Request> for StateService {

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

@@ -143,6 +143,21 @@ impl NonFinalizedState {
         None
     }
 
+    /// Returns the `block` with the given hash in the any chain.
+    pub fn block_by_hash(&self, hash: block::Hash) -> Option<Arc<Block>> {
+        for chain in self.chain_set.iter().rev() {
+            if let Some(block) = chain
+                .height_by_hash
+                .get(&hash)
+                .and_then(|height| chain.blocks.get(height))
+            {
+                return Some(block.clone());
+            }
+        }
+
+        None
+    }
+
     /// Returns the `block` at a given height or hash in the best chain.
     pub fn block(&self, hash_or_height: HashOrHeight) -> Option<Arc<Block>> {
         let best_chain = self.best_chain()?;
@@ -173,6 +188,17 @@ impl NonFinalizedState {
         Some(height)
     }
 
+    /// Returns the height of `hash` in any chain.
+    pub fn 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) {
+                return Some(*height);
+            }
+        }
+
+        None
+    }
+
     /// Returns the given transaction if it exists in the best chain.
     pub fn transaction(&self, hash: transaction::Hash) -> Option<Arc<Transaction>> {
         let best_chain = self.best_chain()?;