use std::sync::Arc;
use brk_types::{Cents, Height, StoredU64, Version};
use vecdb::{
AnyVec, CachedBoxedVec, PrintableIndex, ReadableBoxedVec, ReadableVec, TypedVec, VecIndex,
short_type_name,
};
#[derive(Clone)]
pub(super) struct LazySmaVec {
name: Arc<str>,
version: Version,
window_starts: ReadableBoxedVec<Height, Height>,
prefix_sum: CachedBoxedVec<Height, StoredU64>,
}
impl LazySmaVec {
pub(super) fn new(
name: &str,
version: Version,
window_starts: ReadableBoxedVec<Height, Height>,
prefix_sum: CachedBoxedVec<Height, StoredU64>,
) -> Self {
Self {
name: Arc::from(name),
version,
window_starts,
prefix_sum,
}
}
fn average(prefix_sum: &[StoredU64], index: usize, start: Height) -> Cents {
let start = start.to_usize();
debug_assert!(start <= index, "price SMA window starts after its height");
let current = u64::from(prefix_sum[index]);
let previous = start
.checked_sub(1)
.map(|index| u64::from(prefix_sum[index]))
.unwrap_or_default();
let count = index - start + 1;
Cents::new((current - previous) / count as u64)
}
fn try_for_each_value<E>(
&self,
from: usize,
to: usize,
mut each: impl FnMut(Cents) -> Result<(), E>,
) -> Result<(), E> {
let prefix_sum = self.prefix_sum.cached();
let to = to.min(self.window_starts.len()).min(prefix_sum.len());
if from >= to {
return Ok(());
}
for (offset, start) in self
.window_starts
.collect_range_dyn(from, to)
.into_iter()
.enumerate()
{
each(Self::average(&prefix_sum, from + offset, start))?;
}
Ok(())
}
fn for_each_value(&self, from: usize, to: usize, mut each: impl FnMut(Cents)) {
let prefix_sum = self.prefix_sum.cached();
let to = to.min(self.window_starts.len()).min(prefix_sum.len());
if from >= to {
return;
}
let mut index = from;
self.window_starts
.for_each_range_dyn_at(from, to, &mut |start| {
each(Self::average(&prefix_sum, index, start));
index += 1;
});
}
}
impl AnyVec for LazySmaVec {
fn version(&self) -> Version {
self.version + self.window_starts.version() + self.prefix_sum.version()
}
fn name(&self) -> &str {
&self.name
}
fn len(&self) -> usize {
self.window_starts.len().min(self.prefix_sum.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::<Cents>()
}
fn value_type_to_string(&self) -> &'static str {
short_type_name::<Cents>()
}
}
impl TypedVec for LazySmaVec {
type I = Height;
type T = Cents;
}
impl ReadableVec<Height, Cents> for LazySmaVec {
fn read_into_at(&self, from: usize, to: usize, buf: &mut Vec<Cents>) {
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(Cents)) {
self.for_each_value(from, to, each);
}
fn fold_range_at<B, F: FnMut(B, Cents) -> 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, Cents) -> 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<Cents> {
if index >= self.len() {
return None;
}
let start = self.window_starts.collect_one_at(index)?;
Some(Self::average(&self.prefix_sum.cached(), index, start))
}
fn read_sorted_into_at(&self, indices: &[usize], out: &mut Vec<Cents>) {
let len = self.len();
let indices: Vec<_> = indices
.iter()
.copied()
.filter(|index| *index < len)
.collect();
let starts = self.window_starts.read_sorted_at(&indices);
let prefix_sum = self.prefix_sum.cached();
out.reserve(indices.len());
indices
.into_iter()
.zip(starts)
.for_each(|(index, start)| out.push(Self::average(&prefix_sum, index, start)));
}
}
#[cfg(test)]
mod tests {
use vecdb::{
AnyStoredVec, CachedReadableVec, CachedVec, Database, EagerVec, ImportableVec, PcoVec,
ReadableCloneableVec, WritableVec,
};
use super::*;
use crate::market::moving_average::cached_sma_source::PricePrefixSumVec;
#[test]
fn computes_rolling_average_from_one_shared_prefix_cache() {
let suffix = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos();
let path =
std::env::temp_dir().join(format!("brk-lazy-sma-{}-{suffix}", std::process::id()));
let db = Database::open(&path).unwrap();
let mut prices: EagerVec<PcoVec<Height, Cents>> =
EagerVec::forced_import(&db, "prices", Version::ONE).unwrap();
let mut starts: EagerVec<PcoVec<Height, Height>> =
EagerVec::forced_import(&db, "starts", Version::ONE).unwrap();
for value in [100, 200, 300, 400] {
prices.push(Cents::new(value));
}
for value in [0, 0, 1, 2] {
starts.push(Height::new(value));
}
prices.write().unwrap();
starts.write().unwrap();
let prices = CachedVec::wrap(prices);
let prefix_sum = CachedVec::wrap(PricePrefixSumVec::new(
"prefix",
Version::ONE,
prices.read_only_cached_boxed_clone(),
));
let sma = LazySmaVec::new(
"sma",
Version::ONE,
starts.read_only_boxed_clone(),
prefix_sum.cached_boxed_clone(),
);
assert_eq!(
prefix_sum.cached().as_ref(),
[
StoredU64::new(100),
StoredU64::new(300),
StoredU64::new(600),
StoredU64::new(1_000),
],
);
assert_eq!(
sma.collect_range(Height::ZERO, Height::new(4)),
[
Cents::new(100),
Cents::new(150),
Cents::new(250),
Cents::new(350),
],
);
}
}