Add ord balances to show rune balances (ordinals#2782)

popcnt1 · Dec 1, 2023 · d4d2931 · d4d2931
1 parent 8776a97
commit d4d2931
Show file tree

Hide file tree

Showing 8 changed files with 167 additions and 20 deletions.
diff --git a/src/index.rs b/src/index.rs
@@ -901,37 +901,63 @@ impl Index {
     Ok(runic)
   }
 
-  #[cfg(test)]
-  pub(crate) fn get_rune_balances(&self) -> Vec<(OutPoint, Vec<(RuneId, u128)>)> {
+  pub(crate) fn get_rune_balance_map(&self) -> Result<BTreeMap<Rune, BTreeMap<OutPoint, u128>>> {
+    let outpoint_balances = self.get_rune_balances()?;
+
+    let rtx = self.database.begin_read()?;
+
+    let rune_id_to_rune_entry = rtx.open_table(RUNE_ID_TO_RUNE_ENTRY)?;
+
+    let mut rune_balances: BTreeMap<Rune, BTreeMap<OutPoint, u128>> = BTreeMap::new();
+
+    for (outpoint, balances) in outpoint_balances {
+      for (rune_id, amount) in balances {
+        let rune = RuneEntry::load(
+          rune_id_to_rune_entry
+            .get(&rune_id.store())?
+            .unwrap()
+            .value(),
+        )
+        .rune;
+
+        *rune_balances
+          .entry(rune)
+          .or_default()
+          .entry(outpoint)
+          .or_default() += amount;
+      }
+    }
+
+    Ok(rune_balances)
+  }
+
+  pub(crate) fn get_rune_balances(&self) -> Result<Vec<(OutPoint, Vec<(RuneId, u128)>)>> {
     let mut result = Vec::new();
 
     for entry in self
       .database
-      .begin_read()
-      .unwrap()
-      .open_table(OUTPOINT_TO_RUNE_BALANCES)
-      .unwrap()
-      .iter()
-      .unwrap()
+      .begin_read()?
+      .open_table(OUTPOINT_TO_RUNE_BALANCES)?
+      .iter()?
     {
-      let (outpoint, balances_buffer) = entry.unwrap();
+      let (outpoint, balances_buffer) = entry?;
       let outpoint = OutPoint::load(*outpoint.value());
       let balances_buffer = balances_buffer.value();
 
       let mut balances = Vec::new();
       let mut i = 0;
       while i < balances_buffer.len() {
-        let (id, length) = runes::varint::decode(&balances_buffer[i..]).unwrap();
+        let (id, length) = runes::varint::decode(&balances_buffer[i..])?;
         i += length;
-        let (balance, length) = runes::varint::decode(&balances_buffer[i..]).unwrap();
+        let (balance, length) = runes::varint::decode(&balances_buffer[i..])?;
         i += length;
-        balances.push((RuneId::try_from(id).unwrap(), balance));
+        balances.push((RuneId::try_from(id)?, balance));
       }
 
       result.push((outpoint, balances));
     }
 
-    result
+    Ok(result)
   }
 
   pub(crate) fn block_header(&self, hash: BlockHash) -> Result<Option<Header>> {

diff --git a/src/index/testing.rs b/src/index/testing.rs
@@ -113,7 +113,7 @@ impl Context {
 
     assert_eq!(runes, self.index.runes().unwrap());
 
-    assert_eq!(balances, self.index.get_rune_balances());
+    assert_eq!(balances, self.index.get_rune_balances().unwrap());
 
     let mut outstanding: HashMap<RuneId, u128> = HashMap::new();
 

diff --git a/src/subcommand.rs b/src/subcommand.rs
@@ -1,5 +1,6 @@
 use super::*;
 
+pub mod balances;
 pub mod decode;
 pub mod epochs;
 pub mod find;
@@ -17,6 +18,8 @@ pub mod wallet;
 
 #[derive(Debug, Parser)]
 pub(crate) enum Subcommand {
+  #[command(about = "List all rune balances")]
+  Balances,
   #[command(about = "Decode a transaction")]
   Decode(decode::Decode),
   #[command(about = "List the first satoshis of each reward epoch")]
@@ -50,6 +53,7 @@ pub(crate) enum Subcommand {
 impl Subcommand {
   pub(crate) fn run(self, options: Options) -> SubcommandResult {
     match self {
+      Self::Balances => balances::run(options),
       Self::Decode(decode) => decode.run(),
       Self::Epochs => epochs::run(),
       Self::Find(find) => find.run(options),

diff --git a/src/subcommand/balances.rs b/src/subcommand/balances.rs
@@ -0,0 +1,21 @@
+use super::*;
+
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub struct Output {
+  pub runes: BTreeMap<Rune, BTreeMap<OutPoint, u128>>,
+}
+
+pub(crate) fn run(options: Options) -> SubcommandResult {
+  let index = Index::open(&options)?;
+
+  ensure!(
+    index.has_rune_index(),
+    "`ord balances` requires index created with `--index-runes-pre-alpha-i-agree-to-get-rekt` flag",
+  );
+
+  index.update()?;
+
+  Ok(Box::new(Output {
+    runes: index.get_rune_balance_map()?,
+  }))
+}
diff --git a/src/subcommand/server.rs b/src/subcommand/server.rs
@@ -2151,7 +2151,7 @@ mod tests {
     );
 
     assert_eq!(
-      server.index.get_rune_balances(),
+      server.index.get_rune_balances().unwrap(),
       [(OutPoint { txid, vout: 0 }, vec![(id, u128::max_value())])]
     );
 
@@ -2220,7 +2220,7 @@ mod tests {
     );
 
     assert_eq!(
-      server.index.get_rune_balances(),
+      server.index.get_rune_balances().unwrap(),
       [(OutPoint { txid, vout: 0 }, vec![(id, u128::max_value())])]
     );
 
@@ -2326,7 +2326,7 @@ mod tests {
     );
 
     assert_eq!(
-      server.index.get_rune_balances(),
+      server.index.get_rune_balances().unwrap(),
       [(OutPoint { txid, vout: 0 }, vec![(id, u128::max_value())])]
     );
 
@@ -2396,7 +2396,7 @@ mod tests {
     let output = OutPoint { txid, vout: 0 };
 
     assert_eq!(
-      server.index.get_rune_balances(),
+      server.index.get_rune_balances().unwrap(),
       [(output, vec![(id, u128::max_value())])]
     );
 

diff --git a/test-bitcoincore-rpc/src/server.rs b/test-bitcoincore-rpc/src/server.rs
@@ -253,10 +253,19 @@ impl Api for Server {
       .map(|txout| txout.value)
       .sum::<u64>();
 
-    let (outpoint, input_value) = state
+    let mut utxos = state
       .utxos
+      .clone()
+      .into_iter()
+      .map(|(outpoint, value)| (value, outpoint))
+      .collect::<Vec<(Amount, OutPoint)>>();
+
+    utxos.sort();
+    utxos.reverse();
+
+    let (input_value, outpoint) = utxos
       .iter()
-      .find(|(outpoint, value)| value.to_sat() >= output_value && !state.locked.contains(outpoint))
+      .find(|(value, outpoint)| value.to_sat() >= output_value && !state.locked.contains(outpoint))
       .ok_or_else(Self::not_found)?;
 
     transaction.input.push(TxIn {

diff --git a/tests/balances.rs b/tests/balances.rs
@@ -0,0 +1,86 @@
+use {super::*, ord::subcommand::balances::Output};
+
+#[test]
+fn flag_is_required() {
+  let rpc_server = test_bitcoincore_rpc::builder()
+    .network(Network::Regtest)
+    .build();
+
+  CommandBuilder::new("--regtest balances")
+  .rpc_server(&rpc_server)
+  .expected_exit_code(1)
+  .expected_stderr(
+    "error: `ord balances` requires index created with `--index-runes-pre-alpha-i-agree-to-get-rekt` flag\n",
+  )
+  .run_and_extract_stdout();
+}
+
+#[test]
+fn no_runes() {
+  let rpc_server = test_bitcoincore_rpc::builder()
+    .network(Network::Regtest)
+    .build();
+
+  let output =
+    CommandBuilder::new("--regtest --index-runes-pre-alpha-i-agree-to-get-rekt balances")
+      .rpc_server(&rpc_server)
+      .run_and_deserialize_output::<Output>();
+
+  assert_eq!(
+    output,
+    Output {
+      runes: BTreeMap::new()
+    }
+  );
+}
+
+#[test]
+fn with_runes() {
+  let rpc_server = test_bitcoincore_rpc::builder()
+    .network(Network::Regtest)
+    .build();
+
+  create_wallet(&rpc_server);
+
+  let a = etch(&rpc_server, Rune(RUNE));
+  let b = etch(&rpc_server, Rune(RUNE + 1));
+
+  let output =
+    CommandBuilder::new("--regtest --index-runes-pre-alpha-i-agree-to-get-rekt balances")
+      .rpc_server(&rpc_server)
+      .run_and_deserialize_output::<Output>();
+
+  assert_eq!(
+    output,
+    Output {
+      runes: vec![
+        (
+          Rune(RUNE),
+          vec![(
+            OutPoint {
+              txid: a.transaction,
+              vout: 1
+            },
+            1000
+          )]
+          .into_iter()
+          .collect()
+        ),
+        (
+          Rune(RUNE + 1),
+          vec![(
+            OutPoint {
+              txid: b.transaction,
+              vout: 1
+            },
+            1000
+          )]
+          .into_iter()
+          .collect()
+        ),
+      ]
+      .into_iter()
+      .collect(),
+    }
+  );
+}
diff --git a/tests/lib.rs b/tests/lib.rs
@@ -123,6 +123,7 @@ mod command_builder;
 mod expected;
 mod test_server;
 
+mod balances;
 mod core;
 mod decode;
 mod epochs;