diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs index 1deefb1bc..6b446d990 100644 --- a/zebra-state/src/service.rs +++ b/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 { + 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 { 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> { + 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> { + 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; + + fn next(&mut self) -> Option { + 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) { + 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 for StateService { diff --git a/zebra-state/src/service/memory_state/non_finalized_state.rs b/zebra-state/src/service/memory_state/non_finalized_state.rs index 2384bf070..f4bcbe1e5 100644 --- a/zebra-state/src/service/memory_state/non_finalized_state.rs +++ b/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> { + 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> { 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 { + 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> { let best_chain = self.best_chain()?;