From a0a13eb2a89006c8af7dc9998961063f1c067368 Mon Sep 17 00:00:00 2001 From: nym21 Date: Tue, 1 Jul 2025 13:57:48 +0200 Subject: [PATCH] computer: store part 4 --- crates/brk_computer/src/stores.rs | 121 ++++++++++-- .../stateful/addresstype_to_typeindex_tree.rs | 50 +++++ .../stateful/addresstype_to_typeindex_vec.rs} | 13 +- .../src/vecs/stateful/anyaddressdata.rs | 19 ++ crates/brk_computer/src/vecs/stateful/mod.rs | 173 ++++++++++++++---- .../vecs/stateful/withaddressdatasource.rs | 54 ++++++ crates/brk_core/src/groups/by_address_type.rs | 48 ++++- crates/brk_core/src/groups/mod.rs | 2 - crates/brk_core/src/structs/addressdata.rs | 18 +- crates/brk_core/src/structs/outputtype.rs | 17 ++ 10 files changed, 455 insertions(+), 60 deletions(-) create mode 100644 crates/brk_computer/src/vecs/stateful/addresstype_to_typeindex_tree.rs rename crates/{brk_core/src/groups/addressindex_to_typeindex_to_outputindex.rs => brk_computer/src/vecs/stateful/addresstype_to_typeindex_vec.rs} (69%) create mode 100644 crates/brk_computer/src/vecs/stateful/anyaddressdata.rs create mode 100644 crates/brk_computer/src/vecs/stateful/withaddressdatasource.rs diff --git a/crates/brk_computer/src/stores.rs b/crates/brk_computer/src/stores.rs index 08783268c..fa5f4df34 100644 --- a/crates/brk_computer/src/stores.rs +++ b/crates/brk_computer/src/stores.rs @@ -1,18 +1,19 @@ use std::{path::Path, thread}; use brk_core::{ - AddressData, AddressIndexToTypeIndedToOutputIndex, EmptyAddressData, GroupedByAddressType, - Height, P2AAddressIndex, P2AAddressIndexOutputindex, P2PK33AddressIndex, - P2PK33AddressIndexOutputindex, P2PK65AddressIndex, P2PK65AddressIndexOutputindex, - P2PKHAddressIndex, P2PKHAddressIndexOutputindex, P2SHAddressIndex, P2SHAddressIndexOutputindex, - P2TRAddressIndex, P2TRAddressIndexOutputindex, P2WPKHAddressIndex, - P2WPKHAddressIndexOutputindex, P2WSHAddressIndex, P2WSHAddressIndexOutputindex, Result, Unit, - Version, + AddressData, EmptyAddressData, GroupedByAddressType, Height, OutputIndex, OutputType, + P2AAddressIndex, P2AAddressIndexOutputindex, P2PK33AddressIndex, P2PK33AddressIndexOutputindex, + P2PK65AddressIndex, P2PK65AddressIndexOutputindex, P2PKHAddressIndex, + P2PKHAddressIndexOutputindex, P2SHAddressIndex, P2SHAddressIndexOutputindex, P2TRAddressIndex, + P2TRAddressIndexOutputindex, P2WPKHAddressIndex, P2WPKHAddressIndexOutputindex, + P2WSHAddressIndex, P2WSHAddressIndexOutputindex, Result, TypeIndex, Unit, Version, }; use brk_store::{AnyStore, Store}; use fjall::{PersistMode, TransactionalKeyspace}; use rayon::prelude::*; +use crate::vecs::stateful::AddressTypeToTypeIndexVec; + const VERSION: Version = Version::ZERO; #[derive(Clone)] @@ -487,11 +488,109 @@ impl Stores { .map_err(|e| e.into()) } + pub fn get_addressdata( + &self, + address_type: OutputType, + type_index: TypeIndex, + ) -> Result> { + Ok(match address_type { + OutputType::P2A => self + .p2aaddressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2PK33 => self + .p2pk33addressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2PK65 => self + .p2pk65addressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2PKH => self + .p2pkhaddressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2SH => self + .p2shaddressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2TR => self + .p2traddressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2WPKH => self + .p2wpkhaddressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2WSH => self + .p2wshaddressindex_to_addressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + _ => unreachable!(), + }) + } + + pub fn get_emptyaddressdata( + &self, + address_type: OutputType, + type_index: TypeIndex, + ) -> Result> { + Ok(match address_type { + OutputType::P2A => self + .p2aaddressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2PK33 => self + .p2pk33addressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2PK65 => self + .p2pk65addressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2PKH => self + .p2pkhaddressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2SH => self + .p2shaddressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2TR => self + .p2traddressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2WPKH => self + .p2wpkhaddressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + OutputType::P2WSH => self + .p2wshaddressindex_to_emptyaddressdata + .get(&type_index.into())? + .map(|c| c.into_owned()), + _ => unreachable!(), + }) + } + + // pub fn emptyaddressdata_store_increase_replaced(&mut self, address_type: OutputType) { + // match address_type { + // OutputType::P2A => self.p2aaddressindex_to_addressdata.increase_replaced(), + // OutputType::P2PK33 => self.p2pk33addressindex_to_addressdata.increase_replaced(), + // OutputType::P2PK65 => self.p2pk65addressindex_to_addressdata.increase_replaced(), + // OutputType::P2PKH => self.p2pkhaddressindex_to_addressdata.increase_replaced(), + // OutputType::P2SH => self.p2shaddressindex_to_addressdata.increase_replaced(), + // OutputType::P2TR => self.p2traddressindex_to_addressdata.increase_replaced(), + // OutputType::P2WPKH => self.p2wpkhaddressindex_to_addressdata.increase_replaced(), + // OutputType::P2WSH => self.p2wshaddressindex_to_addressdata.increase_replaced(), + // _ => unreachable!(), + // } + // } + pub fn commit( &mut self, height: Height, - sent: AddressIndexToTypeIndedToOutputIndex, - received: AddressIndexToTypeIndedToOutputIndex, + sent: AddressTypeToTypeIndexVec, + received: AddressTypeToTypeIndexVec, ) -> Result<()> { // &mut self.p2aaddressindex_to_addressdata, // &mut self.p2pk33addressindex_to_addressdata, @@ -521,7 +620,7 @@ impl Stores { p2wsh, p2tr, p2a, - } = received.inner(); + } = received.unwrap(); s.spawn(|| { self.p2aaddressindex_to_outputs_received.commit_( @@ -612,7 +711,7 @@ impl Stores { p2wsh, p2tr, p2a, - } = sent.inner(); + } = sent.unwrap(); s.spawn(|| { self.p2aaddressindex_to_outputs_sent.commit_( diff --git a/crates/brk_computer/src/vecs/stateful/addresstype_to_typeindex_tree.rs b/crates/brk_computer/src/vecs/stateful/addresstype_to_typeindex_tree.rs new file mode 100644 index 000000000..dd5fb9ec6 --- /dev/null +++ b/crates/brk_computer/src/vecs/stateful/addresstype_to_typeindex_tree.rs @@ -0,0 +1,50 @@ +use std::{collections::BTreeMap, mem}; + +use brk_core::TypeIndex; +use derive_deref::{Deref, DerefMut}; + +use super::GroupedByAddressType; + +#[derive(Debug, Deref, DerefMut)] +pub struct AddressTypeToTypeIndexTree(GroupedByAddressType>); + +impl AddressTypeToTypeIndexTree { + pub fn merge(&mut self, mut other: Self) { + Self::merge_(&mut self.p2pk65, &mut other.p2pk65); + Self::merge_(&mut self.p2pk33, &mut other.p2pk33); + Self::merge_(&mut self.p2pkh, &mut other.p2pkh); + Self::merge_(&mut self.p2sh, &mut other.p2sh); + Self::merge_(&mut self.p2wpkh, &mut other.p2wpkh); + Self::merge_(&mut self.p2wsh, &mut other.p2wsh); + Self::merge_(&mut self.p2tr, &mut other.p2tr); + Self::merge_(&mut self.p2a, &mut other.p2a); + } + + fn merge_(own: &mut BTreeMap, other: &mut BTreeMap) { + if own.len() >= other.len() { + own.append(other); + } else { + other.append(own); + mem::swap(own, other); + } + } + + pub fn inner(self) -> GroupedByAddressType> { + self.0 + } +} + +impl Default for AddressTypeToTypeIndexTree { + fn default() -> Self { + Self(GroupedByAddressType { + p2pk65: BTreeMap::default(), + p2pk33: BTreeMap::default(), + p2pkh: BTreeMap::default(), + p2sh: BTreeMap::default(), + p2wpkh: BTreeMap::default(), + p2wsh: BTreeMap::default(), + p2tr: BTreeMap::default(), + p2a: BTreeMap::default(), + }) + } +} diff --git a/crates/brk_core/src/groups/addressindex_to_typeindex_to_outputindex.rs b/crates/brk_computer/src/vecs/stateful/addresstype_to_typeindex_vec.rs similarity index 69% rename from crates/brk_core/src/groups/addressindex_to_typeindex_to_outputindex.rs rename to crates/brk_computer/src/vecs/stateful/addresstype_to_typeindex_vec.rs index 9b0901907..f8c2fef36 100644 --- a/crates/brk_core/src/groups/addressindex_to_typeindex_to_outputindex.rs +++ b/crates/brk_computer/src/vecs/stateful/addresstype_to_typeindex_vec.rs @@ -1,17 +1,14 @@ use std::mem; +use brk_core::TypeIndex; use derive_deref::{Deref, DerefMut}; -use crate::{OutputIndex, TypeIndex}; - use super::GroupedByAddressType; #[derive(Debug, Default, Deref, DerefMut)] -pub struct AddressIndexToTypeIndedToOutputIndex( - GroupedByAddressType>, -); +pub struct AddressTypeToTypeIndexVec(GroupedByAddressType>); -impl AddressIndexToTypeIndedToOutputIndex { +impl AddressTypeToTypeIndexVec { pub fn merge(&mut self, mut other: Self) { Self::merge_(&mut self.p2pk65, &mut other.p2pk65); Self::merge_(&mut self.p2pk33, &mut other.p2pk33); @@ -23,7 +20,7 @@ impl AddressIndexToTypeIndedToOutputIndex { Self::merge_(&mut self.p2a, &mut other.p2a); } - fn merge_(own: &mut Vec<(TypeIndex, OutputIndex)>, other: &mut Vec<(TypeIndex, OutputIndex)>) { + fn merge_(own: &mut Vec<(TypeIndex, T)>, other: &mut Vec<(TypeIndex, T)>) { if own.len() >= other.len() { own.append(other); } else { @@ -32,7 +29,7 @@ impl AddressIndexToTypeIndedToOutputIndex { } } - pub fn inner(self) -> GroupedByAddressType> { + pub fn unwrap(self) -> GroupedByAddressType> { self.0 } } diff --git a/crates/brk_computer/src/vecs/stateful/anyaddressdata.rs b/crates/brk_computer/src/vecs/stateful/anyaddressdata.rs new file mode 100644 index 000000000..9fa18babc --- /dev/null +++ b/crates/brk_computer/src/vecs/stateful/anyaddressdata.rs @@ -0,0 +1,19 @@ +use brk_core::{AddressData, EmptyAddressData}; + +#[derive(Debug)] +pub enum AnyAddressData { + AddressData(AddressData), + EmptyAddressData(EmptyAddressData), +} + +impl From for AnyAddressData { + fn from(value: AddressData) -> Self { + Self::AddressData(value) + } +} + +impl From for AnyAddressData { + fn from(value: EmptyAddressData) -> Self { + Self::EmptyAddressData(value) + } +} diff --git a/crates/brk_computer/src/vecs/stateful/mod.rs b/crates/brk_computer/src/vecs/stateful/mod.rs index 24a2bfb67..b5b632384 100644 --- a/crates/brk_computer/src/vecs/stateful/mod.rs +++ b/crates/brk_computer/src/vecs/stateful/mod.rs @@ -1,7 +1,7 @@ use std::{cmp::Ordering, collections::BTreeMap, mem, path::Path, thread}; use brk_core::{ - AddressIndexToTypeIndedToOutputIndex, DateIndex, GroupedByAddressType, Height, InputIndex, + AddressData, DateIndex, Dollars, EmptyAddressData, GroupedByAddressType, Height, InputIndex, OutputIndex, OutputType, Result, Sats, StoredUsize, Version, }; use brk_exit::Exit; @@ -15,10 +15,7 @@ use rayon::prelude::*; use brk_state::{BlockState, CohortStateTrait, SupplyState, Transacted}; -use crate::{ - stores::Stores, - vecs::{market, stateful::r#trait::CohortVecs}, -}; +use crate::{stores::Stores, vecs::market}; use super::{ Indexes, fetched, @@ -28,10 +25,19 @@ use super::{ mod address_cohort; mod address_cohorts; +mod addresstype_to_typeindex_tree; +mod addresstype_to_typeindex_vec; +mod anyaddressdata; mod common; mod r#trait; mod utxo_cohort; mod utxo_cohorts; +mod withaddressdatasource; + +use addresstype_to_typeindex_tree::*; +pub use addresstype_to_typeindex_vec::*; +use r#trait::CohortVecs; +use withaddressdatasource::WithAddressDataSource; const VERSION: Version = Version::new(5); @@ -353,10 +359,14 @@ impl Vecs { let mut height = starting_height; - let mut addressindex_to_typedindex_to_sent_outputindex = - AddressIndexToTypeIndedToOutputIndex::default(); - let mut addressindex_to_typedindex_to_received_outputindex = - AddressIndexToTypeIndedToOutputIndex::default(); + let mut addresstype_to_typeindex_to_sent_outputindex = + AddressTypeToTypeIndexVec::::default(); + let mut addresstype_to_typeindex_to_received_outputindex = + AddressTypeToTypeIndexVec::::default(); + let mut addresstype_to_typeindex_to_addressdata = + AddressTypeToTypeIndexTree::>::default(); + let mut addresstype_to_typeindex_to_emptyaddressdata = + AddressTypeToTypeIndexTree::>::default(); (height.unwrap_to_usize()..height_to_date_fixed.len()) .map(Height::from) @@ -383,7 +393,7 @@ impl Vecs { let output_count = height_to_output_count_iter.unwrap_get_inner(height); let input_count = height_to_input_count_iter.unwrap_get_inner(height); - let ((mut height_to_sent, new_addressindex_to_typedindex_to_sent_outputindex), (mut received, new_addressindex_to_typedindex_to_received_outputindex)) = thread::scope(|s| { + let ((mut height_to_sent, new_addresstype_to_typedindex_to_sent_outputindex, addresstype_to_typedindex_to_sent_sats_and_addressdata_opt), (mut received, new_addresstype_to_typedindex_to_received_outputindex, addresstype_to_typedindex_to_received_sats_and_addressdata_opt)) = thread::scope(|s| { if chain_state_starting_height <= height { s.spawn(|| { self.utxos_vecs @@ -425,7 +435,11 @@ impl Vecs { unreachable!() } - // stores. + let addressdata_opt = if input_type.is_address() && !addresstype_to_typeindex_to_addressdata.get(input_type).unwrap().contains_key(&typeindex) { + Some(WithAddressDataSource::FromAddressDataStore( stores.get_addressdata(input_type, typeindex).unwrap().unwrap())) + } else { + None + }; let input_txindex = outputindex_to_txindex .get_or_read(outputindex, &outputindex_to_txindex_mmap) @@ -439,29 +453,34 @@ impl Vecs { .unwrap() .into_owned(); - (height, value, input_type, typeindex, outputindex) + (height, value, input_type, typeindex, outputindex, addressdata_opt) }) .fold( || { ( BTreeMap::::default(), - AddressIndexToTypeIndedToOutputIndex::default(), + AddressTypeToTypeIndexVec::::default(), + AddressTypeToTypeIndexVec::<(Sats, Option>)>::default(), ) }, - |(mut tree, mut vecs), (height, value, input_type, typeindex, outputindex)| { + |(mut tree, mut vecs, mut vecs2), (height, value, input_type, typeindex, outputindex, addressdata_opt)| { tree.entry(height).or_default().iterate(value, input_type); - if let Some( vec) = vecs.get_mut(input_type) { + if let Some(vec) = vecs.get_mut(input_type) { vec.push((typeindex, outputindex)); } - (tree, vecs) + if let Some(vec) = vecs2.get_mut(input_type) { + vec.push((typeindex, (value, addressdata_opt))); + } + (tree, vecs, vecs2) }, ) .reduce( || { ( BTreeMap::::default(), - AddressIndexToTypeIndedToOutputIndex::default(), + AddressTypeToTypeIndexVec::::default(), + AddressTypeToTypeIndexVec::<(Sats, Option>)>::default(), ) - }, |(first_tree, mut source_vecs), (second_tree, other_vecs)| { + }, |(first_tree, mut source_vecs,mut source_vecs2), (second_tree, other_vecs, other_vecs2)| { let (mut tree_source, tree_to_consume) = if first_tree.len() > second_tree.len() { (first_tree, second_tree) } else { @@ -471,7 +490,8 @@ impl Vecs { *tree_source.entry(k).or_default() += v; }); source_vecs.merge(other_vecs); - (tree_source, source_vecs) + source_vecs2.merge(other_vecs2); + (tree_source, source_vecs, source_vecs2) }) }); @@ -498,29 +518,56 @@ impl Vecs { .unwrap() .into_owned(); - (value, output_type, typeindex, outputindex) + let addressdata_opt = if output_type.is_address() && !addresstype_to_typeindex_to_addressdata.get(output_type).unwrap().contains_key(&typeindex) && !addresstype_to_typeindex_to_emptyaddressdata.get(output_type).unwrap().contains_key(&typeindex) { + Some(if let Some(addressdata) = stores.get_addressdata(output_type, typeindex).unwrap() { + WithAddressDataSource::FromAddressDataStore(addressdata) + } else if let Some(emptyaddressdata) = stores.get_emptyaddressdata(output_type, typeindex).unwrap() { + WithAddressDataSource::FromEmptyAddressDataStore(emptyaddressdata.into()) + } else { + WithAddressDataSource::New(AddressData::default()) + }) + } else { + None + }; + + (value, output_type, typeindex, outputindex, addressdata_opt) }) .fold( - || (Transacted::default(), AddressIndexToTypeIndedToOutputIndex::default()), - |(mut transacted, mut vecs), (value, output_type, typeindex, outputindex)| { + || (Transacted::default(), AddressTypeToTypeIndexVec::::default(), + AddressTypeToTypeIndexVec::<(Sats, Option>)>::default(), + ), + |(mut transacted, mut vecs, mut vecs2), (value, output_type, typeindex, outputindex, addressdata_opt)| { transacted.iterate(value, output_type); if let Some(vec) = vecs.get_mut(output_type) { vec.push((typeindex, outputindex)); } - (transacted, vecs) + if let Some(vec) = vecs2.get_mut(output_type) { + vec.push((typeindex, (value, addressdata_opt))); + } + (transacted, vecs, vecs2) }, ) - .reduce(|| (Transacted::default(), AddressIndexToTypeIndedToOutputIndex::default()), |(transacted, mut vecs), (other_transacted, other_vecs)| { + .reduce(|| (Transacted::default(), AddressTypeToTypeIndexVec::::default(), + AddressTypeToTypeIndexVec::<(Sats, Option>)>::default(), + ), |(transacted, mut vecs, mut vecs2), (other_transacted, other_vecs, other_vecs2)| { vecs.merge(other_vecs); - (transacted + other_transacted, vecs) + vecs2.merge(other_vecs2); + (transacted + other_transacted, vecs, vecs2) }) }); (sent_handle.join().unwrap(), received_handle.join().unwrap()) }); - addressindex_to_typedindex_to_sent_outputindex.merge(new_addressindex_to_typedindex_to_sent_outputindex); - addressindex_to_typedindex_to_received_outputindex.merge(new_addressindex_to_typedindex_to_received_outputindex); + addresstype_to_typeindex_to_sent_outputindex.merge(new_addresstype_to_typedindex_to_sent_outputindex); + addresstype_to_typeindex_to_received_outputindex.merge(new_addresstype_to_typedindex_to_received_outputindex); + + addresstype_to_typedindex_to_received_sats_and_addressdata_opt.process_received(&mut addresstype_to_typeindex_to_addressdata, &mut addresstype_to_typeindex_to_emptyaddressdata, price); + + // addresstype_to_typedindex_to_sent_sats_and_addressdata_opt; + // take from addressdata store + // apply + // update vecs states if cohort changes unspendable_supply += received .by_type @@ -617,8 +664,8 @@ impl Vecs { )?; stores.commit( height, - mem::take(&mut addressindex_to_typedindex_to_sent_outputindex), - mem::take( &mut addressindex_to_typedindex_to_received_outputindex) + mem::take(&mut addresstype_to_typeindex_to_sent_outputindex), + mem::take( &mut addresstype_to_typeindex_to_received_outputindex) )?; exit.release(); } @@ -633,8 +680,8 @@ impl Vecs { self.flush_states(height, &chain_state, exit)?; stores.commit( height, - mem::take(&mut addressindex_to_typedindex_to_sent_outputindex), - mem::take(&mut addressindex_to_typedindex_to_received_outputindex), + mem::take(&mut addresstype_to_typeindex_to_sent_outputindex), + mem::take(&mut addresstype_to_typeindex_to_received_outputindex), )?; } else { exit.block(); @@ -756,3 +803,67 @@ impl Vecs { .collect::>() } } + +impl AddressTypeToTypeIndexVec<(Sats, Option>)> { + fn process_received( + mut self, + addresstype_to_typeindex_to_addressdata: &mut AddressTypeToTypeIndexTree< + WithAddressDataSource, + >, + addresstype_to_typeindex_to_emptyaddressdata: &mut AddressTypeToTypeIndexTree< + WithAddressDataSource, + >, + price: Option, + ) { + self.into_typed_vec().into_iter().for_each(|(_type, vec)| { + vec.into_iter() + .for_each(|(type_index, (value, addressdata_opt))| { + addresstype_to_typeindex_to_addressdata + .get_mut(_type) + .unwrap() + .entry(type_index) + .or_insert_with(|| { + addressdata_opt.unwrap_or_else(|| { + addresstype_to_typeindex_to_emptyaddressdata + .get_mut(_type) + .unwrap() + .remove(&type_index) + .unwrap() + .into() + }) + }) + .deref_mut() + .receive(value, price); + + // update vecs states if cohort changes + }); + }); + } + + fn process_sent( + mut self, + addresstype_to_typeindex_to_addressdata: &mut AddressTypeToTypeIndexTree< + WithAddressDataSource, + >, + price: Option, + ) -> Result<()> { + self.into_typed_vec() + .into_iter() + .try_for_each(|(_type, vec)| { + vec.into_iter() + .try_for_each(|(type_index, (value, addressdata_opt))| { + addresstype_to_typeindex_to_addressdata + .get_mut(_type) + .unwrap() + .entry(type_index) + .or_insert(addressdata_opt.unwrap()) + .deref_mut() + .send(value, price)?; + + Ok(()) + + // update vecs states if cohort changes + }) + }) + } +} diff --git a/crates/brk_computer/src/vecs/stateful/withaddressdatasource.rs b/crates/brk_computer/src/vecs/stateful/withaddressdatasource.rs new file mode 100644 index 000000000..1c03a9bc1 --- /dev/null +++ b/crates/brk_computer/src/vecs/stateful/withaddressdatasource.rs @@ -0,0 +1,54 @@ +use brk_core::{AddressData, EmptyAddressData}; + +#[derive(Debug)] +pub enum WithAddressDataSource { + New(T), + FromAddressDataStore(T), + FromEmptyAddressDataStore(T), +} + +impl WithAddressDataSource { + pub fn is_new(&self) -> bool { + matches!(self, Self::New(_)) + } + + pub fn deref(&self) -> &T { + match self { + Self::New(v) => v, + Self::FromAddressDataStore(v) => v, + Self::FromEmptyAddressDataStore(v) => v, + } + } + + pub fn deref_mut(&mut self) -> &mut T { + match self { + Self::New(v) => v, + Self::FromAddressDataStore(v) => v, + Self::FromEmptyAddressDataStore(v) => v, + } + } +} + +impl From> for WithAddressDataSource { + fn from(value: WithAddressDataSource) -> Self { + match value { + WithAddressDataSource::New(v) => Self::New(v.into()), + WithAddressDataSource::FromAddressDataStore(v) => Self::FromAddressDataStore(v.into()), + WithAddressDataSource::FromEmptyAddressDataStore(v) => { + Self::FromEmptyAddressDataStore(v.into()) + } + } + } +} + +impl From> for WithAddressDataSource { + fn from(value: WithAddressDataSource) -> Self { + match value { + WithAddressDataSource::New(v) => Self::New(v.into()), + WithAddressDataSource::FromAddressDataStore(v) => Self::FromAddressDataStore(v.into()), + WithAddressDataSource::FromEmptyAddressDataStore(v) => { + Self::FromEmptyAddressDataStore(v.into()) + } + } + } +} diff --git a/crates/brk_core/src/groups/by_address_type.rs b/crates/brk_core/src/groups/by_address_type.rs index 10716a1f0..bd3094ece 100644 --- a/crates/brk_core/src/groups/by_address_type.rs +++ b/crates/brk_core/src/groups/by_address_type.rs @@ -1,4 +1,7 @@ -use std::ops::{Add, AddAssign}; +use std::{ + mem, + ops::{Add, AddAssign}, +}; use crate::OutputType; @@ -17,6 +20,20 @@ pub struct GroupedByAddressType { } impl GroupedByAddressType { + pub fn get(&self, address_type: OutputType) -> Option<&T> { + match address_type { + OutputType::P2PK65 => Some(&self.p2pk65), + OutputType::P2PK33 => Some(&self.p2pk33), + OutputType::P2PKH => Some(&self.p2pkh), + OutputType::P2SH => Some(&self.p2sh), + OutputType::P2WPKH => Some(&self.p2wpkh), + OutputType::P2WSH => Some(&self.p2wsh), + OutputType::P2TR => Some(&self.p2tr), + OutputType::P2A => Some(&self.p2a), + _ => None, + } + } + pub fn get_mut(&mut self, address_type: OutputType) -> Option<&mut T> { match address_type { OutputType::P2PK65 => Some(&mut self.p2pk65), @@ -56,6 +73,35 @@ impl GroupedByAddressType { (OutputType::P2A, &self.p2a), ] } + + pub fn as_mut_typed_vec(&mut self) -> [(OutputType, &mut T); 8] { + [ + (OutputType::P2PK65, &mut self.p2pk65), + (OutputType::P2PK33, &mut self.p2pk33), + (OutputType::P2PKH, &mut self.p2pkh), + (OutputType::P2SH, &mut self.p2sh), + (OutputType::P2WPKH, &mut self.p2wpkh), + (OutputType::P2WSH, &mut self.p2wsh), + (OutputType::P2TR, &mut self.p2tr), + (OutputType::P2A, &mut self.p2a), + ] + } + + pub fn into_typed_vec(&mut self) -> [(OutputType, T); 8] + where + T: Default, + { + [ + (OutputType::P2PK65, mem::take(&mut self.p2pk65)), + (OutputType::P2PK33, mem::take(&mut self.p2pk33)), + (OutputType::P2PKH, mem::take(&mut self.p2pkh)), + (OutputType::P2SH, mem::take(&mut self.p2sh)), + (OutputType::P2WPKH, mem::take(&mut self.p2wpkh)), + (OutputType::P2WSH, mem::take(&mut self.p2wsh)), + (OutputType::P2TR, mem::take(&mut self.p2tr)), + (OutputType::P2A, mem::take(&mut self.p2a)), + ] + } } impl GroupedByAddressType<(GroupFilter, T)> { diff --git a/crates/brk_core/src/groups/mod.rs b/crates/brk_core/src/groups/mod.rs index 268a79586..55cd604af 100644 --- a/crates/brk_core/src/groups/mod.rs +++ b/crates/brk_core/src/groups/mod.rs @@ -1,5 +1,4 @@ mod address; -mod addressindex_to_typeindex_to_outputindex; mod by_address_type; mod by_date_range; mod by_epoch; @@ -16,7 +15,6 @@ mod filter; mod utxo; pub use address::*; -pub use addressindex_to_typeindex_to_outputindex::*; pub use by_address_type::*; pub use by_date_range::*; pub use by_epoch::*; diff --git a/crates/brk_core/src/structs/addressdata.rs b/crates/brk_core/src/structs/addressdata.rs index 03a3260b1..2489ae573 100644 --- a/crates/brk_core/src/structs/addressdata.rs +++ b/crates/brk_core/src/structs/addressdata.rs @@ -31,22 +31,26 @@ impl AddressData { } } - pub fn receive(&mut self, amount: Sats, price: Dollars) { + pub fn receive(&mut self, amount: Sats, price: Option) { self.received += amount; self.outputs_len += 1; - self.realized_cap += price * amount; + if let Some(price) = price { + self.realized_cap += price * amount; + } } - pub fn send(&mut self, amount: Sats, previous_price: Dollars) -> Result<()> { + pub fn send(&mut self, amount: Sats, previous_price: Option) -> Result<()> { if self.amount() < amount { return Err(Error::String("Previous_amount smaller than sent amount")); } self.sent += amount; self.outputs_len -= 1; - self.realized_cap = self - .realized_cap - .checked_sub(previous_price * amount) - .unwrap(); + if let Some(previous_price) = previous_price { + self.realized_cap = self + .realized_cap + .checked_sub(previous_price * amount) + .unwrap(); + } Ok(()) } } diff --git a/crates/brk_core/src/structs/outputtype.rs b/crates/brk_core/src/structs/outputtype.rs index 4a4c70c4f..c1e4b887e 100644 --- a/crates/brk_core/src/structs/outputtype.rs +++ b/crates/brk_core/src/structs/outputtype.rs @@ -50,6 +50,23 @@ impl OutputType { } } + pub fn is_address(&self) -> bool { + match self { + Self::P2PK65 => true, + Self::P2PK33 => true, + Self::P2PKH => true, + Self::P2MS => false, + Self::P2SH => true, + Self::OpReturn => false, + Self::P2WPKH => true, + Self::P2WSH => true, + Self::P2TR => true, + Self::P2A => true, + Self::Empty => false, + Self::Unknown => false, + } + } + pub fn is_unspendable(&self) -> bool { !self.is_spendable() }