use brk_types::{EmptyAddrData, FundedAddrData, OutputType, TypeIndex};
use crate::distribution::addr::AddrTypeToTypeIndexMap;
use super::super::cohort::WithAddrDataSource;
#[derive(Clone, Copy)]
pub enum TrackingStatus {
New,
Tracked,
WasEmpty,
}
pub struct AddrLookup<'a> {
pub funded: &'a mut AddrTypeToTypeIndexMap<WithAddrDataSource<FundedAddrData>>,
pub empty: &'a mut AddrTypeToTypeIndexMap<WithAddrDataSource<EmptyAddrData>>,
}
impl<'a> AddrLookup<'a> {
pub(crate) fn get_or_create_for_receive(
&mut self,
output_type: OutputType,
type_index: TypeIndex,
) -> (&mut WithAddrDataSource<FundedAddrData>, TrackingStatus) {
use std::collections::hash_map::Entry;
let map = self.funded.get_mut(output_type).unwrap();
match map.entry(type_index) {
Entry::Occupied(entry) => {
let status = match entry.get() {
WithAddrDataSource::New(data) => {
if data.funded_txo_count == 0 {
TrackingStatus::New
} else {
TrackingStatus::Tracked
}
}
WithAddrDataSource::FromFunded(..) => TrackingStatus::Tracked,
WithAddrDataSource::FromEmpty(_, data) => {
if data.utxo_count() == 0 {
TrackingStatus::WasEmpty
} else {
TrackingStatus::Tracked
}
}
};
(entry.into_mut(), status)
}
Entry::Vacant(entry) => {
if let Some(empty_data) =
self.empty.get_mut(output_type).unwrap().remove(&type_index)
{
return (entry.insert(empty_data.into()), TrackingStatus::WasEmpty);
}
(
entry.insert(WithAddrDataSource::New(FundedAddrData::default())),
TrackingStatus::New,
)
}
}
}
pub(crate) fn get_for_send(
&mut self,
output_type: OutputType,
type_index: TypeIndex,
) -> &mut WithAddrDataSource<FundedAddrData> {
self.funded
.get_mut(output_type)
.unwrap()
.get_mut(&type_index)
.expect("Addr must exist for send")
}
pub(crate) fn move_to_empty(&mut self, output_type: OutputType, type_index: TypeIndex) {
let data = self
.funded
.get_mut(output_type)
.unwrap()
.remove(&type_index)
.unwrap();
self.empty
.get_mut(output_type)
.unwrap()
.insert(type_index, data.into());
}
}