From 9bcde219094b5a80a77ab9a0a4b95a7540ee7fe0 Mon Sep 17 00:00:00 2001
From: gmart7t2 <49558347+gmart7t2@users.noreply.github.com>
Date: Tue, 19 Sep 2023 17:18:44 -0300
Subject: [PATCH] Allow running `find` on a range of sats (#1992)

---
 src/index.rs           | 54 ++++++++++++++++++++++++++++++++++++++++++
 src/subcommand/find.rs | 21 +++++++++++++---
 tests/find.rs          | 49 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 120 insertions(+), 4 deletions(-)

diff --git a/src/index.rs b/src/index.rs
index 4799630cb1..520b3c6452 100644
--- a/src/index.rs
+++ b/src/index.rs
@@ -8,6 +8,7 @@ use {
     updater::Updater,
   },
   super::*,
+  crate::subcommand::find::FindRangeOutput,
   crate::wallet::Wallet,
   bitcoin::block::Header,
   bitcoincore_rpc::{json::GetBlockHeaderResult, Client},
@@ -811,6 +812,59 @@ impl Index {
     Ok(None)
   }
 
+  pub(crate) fn find_range(
+    &self,
+    range_start: u64,
+    range_end: u64,
+  ) -> Result<Option<Vec<FindRangeOutput>>> {
+    self.require_sat_index("find")?;
+
+    let rtx = self.begin_read()?;
+
+    if rtx.block_count()? < Sat(range_end - 1).height().n() + 1 {
+      return Ok(None);
+    }
+
+    let Some(mut remaining_sats) = range_end.checked_sub(range_start) else {
+      return Err(anyhow!("range end is before range start"));
+    };
+
+    let outpoint_to_sat_ranges = rtx.0.open_table(OUTPOINT_TO_SAT_RANGES)?;
+
+    let mut result = Vec::new();
+    for range in outpoint_to_sat_ranges.range::<&[u8; 36]>(&[0; 36]..)? {
+      let (outpoint_entry, sat_ranges_entry) = range?;
+
+      let mut offset = 0;
+      for sat_range in sat_ranges_entry.value().chunks_exact(11) {
+        let (start, end) = SatRange::load(sat_range.try_into().unwrap());
+
+        if end > range_start && start < range_end {
+          let overlap_start = start.max(range_start);
+          let overlap_end = end.min(range_end);
+
+          result.push(FindRangeOutput {
+            start: overlap_start,
+            size: overlap_end - overlap_start,
+            satpoint: SatPoint {
+              outpoint: Entry::load(*outpoint_entry.value()),
+              offset: offset + overlap_start - start,
+            },
+          });
+
+          remaining_sats -= overlap_end - overlap_start;
+
+          if remaining_sats == 0 {
+            break;
+          }
+        }
+        offset += end - start;
+      }
+    }
+
+    Ok(Some(result))
+  }
+
   fn list_inner(&self, outpoint: OutPointValue) -> Result<Option<Vec<u8>>> {
     Ok(
       self
diff --git a/src/subcommand/find.rs b/src/subcommand/find.rs
index cdc69084b2..b52601a2c6 100644
--- a/src/subcommand/find.rs
+++ b/src/subcommand/find.rs
@@ -4,6 +4,8 @@ use super::*;
 pub(crate) struct Find {
   #[arg(help = "Find output and offset of <SAT>.")]
   sat: Sat,
+  #[clap(help = "Find output and offset of all sats in the range [<SAT>, <END>).")]
+  end: Option<Sat>,
 }
 
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
@@ -11,15 +13,28 @@ pub struct Output {
   pub satpoint: SatPoint,
 }
 
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub struct FindRangeOutput {
+  pub start: u64,
+  pub size: u64,
+  pub satpoint: SatPoint,
+}
+
 impl Find {
   pub(crate) fn run(self, options: Options) -> SubcommandResult {
     let index = Index::open(&options)?;
 
     index.update()?;
 
-    match index.find(self.sat.0)? {
-      Some(satpoint) => Ok(Box::new(Output { satpoint })),
-      None => Err(anyhow!("sat has not been mined as of index height")),
+    match self.end {
+      Some(end) => match index.find_range(self.sat.0, end.0)? {
+        Some(result) => Ok(Box::new(result)),
+        None => Err(anyhow!("range has not been mined as of index height")),
+      },
+      None => match index.find(self.sat.0)? {
+        Some(satpoint) => Ok(Box::new(Output { satpoint })),
+        None => Err(anyhow!("sat has not been mined as of index height")),
+      },
     }
   }
 }
diff --git a/tests/find.rs b/tests/find.rs
index 63d55497f0..70b51f114e 100644
--- a/tests/find.rs
+++ b/tests/find.rs
@@ -1,4 +1,7 @@
-use {super::*, ord::subcommand::find::Output};
+use {
+  super::*,
+  ord::subcommand::find::{FindRangeOutput, Output},
+};
 
 #[test]
 fn find_command_returns_satpoint_for_sat() {
@@ -15,6 +18,50 @@ fn find_command_returns_satpoint_for_sat() {
   );
 }
 
+#[test]
+fn find_range_command_returns_satpoints_and_ranges() {
+  let rpc_server = test_bitcoincore_rpc::spawn();
+
+  rpc_server.mine_blocks(1);
+
+  pretty_assert_eq!(
+    CommandBuilder::new(format!("--index-sats find 0 {}", 55 * COIN_VALUE))
+      .rpc_server(&rpc_server)
+      .run_and_deserialize_output::<Vec<FindRangeOutput>>(),
+    vec![
+      FindRangeOutput {
+        start: 0,
+        size: 50 * COIN_VALUE,
+        satpoint: "4a5e1e4baab89f3a32518a88c31bc87f618f76673e2cc77ab2127b7afdeda33b:0:0"
+          .parse()
+          .unwrap()
+      },
+      FindRangeOutput {
+        start: 50 * COIN_VALUE,
+        size: 5 * COIN_VALUE,
+        satpoint: "30f2f037629c6a21c1f40ed39b9bd6278df39762d68d07f49582b23bcb23386a:0:0"
+          .parse()
+          .unwrap()
+      }
+    ]
+  );
+}
+
+#[test]
+fn find_range_command_fails_for_unmined_sat_ranges() {
+  let rpc_server = test_bitcoincore_rpc::spawn();
+
+  CommandBuilder::new(format!(
+    "--index-sats find {} {}",
+    50 * COIN_VALUE,
+    100 * COIN_VALUE
+  ))
+  .rpc_server(&rpc_server)
+  .expected_exit_code(1)
+  .expected_stderr("error: range has not been mined as of index height\n")
+  .run_and_extract_stdout();
+}
+
 #[test]
 fn unmined_sat() {
   let rpc_server = test_bitcoincore_rpc::spawn();