global: snapshot

crates/brk_iterator/src/lib.rs

@@ -1,73 +1,119 @@
-use brk_reader::Reader;
-use brk_structs::{Block, Height};
+use std::sync::Arc;
 
-mod builder;
+use brk_error::Result;
+use brk_reader::Reader;
+use brk_rpc::Client;
+use brk_structs::{BlockHash, Height};
+
+mod iterator;
 mod range;
 mod source;
+mod state;
 
-use builder::*;
+use iterator::*;
 use range::*;
 use source::*;
+use state::*;
 
-/// Block iterator that can use either RPC or Reader
-pub struct BlockIterator {
-    source: Source,
-}
+///
+/// Block iterator factory
+///
+/// Creates iterators over Bitcoin blocks from various sources (RPC/Reader).
+/// Iterators may end earlier than expected if a chain reorganization occurs.
+///
+/// Thread-safe and free to clone.
+///
+#[derive(Clone)]
+pub struct Blocks(Arc<Source>);
 
-impl From<Source> for BlockIterator {
-    fn from(source: Source) -> Self {
-        Self { source }
+impl Blocks {
+    /// Create with smart mode (auto-select source based on range size)
+    pub fn new(client: Client, reader: Reader) -> Self {
+        Self::new_inner(Source::Smart { client, reader })
     }
-}
 
-impl Iterator for BlockIterator {
-    type Item = Block;
+    /// Create with RPC-only mode
+    pub fn new_rpc(client: Client) -> Self {
+        Self::new_inner(Source::Rpc { client })
+    }
 
-    fn next(&mut self) -> Option<Self::Item> {
-        match &mut self.source {
-            Source::Rpc {
-                client,
-                heights,
-                prev_hash,
-            } => {
-                let height = heights.next()?;
+    /// Create with Reader-only mode
+    pub fn new_reader(reader: Reader) -> Self {
+        Self::new_inner(Source::Reader { reader })
+    }
 
-                let Ok(hash) = client.get_block_hash(height) else {
-                    return None;
-                };
+    fn new_inner(source: Source) -> Self {
+        Self(Arc::new(source))
+    }
 
-                let Ok(block) = client.get_block(&hash) else {
-                    return None;
-                };
+    /// Iterate over a specific range (start..=end)
+    pub fn range(&self, start: Height, end: Height) -> Result<BlockIterator> {
+        self.iter(BlockRange::Span { start, end })
+    }
 
-                if prev_hash
-                    .as_ref()
-                    .is_some_and(|prev_hash| block.header.prev_blockhash != prev_hash.into())
-                {
-                    return None;
+    /// Iterate from start (inclusive) to chain tip
+    pub fn start(&self, start: Height) -> Result<BlockIterator> {
+        self.iter(BlockRange::Start { start })
+    }
+
+    /// Iterate from genesis to end (inclusive)
+    pub fn end(&self, end: Height) -> Result<BlockIterator> {
+        self.iter(BlockRange::End { end })
+    }
+
+    /// Iterate over last n blocks
+    pub fn last(&self, n: u32) -> Result<BlockIterator> {
+        self.iter(BlockRange::Last { n })
+    }
+
+    /// Iterate after hash
+    pub fn after(&self, hash: BlockHash) -> Result<BlockIterator> {
+        self.iter(BlockRange::After { hash })
+    }
+
+    fn iter(&self, range: BlockRange) -> Result<BlockIterator> {
+        let (start, end, hash_opt) = self.resolve_range(range)?;
+
+        let count = end.saturating_sub(*start) + 1;
+
+        let state = match &*self.0 {
+            Source::Smart { client, reader } => {
+                if count <= 10 {
+                    State::new_rpc(client.clone(), start, end, hash_opt)
+                } else {
+                    State::new_reader(reader.clone(), start, end, hash_opt)
                 }
-
-                Some(Block::from((height, hash, block)))
             }
-            Source::Reader { receiver } => receiver.recv().ok().map(|b| b.unwrap()),
+            Source::Rpc { client } => State::new_rpc(client.clone(), start, end, hash_opt),
+            Source::Reader { reader, .. } => {
+                State::new_reader(reader.clone(), start, end, hash_opt)
+            }
+        };
+
+        Ok(BlockIterator::new(state))
+    }
+
+    fn resolve_range(&self, range: BlockRange) -> Result<(Height, Height, Option<BlockHash>)> {
+        let client = self.0.client();
+
+        match range {
+            BlockRange::Span { start, end } => Ok((start, end, None)),
+            BlockRange::Start { start } => {
+                let end = client.get_last_height()?;
+                Ok((start, end, None))
+            }
+            BlockRange::End { end } => Ok((Height::ZERO, end, None)),
+            BlockRange::Last { n } => {
+                let end = client.get_last_height()?;
+                let start = Height::new((*end).saturating_sub(n - 1));
+                Ok((start, end, None))
+            }
+            BlockRange::After { hash } => {
+                let block_info = client.get_block_header_info(&hash)?;
+                let start = (block_info.height + 1).into();
+                let end = client.get_last_height()?;
+                Ok((start, end, Some(hash)))
+            }
         }
     }
 }
-
-impl BlockIterator {
-    pub fn range(start: Height, end: Height) -> BlockIteratorBuilder {
-        BlockIteratorBuilder::from(BlockRange::Span { start, end })
-    }
-
-    pub fn start(start: Height) -> BlockIteratorBuilder {
-        BlockIteratorBuilder::from(BlockRange::Start { start })
-    }
-
-    pub fn end(end: Height) -> BlockIteratorBuilder {
-        BlockIteratorBuilder::from(BlockRange::End { end })
-    }
-
-    pub fn last(n: u32) -> BlockIteratorBuilder {
-        BlockIteratorBuilder::from(BlockRange::Last { n })
-    }
-}