From a2fa233acd73a10ac3864138fbd5df97ccb98d74 Mon Sep 17 00:00:00 2001
From: thunderbiscuit <thunderbiscuit@protonmail.com>
Date: Mon, 23 Oct 2023 14:14:39 -0400
Subject: [PATCH] feat(wallet): add IndexedScript and SpkIterator types, and
 wallet::spks_of_all_keychains()"

---
 bdk-ffi/Cargo.lock                            | 56 +++++------
 bdk-ffi/src/bdk.udl                           | 72 ++++++++------
 bdk-ffi/src/bitcoin.rs                        | 22 +++++
 bdk-ffi/src/lib.rs                            |  5 +-
 bdk-ffi/src/wallet.rs                         | 96 ++++++++++++++++++-
 .../kotlin/org/bitcoindevkit/JvmLibTest.kt    | 17 ++--
 6 files changed, 204 insertions(+), 64 deletions(-)
 create mode 100644 bdk-ffi/src/bitcoin.rs

diff --git a/bdk-ffi/Cargo.lock b/bdk-ffi/Cargo.lock
index 7493e629..48ff58d5 100644
--- a/bdk-ffi/Cargo.lock
+++ b/bdk-ffi/Cargo.lock
@@ -90,9 +90,9 @@ checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
 
 [[package]]
 name = "base64"
-version = "0.21.4"
+version = "0.21.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ba43ea6f343b788c8764558649e08df62f86c6ef251fdaeb1ffd010a9ae50a2"
+checksum = "35636a1494ede3b646cc98f74f8e62c773a38a659ebc777a2cf26b9b74171df9"
 
 [[package]]
 name = "bdk"
