addr: add terminal info

BP-WG · Jul 31, 2023 · e7c7646 · e7c7646
1 parent 5a414a7
commit e7c7646
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 23 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/std/src/chain.rs b/std/src/chain.rs
@@ -27,11 +27,13 @@ use bc::{
     BlockHash, BlockHeader, LockTime, Outpoint, Sats, ScriptPubkey, SeqNo, SigScript, Txid, Witness,
 };
 
-use crate::{Address, DeriveSpk, Idx, NormalIndex, WalletCache, WalletDescr};
+use crate::{
+    Address, DeriveSpk, DerivedAddr, Idx, NormalIndex, Terminal, WalletCache, WalletDescr,
+};
 
 pub type BlockHeight = NonZeroU32;
 
-#[derive(Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct BlockInfo {
     pub header: BlockHeader,
     pub difficulty: u8,
@@ -41,14 +43,14 @@ pub struct BlockInfo {
     pub mediantime: u32,
 }
 
-#[derive(Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct MiningInfo {
     pub height: BlockHeight,
     pub time: u64,
     pub block_hash: BlockHash,
 }
 
-#[derive(Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub enum TxStatus {
     Mined(MiningInfo),
     Mempool,
@@ -79,29 +81,40 @@ pub struct TxInInfo {
     pub value: Option<Sats>,
 }
 
-#[derive(Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct TxOutInfo {
     pub outpoint: Outpoint,
     pub value: Sats,
-    pub derivation: Option<(NormalIndex, NormalIndex)>,
+    pub derivation: Option<Terminal>,
 }
 
-#[derive(Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct UtxoInfo {
     pub outpoint: Outpoint,
     pub value: Sats,
     pub address: Address,
-    pub derivation: (NormalIndex, NormalIndex),
+    pub derivation: Terminal,
 }
 
-#[derive(Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct AddrInfo {
-    pub derivation: (NormalIndex, NormalIndex),
+    pub derivation: Terminal,
     pub used: u32,
     pub volume: Sats,
     pub balance: Sats,
 }
 
