use std::cell::RefCell;
pub struct Arena {
buf: RefCell<Vec<f64>>,
capacity: usize,
offset: RefCell<usize>,
}
impl Arena {
pub fn new(capacity: usize) -> Self {
Self {
buf: RefCell::new(vec![0.0; capacity]),
capacity,
offset: RefCell::new(0),
}
}
pub fn alloc_slice(&self, n: usize) -> Option<ArenaSlice> {
let mut off = self.offset.borrow_mut();
if *off + n > self.capacity {
return None;
}
let start = *off;
*off += n;
Some(ArenaSlice { start, len: n })
}
pub fn write(&self, slice: &ArenaSlice, data: &[f64]) {
let mut buf = self.buf.borrow_mut();
let end = slice.start + data.len().min(slice.len);
buf[slice.start..end].copy_from_slice(&data[..end - slice.start]);
}
pub fn read(&self, slice: &ArenaSlice) -> Vec<f64> {
let buf = self.buf.borrow();
buf[slice.start..slice.start + slice.len].to_vec()
}
pub fn get(&self, slice: &ArenaSlice, index: usize) -> Option<f64> {
if index >= slice.len {
return None;
}
Some(self.buf.borrow()[slice.start + index])
}
pub fn set(&self, slice: &ArenaSlice, index: usize, value: f64) {
if index < slice.len {
self.buf.borrow_mut()[slice.start + index] = value;
}
}
pub fn reset(&self) {
*self.offset.borrow_mut() = 0;
}
pub fn used(&self) -> usize {
*self.offset.borrow()
}
pub fn remaining(&self) -> usize {
self.capacity - self.used()
}
pub fn capacity(&self) -> usize {
self.capacity
}
}
#[derive(Debug, Clone, Copy)]
pub struct ArenaSlice {
start: usize,
len: usize,
}
impl ArenaSlice {
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
}
pub struct ArenaMatrix {
slice: ArenaSlice,
pub rows: usize,
pub cols: usize,
}
impl ArenaMatrix {
pub fn new(arena: &Arena, rows: usize, cols: usize) -> Option<Self> {
let slice = arena.alloc_slice(rows * cols)?;
Some(Self { slice, rows, cols })
}
pub fn get(&self, arena: &Arena, r: usize, c: usize) -> Option<f64> {
if r >= self.rows || c >= self.cols {
return None;
}
arena.get(&self.slice, r * self.cols + c)
}
pub fn set(&self, arena: &Arena, r: usize, c: usize, value: f64) {
if r < self.rows && c < self.cols {
arena.set(&self.slice, r * self.cols + c, value);
}
}
pub fn to_vec(&self, arena: &Arena) -> Vec<Vec<f64>> {
let flat = arena.read(&self.slice);
flat.chunks(self.cols).map(|row| row.to_vec()).collect()
}
}
pub struct ScratchPool {
arena: Arena,
scratch_size: usize,
count: usize,
}
impl ScratchPool {
pub fn new(scratch_size: usize, count: usize) -> Self {
Self {
arena: Arena::new(scratch_size * count),
scratch_size,
count,
}
}
pub fn get(&self, index: usize) -> Option<ArenaSlice> {
if index >= self.count {
return None;
}
Some(ArenaSlice {
start: index * self.scratch_size,
len: self.scratch_size,
})
}
pub fn write(&self, index: usize, data: &[f64]) {
if let Some(slice) = self.get(index) {
self.arena.write(&slice, data);
}
}
pub fn read(&self, index: usize) -> Vec<f64> {
self.get(index)
.map(|s| self.arena.read(&s))
.unwrap_or_default()
}
pub fn scratch_size(&self) -> usize {
self.scratch_size
}
pub fn count(&self) -> usize {
self.count
}
pub fn arena(&self) -> &Arena {
&self.arena
}
}