computer: trying the new stateful

crates/brk_computer/src/stateful/compute/block_loop.rs

@@ -0,0 +1,567 @@
+//! Main block processing loop.
+//!
+//! Iterates through blocks and processes each one:
+//! 1. Reset per-block state values
+//! 2. Tick-tock age transitions
+//! 3. Process outputs (receive) in parallel
+//! 4. Process inputs (send) in parallel
+//! 5. Push to height-indexed vectors
+//! 6. Periodically flush checkpoints
+
+use std::{mem, thread};
+
+use brk_error::Result;
+use brk_grouper::ByAddressType;
+use brk_indexer::Indexer;
+use brk_types::{
+    DateIndex, Dollars, Height, OutputType, Sats, Timestamp, TypeIndex,
+};
+use log::info;
+use rayon::prelude::*;
+use vecdb::{AnyStoredVec, Exit, GenericStoredVec, IterableVec, TypedVecIterator, VecIndex};
+
+use crate::states::{BlockState, Transacted};
+use crate::utils::OptionExt;
+use crate::{chain, indexes, price};
+
+use super::super::address::AddressTypeToTypeIndexMap;
+use super::super::cohorts::{AddressCohorts, DynCohortVecs, UTXOCohorts};
+use super::super::vecs::Vecs;
+use super::{
+    BIP30_DUPLICATE_HEIGHT_1, BIP30_DUPLICATE_HEIGHT_2, BIP30_ORIGINAL_HEIGHT_1,
+    BIP30_ORIGINAL_HEIGHT_2, FLUSH_INTERVAL, IndexerReaders, VecsReaders,
+    build_txinindex_to_txindex, build_txoutindex_to_txindex,
+    flush::flush_checkpoint as flush_checkpoint_full,
+};
+use crate::stateful::address::AddressTypeToAddressCount;
+use crate::stateful::process::{
+    AddressLookup, EmptyAddressDataWithSource, InputsResult, LoadedAddressDataWithSource,
+    build_txoutindex_to_height_map, process_inputs, process_outputs, process_received,
+    process_sent, update_tx_counts,
+};
+
+/// Process all blocks from starting_height to last_height.
+#[allow(clippy::too_many_arguments)]
+pub fn process_blocks(
+    vecs: &mut Vecs,
+    indexer: &Indexer,
+    indexes: &indexes::Vecs,
+    chain: &chain::Vecs,
+    price: Option<&price::Vecs>,
+    starting_height: Height,
+    last_height: Height,
+    chain_state: &mut Vec<BlockState>,
+    exit: &Exit,
+) -> Result<()> {
+    if starting_height > last_height {
+        return Ok(());
+    }
+
+    info!(
+        "Processing blocks {} to {}...",
+        starting_height, last_height
+    );
+
+    // References to vectors using correct field paths
+    // From indexer.vecs:
+    let height_to_first_txindex = &indexer.vecs.tx.height_to_first_txindex;
+    let height_to_first_txoutindex = &indexer.vecs.txout.height_to_first_txoutindex;
+    let height_to_first_txinindex = &indexer.vecs.txin.height_to_first_txinindex;
+
+    // From chain (via .height.u() or .height.unwrap_sum() patterns):
+    let height_to_tx_count = chain.indexes_to_tx_count.height.u();
+    let height_to_output_count = chain.indexes_to_output_count.height.unwrap_sum();
+    let height_to_input_count = chain.indexes_to_input_count.height.unwrap_sum();
+    let height_to_unclaimed_rewards = chain
+        .indexes_to_unclaimed_rewards
+        .sats
+        .height
+        .as_ref()
+        .unwrap();
+
+    // From indexes:
+    let height_to_timestamp = &indexes.height_to_timestamp_fixed;
+    let height_to_date = &indexes.height_to_date_fixed;
+    let dateindex_to_first_height = &indexes.dateindex_to_first_height;
+    let dateindex_to_height_count = &indexes.dateindex_to_height_count;
+    let txindex_to_output_count = &indexes.txindex_to_output_count;
+    let txindex_to_input_count = &indexes.txindex_to_input_count;
+
+    // From price (optional):
+    let height_to_price = price.map(|p| &p.chainindexes_to_price_close.height);
+    let dateindex_to_price = price.map(|p| p.timeindexes_to_price_close.dateindex.u());
+
+    // Collect price and timestamp vectors for process_sent (needs slice access, not iterators)
+    // These are used in the spawned thread and need to be separate from chain_state
+    let height_to_price_vec: Option<Vec<Dollars>> =
+        height_to_price.map(|v| v.into_iter().map(|d| *d).collect());
+    let height_to_timestamp_vec: Vec<Timestamp> = height_to_timestamp.into_iter().collect();
+
+    // Create iterators for sequential access
+    let mut height_to_first_txindex_iter = height_to_first_txindex.into_iter();
+    let mut height_to_first_txoutindex_iter = height_to_first_txoutindex.into_iter();
+    let mut height_to_first_txinindex_iter = height_to_first_txinindex.into_iter();
+    let mut height_to_tx_count_iter = height_to_tx_count.into_iter();
+    let mut height_to_output_count_iter = height_to_output_count.into_iter();
+    let mut height_to_input_count_iter = height_to_input_count.into_iter();
+    let mut height_to_unclaimed_rewards_iter = height_to_unclaimed_rewards.into_iter();
+    let mut height_to_timestamp_iter = height_to_timestamp.into_iter();
+    let mut height_to_date_iter = height_to_date.into_iter();
+    let mut dateindex_to_first_height_iter = dateindex_to_first_height.into_iter();
+    let mut dateindex_to_height_count_iter = dateindex_to_height_count.into_iter();
+    let mut txindex_to_output_count_iter = txindex_to_output_count.iter();
+    let mut txindex_to_input_count_iter = txindex_to_input_count.iter();
+    let mut height_to_price_iter = height_to_price.map(|v| v.into_iter());
+    let mut dateindex_to_price_iter = dateindex_to_price.map(|v| v.into_iter());
+
+    // Build txoutindex -> height map for input processing
+    let txoutindex_to_height = build_txoutindex_to_height_map(height_to_first_txoutindex);
+
+    // Create readers for parallel data access
+    let ir = IndexerReaders::new(indexer);
+    let mut vr = VecsReaders::new(&vecs.any_address_indexes, &vecs.addresses_data);
+
+    // Create iterators for first address indexes per type
+    let mut first_p2a_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2aaddressindex
+        .into_iter();
+    let mut first_p2pk33_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2pk33addressindex
+        .into_iter();
+    let mut first_p2pk65_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2pk65addressindex
+        .into_iter();
+    let mut first_p2pkh_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2pkhaddressindex
+        .into_iter();
+    let mut first_p2sh_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2shaddressindex
+        .into_iter();
+    let mut first_p2tr_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2traddressindex
+        .into_iter();
+    let mut first_p2wpkh_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2wpkhaddressindex
+        .into_iter();
+    let mut first_p2wsh_iter = indexer
+        .vecs
+        .address
+        .height_to_first_p2wshaddressindex
+        .into_iter();
+
+    // Track running totals - recover from previous height if resuming
+    let (
+        mut unspendable_supply,
+        mut opreturn_supply,
+        mut addresstype_to_addr_count,
+        mut addresstype_to_empty_addr_count,
+    ) = if starting_height > Height::ZERO {
+        let prev_height = starting_height.decremented().unwrap();
+        (
+            vecs.height_to_unspendable_supply
+                .into_iter()
+                .get_unwrap(prev_height),
+            vecs.height_to_opreturn_supply
+                .into_iter()
+                .get_unwrap(prev_height),
+            AddressTypeToAddressCount::from((
+                &vecs.addresstype_to_height_to_addr_count,
+                starting_height,
+            )),
+            AddressTypeToAddressCount::from((
+                &vecs.addresstype_to_height_to_empty_addr_count,
+                starting_height,
+            )),
+        )
+    } else {
+        (
+            Sats::ZERO,
+            Sats::ZERO,
+            AddressTypeToAddressCount::default(),
+            AddressTypeToAddressCount::default(),
+        )
+    };
+
+    // Persistent address data caches (accumulate across blocks, flushed at checkpoints)
+    let mut loaded_cache: AddressTypeToTypeIndexMap<LoadedAddressDataWithSource> =
+        AddressTypeToTypeIndexMap::default();
+    let mut empty_cache: AddressTypeToTypeIndexMap<EmptyAddressDataWithSource> =
+        AddressTypeToTypeIndexMap::default();
+
+    // Main block iteration
+    for height in starting_height.to_usize()..=last_height.to_usize() {
+        let height = Height::from(height);
+
+        if height.to_usize() % 10000 == 0 {
+            info!("Processing chain at {}...", height);
+        }
+
+        // Get block metadata
+        let first_txindex = height_to_first_txindex_iter.get_unwrap(height);
+        let tx_count = u64::from(height_to_tx_count_iter.get_unwrap(height));
+        let first_txoutindex = height_to_first_txoutindex_iter
+            .get_unwrap(height)
+            .to_usize();
+        let output_count = u64::from(height_to_output_count_iter.get_unwrap(height)) as usize;
+        let first_txinindex = height_to_first_txinindex_iter.get_unwrap(height).to_usize();
+        let input_count = u64::from(height_to_input_count_iter.get_unwrap(height)) as usize;
+        let timestamp = height_to_timestamp_iter.get_unwrap(height);
+        let block_price = height_to_price_iter.as_mut().map(|v| *v.get_unwrap(height));
+
+        // Build txindex mappings for this block
+        let txoutindex_to_txindex =
+            build_txoutindex_to_txindex(first_txindex, tx_count, &mut txindex_to_output_count_iter);
+        let txinindex_to_txindex =
+            build_txinindex_to_txindex(first_txindex, tx_count, &mut txindex_to_input_count_iter);
+
+        // Get first address indexes for this height
+        let first_addressindexes = ByAddressType {
+            p2a: TypeIndex::from(first_p2a_iter.get_unwrap(height).to_usize()),
+            p2pk33: TypeIndex::from(first_p2pk33_iter.get_unwrap(height).to_usize()),
+            p2pk65: TypeIndex::from(first_p2pk65_iter.get_unwrap(height).to_usize()),
+            p2pkh: TypeIndex::from(first_p2pkh_iter.get_unwrap(height).to_usize()),
+            p2sh: TypeIndex::from(first_p2sh_iter.get_unwrap(height).to_usize()),
+            p2tr: TypeIndex::from(first_p2tr_iter.get_unwrap(height).to_usize()),
+            p2wpkh: TypeIndex::from(first_p2wpkh_iter.get_unwrap(height).to_usize()),
+            p2wsh: TypeIndex::from(first_p2wsh_iter.get_unwrap(height).to_usize()),
+        };
+
+        // Reset per-block values for all separate cohorts
+        reset_block_values(&mut vecs.utxo_cohorts, &mut vecs.address_cohorts);
+
+        // Process outputs and inputs in parallel with tick-tock
+        let (outputs_result, inputs_result) = thread::scope(|scope| {
+            // Tick-tock age transitions in background
+            scope.spawn(|| {
+                vecs.utxo_cohorts
+                    .tick_tock_next_block(chain_state, timestamp);
+            });
+
+            // Process outputs (receive)
+            let outputs_result = process_outputs(
+                first_txoutindex,
+                output_count,
+                &txoutindex_to_txindex,
+                &indexer.vecs.txout.txoutindex_to_value,
+                &indexer.vecs.txout.txoutindex_to_outputtype,
+                &indexer.vecs.txout.txoutindex_to_typeindex,
+                &ir,
+                &first_addressindexes,
+                &loaded_cache,
+                &empty_cache,
+                &vr,
+                &vecs.any_address_indexes,
+                &vecs.addresses_data,
+            );
+
+            // Process inputs (send) - skip coinbase input
+            let inputs_result = if input_count > 1 {
+                process_inputs(
+                    first_txinindex + 1, // Skip coinbase
+                    input_count - 1,
+                    &txinindex_to_txindex[1..], // Skip coinbase
+                    &indexer.vecs.txin.txinindex_to_outpoint,
+                    &indexer.vecs.tx.txindex_to_first_txoutindex,
+                    &indexer.vecs.txout.txoutindex_to_value,
+                    &indexer.vecs.txout.txoutindex_to_outputtype,
+                    &indexer.vecs.txout.txoutindex_to_typeindex,
+                    &txoutindex_to_height,
+                    &ir,
+                    &first_addressindexes,
+                    &loaded_cache,
+                    &empty_cache,
+                    &vr,
+                    &vecs.any_address_indexes,
+                    &vecs.addresses_data,
+                )
+            } else {
+                InputsResult {
+                    height_to_sent: Default::default(),
+                    sent_data: Default::default(),
+                    address_data: Default::default(),
+                    txindex_vecs: Default::default(),
+                    txoutindex_to_txinindex_updates: Default::default(),
+                }
+            };
+
+            (outputs_result, inputs_result)
+        });
+
+        // Merge new address data into caches
+        loaded_cache.merge_mut(outputs_result.address_data);
+        loaded_cache.merge_mut(inputs_result.address_data);
+
+        // Combine txindex_vecs from outputs and inputs, then update tx_count
+        let combined_txindex_vecs = outputs_result
+            .txindex_vecs
+            .merge_vec(inputs_result.txindex_vecs);
+        update_tx_counts(&mut loaded_cache, &mut empty_cache, combined_txindex_vecs);
+
+        let mut transacted = outputs_result.transacted;
+        let mut height_to_sent = inputs_result.height_to_sent;
+
+        // Update supply tracking
+        unspendable_supply += transacted.by_type.unspendable.opreturn.value
+            + height_to_unclaimed_rewards_iter.get_unwrap(height);
+        opreturn_supply += transacted.by_type.unspendable.opreturn.value;
+
+        // Handle special cases
+        if height == Height::ZERO {
+            // Genesis block - reset transacted, add 50 BTC to unspendable
+            transacted = Transacted::default();
+            unspendable_supply += Sats::FIFTY_BTC;
+        } else if height == Height::new(BIP30_DUPLICATE_HEIGHT_1)
+            || height == Height::new(BIP30_DUPLICATE_HEIGHT_2)
+        {
+            // BIP30: Add 50 BTC to spent from original height
+            let original_height = if height == Height::new(BIP30_DUPLICATE_HEIGHT_1) {
+                Height::new(BIP30_ORIGINAL_HEIGHT_1)
+            } else {
+                Height::new(BIP30_ORIGINAL_HEIGHT_2)
+            };
+            height_to_sent
+                .entry(original_height)
+                .or_default()
+                .iterate(Sats::FIFTY_BTC, OutputType::P2PK65);
+        }
+
+        // Push current block state before processing cohort updates
+        chain_state.push(BlockState {
+            supply: transacted.spendable_supply.clone(),
+            price: block_price,
+            timestamp,
+        });
+
+        // Process UTXO cohorts and Address cohorts in parallel
+        // - Main thread: UTXO cohorts receive/send
+        // - Spawned thread: Address cohorts process_received/process_sent
+        thread::scope(|scope| {
+            // Spawn address cohort processing in background thread
+            scope.spawn(|| {
+                // Create lookup closure that returns None (data was pre-fetched in parallel phase)
+                let get_address_data =
+                    |_output_type, _type_index| -> Option<LoadedAddressDataWithSource> { None };
+
+                let mut lookup = AddressLookup {
+                    get_address_data,
+                    loaded: &mut loaded_cache,
+                    empty: &mut empty_cache,
+                };
+
+                // Process received outputs (addresses receiving funds)
+                process_received(
+                    outputs_result.received_data,
+                    &mut vecs.address_cohorts,
+                    &mut lookup,
+                    block_price,
+                    &mut addresstype_to_addr_count,
+                    &mut addresstype_to_empty_addr_count,
+                );
+
+                // Process sent inputs (addresses sending funds)
+                // Uses separate price/timestamp vecs to avoid borrowing chain_state
+                process_sent(
+                    inputs_result.sent_data,
+                    &mut vecs.address_cohorts,
+                    &mut lookup,
+                    block_price,
+                    &mut addresstype_to_addr_count,
+                    &mut addresstype_to_empty_addr_count,
+                    height_to_price_vec.as_deref(),
+                    &height_to_timestamp_vec,
+                    height,
+                    timestamp,
+                )
+                .unwrap();
+            });
+
+            // Main thread: Update UTXO cohorts
+            vecs.utxo_cohorts.receive(transacted, height, block_price);
+            vecs.utxo_cohorts.send(height_to_sent, chain_state);
+        });
+
+        // Update txoutindex_to_txinindex
+        vecs.update_txoutindex_to_txinindex(
+            output_count,
+            inputs_result.txoutindex_to_txinindex_updates,
+        )?;
+
+        // Push to height-indexed vectors
+        vecs.height_to_unspendable_supply
+            .truncate_push(height, unspendable_supply)?;
+        vecs.height_to_opreturn_supply
+            .truncate_push(height, opreturn_supply)?;
+        vecs.addresstype_to_height_to_addr_count
+            .truncate_push(height, &addresstype_to_addr_count)?;
+        vecs.addresstype_to_height_to_empty_addr_count
+            .truncate_push(height, &addresstype_to_empty_addr_count)?;
+
+        // Get date info for unrealized state computation
+        let date = height_to_date_iter.get_unwrap(height);
+        let dateindex = DateIndex::try_from(date).unwrap();
+        let date_first_height = dateindex_to_first_height_iter.get_unwrap(dateindex);
+        let date_height_count = dateindex_to_height_count_iter.get_unwrap(dateindex);
+        let is_date_last_height =
+            date_first_height + Height::from(date_height_count).decremented().unwrap() == height;
+        let date_price = dateindex_to_price_iter
+            .as_mut()
+            .map(|v| is_date_last_height.then(|| *v.get_unwrap(dateindex)));
+        let dateindex_opt = is_date_last_height.then_some(dateindex);
+
+        // Push cohort states and compute unrealized
+        push_cohort_states(
+            &mut vecs.utxo_cohorts,
+            &mut vecs.address_cohorts,
+            height,
+            block_price,
+            dateindex_opt,
+            date_price,
+        )?;
+
+        // Compute and push percentiles for aggregate cohorts (all, sth, lth)
+        vecs.utxo_cohorts
+            .truncate_push_aggregate_percentiles(height)?;
+
+        // Periodic checkpoint flush
+        if height != last_height
+            && height != Height::ZERO
+            && height.to_usize() % FLUSH_INTERVAL == 0
+        {
+            let _lock = exit.lock();
+
+            // Drop readers before flush to release mmap handles
+            drop(vr);
+
+            flush_checkpoint(vecs, height, &mut loaded_cache, &mut empty_cache, exit)?;
+
+            // Recreate readers after flush to pick up new data
+            vr = VecsReaders::new(&vecs.any_address_indexes, &vecs.addresses_data);
+        }
+    }
+
+    // Final flush
+    let _lock = exit.lock();
+    drop(vr);
+    flush_checkpoint(vecs, last_height, &mut loaded_cache, &mut empty_cache, exit)?;
+
+    Ok(())
+}
+
+/// Reset per-block values for all separate cohorts.
+fn reset_block_values(utxo_cohorts: &mut UTXOCohorts, address_cohorts: &mut AddressCohorts) {
+    utxo_cohorts.par_iter_separate_mut().for_each(|v| {
+        if let Some(state) = v.state.as_mut() {
+            state.reset_single_iteration_values();
+        }
+    });
+
+    address_cohorts.par_iter_separate_mut().for_each(|v| {
+        if let Some(state) = v.state.as_mut() {
+            state.inner.reset_single_iteration_values();
+        }
+    });
+}
+
+/// Push cohort states to height-indexed vectors.
+fn push_cohort_states(
+    utxo_cohorts: &mut UTXOCohorts,
+    address_cohorts: &mut AddressCohorts,
+    height: Height,
+    height_price: Option<brk_types::Dollars>,
+    dateindex: Option<DateIndex>,
+    date_price: Option<Option<brk_types::Dollars>>,
+) -> Result<()> {
+    utxo_cohorts
+        .par_iter_separate_mut()
+        .map(|v| v as &mut dyn DynCohortVecs)
+        .chain(
+            address_cohorts
+                .par_iter_separate_mut()
+                .map(|v| v as &mut dyn DynCohortVecs),
+        )
+        .try_for_each(|v| {
+            v.truncate_push(height)?;
+            v.compute_then_truncate_push_unrealized_states(
+                height,
+                height_price,
+                dateindex,
+                date_price,
+            )
+        })?;
+
+    Ok(())
+}
+
+/// Flush checkpoint to disk.
+///
+/// Flushes all accumulated data including:
+/// - Cohort stateful vectors
+/// - Height-indexed vectors
+/// - Address data caches (loaded and empty)
+/// - Chain state
+#[allow(clippy::too_many_arguments)]
+fn flush_checkpoint(
+    vecs: &mut Vecs,
+    height: Height,
+    loaded_cache: &mut AddressTypeToTypeIndexMap<LoadedAddressDataWithSource>,
+    empty_cache: &mut AddressTypeToTypeIndexMap<EmptyAddressDataWithSource>,
+    exit: &Exit,
+) -> Result<()> {
+    info!("Flushing checkpoint at height {}...", height);
+
+    // Flush cohort states
+    vecs.utxo_cohorts.safe_flush_stateful_vecs(height, exit)?;
+    vecs.address_cohorts
+        .safe_flush_stateful_vecs(height, exit)?;
+
+    // Flush height-indexed vectors
+    vecs.height_to_unspendable_supply.safe_write(exit)?;
+    vecs.height_to_opreturn_supply.safe_write(exit)?;
+    vecs.addresstype_to_height_to_addr_count.safe_flush(exit)?;
+    vecs.addresstype_to_height_to_empty_addr_count
+        .safe_flush(exit)?;
+
+    // Process and flush address data updates
+    let empty_updates = mem::take(empty_cache);
+    let loaded_updates = mem::take(loaded_cache);
+    flush_checkpoint_full(
+        height,
+        &mut vecs
+            .utxo_cohorts
+            .par_iter_separate_mut()
+            .map(|v| v as &mut dyn DynCohortVecs)
+            .collect::<Vec<_>>()[..],
+        &mut vecs
+            .address_cohorts
+            .par_iter_separate_mut()
+            .map(|v| v as &mut dyn DynCohortVecs)
+            .collect::<Vec<_>>()[..],
+        &mut vecs.any_address_indexes,
+        &mut vecs.addresses_data,
+        empty_updates,
+        loaded_updates,
+        true,
+        exit,
+    )?;
+
+    // Flush chain state and txoutindex_to_txinindex with stamp
+    vecs.txoutindex_to_txinindex
+        .stamped_flush_with_changes(height.into())?;
+    vecs.chain_state.stamped_flush_with_changes(height.into())?;
+
+    Ok(())
+}