use std::sync::Arc;
use brk_types::{Height, PoolSlug, StoredU64};
use parking_lot::RwLock;
use rustc_hash::FxHashMap;
use vecdb::{
AnyVec, BytesVec, PrintableIndex, ReadOnlyClone, ReadableVec, TypedVec, VecIndex, Version,
short_type_name,
};
#[derive(Default)]
struct State {
by_pool: FxHashMap<PoolSlug, Vec<Height>>,
len: usize,
version: Version,
}
#[derive(Clone, Default)]
pub struct PoolHeights(Arc<RwLock<State>>);
impl PoolHeights {
pub fn build(pool: &BytesVec<Height, PoolSlug>) -> Self {
let len = pool.len();
let mut by_pool: FxHashMap<PoolSlug, Vec<Height>> = FxHashMap::default();
let reader = pool.reader();
for h in 0..len {
by_pool
.entry(reader.get_at(h))
.or_default()
.push(Height::from(h));
}
Self(Arc::new(RwLock::new(State {
by_pool,
len,
version: pool.version(),
})))
}
pub fn truncate(&self, min: usize) {
let mut state = self.0.write();
for heights in state.by_pool.values_mut() {
let cut = heights.partition_point(|h| h.to_usize() < min);
heights.truncate(cut);
}
state.len = min;
}
pub fn push(&self, slug: PoolSlug, height: Height) {
let mut state = self.0.write();
debug_assert_eq!(height.to_usize(), state.len);
state.by_pool.entry(slug).or_default().push(height);
state.len = height.to_usize() + 1;
}
pub fn block_numbers(&self, slugs: &[PoolSlug], first_height: Height) -> Vec<u64> {
let state = self.0.read();
let first_height = first_height.to_usize();
slugs
.iter()
.enumerate()
.map(|(offset, slug)| {
state.by_pool.get(slug).map_or(0, |heights| {
Self::cumulative_count(heights, first_height + offset) as u64
})
})
.collect()
}
pub fn latest_heights(
&self,
slug: PoolSlug,
through_height: Height,
limit: usize,
) -> Vec<Height> {
let state = self.0.read();
let Some(heights) = state.by_pool.get(&slug) else {
return Vec::new();
};
let end = Self::cumulative_count(heights, through_height.to_usize());
let start = end.saturating_sub(limit);
heights[start..end].iter().rev().copied().collect()
}
pub fn latest_height(&self, slug: PoolSlug, through_height: Height) -> Option<Height> {
let state = self.0.read();
let heights = state.by_pool.get(&slug)?;
let end = Self::cumulative_count(heights, through_height.to_usize());
end.checked_sub(1).map(|index| heights[index])
}
fn cumulative_count(heights: &[Height], through_height: usize) -> usize {
heights.partition_point(|height| height.to_usize() <= through_height)
}
fn len(&self) -> usize {
self.0.read().len
}
fn version(&self) -> Version {
self.0.read().version
}
}
#[derive(Clone)]
pub(super) struct PoolCumulativeVec {
name: Arc<str>,
slug: PoolSlug,
pool_heights: PoolHeights,
}
impl PoolCumulativeVec {
pub(super) fn new(name: &str, slug: PoolSlug, pool_heights: PoolHeights) -> Self {
Self {
name: Arc::from(name),
slug,
pool_heights,
}
}
fn for_each_value(&self, from: usize, to: usize, mut each: impl FnMut(StoredU64)) {
let result = self.try_for_each_value(from, to, |value| {
each(value);
Ok::<_, std::convert::Infallible>(())
});
match result {
Ok(()) => {}
Err(error) => match error {},
}
}
fn try_for_each_value<E>(
&self,
from: usize,
to: usize,
mut each: impl FnMut(StoredU64) -> Result<(), E>,
) -> Result<(), E> {
let to = to.min(self.len());
if from >= to {
return Ok(());
}
let state = self.pool_heights.0.read();
let heights = state.by_pool.get(&self.slug).map_or(&[][..], Vec::as_slice);
let mut position = heights.partition_point(|height| height.to_usize() < from);
let mut cumulative = position as u64;
for height in from..to {
while heights
.get(position)
.is_some_and(|pool_height| pool_height.to_usize() == height)
{
position += 1;
cumulative += 1;
}
each(StoredU64::from(cumulative))?;
}
Ok(())
}
}
impl AnyVec for PoolCumulativeVec {
fn version(&self) -> Version {
self.pool_heights.version()
}
fn name(&self) -> &str {
&self.name
}
fn len(&self) -> usize {
self.pool_heights.len()
}
fn index_type_to_string(&self) -> &'static str {
<Height as PrintableIndex>::to_string()
}
fn region_names(&self) -> Vec<String> {
Vec::new()
}
fn value_type_to_size_of(&self) -> usize {
size_of::<StoredU64>()
}
fn value_type_to_string(&self) -> &'static str {
short_type_name::<StoredU64>()
}
}
impl TypedVec for PoolCumulativeVec {
type I = Height;
type T = StoredU64;
}
impl ReadableVec<Height, StoredU64> for PoolCumulativeVec {
fn read_into_at(&self, from: usize, to: usize, buf: &mut Vec<StoredU64>) {
buf.reserve(to.min(self.len()).saturating_sub(from));
self.for_each_value(from, to, |value| buf.push(value));
}
fn for_each_range_dyn_at(&self, from: usize, to: usize, each: &mut dyn FnMut(StoredU64)) {
self.for_each_value(from, to, each);
}
fn fold_range_at<B, F: FnMut(B, StoredU64) -> B>(
&self,
from: usize,
to: usize,
init: B,
mut fold: F,
) -> B {
let mut acc = Some(init);
self.for_each_value(from, to, |value| {
acc = Some(fold(acc.take().unwrap(), value));
});
acc.unwrap()
}
fn try_fold_range_at<B, E, F: FnMut(B, StoredU64) -> Result<B, E>>(
&self,
from: usize,
to: usize,
init: B,
mut fold: F,
) -> Result<B, E> {
let mut acc = Some(init);
self.try_for_each_value(from, to, |value| {
acc = Some(fold(acc.take().unwrap(), value)?);
Ok(())
})?;
Ok(acc.unwrap())
}
fn collect_one_at(&self, index: usize) -> Option<StoredU64> {
if index >= self.len() {
return None;
}
let state = self.pool_heights.0.read();
let heights = state.by_pool.get(&self.slug).map_or(&[][..], Vec::as_slice);
Some(StoredU64::from(
PoolHeights::cumulative_count(heights, index) as u64,
))
}
fn read_sorted_into_at(&self, indices: &[usize], out: &mut Vec<StoredU64>) {
let Some(&first) = indices.first() else {
return;
};
let len = self.len();
if first >= len {
return;
}
let state = self.pool_heights.0.read();
let heights = state.by_pool.get(&self.slug).map_or(&[][..], Vec::as_slice);
let mut position = PoolHeights::cumulative_count(heights, first);
out.reserve(indices.len());
out.push(StoredU64::from(position as u64));
for &index in &indices[1..] {
if index >= len {
break;
}
while heights
.get(position)
.is_some_and(|height| height.to_usize() <= index)
{
position += 1;
}
out.push(StoredU64::from(position as u64));
}
}
}
impl ReadOnlyClone for PoolCumulativeVec {
type ReadOnly = Self;
fn read_only_clone(&self) -> Self {
self.clone()
}
}
#[cfg(test)]
mod tests {
use vecdb::ReadableVec;
use super::*;
fn fixture() -> (PoolHeights, PoolCumulativeVec) {
let pool_heights = PoolHeights::default();
for (height, slug) in [
PoolSlug::F2Pool,
PoolSlug::Unknown,
PoolSlug::F2Pool,
PoolSlug::F2Pool,
PoolSlug::Unknown,
]
.into_iter()
.enumerate()
{
pool_heights.push(slug, Height::from(height));
}
let cumulative = PoolCumulativeVec::new(
"f2pool_blocks_mined_cumulative",
PoolSlug::F2Pool,
pool_heights.clone(),
);
(pool_heights, cumulative)
}
#[test]
fn cumulative_counts_include_the_current_height() {
let (_, cumulative) = fixture();
assert_eq!(
cumulative.collect(),
[1_u64, 1, 2, 3, 3].map(StoredU64::from)
);
assert_eq!(
cumulative.collect_range_at(1, 4),
[1_u64, 2, 3].map(StoredU64::from)
);
assert_eq!(cumulative.collect_one_at(4), Some(StoredU64::from(3_u64)));
assert_eq!(cumulative.collect_one_at(5), None);
}
#[test]
fn sorted_and_appending_reads_use_cumulative_ranks() {
let (_, cumulative) = fixture();
let mut sorted = vec![StoredU64::from(99_u64)];
cumulative.read_sorted_into_at(&[0, 2, 4], &mut sorted);
assert_eq!(sorted, [99_u64, 1, 2, 3].map(StoredU64::from));
let mut range = vec![StoredU64::from(99_u64)];
cumulative.read_into_at(2, 5, &mut range);
assert_eq!(range, [99_u64, 2, 3, 3].map(StoredU64::from));
}
#[test]
fn truncate_and_push_keep_length_and_counts_in_sync() {
let (pool_heights, cumulative) = fixture();
pool_heights.truncate(3);
assert_eq!(cumulative.collect(), [1_u64, 1, 2].map(StoredU64::from));
pool_heights.push(PoolSlug::Unknown, Height::from(3_u32));
pool_heights.push(PoolSlug::F2Pool, Height::from(4_u32));
assert_eq!(
cumulative.collect(),
[1_u64, 1, 2, 2, 3].map(StoredU64::from)
);
}
#[test]
fn absent_pool_is_zero_for_every_height() {
let (pool_heights, _) = fixture();
let cumulative = PoolCumulativeVec::new("absent", PoolSlug::Luxor, pool_heights);
assert_eq!(cumulative.collect(), [StoredU64::from(0_u64); 5]);
}
#[test]
fn bulk_block_numbers_use_one_height_per_slug() {
let (pool_heights, _) = fixture();
assert_eq!(
pool_heights.block_numbers(
&[PoolSlug::F2Pool, PoolSlug::Unknown, PoolSlug::F2Pool],
Height::from(0_u32),
),
[1, 1, 2]
);
}
#[test]
fn latest_heights_are_limited_and_descending() {
let (pool_heights, _) = fixture();
assert_eq!(
pool_heights.latest_heights(PoolSlug::F2Pool, Height::from(4_u32), 2),
[Height::from(3_u32), Height::from(2_u32)]
);
assert_eq!(
pool_heights.latest_heights(PoolSlug::F2Pool, Height::from(2_u32), 10),
[Height::from(2_u32), Height::from(0_u32)]
);
}
#[test]
fn latest_height_is_bounded_by_height() {
let (pool_heights, _) = fixture();
assert_eq!(
pool_heights.latest_height(PoolSlug::F2Pool, Height::from(4_u32)),
Some(Height::from(3_u32))
);
assert_eq!(
pool_heights.latest_height(PoolSlug::F2Pool, Height::from(1_u32)),
Some(Height::from(0_u32))
);
assert_eq!(
pool_heights.latest_height(PoolSlug::Luxor, Height::from(4_u32)),
None
);
}
}