allow breaking inside for_blocks closures using ControlFlow

Author: antonilol

src/daemon.rs

@@ -8,6 +8,7 @@ use serde_json::{json, Value};
 
 use std::fs::File;
 use std::io::Read;
+use std::ops::ControlFlow;
 use std::path::Path;
 
 use crate::{
@@ -228,10 +229,10 @@ impl Daemon {
         self.p2p.lock().get_new_headers(chain)
     }
 
-    pub(crate) fn for_blocks<B, F>(&self, blockhashes: B, func: F) -> Result<()>
+    pub(crate) fn for_blocks<B, F, R>(&self, blockhashes: B, func: F) -> Result<ControlFlow<R>>
     where
         B: IntoIterator<Item = BlockHash>,
-        F: FnMut(BlockHash, Block),
+        F: FnMut(BlockHash, Block) -> ControlFlow<R>,
     {
         self.p2p.lock().for_blocks(blockhashes, func)
     }

src/index.rs

@@ -1,6 +1,7 @@
 use anyhow::{Context, Result};
 use bitcoin::consensus::{deserialize, serialize};
 use bitcoin::{Block, BlockHash, OutPoint, Txid};
+use std::ops::ControlFlow;
 
 use crate::{
     chain::{Chain, NewHeader},
@@ -205,6 +206,7 @@ impl Index {
                 index_single_block(blockhash, block, height, &mut batch);
             });
             self.stats.height.set("tip", height as f64);
+            ControlFlow::Continue::<()>(())
         })?;
         let heights: Vec<_> = heights.collect();
         assert!(

src/p2p.rs

@@ -21,6 +21,7 @@ use crossbeam_channel::{bounded, select, Receiver, Sender};
 
 use std::io::{self, ErrorKind, Write};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpStream};
+use std::ops::ControlFlow;
 use std::sync::Arc;
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
@@ -92,21 +93,26 @@ impl Connection {
     /// Request and process the specified blocks (in the specified order).
     /// See https://en.bitcoin.it/wiki/Protocol_documentation#getblocks for details.
     /// Defined as `&mut self` to prevent concurrent invocations (https://github.com/romanz/electrs/pull/526#issuecomment-934685515).
-    pub(crate) fn for_blocks<B, F>(&mut self, blockhashes: B, mut func: F) -> Result<()>
+    pub(crate) fn for_blocks<B, F, R>(
+        &mut self,
+        blockhashes: B,
+        mut func: F,
+    ) -> Result<ControlFlow<R>>
     where
         B: IntoIterator<Item = BlockHash>,
-        F: FnMut(BlockHash, Block),
+        F: FnMut(BlockHash, Block) -> ControlFlow<R>,
     {
         self.blocks_duration.observe_duration("total", || {
             let blockhashes: Vec<BlockHash> = blockhashes.into_iter().collect();
             if blockhashes.is_empty() {
-                return Ok(());
+                return Ok(ControlFlow::Continue(()));
             }
             self.blocks_duration.observe_duration("request", || {
                 debug!("loading {} blocks", blockhashes.len());
                 self.req_send.send(Request::get_blocks(&blockhashes))
             })?;
 
+            let mut ret = ControlFlow::Continue(());
             for hash in blockhashes {
                 let block = self.blocks_duration.observe_duration("response", || {
                     let block = self
@@ -116,10 +122,13 @@ impl Connection {
                     ensure!(block.block_hash() == hash, "got unexpected block");
                     Ok(block)
                 })?;
-                self.blocks_duration
-                    .observe_duration("process", || func(hash, block));
+                if ret.is_continue() {
+                    ret = self
+                        .blocks_duration
+                        .observe_duration("process", || func(hash, block));
+                }
             }
-            Ok(())
+            Ok(ret)
         })
     }
 

src/status.rs

@@ -8,6 +8,7 @@ use serde::ser::{Serialize, Serializer};
 
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::convert::TryFrom;
+use std::ops::ControlFlow;
 
 use crate::{
     cache::Cache,
@@ -303,10 +304,15 @@ impl ScriptHashStatus {
     }
 
     /// Apply func only on the new blocks (fetched from daemon).
-    fn for_new_blocks<B, F>(&self, blockhashes: B, daemon: &Daemon, func: F) -> Result<()>
+    fn for_new_blocks<B, F, R>(
+        &self,
+        blockhashes: B,
+        daemon: &Daemon,
+        func: F,
+    ) -> Result<ControlFlow<R>>
     where
         B: IntoIterator<Item = BlockHash>,
-        F: FnMut(BlockHash, Block),
+        F: FnMut(BlockHash, Block) -> ControlFlow<R>,
     {
         daemon.for_blocks(
             blockhashes
@@ -346,6 +352,7 @@ impl ScriptHashStatus {
                         .outputs = funding_outputs;
                 },
             );
+            ControlFlow::Continue::<()>(())
         })?;
         let spending_blockhashes: HashSet<BlockHash> = outpoints
             .par_iter()
@@ -367,6 +374,7 @@ impl ScriptHashStatus {
                         .spent = spent_outpoints;
                 },
             );
+            ControlFlow::Continue::<()>(())
         })?;
 
         Ok(result

src/tracker.rs

@@ -1,5 +1,6 @@
 use anyhow::{Context, Result};
 use bitcoin::{BlockHash, Transaction, Txid};
+use std::ops::ControlFlow;
 
 use crate::{
     cache::Cache,
@@ -100,18 +101,22 @@ impl Tracker {
     ) -> Result<Option<(BlockHash, Transaction)>> {
         // Note: there are two blocks with coinbase transactions having same txid (see BIP-30)
         let blockhashes = self.index.filter_by_txid(txid);
-        let mut result = None;
-        daemon.for_blocks(blockhashes, |blockhash, block| {
+        let result = daemon.for_blocks(blockhashes, |blockhash, block| {
             for tx in block.txdata {
-                if result.is_some() {
-                    return;
-                }
                 if tx.txid() == txid {
-                    result = Some((blockhash, tx));
-                    return;
+                    return ControlFlow::Break((blockhash, tx));
                 }
             }
+            ControlFlow::Continue(())
         })?;
-        Ok(result)
+        Ok(control_flow_break_value(result))
+    }
+}
+
+/// See unstable ControlFlow::break_value (https://github.com/rust-lang/rust/issues/75744)
+fn control_flow_break_value<B, C>(value: ControlFlow<B, C>) -> Option<B> {
+    match value {
+        ControlFlow::Continue(..) => None,
+        ControlFlow::Break(x) => Some(x),
     }
 }