#![allow(dead_code)]
use parking_lot::RwLock;
pub struct VectorStore {
buffer: RwLock<Vec<f32>>,
dimension: usize,
count: RwLock<usize>,
free_slots: RwLock<Vec<usize>>,
}
impl VectorStore {
#[must_use]
pub fn new(dimension: usize, initial_capacity: usize) -> Self {
Self {
buffer: RwLock::new(Vec::with_capacity(dimension * initial_capacity)),
dimension,
count: RwLock::new(0),
free_slots: RwLock::new(Vec::new()),
}
}
#[must_use]
pub fn dimension(&self) -> usize {
self.dimension
}
#[must_use]
pub fn len(&self) -> usize {
*self.count.read()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn insert(&self, vector: &[f32]) -> usize {
assert_eq!(
vector.len(),
self.dimension,
"Vector dimension mismatch: expected {}, got {}",
self.dimension,
vector.len()
);
let mut free_slots = self.free_slots.write();
let mut buffer = self.buffer.write();
let mut count = self.count.write();
if let Some(idx) = free_slots.pop() {
let offset = idx * self.dimension;
buffer[offset..offset + self.dimension].copy_from_slice(vector);
return idx;
}
let idx = *count;
buffer.extend_from_slice(vector);
*count += 1;
idx
}
#[must_use]
pub fn get(&self, idx: usize) -> Option<Vec<f32>> {
let buffer = self.buffer.read();
let offset = idx * self.dimension;
if offset + self.dimension <= buffer.len() {
Some(buffer[offset..offset + self.dimension].to_vec())
} else {
None
}
}
#[must_use]
pub fn get_slice(&self, idx: usize) -> Option<VectorRef<'_>> {
let buffer = self.buffer.read();
let offset = idx * self.dimension;
if offset + self.dimension <= buffer.len() {
Some(VectorRef {
guard: buffer,
offset,
dimension: self.dimension,
})
} else {
None
}
}
pub fn remove(&self, idx: usize) -> bool {
let buffer = self.buffer.read();
let offset = idx * self.dimension;
if offset + self.dimension <= buffer.len() {
drop(buffer);
let mut free_slots = self.free_slots.write();
free_slots.push(idx);
true
} else {
false
}
}
pub fn update(&self, idx: usize, vector: &[f32]) -> bool {
assert_eq!(
vector.len(),
self.dimension,
"Vector dimension mismatch: expected {}, got {}",
self.dimension,
vector.len()
);
let mut buffer = self.buffer.write();
let offset = idx * self.dimension;
if offset + self.dimension <= buffer.len() {
buffer[offset..offset + self.dimension].copy_from_slice(vector);
true
} else {
false
}
}
#[must_use]
pub fn memory_usage(&self) -> usize {
let buffer = self.buffer.read();
buffer.capacity() * std::mem::size_of::<f32>()
}
#[inline]
pub fn prefetch(&self, idx: usize) {
let buffer = self.buffer.read();
let offset = idx * self.dimension;
if offset < buffer.len() {
#[cfg(target_arch = "x86_64")]
unsafe {
use std::arch::x86_64::{_mm_prefetch, _MM_HINT_T0};
let ptr = buffer.as_ptr().add(offset);
_mm_prefetch(ptr.cast::<i8>(), _MM_HINT_T0);
}
#[cfg(target_arch = "aarch64")]
{
}
}
}
}
pub struct VectorRef<'a> {
guard: parking_lot::RwLockReadGuard<'a, Vec<f32>>,
offset: usize,
dimension: usize,
}
impl VectorRef<'_> {
#[must_use]
pub fn as_slice(&self) -> &[f32] {
unsafe {
let ptr = self.guard.as_ptr().add(self.offset);
std::slice::from_raw_parts(ptr, self.dimension)
}
}
}
impl std::ops::Deref for VectorRef<'_> {
type Target = [f32];
fn deref(&self) -> &Self::Target {
self.as_slice()
}
}