brk/crates/brk_computer/src/distribution/address/address_count.rs

use brk_cohort::ByAddressType;
use brk_error::Result;
use brk_traversable::Traversable;
use brk_types::{Height, StoredU64, Version};
use derive_deref::{Deref, DerefMut};
use rayon::prelude::*;
use vecdb::{
    AnyStoredVec, AnyVec, Database, EagerVec, Exit, GenericStoredVec, ImportableVec,
    IterableCloneableVec, PcoVec, TypedVecIterator,
};

use crate::{
    ComputeIndexes, indexes,
    internal::{ComputedVecsFromHeight, Source, VecBuilderOptions},
};

/// Address count per address type (runtime state).
#[derive(Debug, Default, Deref, DerefMut)]
pub struct AddressTypeToAddressCount(ByAddressType<u64>);

impl From<(&AddressTypeToHeightToAddressCount, Height)> for AddressTypeToAddressCount {
    #[inline]
    fn from((groups, starting_height): (&AddressTypeToHeightToAddressCount, Height)) -> Self {
        if let Some(prev_height) = starting_height.decremented() {
            Self(ByAddressType {
                p2pk65: groups.p2pk65.into_iter().get_unwrap(prev_height).into(),
                p2pk33: groups.p2pk33.into_iter().get_unwrap(prev_height).into(),
                p2pkh: groups.p2pkh.into_iter().get_unwrap(prev_height).into(),
                p2sh: groups.p2sh.into_iter().get_unwrap(prev_height).into(),
                p2wpkh: groups.p2wpkh.into_iter().get_unwrap(prev_height).into(),
                p2wsh: groups.p2wsh.into_iter().get_unwrap(prev_height).into(),
                p2tr: groups.p2tr.into_iter().get_unwrap(prev_height).into(),
                p2a: groups.p2a.into_iter().get_unwrap(prev_height).into(),
            })
        } else {
            Default::default()
        }
    }
}

/// Address count per address type, indexed by height.
#[derive(Debug, Clone, Deref, DerefMut, Traversable)]
pub struct AddressTypeToHeightToAddressCount(ByAddressType<EagerVec<PcoVec<Height, StoredU64>>>);

impl From<ByAddressType<EagerVec<PcoVec<Height, StoredU64>>>>
    for AddressTypeToHeightToAddressCount
{
    #[inline]
    fn from(value: ByAddressType<EagerVec<PcoVec<Height, StoredU64>>>) -> Self {
        Self(value)
    }
}

impl AddressTypeToHeightToAddressCount {
    pub fn min_len(&self) -> usize {
        self.p2pk65
            .len()
            .min(self.p2pk33.len())
            .min(self.p2pkh.len())
            .min(self.p2sh.len())
            .min(self.p2wpkh.len())
            .min(self.p2wsh.len())
            .min(self.p2tr.len())
            .min(self.p2a.len())
    }

    pub fn forced_import(db: &Database, name: &str, version: Version) -> Result<Self> {
        Ok(Self::from(ByAddressType::new_with_name(|type_name| {
            Ok(EagerVec::forced_import(
                db,
                &format!("{type_name}_{name}"),
                version,
            )?)
        })?))
    }

    /// Returns a parallel iterator over all vecs for parallel writing.
    pub fn par_iter_mut(&mut self) -> impl ParallelIterator<Item = &mut dyn AnyStoredVec> {
        let inner = &mut self.0;
        [
            &mut inner.p2pk65 as &mut dyn AnyStoredVec,
            &mut inner.p2pk33 as &mut dyn AnyStoredVec,
            &mut inner.p2pkh as &mut dyn AnyStoredVec,
            &mut inner.p2sh as &mut dyn AnyStoredVec,
            &mut inner.p2wpkh as &mut dyn AnyStoredVec,
            &mut inner.p2wsh as &mut dyn AnyStoredVec,
            &mut inner.p2tr as &mut dyn AnyStoredVec,
            &mut inner.p2a as &mut dyn AnyStoredVec,
        ]
        .into_par_iter()
    }

