use crate::array2::Array2;
use crate::numeric::Float;
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ScratchBuffer<T> {
data: Vec<T>,
}
impl<T> ScratchBuffer<T> {
pub fn new() -> Self {
Self { data: Vec::new() }
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn capacity(&self) -> usize {
self.data.capacity()
}
pub fn as_slice(&self) -> &[T] {
&self.data
}
pub fn as_mut_slice(&mut self) -> &mut [T] {
&mut self.data
}
pub fn clear(&mut self) {
self.data.clear();
}
}
impl<T: Clone> ScratchBuffer<T> {
pub fn resize(&mut self, len: usize, value: T) -> &mut [T] {
self.data.resize(len, value);
&mut self.data
}
}
impl<T: Float> ScratchBuffer<T> {
pub fn zeros(&mut self, len: usize) -> &mut [T] {
self.resize(len, T::zero())
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Workspace<T> {
buffers: Vec<ScratchBuffer<T>>,
}
impl<T> Workspace<T> {
pub fn new() -> Self {
Self {
buffers: Vec::new(),
}
}
pub fn buffers(&self) -> usize {
self.buffers.len()
}
pub fn buffer(&mut self, index: usize) -> &mut ScratchBuffer<T> {
while self.buffers.len() <= index {
self.buffers.push(ScratchBuffer::new());
}
&mut self.buffers[index]
}
pub fn two_buffers_mut(
&mut self,
left: usize,
right: usize,
) -> (&mut ScratchBuffer<T>, &mut ScratchBuffer<T>) {
assert_ne!(left, right, "scratch buffer indices must be distinct");
let max = left.max(right);
while self.buffers.len() <= max {
self.buffers.push(ScratchBuffer::new());
}
if left < right {
let (head, tail) = self.buffers.split_at_mut(right);
(&mut head[left], &mut tail[0])
} else {
let (head, tail) = self.buffers.split_at_mut(left);
(&mut tail[0], &mut head[right])
}
}
pub fn clear(&mut self) {
for buffer in &mut self.buffers {
buffer.clear();
}
}
}
impl<T: Float> Workspace<T> {
pub fn matrix(&mut self, shape: [usize; 2]) -> Array2<T> {
Array2::zeros(shape)
}
}