brk/crates/brk_computer/src/storage/vecs/blocks.rs

use std::{fs, path::Path};

use brk_core::{CheckedSub, CounterU32, Timestamp};
use brk_exit::Exit;
use brk_indexer::Indexer;
use brk_vec::{AnyStorableVec, Compressed, Version};

use super::{
    Indexes,
    grouped::{ComputedVecsFromHeight, StorableVecGeneatorOptions},
    indexes,
};

#[derive(Clone)]
pub struct Vecs {
    pub indexes_to_block_interval: ComputedVecsFromHeight<Timestamp>,
    pub indexes_to_block_count: ComputedVecsFromHeight<CounterU32>,
}

impl Vecs {
    pub fn forced_import(path: &Path, compressed: Compressed) -> color_eyre::Result<Self> {
        fs::create_dir_all(path)?;

        Ok(Self {
            indexes_to_block_interval: ComputedVecsFromHeight::forced_import(
                path,
                "block_interval",
                Version::ONE,
                compressed,
                StorableVecGeneatorOptions::default()
                    .add_percentiles()
                    .add_minmax()
                    .add_average(),
            )?,
            indexes_to_block_count: ComputedVecsFromHeight::forced_import(
                path,
                "block_count",
                Version::ONE,
                compressed,
                StorableVecGeneatorOptions::default().add_sum().add_total(),
            )?,
        })
    }

    pub fn compute(
        &mut self,
        indexer: &mut Indexer,
        indexes: &mut indexes::Vecs,
        starting_indexes: &Indexes,
        exit: &Exit,
    ) -> color_eyre::Result<()> {
        self.indexes_to_block_interval.compute(
            indexer,
            indexes,
            starting_indexes,
            exit,
            |v, indexer, _, starting_indexes, exit| {
                let indexer_vecs = indexer.mut_vecs();

                v.compute_transform(
                    starting_indexes.height,
                    indexer_vecs.height_to_timestamp.mut_vec(),
                    |(height, timestamp, _, height_to_timestamp)| {
                        let interval = height.decremented().map_or(Timestamp::ZERO, |prev_h| {
                            let prev_timestamp = *height_to_timestamp.get(prev_h).unwrap().unwrap();
                            timestamp
                                .checked_sub(prev_timestamp)
                                .unwrap_or(Timestamp::ZERO)
                        });
                        (height, interval)
                    },
                    exit,
                )
            },
        )?;

        self.indexes_to_block_count.compute(
            indexer,
            indexes,
            starting_indexes,
            exit,
            |v, indexer, _, starting_indexes, exit| {
                let indexer_vecs = indexer.mut_vecs();

                v.compute_transform(
                    starting_indexes.height,
                    indexer_vecs.height_to_weight.mut_vec(),
                    |(h, ..)| (h, CounterU32::from(1_u32)),
                    exit,
                )
            },
        )?;

        Ok(())
    }

    pub fn as_any_vecs(&self) -> Vec<&dyn AnyStorableVec> {
        [
            self.indexes_to_block_interval.any_vecs(),
            self.indexes_to_block_count.any_vecs(),
        ]
        .concat()
    }
}