    pub fn write(&mut self) -> Result<()> {
        self.p2pk65.write()?;
        self.p2pk33.write()?;
        self.p2pkh.write()?;
        self.p2sh.write()?;
        self.p2wpkh.write()?;
        self.p2wsh.write()?;
        self.p2tr.write()?;
        self.p2a.write()?;
        Ok(())
    }

    pub fn truncate_push(
        &mut self,
        height: Height,
        addresstype_to_usize: &AddressTypeToAddressCount,
    ) -> Result<()> {
        self.p2pk65
            .truncate_push(height, addresstype_to_usize.p2pk65.into())?;
        self.p2pk33
            .truncate_push(height, addresstype_to_usize.p2pk33.into())?;
        self.p2pkh
            .truncate_push(height, addresstype_to_usize.p2pkh.into())?;
        self.p2sh
            .truncate_push(height, addresstype_to_usize.p2sh.into())?;
        self.p2wpkh
            .truncate_push(height, addresstype_to_usize.p2wpkh.into())?;
        self.p2wsh
            .truncate_push(height, addresstype_to_usize.p2wsh.into())?;
        self.p2tr
            .truncate_push(height, addresstype_to_usize.p2tr.into())?;
        self.p2a
            .truncate_push(height, addresstype_to_usize.p2a.into())?;
        Ok(())
    }

    pub fn reset(&mut self) -> Result<()> {
        use vecdb::GenericStoredVec;
        self.p2pk65.reset()?;
        self.p2pk33.reset()?;
        self.p2pkh.reset()?;
        self.p2sh.reset()?;
        self.p2wpkh.reset()?;
        self.p2wsh.reset()?;
        self.p2tr.reset()?;
        self.p2a.reset()?;
        Ok(())
    }
}

/// Address count per address type, indexed by various indexes (dateindex, etc.).
#[derive(Clone, Deref, DerefMut, Traversable)]
pub struct AddressTypeToIndexesToAddressCount(ByAddressType<ComputedVecsFromHeight<StoredU64>>);

impl From<ByAddressType<ComputedVecsFromHeight<StoredU64>>> for AddressTypeToIndexesToAddressCount {
    #[inline]
    fn from(value: ByAddressType<ComputedVecsFromHeight<StoredU64>>) -> Self {
        Self(value)
    }
}

impl AddressTypeToIndexesToAddressCount {
    pub fn forced_import(
        db: &Database,
        name: &str,
        version: Version,
        indexes: &indexes::Vecs,
        sources: &AddressTypeToHeightToAddressCount,
    ) -> Result<Self> {
        Ok(Self::from(ByAddressType::<
            ComputedVecsFromHeight<StoredU64>,
        >::try_zip_with_name(
            sources,
            |type_name, source| {
                ComputedVecsFromHeight::forced_import(
                    db,
                    &format!("{type_name}_{name}"),
                    Source::Vec(source.boxed_clone()),
                    version,
                    indexes,
                    VecBuilderOptions::default().add_last(),
                )
            },
        )?))
    }

    pub fn compute(
        &mut self,
        indexes: &indexes::Vecs,
        starting_indexes: &ComputeIndexes,
        exit: &Exit,
        addresstype_to_height_to_addresscount: &AddressTypeToHeightToAddressCount,
    ) -> Result<()> {
        self.p2pk65.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2pk65),
        )?;
        self.p2pk33.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2pk33),
        )?;
        self.p2pkh.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2pkh),
        )?;
        self.p2sh.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2sh),
        )?;
        self.p2wpkh.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2wpkh),
        )?;
        self.p2wsh.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2wsh),
        )?;
        self.p2tr.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2tr),
        )?;
        self.p2a.compute_rest(
            indexes,
            starting_indexes,
            exit,
            Some(&addresstype_to_height_to_addresscount.p2a),
        )?;
        Ok(())
    }
}