+impl From<DerivedAddr> for AddrInfo {
+    fn from(derived: DerivedAddr) -> Self {
+        AddrInfo {
+            derivation: derived.terminal,
+            used: 0,
+            volume: Sats::ZERO,
+            balance: Sats::ZERO,
+        }
+    }
+}
+
 pub trait Blockchain {
     type Error;
 
@@ -137,13 +150,13 @@ impl WalletCache {
                 txids.extend(r.iter().map(|utxo| utxo.outpoint.txid));
                 let max_known = cache.max_known.entry(*keychain).or_default();
                 *max_known = max(
-                    r.iter().map(|utxo| utxo.derivation.1).max().unwrap_or_default(),
+                    r.iter().map(|utxo| utxo.derivation.index).max().unwrap_or_default(),
                     *max_known,
                 );
                 if r.is_empty() {
                     break;
                 }
-                cache.utxo.extend(r.into_iter().map(|utxo| (utxo.outpoint, utxo)));
+                cache.utxo.extend(r.into_iter().map(|utxo| (utxo.address, set! {utxo})));
                 if !index.saturating_add_assign(BATCH_SIZE) {
                     break;
                 }

diff --git a/std/src/derive.rs b/std/src/derive.rs
@@ -20,23 +20,42 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::cmp::Ordering;
+
 use bc::{InternalPk, ScriptPubkey};
 
 use crate::{Address, AddressNetwork, ComprPubkey, Idx, NormalIndex, WalletDescr, XpubDescriptor};
 
+#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug, Display)]
+#[display("/{keychain}/{index}")]
+pub struct Terminal {
+    pub keychain: NormalIndex,
+    pub index: NormalIndex,
+}
+
+impl Terminal {
+    pub fn new(keychain: NormalIndex, index: NormalIndex) -> Self { Terminal { keychain, index } }
+}
+
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct DerivedAddr {
     pub addr: Address,
-    pub keychain: NormalIndex,
-    pub index: NormalIndex,
+    pub terminal: Terminal,
+}
+
+impl Ord for DerivedAddr {
+    fn cmp(&self, other: &Self) -> Ordering { self.terminal.cmp(&other.terminal) }
+}
+
+impl PartialOrd for DerivedAddr {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> { Some(self.cmp(other)) }
 }
 
 impl DerivedAddr {
     pub fn new(addr: Address, keychain: NormalIndex, index: NormalIndex) -> Self {
         DerivedAddr {
             addr,
-            keychain,
-            index,
+            terminal: Terminal::new(keychain, index),
         }
     }
 }
@@ -59,7 +78,7 @@ impl<'descr, D: DeriveSpk> Iterator for AddrIter<'descr, D> {
 }
 
 impl<D: DeriveSpk> WalletDescr<D> {
-    pub fn addresses<'descr>(&'descr self) -> AddrIter<'descr, D> {
+    pub fn addresses(&self) -> AddrIter<D> {
         AddrIter {
             script_pubkey: &self.script_pubkey,
             network: self.chain.into(),

diff --git a/std/src/lib.rs b/std/src/lib.rs
@@ -44,7 +44,9 @@ pub use address::{
 };
 pub use bc::{secp256k1, *};
 pub use chain::{AddrInfo, BlockInfo, MiningInfo, TxInInfo, TxInfo, TxOutInfo, TxStatus, UtxoInfo};
-pub use derive::{AddrIter, Derive, DeriveCompr, DeriveSet, DeriveSpk, DeriveXOnly, DerivedAddr};
+pub use derive::{
+    AddrIter, Derive, DeriveCompr, DeriveSet, DeriveSpk, DeriveXOnly, DerivedAddr, Terminal,
+};
 pub use descriptors::{DescriptorStd, TrKey};
 pub use index::{
     DerivationIndex, HardenedIndex, Idx, IndexError, IndexParseError, NormalIndex,

diff --git a/std/src/wallet.rs b/std/src/wallet.rs
@@ -20,14 +20,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::num::NonZeroU32;
 use std::ops::Deref;
 
 use bc::{Chain, Outpoint, Txid};
 
 use crate::derive::DeriveSpk;
-use crate::{AddrInfo, Address, BlockInfo, Idx, NormalIndex, TxInfo, UtxoInfo};
+use crate::{AddrInfo, Address, BlockInfo, Idx, NormalIndex, Terminal, TxInfo, UtxoInfo};
 
 #[derive(Getters, Clone, Eq, PartialEq, Debug)]
 pub struct WalletDescr<D>
@@ -70,8 +70,8 @@ pub struct WalletCache {
     pub(crate) tip: u32,
     pub(crate) headers: HashMap<NonZeroU32, BlockInfo>,
     pub(crate) tx: HashMap<Txid, TxInfo>,
-    pub(crate) utxo: HashMap<Outpoint, UtxoInfo>,
-    pub(crate) addr: HashMap<(NormalIndex, NormalIndex), AddrInfo>,
+    pub(crate) utxo: HashMap<Address, HashSet<UtxoInfo>>,
+    pub(crate) addr: HashMap<Terminal, AddrInfo>,
     pub(crate) max_known: HashMap<NormalIndex, NormalIndex>,
 }
 
@@ -98,6 +98,23 @@ impl<D: DeriveSpk, L2: Default> Wallet<D, L2> {
             layer2: default!(),
         }
     }
+
+    pub fn coins(&self) -> impl Iterator<Item = UtxoInfo> + '_ {
+        self.cache.utxo.values().flatten().copied()
+    }
+
+    pub fn address_coins(
+        &self,
+    ) -> impl Iterator<Item = (Address, impl Iterator<Item = UtxoInfo> + '_)> + '_ {
+        self.cache.utxo.iter().map(|(k, v)| (*k, v.iter().copied()))
+    }
+
+    pub fn address_all(&self) -> impl Iterator<Item = AddrInfo> + '_ {
+        self.descr.addresses().map(|derived| match self.cache.addr.get(&derived.terminal) {
+            None => AddrInfo::from(derived),
+            Some(info) => *info,
+        })
+    }
 }
 
 impl WalletCache {