@@ -248,9 +248,9 @@ dependencies = [
 
 [[package]]
 name = "cargo-platform"
-version = "0.1.3"
+version = "0.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2cfa25e60aea747ec7e1124f238816749faa93759c6ff5b31f1ccdda137f4479"
+checksum = "12024c4645c97566567129c204f65d5815a8c9aecf30fcbe682b2fe034996d36"
 dependencies = [
  "serde",
 ]
@@ -463,9 +463,9 @@ dependencies = [
 
 [[package]]
 name = "libc"
-version = "0.2.148"
+version = "0.2.149"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9cdc71e17332e86d2e1d38c1f99edcb6288ee11b815fb1a4b049eaa2114d369b"
+checksum = "a08173bc88b7955d1b3145aa561539096c421ac8debde8cbc3612ec635fee29b"
 
 [[package]]
 name = "log"
@@ -551,9 +551,9 @@ checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
 
 [[package]]
 name = "os_str_bytes"
-version = "6.5.1"
+version = "6.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4d5d9eb14b174ee9aa2ef96dc2b94637a2d4b6e7cb873c7e171f0c20c6cf3eac"
+checksum = "e2355d85b9a3786f481747ced0e0ff2ba35213a1f9bd406ed906554d7af805a1"
 
 [[package]]
 name = "paste"
@@ -605,9 +605,9 @@ dependencies = [
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.67"
+version = "1.0.69"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3d433d9f1a3e8c1263d9456598b16fec66f4acc9a74dacffd35c7bb09b3a1328"
+checksum = "134c189feb4956b20f6f547d2cf727d4c0fe06722b20a0eec87ed445a97f92da"
 dependencies = [
  "unicode-ident",
 ]
@@ -717,7 +717,7 @@ checksum = "1db149f81d46d2deba7cd3c50772474707729550221e69588478ebf9ada425ae"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.37",
+ "syn 2.0.38",
 ]
 
 [[package]]
@@ -753,31 +753,31 @@ dependencies = [
 
 [[package]]
 name = "semver"
-version = "1.0.19"
+version = "1.0.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ad977052201c6de01a8ef2aa3378c4bd23217a056337d1d6da40468d267a4fb0"
+checksum = "836fa6a3e1e547f9a2c4040802ec865b5d85f4014efe00555d7090a3dcaa1090"
 dependencies = [
  "serde",
 ]
 
 [[package]]
 name = "serde"
-version = "1.0.188"
+version = "1.0.189"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cf9e0fcba69a370eed61bcf2b728575f726b50b55cba78064753d708ddc7549e"
+checksum = "8e422a44e74ad4001bdc8eede9a4570ab52f71190e9c076d14369f38b9200537"
 dependencies = [
  "serde_derive",
 ]
 
 [[package]]
 name = "serde_derive"
-version = "1.0.188"
+version = "1.0.189"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4eca7ac642d82aa35b60049a6eccb4be6be75e599bd2e9adb5f875a737654af2"
+checksum = "1e48d1f918009ce3145511378cf68d613e3b3d9137d67272562080d68a2b32d5"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.37",
+ "syn 2.0.38",
 ]
 
 [[package]]
@@ -848,9 +848,9 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "2.0.37"
+version = "2.0.38"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7303ef2c05cd654186cb250d29049a24840ca25d2747c25c0381c8d9e2f582e8"
+checksum = "e96b79aaa137db8f61e26363a0c9b47d8b4ec75da28b7d1d614c2303e232408b"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -874,22 +874,22 @@ checksum = "222a222a5bfe1bba4a77b45ec488a741b3cb8872e5e499451fd7d0129c9c7c3d"
 
 [[package]]
 name = "thiserror"
-version = "1.0.49"
+version = "1.0.50"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1177e8c6d7ede7afde3585fd2513e611227efd6481bd78d2e82ba1ce16557ed4"
+checksum = "f9a7210f5c9a7156bb50aa36aed4c95afb51df0df00713949448cf9e97d382d2"
 dependencies = [
  "thiserror-impl",
 ]
 
 [[package]]
 name = "thiserror-impl"
-version = "1.0.49"
+version = "1.0.50"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "10712f02019e9288794769fba95cd6847df9874d49d871d062172f9dd41bc4cc"
+checksum = "266b2e40bc00e5a6c09c3584011e08b06f123c00362c92b975ba9843aaaa14b8"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.37",
+ "syn 2.0.38",
 ]
 
 [[package]]
@@ -1064,7 +1064,7 @@ version = "2.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f5ccd538d4a604753ebc2f17cd9946e89b77bf87f6a8e2309667c6f2e87855e3"
 dependencies = [
- "base64 0.21.4",
+ "base64 0.21.5",
  "flate2",
  "log",
  "once_cell",
@@ -1121,7 +1121,7 @@ dependencies = [
  "once_cell",
  "proc-macro2",
  "quote",
- "syn 2.0.37",
+ "syn 2.0.38",
  "wasm-bindgen-shared",
 ]
 
@@ -1143,7 +1143,7 @@ checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.37",
+ "syn 2.0.38",
  "wasm-bindgen-backend",
  "wasm-bindgen-shared",
 ]
diff --git a/bdk-ffi/src/bdk.udl b/bdk-ffi/src/bdk.udl
index f5830bb3..49096853 100644
--- a/bdk-ffi/src/bdk.udl
+++ b/bdk-ffi/src/bdk.udl
@@ -82,40 +82,22 @@ interface Wallet {
 
   [Throws=BdkError]
   void apply_update(Update update);
-};
-
-interface Update {};
-
-// ------------------------------------------------------------------------
-// bdk crate - bitcoin reexports
-// ------------------------------------------------------------------------
 
-enum Network {
-  "Bitcoin",
-  "Testnet",
-  "Signet",
-  "Regtest",
+  record<KeychainKind, SpkIterator> spks_of_all_keychains();
 };
 
-enum WordCount {
-  "Words12",
-  "Words15",
-  "Words18",
-  "Words21",
-  "Words24",
+dictionary IndexedScript {
+  u32 index;
+  Script script;
 };
 
-interface Address {
-  [Throws=BdkError]
-  constructor(string address, Network network);
-
-  Network network();
-
-  string to_qr_uri();
-
-  string as_string();
+interface SpkIterator {
+    constructor(Descriptor descriptor);
+    IndexedScript? next();
 };
 
+interface Update {};
+
 // ------------------------------------------------------------------------
 // bdk crate - descriptor module
 // ------------------------------------------------------------------------
@@ -209,3 +191,39 @@ interface Descriptor {
 interface EsploraClient {
   constructor(string url);
 };
+
+// ------------------------------------------------------------------------
+// bdk crate - bitcoin reexports
+// ------------------------------------------------------------------------
+
+interface Script {
+  constructor(sequence<u8> raw_output_script);
+
+  sequence<u8> to_bytes();
+};
+
+enum Network {
+  "Bitcoin",
+  "Testnet",
+  "Signet",
+  "Regtest",
+};
+
+enum WordCount {
+  "Words12",
+  "Words15",
+  "Words18",
+  "Words21",
+  "Words24",
+};
+
+interface Address {
+  [Throws=BdkError]
+  constructor(string address, Network network);
+
+  Network network();
+
+  string to_qr_uri();
+
+  string as_string();
+};
diff --git a/bdk-ffi/src/bitcoin.rs b/bdk-ffi/src/bitcoin.rs
new file mode 100644
index 00000000..c3b56fa5
--- /dev/null
+++ b/bdk-ffi/src/bitcoin.rs
@@ -0,0 +1,22 @@
+use bdk::bitcoin::blockdata::script::ScriptBuf as BdkScriptBuf;
+
+/// A Bitcoin script.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Script(pub(crate) BdkScriptBuf);
+
+impl Script {
+    pub fn new(raw_output_script: Vec<u8>) -> Self {
+        let script: BdkScriptBuf = BdkScriptBuf::from(raw_output_script);
+        Script(script)
+    }
+
+    pub fn to_bytes(&self) -> Vec<u8> {
+        self.0.to_bytes()
+    }
+}
+
+impl From<BdkScriptBuf> for Script {
+    fn from(script: BdkScriptBuf) -> Self {
+        Script(script)
+    }
+}
diff --git a/bdk-ffi/src/lib.rs b/bdk-ffi/src/lib.rs
index 02cd0d35..3b1584b1 100644
--- a/bdk-ffi/src/lib.rs
+++ b/bdk-ffi/src/lib.rs
@@ -1,3 +1,4 @@
+mod bitcoin;
 mod descriptor;
 mod esplora;
 mod keys;
@@ -15,16 +16,18 @@ use bdk::KeychainKind;
 use std::sync::Arc;
 
 // TODO 6: Why are these imports required?
+use crate::bitcoin::Script;
 use crate::descriptor::Descriptor;
 use crate::esplora::EsploraClient;
 use crate::keys::DerivationPath;
 use crate::keys::DescriptorPublicKey;
 use crate::keys::DescriptorSecretKey;
 use crate::keys::Mnemonic;
+use crate::wallet::IndexedScript;
+use crate::wallet::SpkIterator;
 use crate::wallet::Update;
 use crate::wallet::Wallet;
 use bdk::keys::bip39::WordCount;
-// use bdk_esplora::EsploraExt;
 
 uniffi::include_scaffolding!("bdk");
 
diff --git a/bdk-ffi/src/wallet.rs b/bdk-ffi/src/wallet.rs
index e3c806b9..ff83a3a0 100644
--- a/bdk-ffi/src/wallet.rs
+++ b/bdk-ffi/src/wallet.rs
@@ -1,9 +1,14 @@
+use crate::bitcoin::Script;
 use crate::descriptor::Descriptor;
 use crate::Balance;
 use crate::{AddressIndex, AddressInfo, Network};
+use bdk::chain::SpkIterator as BdkSpkIterator;
+use bdk::descriptor::Descriptor as BdkDescriptor;
+use bdk::miniscript::DescriptorPublicKey;
 use bdk::wallet::Update as BdkUpdate;
-use bdk::Error as BdkError;
 use bdk::Wallet as BdkWallet;
+use bdk::{Error as BdkError, KeychainKind};
+use std::collections::HashMap;
 use std::sync::{Arc, Mutex, MutexGuard};
 
 #[derive(Debug)]
@@ -68,6 +73,22 @@ impl Wallet {
             .map_err(|e| BdkError::Generic(e.to_string()))
     }
 
+    pub fn spks_of_all_keychains(&self) -> HashMap<KeychainKind, Arc<SpkIterator>> {
+        let all_spks = self
+            .get_wallet()
+            .spk_index()
+            .spks_of_all_keychains()
+            .into_iter()
+            .fold(HashMap::new(), |mut map, (keychain, iter)| {
+                let spk_iter = SpkIterator(Mutex::new(
+                    iter as BdkSpkIterator<BdkDescriptor<DescriptorPublicKey>>,
+                ));
+                map.insert(keychain, Arc::new(spk_iter));
+                map
+            });
+        all_spks
+    }
+
     // pub fn commit(&self) -> Result<(), BdkError> {}
 
     // fn is_mine(&self, script: Arc<Script>) -> bool {
@@ -75,6 +96,32 @@ impl Wallet {
     // }
 }
 
+#[derive(Debug)]
+pub struct IndexedScript {
+    pub index: u32,
+    pub script: Arc<Script>,
+}
+
+pub struct SpkIterator(pub(crate) Mutex<BdkSpkIterator<BdkDescriptor<DescriptorPublicKey>>>);
+
+impl SpkIterator {
+    pub fn new(descriptor: Arc<Descriptor>) -> SpkIterator {
+        Self(Mutex::new(BdkSpkIterator::new(
+            descriptor.extended_descriptor.clone(),
+        )))
+    }
+
+    pub fn next(&self) -> Option<IndexedScript> {
+        let mut bdk_spk_iterator = self.0.lock().unwrap();
+        bdk_spk_iterator
+            .next()
+            .map(|(index, script)| IndexedScript {
+                index,
+                script: Arc::new(script.into()),
+            })
+    }
+}
+
 pub struct Update(pub(crate) BdkUpdate);
 
 // /// A Bitcoin wallet.
@@ -639,7 +686,52 @@ pub struct Update(pub(crate) BdkUpdate);
 //             .map(Arc::new)
 //     }
 // }
-//
+
+#[cfg(test)]
+mod test {
+    use crate::descriptor::Descriptor;
+    use crate::wallet::Wallet;
+    use crate::Network;
+    use bdk::bitcoin::Address;
+    use bdk::{bitcoin, KeychainKind};
+    use std::sync::Arc;
+
+    #[test]
+    fn test_spks_of_all_keychains() {
+        let descriptor = Descriptor::new(
+            "wpkh(tprv8ZgxMBicQKsPf2qfrEygW6fdYseJDDrVnDv26PH5BHdvSuG6ecCbHqLVof9yZcMoM31z9ur3tTYbSnr1WBqbGX97CbXcmp5H6qeMpyvx35B/84h/1h/0h/0/*)".to_string(),
+            Network::Testnet
+        ).unwrap();
+        let wallet = Wallet::new_no_persist(Arc::new(descriptor), None, Network::Testnet).unwrap();
+        let spks = wallet.spks_of_all_keychains();
+        spks.iter().for_each(|(keychain, spk_iter)| {
+            dbg!(keychain);
+            assert_eq!(*keychain, KeychainKind::External);
+            let mut spks_iter = spk_iter.0.lock().unwrap().clone();
+            let next_addr = spks_iter.next().map(|(i, s)| {
+                (
+                    i,
+                    Address::from_script(s.as_script(), bitcoin::Network::Testnet)
+                        .unwrap()
+                        .to_string(),
+                )
+            });
+            let expected_addr = "tb1qrnfslnrve9uncz9pzpvf83k3ukz22ljgees989".to_string();
+            assert!(matches!(next_addr, Some((i, addr)) if i == 0 && addr == expected_addr));
+            let next_addr = spks_iter.next().map(|(i, s)| {
+                (
+                    i,
+                    Address::from_script(s.as_script(), bitcoin::Network::Testnet)
+                        .unwrap()
+                        .to_string(),
+                )
+            });
+            let expected_addr = "tb1qltq3nuep2fghvnytukac5mq5hxhascs0v6hm3x".to_string();
+            assert!(matches!(next_addr, Some((i, addr)) if i == 1 && addr == expected_addr));
+        })
+    }
+}
+
 // // The goal of these tests to to ensure `bdk-ffi` intermediate code correctly calls `bdk` APIs.
 // // These tests should not be used to verify `bdk` behavior that is already tested in the `bdk`
 // // crate.
diff --git a/bdk-jvm/lib/src/test/kotlin/org/bitcoindevkit/JvmLibTest.kt b/bdk-jvm/lib/src/test/kotlin/org/bitcoindevkit/JvmLibTest.kt
index 049cbc48..18d2231e 100644
--- a/bdk-jvm/lib/src/test/kotlin/org/bitcoindevkit/JvmLibTest.kt
+++ b/bdk-jvm/lib/src/test/kotlin/org/bitcoindevkit/JvmLibTest.kt
@@ -1,11 +1,6 @@
 package org.bitcoindevkit
 
-import org.junit.Assert.*
 import org.junit.Test
-import org.slf4j.Logger
-import org.slf4j.LoggerFactory
-import java.io.File
-import java.nio.file.Files
 
 class WalletTest {
     @Test
@@ -24,7 +19,7 @@ class WalletTest {
     fun testUsedWallet() {
         val descriptor = Descriptor("wpkh(tprv8ZgxMBicQKsPf2qfrEygW6fdYseJDDrVnDv26PH5BHdvSuG6ecCbHqLVof9yZcMoM31z9ur3tTYbSnr1WBqbGX97CbXcmp5H6qeMpyvx35B/84h/1h/0h/0/*)", Network.TESTNET)
         val wallet = Wallet.newNoPersist(descriptor, null, Network.TESTNET)
-        val (index, address, keychain)  = wallet.getAddress(AddressIndex.LastUnused)
+        val (index, address, keychain) = wallet.getAddress(AddressIndex.LastUnused)
         println("Address ${address.asString()} at index $index")
     }
 
@@ -36,6 +31,16 @@ class WalletTest {
         assert(wallet.getBalance().total() == 0uL)
     }
 
+    @Test
+    fun testSpks() {
+        val descriptor = Descriptor("wpkh(tprv8ZgxMBicQKsPf2qfrEygW6fdYseJDDrVnDv26PH5BHdvSuG6ecCbHqLVof9yZcMoM31z9ur3tTYbSnr1WBqbGX97CbXcmp5H6qeMpyvx35B/84h/1h/0h/0/*)", Network.TESTNET)
+        val wallet = Wallet.newNoPersist(descriptor, null, Network.TESTNET)
+        val (index, address, keychain) = wallet.getAddress(AddressIndex.LastUnused)
+        //
+        val spks: Map<KeychainKind, List<IndexedScript>> = wallet.getAllSpks()
+        // println(spks)
+    }
+
     // @Test
     // fun testSyncedBalance() {
     //     val descriptor = Descriptor("wpkh(tprv8ZgxMBicQKsPf2qfrEygW6fdYseJDDrVnDv26PH5BHdvSuG6ecCbHqLVof9yZcMoM31z9ur3tTYbSnr1WBqbGX97CbXcmp5H6qeMpyvx35B/84h/1h/0h/0/*)", Network.TESTNET)