use crate::ring_buffer::multi_buffer::multi_buffer::{BufferElement, MultiBuffer};
use crate::ring_buffer::single_buffer::generic_buffer::Buffer as SingleBuffer;
use crate::indicators::simd_indicators::{
max_simd::{SimdState as MaxState, find_max_scalar}, min_simd::{SimdState as MinState, find_min_scalar}, simd_types::UsizeConstants,
};
use crate::indicators::{
max::{find_max_simd},
min::{find_min_simd},
};
use std::simd::{
cmp::{SimdPartialEq, SimdPartialOrd},
Select, Simd
};
pub trait MirrorBuffer<const B: usize, T: BufferElement = f64> {
fn new(capacity: usize) -> Self;
unsafe fn push_unchecked(&mut self, values: [T; B]);
fn push(&mut self, values: [T; B]);
fn push_with_info(&mut self, values: [T; B]) -> Option<[T; B]>;
unsafe fn push_with_info_unchecked(&mut self, values: [T; B]) -> [T; B];
fn get_slice(&self, lane: usize, offset: usize) -> &[T];
fn get_slices(&self, offset: usize) -> [&[T]; B];
fn window_index_to_bars_ago(&self, idx: usize) -> usize;
fn window_index_to_bars_ago_simd(&self, window_index: Simd<usize, B>) -> Simd<usize, B>;
fn from_slice(vals: [&[T]; B], capacity: usize) -> Self;
fn to_single_buffers(&self) -> [SingleBuffer<T>; B];
}
impl<const B: usize, T: BufferElement> MirrorBuffer<B, T> for MultiBuffer<B, T> {
fn new(capacity: usize) -> Self {
Self {
vals: core::array::from_fn(|_| vec![T::default(); capacity * 2]),
index: 0,
prev_idx: 0,
capacity,
count: 0,
}
}
fn from_slice(vals: [&[T]; B], capacity: usize) -> Self {
let count = vals[0].len().min(capacity);
let buffer_vals: [Vec<T>; B] = core::array::from_fn(|lane| {
let mut vec = vals[lane].to_vec();
if count < capacity {
vec.resize(capacity, T::default());
}
vec.extend_from_within(..);
vec
});
let index = count % capacity;
Self {
vals: buffer_vals,
index: index,
prev_idx: index.wrapping_sub(1) % capacity,
capacity,
count,
}
}
#[inline(always)]
fn push(&mut self, values: [T; B]) {
write_values(self, values);
self.update_internals();
}
#[inline(always)]
unsafe fn push_unchecked(&mut self, values: [T; B]) {
write_values(self, values);
self.update_internals_unchecked();
}
#[inline(always)]
fn push_with_info(&mut self, values: [T; B]) -> Option<[T; B]> {
if self.count == self.capacity {
let replaced = write_values_pop(self, values);
return Some(replaced);
}
write_values(self, values);
self.update_internals();
None
}
#[inline(always)]
unsafe fn push_with_info_unchecked(&mut self, values: [T; B]) -> [T; B] {
write_values_pop(self, values)
}
#[inline(always)]
fn get_slice(&self, lane: usize, offset: usize) -> &[T] {
debug_assert!(
lane < B,
"Lane index {} out of bounds for buffer with {} lanes",
lane,
B
);
if self.count == 0 {
return &[];
} else if self.count == self.capacity {
return unsafe { self.vals[lane].get_unchecked(self.index..self.index + self.count - offset) };
}
unsafe { self.vals[lane].get_unchecked(0..self.count - offset) }
}
#[inline(always)]
fn get_slices(&self, offset: usize) -> [&[T]; B] {
if self.count == 0 {
return core::array::from_fn(|_| [].as_slice());
} else if self.count == self.capacity {
return core::array::from_fn(|lane| unsafe {
self.vals[lane].get_unchecked(self.index..self.index + self.count - offset)
});
}
core::array::from_fn(|lane| unsafe { self.vals[lane].get_unchecked(0..self.count - offset) })
}
#[inline(always)]
fn window_index_to_bars_ago(&self, window_index: usize) -> usize {
self.count - 1 - window_index
}
#[inline(always)]
fn window_index_to_bars_ago_simd(&self, window_index: Simd<usize, B>) -> Simd<usize, B> {
Simd::splat(self.count - 1) - window_index
}
fn to_single_buffers(&self) -> [SingleBuffer<T>; B] {
std::array::from_fn(|i|
SingleBuffer {
index: self.index,
count: self.count,
prev_idx: self.prev_idx,
capacity: self.capacity,
vals: self.vals[i].clone()
}
)
}
}
#[inline(always)]
pub(crate) fn write_values<const B: usize, T: BufferElement>(
buffer: &mut MultiBuffer<B, T>,
values: [T; B],
) {
for (buff, &vals) in buffer.vals.iter_mut().zip(values.iter()) {
unsafe { *buff.get_unchecked_mut(buffer.index) = vals };
unsafe { *buff.get_unchecked_mut(buffer.index + buffer.capacity) = vals };
}
}
#[inline(always)]
pub(crate) fn write_values_pop<const N: usize, T: BufferElement>(
buffer: &mut MultiBuffer<N, T>,
values: [T; N],
) -> [T; N] {
let mut results = [T::default(); N];
for ((buff, &vals), result) in buffer
.vals
.iter_mut()
.zip(values.iter())
.zip(results.iter_mut())
{
*result = unsafe { *buff.get_unchecked(buffer.index) };
unsafe { *buff.get_unchecked_mut(buffer.index) = vals };
unsafe { *buff.get_unchecked_mut(buffer.index + buffer.capacity) = vals };
}
results
}
pub trait MinMaxBuffer<const B: usize>: MirrorBuffer<B, f64> {
fn max<const CHUNK_SIZE: usize>(
&self,
state: &mut MaxState<B>,
bar: Simd<f64, B>,
look_back: usize,
) -> (Simd<f64, B>, Simd<usize, B>);
fn min<const CHUNK_SIZE: usize>(
&self,
state: &mut MinState<B>,
bar: Simd<f64, B>,
look_back: usize,
) -> (Simd<f64, B>, Simd<usize, B>);
}
impl<const B: usize> MinMaxBuffer<B> for MultiBuffer<B, f64> {
fn max<const CHUNK_SIZE: usize>(
&self,
state: &mut MaxState<B>,
bar: Simd<f64, B>,
look_back: usize,
) -> (Simd<f64, B>, Simd<usize, B>) {
let (mut max, mut trail) = (state.max, state.trail);
trail += UsizeConstants::ONE;
let lookback_simd = Simd::splat(look_back);
let needs_search = lookback_simd.simd_eq(trail);
let search_mask = needs_search.to_bitmask();
let current_is_new_max = bar.simd_ge(max);
max = current_is_new_max.select(bar, max);
trail = current_is_new_max.select(UsizeConstants::ZERO, trail);
if search_mask != 0 {
let max_array = max.as_mut_array();
let trail_array = trail.as_mut_array();
let mut lane = 0;
while lane < B {
if search_mask & (1 << lane) != 0 {
let (max_val, max_idx) = if CHUNK_SIZE == 1 {
find_max_scalar(self.get_slice(lane, 1), bar[lane])
} else {
find_max_simd::<CHUNK_SIZE>(self.get_slice(lane, 0))
};
max_array[lane] = max_val;
trail_array[lane] = self.window_index_to_bars_ago(max_idx);
}
lane += 1;
}
}
(state.max, state.trail) = (max, trail);
(max, trail)
}
fn min<const CHUNK_SIZE: usize>(
&self,
state: &mut MinState<B>,
bar: Simd<f64, B>,
look_back: usize,
) -> (Simd<f64, B>, Simd<usize, B>) {
let (mut min, mut trail) = (state.min, state.trail);
trail += UsizeConstants::ONE;
let lookback_simd = Simd::splat(look_back);
let needs_search = lookback_simd.simd_eq(trail);
let search_mask = needs_search.to_bitmask();
let current_is_new_min = bar.simd_le(min);
min = current_is_new_min.select(bar, min);
trail = current_is_new_min.select(UsizeConstants::ZERO, trail);
if search_mask != 0 {
let min_array = min.as_mut_array();
let trail_array = trail.as_mut_array();
let mut lane = 0;
while lane < B {
if search_mask & (1 << lane) != 0 {
let (min_val, min_idx) = if CHUNK_SIZE == 1 {
find_min_scalar(self.get_slice(lane, 1), bar[lane])
} else {
find_min_simd::<CHUNK_SIZE>(self.get_slice(lane, 0))
};
min_array[lane] = min_val;
trail_array[lane] = self.window_index_to_bars_ago(min_idx);
}
lane += 1;
}
}
(state.min, state.trail) = (min, trail);
(min, trail)
}
}