use pgx::PGXSharedMemory;
use thiserror::Error;
trait F32VecBlockData {}
trait I64BlockData {}
#[derive(Default)]
pub struct F32Block<const SIZE: usize>(heapless::Vec<f32, SIZE>);
unsafe impl<const SIZE: usize> PGXSharedMemory for F32Block<SIZE> {}
#[derive(Error, Debug, PartialEq, Eq)]
pub enum F32BlockError {
#[error("Trying to write over end of block: Start {write_start}, Length {write_length}, block size {block_size}")]
Overflow {
write_start: usize,
write_length: usize,
block_size: usize,
},
#[error("Trying to free unallocated memory. Only allocated {data_len}. Segment to delete starts at {start}, and is of length {len}")]
DeleteEmpty {
data_len: usize,
start: usize,
len: usize,
},
#[error("Placeholder")]
PlaceHolder,
}
impl F32BlockError {
fn overflow(start: usize, length: usize, block: usize) -> F32BlockError {
F32BlockError::Overflow {
write_start: start,
write_length: length,
block_size: block,
}
}
}
impl<const Size: usize> F32Block<Size> {
fn write_vec(&mut self, vec: &Vec<f32>, start: usize) -> Result<(), F32BlockError> {
let end = start + vec.len();
if end < start || (end) >= Size {
Err(F32BlockError::overflow(start, vec.len(), Size))
} else {
let vec_data = &mut self.0;
let data_len = vec_data.len();
if start > data_len {
(data_len..start)
.map(|_| &0.0)
.chain(vec.iter())
.map(|x| vec_data.push(*x))
.collect::<Result<Vec<()>, f32>>()
.map(|_| ())
.map_err(|_| F32BlockError::overflow(start, vec.len(), Size))
} else {
let intersect = data_len.min(end);
let intersect_shift = intersect - start;
vec_data[start..intersect]
.iter_mut()
.zip(vec.iter())
.for_each(|(buffer, new)| *buffer = *new);
vec[intersect_shift..]
.iter()
.map(|x| vec_data.push(*x))
.collect::<Result<Vec<()>, f32>>()
.map_err(|_| F32BlockError::overflow(start, vec.len(), Size))?;
Ok(())
}
}
}
fn free(&mut self, start: usize, len: usize) -> Result<(), F32BlockError> {
const SEGMENT_SIZE: usize = 4096;
let end = start + len;
let data_len = self.0.len();
if end < start || end > data_len {
Err(F32BlockError::DeleteEmpty {
data_len,
start,
len,
})
} else {
let mut tmp = [0.0; SEGMENT_SIZE];
let to_move = data_len - end;
(end..data_len).step_by(SEGMENT_SIZE).for_each(|i| {
let vec_data = &self.0;
tmp.iter_mut()
.zip(vec_data[i..(i + SEGMENT_SIZE)].iter())
.for_each(|(tmp_val, vec_data_val)| *tmp_val = *vec_data_val);
let vec_data = &mut self.0;
vec_data[(i - len)..(i - len + SEGMENT_SIZE)]
.iter_mut()
.zip(tmp.iter())
.for_each(|(vec_data_val, tmp_val)| *vec_data_val = *tmp_val);
});
(0..len)
.map(|_| {
self.0.pop().ok_or(F32BlockError::DeleteEmpty {
data_len,
start,
len,
})
})
.collect::<Result<Vec<f32>, F32BlockError>>()
.map(|_| ())
}
}
fn read_data(
&self,
start: usize,
len: usize,
) -> Result<impl Iterator<Item = &f32>, F32BlockError> {
let data_len = self.0.len();
let end = start + len;
if (end < start) || (end >= data_len) {
Err(F32BlockError::overflow(start, len, Size))
} else {
Ok(self.0[start..end].iter())
}
}
}
#[cfg(test)]
mod f32_block_tests {
use crate::vec_block_alloc::F32Block;
use crate::vec_block_alloc::F32BlockError;
use core::cell::RefCell;
/// test_push: given block of finite size, test that we can push a vector of finite dim
/// test_drop:
/// test_compact:
#[test]
fn test_write() {
const BLOCK_SIZE: usize = 10000;
let good_start = 1000;
let vec_size = 512;
let bad_start = 10000 - (vec_size / 2);
let vec = (0..vec_size).map(|x| (x as f32).sin()).collect::<Vec<_>>();
let mem = RefCell::new(Box::new(F32Block::<BLOCK_SIZE>::default()));
assert!(mem
.borrow_mut()
.write_vec(&vec, good_start)
.map(|_| true)
.unwrap());
assert_eq!(
mem.borrow_mut().write_vec(&vec, bad_start).unwrap_err(),
F32BlockError::overflow(bad_start, vec_size, BLOCK_SIZE)
);
}
#[test]
fn test_delete() -> Result<(), F32BlockError> {
const BLOCK_SIZE: usize = 100000;
let good_start = 10000; // something less than BLOCK_SIZE;
let len = 10000; // something less than BLOCK_SIZE - good_start;
let mem = RefCell::new(Box::new(F32Block::<BLOCK_SIZE>::default()));
assert_eq!(
mem.borrow_mut().free(good_start, len).unwrap_err(),
F32BlockError::DeleteEmpty {
data_len: 0,
start: good_start,
len
}
);
mem.borrow_mut().write_vec(
&(0..(good_start + len))
.map(|x| (x as f32).cos())
.collect::<Vec<_>>(),
0,
)?;
assert!(mem
.borrow_mut()
.free(good_start, len)
.map(|_| true)
.unwrap());
Ok(())
}
}
/// Later we can replace the individual fields with a seq of blocks
#[derive(Default)]
pub struct FiniteDimAllocator<const BlockSize: usize> {
dimension: u16,
start: usize,
len: usize,
f32_data: F32Block<BlockSize>,
}
unsafe impl<const Size: usize> PGXSharedMemory for FiniteDimAllocator<Size> {}
impl<const BlockSize: usize> FiniteDimAllocator<BlockSize> {
fn new(dimension: u16) -> FiniteDimAllocator<BlockSize> {
FiniteDimAllocator {
dimension,
start: 0,
len: 0,
f32_data: F32Block::<BlockSize>::default(),
}
}
pub fn set_dim(&mut self, dimension: u16) {
self.dimension = dimension
}
fn dim(&self) -> usize {
self.dimension as usize
}
fn end(&self) -> usize {
self.start + self.len * (self.dimension as usize)
}
pub fn add_vec(&mut self, vec: &Vec<f32>) -> Result<(), F32BlockError> {
// Replace with custom error type
let end = self.end();
let write_vec = vec[..self.dim()].iter().map(|x| *x).collect::<Vec<_>>(); // Unecessary allocation ! need to implement as iters
self.f32_data.write_vec(&write_vec, end)?;
self.len += 1;
Ok(())
}
fn remove_vec(&mut self, index: usize) -> Result<(), F32BlockError> {
if index >= self.len {
Err(F32BlockError::PlaceHolder)
} else {
let dimension = self.dim();
let slice_start = self.start + index * dimension;
let slice_end = slice_start + dimension;
self.f32_data.free(slice_start, slice_end)?;
self.len -= 1;
Ok(())
}
}
pub fn get_iter(&self, index: usize) -> Result<impl Iterator<Item = &f32>, F32BlockError> {
let dimension = self.dim();
let slice_start = self.start + (dimension) * index;
let slice_end = slice_start + (dimension);
self.f32_data.read_data(slice_start, slice_end)
}
pub fn ip_many(&self, vec: &Vec<f32>, indices: &Vec<usize>) -> Result<Vec<f32>, F32BlockError> {
indices
.iter()
.map(|&index| {
let data_iter = self.get_iter(index)?;
Ok(data_iter.zip(vec.iter()).map(|(&a, &b)| a * b).sum())
})
.collect::<Result<Vec<f32>, F32BlockError>>()
}
}
#[cfg(test)]
mod test_finite_dim_allocator {
use crate::util::random_vecs;
use crate::vec_block_alloc::F32BlockError;
use crate::vec_block_alloc::FiniteDimAllocator;
#[test]
fn test_insert() {
const SIZE: usize = 100000;
let dim: u16 = 32;
let mut allocator = FiniteDimAllocator::<SIZE>::new(dim);
let allocator_borr = &mut allocator;
random_vecs(dim, 100)
.map(|vec| allocator_borr.add_vec(&vec))
.collect::<Result<Vec<()>, F32BlockError>>()
.unwrap();
}
#[test]
fn print_size() {
const SIZE: usize = 523934;
let size = std::mem::size_of::<FiniteDimAllocator<SIZE>>();
print!("Size of allocator: {}", size);
assert!(false);
}
}