pub mod builder;
mod hnsw;
mod ivf;
use crate::platform::{memcpy, memset};
use crate::types::{
DataType, DistanceType, IndexType, RemDbError, Result, TableDef, VectorIndexType,
VectorMetadata,
};
use core::ptr::NonNull;
pub use builder::init_index_build_thread_pool;
#[cfg(feature = "log")]
use crate::log::error;
const BTREE_ORDER: usize = 4;
const TTREE_ORDER: usize = 3;
#[derive(Default)]
pub struct IndexStats {
pub access_count: usize,
pub hit_count: usize,
pub size: usize,
pub item_count: usize,
}
#[repr(C)]
pub struct PrimaryIndexItem {
pub next: Option<NonNull<PrimaryIndexItem>>,
pub record_id: u16,
pub key_size: u8,
pub key_data: [u8; 128], }
#[derive(Copy, Clone)]
#[repr(C)]
pub struct SecondaryIndexItem {
pub key_size: u8,
pub record_id: u16,
pub key_data: [u8; 128], }
impl Default for SecondaryIndexItem {
fn default() -> Self {
SecondaryIndexItem {
key_size: 0,
record_id: 0,
key_data: [0u8; 128],
}
}
}
#[repr(C)]
pub struct BTreeNode {
pub is_leaf: bool,
pub key_count: u8,
pub keys: [SecondaryIndexItem; BTREE_ORDER],
pub children: [Option<NonNull<BTreeNode>>; BTREE_ORDER + 1],
}
pub struct BTreeIndex {
pub def: alloc::sync::Arc<TableDef>,
pub root: Option<NonNull<BTreeNode>>,
pub nodes: NonNull<BTreeNode>,
pub free_nodes: Option<NonNull<BTreeNode>>,
pub max_nodes: usize,
pub stats: IndexStats,
pub lock: u32,
}
#[repr(C)]
pub struct TTreeNode {
pub key_count: u8,
pub keys: [SecondaryIndexItem; TTREE_ORDER],
pub left: Option<NonNull<TTreeNode>>,
pub middle: Option<NonNull<TTreeNode>>,
pub right: Option<NonNull<TTreeNode>>,
}
pub struct TTreeIndex {
pub def: alloc::sync::Arc<TableDef>,
pub root: Option<NonNull<TTreeNode>>,
pub nodes: NonNull<TTreeNode>,
pub free_nodes: Option<NonNull<TTreeNode>>,
pub max_nodes: usize,
pub stats: IndexStats,
pub lock: u32,
}
pub struct PrimaryIndex {
def: alloc::sync::Arc<TableDef>,
hash_table: NonNull<Option<NonNull<PrimaryIndexItem>>>,
hash_table_size: usize,
items: NonNull<PrimaryIndexItem>,
free_items: Option<NonNull<PrimaryIndexItem>>,
stats: IndexStats,
lock: u32,
}
unsafe fn encode_composite_key(
record_ptr: *const u8,
primary_key_fields: &[&crate::types::FieldDef],
) -> (Vec<u8>, usize) {
let mut encoded_key = Vec::new();
for field in primary_key_fields {
let field_ptr = record_ptr.add(field.offset);
match field.data_type {
crate::types::DataType::UInt8 => {
let value = core::ptr::read_unaligned(field_ptr as *const u8);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::UInt16 => {
let value = core::ptr::read_unaligned(field_ptr as *const u16);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::UInt32 => {
let value = core::ptr::read_unaligned(field_ptr as *const u32);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::UInt64 => {
let value = core::ptr::read_unaligned(field_ptr as *const u64);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::Int8 => {
let value = core::ptr::read_unaligned(field_ptr as *const i8);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::Int16 => {
let value = core::ptr::read_unaligned(field_ptr as *const i16);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::Int32 => {
let value = core::ptr::read_unaligned(field_ptr as *const i32);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::Int64 => {
let value = core::ptr::read_unaligned(field_ptr as *const i64);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::Float32 => {
let value = core::ptr::read_unaligned(field_ptr as *const f32);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::Float64 => {
let value = core::ptr::read_unaligned(field_ptr as *const f64);
encoded_key.extend_from_slice(&value.to_le_bytes());
}
crate::types::DataType::Bool => {
let value = core::ptr::read_unaligned(field_ptr as *const bool);
encoded_key.push(value as u8);
}
crate::types::DataType::Timestamp => {
let value =
core::ptr::read_unaligned(field_ptr as *const crate::types::db_timestamp);
encoded_key.extend_from_slice(&value.value.to_le_bytes());
}
crate::types::DataType::VarChar | crate::types::DataType::Char => {
let str_slice = core::slice::from_raw_parts(field_ptr, field.size);
let str_len = str_slice.iter().position(|&c| c == 0).unwrap_or(field.size);
encoded_key.push(str_len as u8);
encoded_key.extend_from_slice(&str_slice[0..str_len]);
}
crate::types::DataType::Text => {
let text_storage =
core::ptr::read_unaligned(field_ptr as *const crate::types::TextStorage);
if text_storage.is_inline() {
if let Some(data) = text_storage.as_inline() {
let str_len = data.iter().position(|&c| c == 0).unwrap_or(data.len());
let capped_len = core::cmp::min(str_len, 255);
encoded_key.push(capped_len as u8);
encoded_key.extend_from_slice(&data[..capped_len]);
} else {
encoded_key.push(0);
}
} else if text_storage.is_external() {
if let Some(ext) = text_storage.as_external() {
if !ext.data_ptr.is_null() {
let str_len = ext.length as usize;
let capped_len = core::cmp::min(str_len, 255);
let bytes =
unsafe { core::slice::from_raw_parts(ext.data_ptr, capped_len) };
encoded_key.push(capped_len as u8);
encoded_key.extend_from_slice(bytes);
} else {
encoded_key.push(0);
}
} else {
encoded_key.push(0);
}
} else {
encoded_key.push(0);
}
}
_ => {
}
}
}
let len = encoded_key.len();
(encoded_key, len)
}
impl PrimaryIndex {
pub unsafe fn new(
def: alloc::sync::Arc<TableDef>,
hash_table_start: *mut Option<NonNull<PrimaryIndexItem>>,
items_start: *mut PrimaryIndexItem,
hash_table_size: usize,
max_items: usize,
) -> Self {
let hash_table = NonNull::new_unchecked(hash_table_start);
for i in 0..hash_table_size {
let slot_ptr = hash_table.as_ptr().add(i);
*slot_ptr = None;
}
let items = NonNull::new_unchecked(items_start);
let mut free_items = None;
for i in (0..max_items).rev() {
let item_ptr = items.as_ptr().add(i);
(*item_ptr).next = free_items;
(*item_ptr).record_id = 0;
(*item_ptr).key_size = 0;
memset((*item_ptr).key_data.as_mut_ptr(), 0, 128);
free_items = Some(NonNull::new_unchecked(item_ptr));
}
PrimaryIndex {
def,
hash_table,
hash_table_size,
items,
free_items,
stats: IndexStats {
access_count: 0,
hit_count: 0,
size: hash_table_size * core::mem::size_of::<Option<NonNull<PrimaryIndexItem>>>()
+ max_items * core::mem::size_of::<PrimaryIndexItem>(),
item_count: 0,
},
lock: 0,
}
}
pub const fn calculate_memory_size(
_def: &TableDef,
hash_table_size: usize,
max_items: usize,
) -> usize {
let hash_table_size_bytes =
hash_table_size * core::mem::size_of::<Option<NonNull<PrimaryIndexItem>>>();
let items_size_bytes = max_items * core::mem::size_of::<PrimaryIndexItem>();
hash_table_size_bytes + items_size_bytes
}
fn hash_key(&self, key: *const u8, key_size: usize) -> usize {
let mut hash = 0u64;
let seed = 0x5bd1e995u64;
for i in 0..key_size {
let byte = unsafe { *key.add(i) };
hash ^= byte as u64;
hash = hash.wrapping_mul(seed);
hash ^= hash >> 47;
}
(hash as usize) % self.hash_table_size
}
pub unsafe fn insert(&mut self, key: *const u8, key_size: usize, record_id: u16) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_inserts()
}
crate::platform::spin_lock(&mut self.lock);
if key_size > 128 {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::UnsupportedOperation);
}
let mut item = match self.free_items {
Some(item_ptr) => {
let next = (*item_ptr.as_ptr()).next;
self.free_items = next;
item_ptr
}
None => {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::OutOfMemory);
}
};
let item_mut = item.as_mut();
item_mut.record_id = record_id;
item_mut.key_size = key_size as u8;
memcpy(item_mut.key_data.as_mut_ptr(), key, key_size);
let hash = self.hash_key(key, key_size);
let slot_ptr = self.hash_table.as_ptr().add(hash);
item_mut.next = *slot_ptr;
*slot_ptr = Some(item);
self.stats.item_count += 1;
crate::platform::spin_unlock(&mut self.lock);
Ok(())
}
pub unsafe fn insert_composite(&mut self, record_ptr: *const u8, record_id: u16) -> Result<()> {
let primary_key_fields: Vec<&crate::types::FieldDef> = self
.def
.primary_key
.iter()
.map(|&idx| &self.def.fields[idx])
.collect();
let (encoded_key, key_size) = encode_composite_key(record_ptr, &primary_key_fields);
self.insert(encoded_key.as_ptr(), key_size, record_id)
}
pub unsafe fn find(&mut self, key: *const u8, key_size: usize) -> Result<u16> {
self.stats.access_count += 1;
let hash = self.hash_key(key, key_size);
let slot_ptr = self.hash_table.as_ptr().add(hash);
let mut current = *slot_ptr;
while let Some(item) = current {
let item_ref = item.as_ref();
if item_ref.key_size == key_size as u8 {
let mut match_found = true;
for i in 0..key_size {
if item_ref.key_data[i] != *key.add(i) {
match_found = false;
break;
}
}
if match_found {
self.stats.hit_count += 1;
return Ok(item_ref.record_id);
}
}
current = item_ref.next;
}
Err(RemDbError::RecordNotFound)
}
pub unsafe fn find_composite(&mut self, record_ptr: *const u8) -> Result<u16> {
let primary_key_fields: Vec<&crate::types::FieldDef> = self
.def
.primary_key
.iter()
.map(|&idx| &self.def.fields[idx])
.collect();
let (encoded_key, key_size) = encode_composite_key(record_ptr, &primary_key_fields);
self.find(encoded_key.as_ptr(), key_size)
}
pub unsafe fn delete(&mut self, key: *const u8, key_size: usize) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_deletes()
}
crate::platform::spin_lock(&mut self.lock);
let hash = self.hash_key(key, key_size);
let slot_ptr = self.hash_table.as_ptr().add(hash);
let mut current = *slot_ptr;
let mut prev: Option<NonNull<PrimaryIndexItem>> = None;
while let Some(mut item) = current {
let item_ref = item.as_ref();
if item_ref.key_size == key_size as u8 {
let mut match_found = true;
for i in 0..key_size {
if item_ref.key_data[i] != *key.add(i) {
match_found = false;
break;
}
}
if match_found {
if let Some(mut prev_item) = prev {
prev_item.as_mut().next = item_ref.next;
} else {
*slot_ptr = item_ref.next;
}
let item_mut = item.as_mut();
item_mut.next = self.free_items;
self.free_items = Some(item);
self.stats.item_count -= 1;
crate::platform::spin_unlock(&mut self.lock);
return Ok(());
}
}
prev = Some(item);
current = item_ref.next;
}
crate::platform::spin_unlock(&mut self.lock);
Err(RemDbError::RecordNotFound)
}
pub unsafe fn delete_composite(&mut self, record_ptr: *const u8) -> Result<()> {
let primary_key_fields: Vec<&crate::types::FieldDef> = self
.def
.primary_key
.iter()
.map(|&idx| &self.def.fields[idx])
.collect();
let (encoded_key, key_size) = encode_composite_key(record_ptr, &primary_key_fields);
self.delete(encoded_key.as_ptr(), key_size)
}
pub fn stats(&self) -> &IndexStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats.access_count = 0;
self.stats.hit_count = 0;
}
}
impl Drop for TTreeIndex {
fn drop(&mut self) {
unsafe {
crate::memory::allocator::free(self.nodes.cast::<u8>());
}
}
}
impl Drop for VectorIndex {
fn drop(&mut self) {
unsafe {
if !self.items.is_null() {
if let Some(ptr) = NonNull::new(self.items as *mut u8) {
crate::memory::allocator::free(ptr);
}
}
}
}
}
#[derive(Copy, Clone)]
struct VectorIndexItem {
vector_offset: usize,
record_id: u16,
}
enum VectorIndexImpl {
LinearSearch,
HNSW(Option<hnsw::HNSWIndex>),
IVFFlat(Option<ivf::IVFIndex>),
}
pub struct VectorIndex {
def: alloc::sync::Arc<TableDef>,
stats: IndexStats,
lock: u32,
distance_type: DistanceType,
dimension: u16,
vector_index_type: VectorIndexType,
items: *mut VectorIndexItem,
vectors: *mut f32,
max_items: usize,
item_count: usize,
vector_count: usize,
index_impl: VectorIndexImpl,
}
impl VectorIndex {
pub unsafe fn new(
def: alloc::sync::Arc<TableDef>,
memory_start: *mut u8,
max_items: usize,
) -> Result<Self> {
let dimension;
let distance_type;
let vector_index_type;
let secondary_index = def
.secondary_index
.as_ref()
.ok_or(RemDbError::TypeMismatch)?;
if secondary_index.len() != 1 {
return Err(RemDbError::TypeMismatch);
}
let secondary_index = secondary_index[0];
if secondary_index >= def.fields.len() {
return Err(RemDbError::TypeMismatch);
}
let field = &def.fields[secondary_index];
if field.data_type != DataType::Vector {
#[cfg(feature = "log")]
error!(
"TypeMismatch: field.data_type != DataType::Vector, actual: {:?}",
field.data_type
);
return Err(RemDbError::TypeMismatch);
}
let vector_meta = match &field.vector_metadata {
Some(meta) => {
dimension = meta.dimension;
distance_type = meta.distance_type;
vector_index_type = meta.index_type;
meta
}
None => {
#[cfg(feature = "log")]
error!("TypeMismatch: field.vector_metadata is None");
return Err(RemDbError::TypeMismatch);
}
};
if dimension == 0 || dimension > 1024 {
#[cfg(feature = "log")]
error!("TypeMismatch: invalid dimension: {}", dimension);
return Err(RemDbError::TypeMismatch);
}
let items_size = max_items * core::mem::size_of::<VectorIndexItem>();
let vectors_size = max_items * dimension as usize * core::mem::size_of::<f32>();
let items = memory_start as *mut VectorIndexItem;
for i in 0..max_items {
let item_ptr = unsafe { items.add(i) };
*item_ptr = VectorIndexItem {
vector_offset: 0,
record_id: 0,
};
}
let vectors = (memory_start.add(items_size)) as *mut f32;
for i in 0..(max_items * dimension as usize) {
let vec_ptr = unsafe { vectors.add(i) };
*vec_ptr = 0.0;
}
let index_impl = match vector_index_type {
VectorIndexType::HNSW | VectorIndexType::HNSW_SQ | VectorIndexType::HNSW_BQ => {
let hnsw_memory = (memory_start.add(items_size + vectors_size)) as *mut u8;
let hnsw_index =
hnsw::HNSWIndex::new(*vector_meta, vectors, hnsw_memory, max_items)?;
VectorIndexImpl::HNSW(Some(hnsw_index))
}
VectorIndexType::IVF | VectorIndexType::IVF_PQ => {
let ivf_index = ivf::IVFIndex::new(
*vector_meta,
vectors,
vector_meta.ivf_nlist,
vector_meta.ivf_nprobe,
)?;
VectorIndexImpl::IVFFlat(Some(ivf_index))
}
};
Ok(VectorIndex {
def,
stats: IndexStats::default(),
lock: 0,
distance_type,
dimension,
vector_index_type,
items,
vectors,
max_items,
item_count: 0,
vector_count: 0,
index_impl,
})
}
pub fn calculate_memory_size(def: &TableDef, max_items: usize) -> usize {
let dimension = match def.secondary_index.as_ref() {
Some(secondary_index) if !secondary_index.is_empty() => {
let secondary_index = secondary_index[0];
if secondary_index < def.fields.len() {
let field = &def.fields[secondary_index];
if field.data_type == DataType::Vector {
if let Some(vector_meta) = &field.vector_metadata {
vector_meta.dimension
} else {
128 }
} else {
128 }
} else {
128 }
}
_ => 128, };
let dimension = core::cmp::max(dimension, 1);
let items_size = max_items * core::mem::size_of::<VectorIndexItem>();
let vectors_size = max_items * dimension as usize * core::mem::size_of::<f32>();
let base_size = items_size + vectors_size;
let index_impl_size = if let Some(secondary_index) = def.secondary_index.as_ref() {
if !secondary_index.is_empty() {
let secondary_index = secondary_index[0];
if secondary_index < def.fields.len() {
let field = &def.fields[secondary_index];
if field.data_type == DataType::Vector {
if let Some(vector_meta) = &field.vector_metadata {
Self::calculate_index_impl_memory_size(vector_meta, max_items)
} else {
0
}
} else {
0
}
} else {
0
}
} else {
0
}
} else {
0
};
core::cmp::max(base_size + index_impl_size, 1)
}
fn calculate_index_impl_memory_size(meta: &VectorMetadata, max_items: usize) -> usize {
match meta.index_type {
VectorIndexType::HNSW | VectorIndexType::HNSW_SQ | VectorIndexType::HNSW_BQ => {
let max_level = (max_items as f64).ln() as usize;
let max_level = core::cmp::max(max_level, 1);
let base_node_size = core::mem::size_of::<usize>() + core::mem::size_of::<u16>();
let neighbor_counts_capacity = max_level + 1;
let neighbors_capacity = (max_level + 1) * 32;
let vec_overhead = 3 * core::mem::size_of::<usize>();
let neighbor_counts_size = vec_overhead + neighbor_counts_capacity;
let neighbors_size =
vec_overhead + neighbors_capacity * core::mem::size_of::<NonNull<()>>();
let node_total_size = base_node_size + neighbor_counts_size + neighbors_size;
max_items * node_total_size
}
VectorIndexType::IVF | VectorIndexType::IVF_PQ => {
let nlist = meta.ivf_nlist as usize;
let nlist = core::cmp::max(nlist, 1);
let centroid_size = meta.dimension as usize * core::mem::size_of::<f32>();
nlist * centroid_size + nlist * core::mem::size_of::<usize>()
}
}
}
unsafe fn calculate_distance(&self, vec1: *const f32, vec2: *const f32) -> f32 {
match self.distance_type {
DistanceType::L2 => {
let mut sum = 0.0;
for i in 0..self.dimension {
let diff = *vec1.add(i as usize) - *vec2.add(i as usize);
sum += diff * diff;
}
sum.sqrt()
}
DistanceType::InnerProduct => {
let mut sum = 0.0;
for i in 0..self.dimension {
sum += *vec1.add(i as usize) * *vec2.add(i as usize);
}
sum
}
DistanceType::Cosine => {
let mut dot = 0.0;
let mut norm1 = 0.0;
let mut norm2 = 0.0;
for i in 0..self.dimension {
let v1 = *vec1.add(i as usize);
let v2 = *vec2.add(i as usize);
dot += v1 * v2;
norm1 += v1 * v1;
norm2 += v2 * v2;
}
let norm1 = norm1.sqrt();
let norm2 = norm2.sqrt();
if norm1 == 0.0 || norm2 == 0.0 {
0.0
} else {
dot / (norm1 * norm2)
}
}
}
}
#[cfg(feature = "std")]
pub fn save(&self, file_path: &str) -> Result<()> {
use std::fs::File;
use std::io::Write;
let mut file = File::create(file_path).map_err(|_| RemDbError::FileIoError)?;
file.write_all(&[self.vector_index_type as u8])
.map_err(|_| RemDbError::FileIoError)?;
file.write_all(&[self.distance_type as u8])
.map_err(|_| RemDbError::FileIoError)?;
file.write_all(&self.dimension.to_le_bytes())
.map_err(|_| RemDbError::FileIoError)?;
file.write_all(&self.max_items.to_le_bytes())
.map_err(|_| RemDbError::FileIoError)?;
file.write_all(&self.item_count.to_le_bytes())
.map_err(|_| RemDbError::FileIoError)?;
file.write_all(&self.vector_count.to_le_bytes())
.map_err(|_| RemDbError::FileIoError)?;
unsafe {
for i in 0..self.item_count {
let item = self.items.add(i).read();
file.write_all(&item.vector_offset.to_le_bytes())
.map_err(|_| RemDbError::FileIoError)?;
file.write_all(&item.record_id.to_le_bytes())
.map_err(|_| RemDbError::FileIoError)?;
}
}
unsafe {
let vector_size = self.vector_count * self.dimension as usize;
let vec_slice = core::slice::from_raw_parts(self.vectors, vector_size);
let vec_bytes = unsafe {
core::slice::from_raw_parts(
vec_slice.as_ptr() as *const u8,
vector_size * core::mem::size_of::<f32>(),
)
};
file.write_all(vec_bytes)
.map_err(|_| RemDbError::FileIoError)?;
}
match &self.index_impl {
VectorIndexImpl::HNSW(Some(hnsw_index)) => {
hnsw_index.save(&mut file)?;
}
VectorIndexImpl::IVFFlat(Some(ivf_index)) => {
ivf_index.save(&mut file)?;
}
_ => {
}
}
Ok(())
}
#[cfg(feature = "std")]
pub unsafe fn load(
def: alloc::sync::Arc<TableDef>,
file_path: &str,
memory_start: *mut u8,
) -> Result<Self> {
use std::fs::File;
use std::io::Read;
let mut file = File::open(file_path).map_err(|_| RemDbError::FileIoError)?;
let mut index_type_byte = [0u8; 1];
file.read_exact(&mut index_type_byte)
.map_err(|_| RemDbError::FileIoError)?;
let vector_index_type = match index_type_byte[0] {
0 => VectorIndexType::HNSW,
1 => VectorIndexType::HNSW_SQ,
2 => VectorIndexType::HNSW_BQ,
3 => VectorIndexType::IVF,
4 => VectorIndexType::IVF_PQ,
_ => VectorIndexType::HNSW, };
let mut distance_type_byte = [0u8; 1];
file.read_exact(&mut distance_type_byte)
.map_err(|_| RemDbError::FileIoError)?;
let distance_type = match distance_type_byte[0] {
0 => DistanceType::L2,
1 => DistanceType::InnerProduct,
2 => DistanceType::Cosine,
_ => DistanceType::L2, };
let mut dimension_bytes = [0u8; 2];
file.read_exact(&mut dimension_bytes)
.map_err(|_| RemDbError::FileIoError)?;
let dimension = u16::from_le_bytes(dimension_bytes);
let mut max_items_bytes = [0u8; 8];
file.read_exact(&mut max_items_bytes)
.map_err(|_| RemDbError::FileIoError)?;
let max_items = usize::from_le_bytes(max_items_bytes);
let mut item_count_bytes = [0u8; 8];
file.read_exact(&mut item_count_bytes)
.map_err(|_| RemDbError::FileIoError)?;
let item_count = usize::from_le_bytes(item_count_bytes);
let mut vector_count_bytes = [0u8; 8];
file.read_exact(&mut vector_count_bytes)
.map_err(|_| RemDbError::FileIoError)?;
let vector_count = usize::from_le_bytes(vector_count_bytes);
let items_size = max_items * core::mem::size_of::<VectorIndexItem>();
let vectors_size = max_items * dimension as usize * core::mem::size_of::<f32>();
let items = memory_start as *mut VectorIndexItem;
for i in 0..item_count {
let mut item = VectorIndexItem {
vector_offset: 0,
record_id: 0,
};
let mut offset_bytes = [0u8; 8];
file.read_exact(&mut offset_bytes)
.map_err(|_| RemDbError::FileIoError)?;
item.vector_offset = usize::from_le_bytes(offset_bytes);
let mut record_id_bytes = [0u8; 2];
file.read_exact(&mut record_id_bytes)
.map_err(|_| RemDbError::FileIoError)?;
item.record_id = u16::from_le_bytes(record_id_bytes);
*items.add(i) = item;
}
let vectors = (memory_start.add(items_size)) as *mut f32;
let vector_size = vector_count * dimension as usize;
let vec_slice = core::slice::from_raw_parts_mut(vectors, vector_size);
let vec_bytes = unsafe {
core::slice::from_raw_parts_mut(
vec_slice.as_mut_ptr() as *mut u8,
vector_size * core::mem::size_of::<f32>(),
)
};
file.read_exact(vec_bytes)
.map_err(|_| RemDbError::FileIoError)?;
let index_impl = match vector_index_type {
VectorIndexType::HNSW | VectorIndexType::HNSW_SQ | VectorIndexType::HNSW_BQ => {
let hnsw_memory = (memory_start.add(items_size + vectors_size)) as *mut u8;
let vector_meta = crate::types::VectorMetadata {
dimension,
distance_type,
index_type: vector_index_type,
compression_enabled: false,
compression_scheme: 0,
compression_level: 3,
hnsw_m: 16, hnsw_ef_construction: 200, hnsw_ef_search: 100, ivf_nlist: 100, ivf_nprobe: 10, };
let hnsw_index =
hnsw::HNSWIndex::load(vector_meta, vectors, hnsw_memory, max_items, &mut file)
.map_err(|_| RemDbError::FileIoError)?;
VectorIndexImpl::HNSW(Some(hnsw_index))
}
VectorIndexType::IVF | VectorIndexType::IVF_PQ => {
let vector_meta = crate::types::VectorMetadata {
dimension,
distance_type,
index_type: vector_index_type,
compression_enabled: false,
compression_scheme: 0,
compression_level: 3,
hnsw_m: 16, hnsw_ef_construction: 200, hnsw_ef_search: 100, ivf_nlist: 100, ivf_nprobe: 10, };
let ivf_index = ivf::IVFIndex::load(vector_meta, vectors, &mut file)
.map_err(|_| RemDbError::FileIoError)?;
VectorIndexImpl::IVFFlat(Some(ivf_index))
}
};
Ok(VectorIndex {
def,
stats: IndexStats::default(),
lock: 0,
distance_type,
dimension,
vector_index_type,
items,
vectors,
max_items,
item_count,
vector_count,
index_impl,
})
}
unsafe fn calculate_distance_old(&self, vec1: *const f32, vec2: *const f32) -> f32 {
match self.distance_type {
DistanceType::L2 => {
let mut sum = 0.0;
for i in 0..self.dimension {
let diff = *vec1.add(i as usize) - *vec2.add(i as usize);
sum += diff * diff;
}
sum.sqrt()
}
DistanceType::InnerProduct => {
let mut sum = 0.0;
for i in 0..self.dimension {
sum += *vec1.add(i as usize) * *vec2.add(i as usize);
}
-sum }
DistanceType::Cosine => {
let mut dot = 0.0;
let mut norm1 = 0.0;
let mut norm2 = 0.0;
for i in 0..self.dimension {
let v1 = *vec1.add(i as usize);
let v2 = *vec2.add(i as usize);
dot += v1 * v2;
norm1 += v1 * v1;
norm2 += v2 * v2;
}
let norm1 = norm1.sqrt();
let norm2 = norm2.sqrt();
if norm1 == 0.0 || norm2 == 0.0 {
-1.0 } else {
-(dot / (norm1 * norm2)) }
}
}
}
pub unsafe fn insert(&mut self, key: *const u8, key_size: usize, record_id: u16) -> Result<()> {
crate::platform::spin_lock(&mut self.lock);
if self.item_count >= self.max_items {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::OutOfMemory);
}
let vec_len = self.dimension as usize;
let needed = vec_len.saturating_mul(core::mem::size_of::<f32>());
if key_size < needed {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::TypeMismatch);
}
let vec_ptr = key as *const f32;
let start_offset = self.vector_count * vec_len;
for i in 0..vec_len {
*self.vectors.add(start_offset + i) = *vec_ptr.add(i);
}
let impl_result = match &mut self.index_impl {
VectorIndexImpl::HNSW(Some(hnsw_index)) => hnsw_index.insert(start_offset, record_id),
VectorIndexImpl::IVFFlat(Some(ivf_index)) => ivf_index.insert(start_offset, record_id),
_ => {
Ok(())
}
};
match impl_result {
Ok(()) => {}
Err(e) => {
crate::platform::spin_unlock(&mut self.lock);
return Err(e);
}
}
let item_ptr = self.items.add(self.item_count);
*item_ptr = VectorIndexItem {
vector_offset: start_offset,
record_id,
};
self.item_count += 1;
self.vector_count = self.item_count;
self.stats.item_count += 1;
self.stats.size +=
core::mem::size_of::<VectorIndexItem>() + vec_len * core::mem::size_of::<f32>();
crate::platform::spin_unlock(&mut self.lock);
Ok(())
}
pub unsafe fn find(&mut self, key: *const u8, key_size: usize) -> Result<u16> {
self.stats.access_count += 1;
let query_vec: *const f32;
let mut query_vec_buf = [0.0f32; 1024];
let vec_len = self.dimension as usize;
if key_size > 4 && *key == b'[' {
let vec_str = match core::str::from_utf8(core::slice::from_raw_parts(key, key_size)) {
Ok(s) => s,
Err(_) => return Err(RemDbError::TypeMismatch),
};
let vec_str = vec_str.trim_start_matches('[').trim_end_matches(']');
let elements: Vec<&str> = vec_str.split(',').map(|s| s.trim()).collect();
if elements.len() != vec_len {
return Err(RemDbError::TypeMismatch);
}
for (i, elem) in elements.iter().enumerate() {
let parsed = match elem.parse::<f32>() {
Ok(v) => v,
Err(_) => return Err(RemDbError::TypeMismatch),
};
match query_vec_buf.get_mut(i) {
Some(slot) => *slot = parsed,
None => return Err(RemDbError::TypeMismatch),
}
}
query_vec = query_vec_buf.as_ptr();
} else {
let needed = vec_len.saturating_mul(core::mem::size_of::<f32>());
if key_size < needed {
return Err(RemDbError::TypeMismatch);
}
query_vec = key as *const f32;
}
crate::platform::spin_lock(&mut self.lock);
let result = match &self.index_impl {
VectorIndexImpl::HNSW(Some(hnsw_index)) => match hnsw_index.search(query_vec, 1) {
Ok(results) => match results.first() {
Some(&(_, record_id)) => Ok(record_id),
None => Err(RemDbError::RecordNotFound),
},
Err(e) => Err(e),
},
VectorIndexImpl::IVFFlat(Some(ivf_index)) => match ivf_index.search(query_vec, 1) {
Ok(results) => match results.first() {
Some(&(_, record_id)) => Ok(record_id),
None => Err(RemDbError::RecordNotFound),
},
Err(e) => Err(e),
},
_ => {
let mut min_distance = f32::MAX;
let mut best_record_id = 0;
let mut found = false;
for i in 0..self.item_count {
let item_ptr = self.items.add(i);
let vec_ptr = self.vectors.add((*item_ptr).vector_offset);
let distance = self.calculate_distance(query_vec, vec_ptr);
if distance < min_distance {
min_distance = distance;
best_record_id = (*item_ptr).record_id;
found = true;
}
}
if found {
Ok(best_record_id)
} else {
Err(RemDbError::RecordNotFound)
}
}
};
crate::platform::spin_unlock(&mut self.lock);
if result.is_ok() {
self.stats.hit_count += 1;
}
result
}
pub unsafe fn search_knn(
&mut self,
query_vec: *const f32,
k: usize,
) -> Result<Vec<(f32, u16)>> {
self.stats.access_count += 1;
if self.item_count == 0 {
return Ok(Vec::new());
}
crate::platform::spin_lock(&mut self.lock);
let k = core::cmp::min(k, self.item_count);
let results = match &self.index_impl {
VectorIndexImpl::HNSW(Some(hnsw_index)) => {
let r = hnsw_index.search(query_vec, k);
crate::platform::spin_unlock(&mut self.lock);
r?
}
VectorIndexImpl::IVFFlat(Some(ivf_index)) => {
let r = ivf_index.search(query_vec, k);
crate::platform::spin_unlock(&mut self.lock);
r?
}
_ => {
let mut all_results: Vec<(f32, u16)> = Vec::new();
for i in 0..self.item_count {
let item_ptr = self.items.add(i);
let vec_ptr = self.vectors.add((*item_ptr).vector_offset);
let distance = self.calculate_distance(query_vec, vec_ptr);
all_results.push((distance, (*item_ptr).record_id));
}
all_results
.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(core::cmp::Ordering::Equal));
all_results.truncate(k);
crate::platform::spin_unlock(&mut self.lock);
all_results
}
};
self.stats.hit_count += 1;
Ok(results)
}
pub unsafe fn find_range(
&mut self,
start_key: *const u8,
_start_key_size: usize,
end_key: *const u8,
_end_key_size: usize,
) -> Result<u16> {
self.stats.access_count += 1;
let query_vec: *const f32;
let mut query_vec_buf = [0.0f32; 1024];
let vec_len = self.dimension as usize;
if _start_key_size > 4 && *start_key == b'[' {
let vec_str =
match core::str::from_utf8(core::slice::from_raw_parts(start_key, _start_key_size))
{
Ok(s) => s,
Err(_) => return Err(RemDbError::TypeMismatch),
};
let vec_str = vec_str.trim_start_matches('[').trim_end_matches(']');
let elements: Vec<&str> = vec_str.split(',').map(|s| s.trim()).collect();
if elements.len() != vec_len {
return Err(RemDbError::TypeMismatch);
}
for (i, elem) in elements.iter().enumerate() {
let parsed = match elem.parse::<f32>() {
Ok(v) => v,
Err(_) => return Err(RemDbError::TypeMismatch),
};
match query_vec_buf.get_mut(i) {
Some(slot) => *slot = parsed,
None => return Err(RemDbError::TypeMismatch),
}
}
query_vec = query_vec_buf.as_ptr();
} else {
let needed = vec_len.saturating_mul(core::mem::size_of::<f32>());
if _start_key_size < needed {
return Err(RemDbError::TypeMismatch);
}
query_vec = start_key as *const f32;
}
let range_value: f32 = if _end_key_size > 4 && *end_key == b'[' {
1000.0 } else {
core::ptr::read_unaligned(end_key as *const f32)
};
crate::platform::spin_lock(&mut self.lock);
let mut found_record: Option<u16> = None;
for i in 0..self.item_count {
let item_ptr = self.items.add(i);
let vec_ptr = self.vectors.add((*item_ptr).vector_offset);
let distance = self.calculate_distance(vec_ptr, query_vec);
if distance <= range_value {
found_record = Some((*item_ptr).record_id);
break;
}
}
crate::platform::spin_unlock(&mut self.lock);
match found_record {
Some(record_id) => {
self.stats.hit_count += 1;
Ok(record_id)
}
None => Err(RemDbError::RecordNotFound),
}
}
pub unsafe fn find_range_all(
&mut self,
start_key: *const u8,
_start_key_size: usize,
end_key: *const u8,
_end_key_size: usize,
out_record_ids: *mut u16,
max_records: usize,
) -> Result<usize> {
self.stats.access_count += 1;
if out_record_ids.is_null() {
return Err(RemDbError::UnsupportedOperation);
}
let query_vec: *const f32;
let mut query_vec_buf = [0.0f32; 1024];
let vec_len = self.dimension as usize;
if _start_key_size > 4 && *start_key == b'[' {
let vec_str =
match core::str::from_utf8(core::slice::from_raw_parts(start_key, _start_key_size))
{
Ok(s) => s,
Err(_) => return Err(RemDbError::TypeMismatch),
};
let vec_str = vec_str.trim_start_matches('[').trim_end_matches(']');
let elements: Vec<&str> = vec_str.split(',').map(|s| s.trim()).collect();
if elements.len() != vec_len {
return Err(RemDbError::TypeMismatch);
}
for (i, elem) in elements.iter().enumerate() {
let parsed = match elem.parse::<f32>() {
Ok(v) => v,
Err(_) => return Err(RemDbError::TypeMismatch),
};
match query_vec_buf.get_mut(i) {
Some(slot) => *slot = parsed,
None => return Err(RemDbError::TypeMismatch),
}
}
query_vec = query_vec_buf.as_ptr();
} else {
let needed = vec_len.saturating_mul(core::mem::size_of::<f32>());
if _start_key_size < needed {
return Err(RemDbError::TypeMismatch);
}
query_vec = start_key as *const f32;
}
let range_value: f32 = if _end_key_size > 4 && *end_key == b'[' {
1000.0 } else {
core::ptr::read_unaligned(end_key as *const f32)
};
crate::platform::spin_lock(&mut self.lock);
let mut match_count = 0;
for i in 0..self.item_count {
if match_count >= max_records {
break;
}
let item_ptr = self.items.add(i);
let vec_ptr = self.vectors.add((*item_ptr).vector_offset);
let distance = self.calculate_distance(query_vec, vec_ptr);
if distance <= range_value {
*out_record_ids.add(match_count) = (*item_ptr).record_id;
match_count += 1;
}
}
crate::platform::spin_unlock(&mut self.lock);
if match_count > 0 {
self.stats.hit_count += match_count;
}
Ok(match_count)
}
pub unsafe fn delete(&mut self, key: *const u8, _key_size: usize) -> Result<()> {
crate::platform::spin_lock(&mut self.lock);
let query_vec = key as *const f32;
let vec_len = self.dimension as usize;
let mut found_idx = None;
for i in 0..self.item_count {
let item_ptr = self.items.add(i);
let vec_ptr = self.vectors.add((*item_ptr).vector_offset);
let mut match_found = true;
for j in 0..vec_len {
if *query_vec.add(j) != *vec_ptr.add(j) {
match_found = false;
break;
}
}
if match_found {
found_idx = Some(i);
break;
}
}
if let Some(idx) = found_idx {
let deleted_item_ptr = self.items.add(idx);
let deleted_offset = (*deleted_item_ptr).vector_offset;
if idx < self.item_count - 1 {
let last_item_ptr = self.items.add(self.item_count - 1);
let last_offset = (*last_item_ptr).vector_offset;
let src_vec_ptr = self.vectors.add(last_offset);
let dst_vec_ptr = self.vectors.add(deleted_offset);
for i in 0..vec_len {
*dst_vec_ptr.add(i) = *src_vec_ptr.add(i);
}
*deleted_item_ptr = *last_item_ptr;
(*deleted_item_ptr).vector_offset = deleted_offset;
}
self.item_count -= 1;
self.vector_count = self.item_count;
self.stats.item_count -= 1;
self.stats.size -=
core::mem::size_of::<VectorIndexItem>() + vec_len * core::mem::size_of::<f32>();
}
crate::platform::spin_unlock(&mut self.lock);
Ok(())
}
pub fn stats(&self) -> &IndexStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats = IndexStats::default();
}
}
pub struct JsonIndex {
def: alloc::sync::Arc<TableDef>,
btree: BTreeIndex,
json_path: Option<String>,
stats: IndexStats,
lock: u32,
}
impl JsonIndex {
pub unsafe fn new(
def: alloc::sync::Arc<TableDef>,
memory_start: *mut u8,
max_items: usize,
) -> Result<Self> {
let btree = BTreeIndex::new(def.clone(), memory_start as *mut BTreeNode, max_items);
let json_path = None;
Ok(JsonIndex {
def,
btree,
json_path,
stats: IndexStats::default(),
lock: 0,
})
}
pub fn calculate_memory_size(max_items: usize) -> usize {
BTreeIndex::calculate_memory_size(max_items)
}
fn encode_json_key(
&self,
_json_doc: &crate::json::JsonDocument,
_path: &str,
) -> Result<(Vec<u8>, usize)> {
Err(RemDbError::UnsupportedOperation)
}
pub unsafe fn insert(&mut self, key: *const u8, key_size: usize, record_id: u16) -> Result<()> {
crate::platform::spin_lock(&mut self.lock);
let result = self.btree.insert(key, key_size, record_id);
crate::platform::spin_unlock(&mut self.lock);
result
}
pub unsafe fn find(&mut self, key: *const u8, key_size: usize) -> Result<u16> {
self.stats.access_count += 1;
let result = self.btree.find(key, key_size);
if result.is_ok() {
self.stats.hit_count += 1;
}
result
}
pub unsafe fn find_range(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
) -> Result<u16> {
self.stats.access_count += 1;
let result = self
.btree
.find_range(start_key, start_key_size, end_key, end_key_size);
if result.is_ok() {
self.stats.hit_count += 1;
}
result
}
pub unsafe fn find_range_all(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
out_record_ids: *mut u16,
max_records: usize,
) -> Result<usize> {
self.stats.access_count += 1;
let result = self.btree.find_range_all(
start_key,
start_key_size,
end_key,
end_key_size,
out_record_ids,
max_records,
);
if let Ok(count) = result {
if count > 0 {
self.stats.hit_count += count;
}
}
result
}
pub unsafe fn delete(&mut self, key: *const u8, key_size: usize) -> Result<()> {
crate::platform::spin_lock(&mut self.lock);
let result = self.btree.delete(key, key_size);
crate::platform::spin_unlock(&mut self.lock);
result
}
pub fn stats(&self) -> &IndexStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats = IndexStats::default();
}
}
pub enum AnySecondaryIndex {
SortedArray(SecondaryIndex),
BTree(BTreeIndex),
TTree(TTreeIndex),
Vector(VectorIndex),
Json(JsonIndex),
}
impl AnySecondaryIndex {
pub unsafe fn new(
def: alloc::sync::Arc<TableDef>,
memory_start: *mut u8,
max_items: usize,
) -> Result<Self> {
match def.secondary_index_type {
IndexType::SortedArray => {
let index =
SecondaryIndex::new(def, memory_start as *mut SecondaryIndexItem, max_items);
Ok(AnySecondaryIndex::SortedArray(index))
}
IndexType::BTree => {
let index = BTreeIndex::new(def, memory_start as *mut BTreeNode, max_items);
Ok(AnySecondaryIndex::BTree(index))
}
IndexType::TTree => {
let index = TTreeIndex::new(def, memory_start as *mut TTreeNode, max_items);
Ok(AnySecondaryIndex::TTree(index))
}
IndexType::Vector => {
let index = VectorIndex::new(def, memory_start, max_items)?;
Ok(AnySecondaryIndex::Vector(index))
}
IndexType::Json => {
let index = JsonIndex::new(def, memory_start, max_items)?;
Ok(AnySecondaryIndex::Json(index))
}
_ => Err(RemDbError::UnsupportedOperation),
}
}
pub fn calculate_memory_size(def: &TableDef, max_items: usize) -> usize {
match def.secondary_index_type {
IndexType::SortedArray => SecondaryIndex::calculate_memory_size(max_items),
IndexType::BTree => BTreeIndex::calculate_memory_size(max_items),
IndexType::TTree => TTreeIndex::calculate_memory_size(max_items),
IndexType::Vector => VectorIndex::calculate_memory_size(def, max_items),
IndexType::Json => JsonIndex::calculate_memory_size(max_items),
_ => 0,
}
}
pub unsafe fn insert(&mut self, key: *const u8, key_size: usize, record_id: u16) -> Result<()> {
match self {
AnySecondaryIndex::SortedArray(index) => index.insert(key, key_size, record_id),
AnySecondaryIndex::BTree(index) => index.insert(key, key_size, record_id),
AnySecondaryIndex::TTree(index) => index.insert(key, key_size, record_id),
AnySecondaryIndex::Vector(index) => index.insert(key, key_size, record_id),
AnySecondaryIndex::Json(index) => index.insert(key, key_size, record_id),
}
}
pub unsafe fn find(&mut self, key: *const u8, key_size: usize) -> Result<u16> {
match self {
AnySecondaryIndex::SortedArray(index) => index.find(key, key_size),
AnySecondaryIndex::BTree(index) => index.find(key, key_size),
AnySecondaryIndex::TTree(index) => index.find(key, key_size),
AnySecondaryIndex::Vector(index) => index.find(key, key_size),
AnySecondaryIndex::Json(index) => index.find(key, key_size),
}
}
pub unsafe fn find_range(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
) -> Result<u16> {
match self {
AnySecondaryIndex::SortedArray(index) => {
index.find_range(start_key, start_key_size, end_key, end_key_size)
}
AnySecondaryIndex::BTree(index) => {
index.find_range(start_key, start_key_size, end_key, end_key_size)
}
AnySecondaryIndex::TTree(index) => {
index.find_range(start_key, start_key_size, end_key, end_key_size)
}
AnySecondaryIndex::Vector(index) => {
index.find_range(start_key, start_key_size, end_key, end_key_size)
}
AnySecondaryIndex::Json(index) => {
index.find_range(start_key, start_key_size, end_key, end_key_size)
}
}
}
pub unsafe fn find_range_all(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
out_record_ids: *mut u16,
max_records: usize,
) -> Result<usize> {
match self {
AnySecondaryIndex::SortedArray(index) => index.find_range_all(
start_key,
start_key_size,
end_key,
end_key_size,
out_record_ids,
max_records,
),
AnySecondaryIndex::BTree(index) => index.find_range_all(
start_key,
start_key_size,
end_key,
end_key_size,
out_record_ids,
max_records,
),
AnySecondaryIndex::TTree(index) => index.find_range_all(
start_key,
start_key_size,
end_key,
end_key_size,
out_record_ids,
max_records,
),
AnySecondaryIndex::Vector(index) => index.find_range_all(
start_key,
start_key_size,
end_key,
end_key_size,
out_record_ids,
max_records,
),
AnySecondaryIndex::Json(index) => index.find_range_all(
start_key,
start_key_size,
end_key,
end_key_size,
out_record_ids,
max_records,
),
}
}
pub unsafe fn delete(&mut self, key: *const u8, key_size: usize) -> Result<()> {
match self {
AnySecondaryIndex::SortedArray(index) => index.delete(key, key_size),
AnySecondaryIndex::BTree(index) => index.delete(key, key_size),
AnySecondaryIndex::TTree(index) => index.delete(key, key_size),
AnySecondaryIndex::Vector(index) => index.delete(key, key_size),
AnySecondaryIndex::Json(index) => index.delete(key, key_size),
}
}
pub fn stats(&self) -> &IndexStats {
match self {
AnySecondaryIndex::SortedArray(index) => index.stats(),
AnySecondaryIndex::BTree(index) => index.stats(),
AnySecondaryIndex::TTree(index) => index.stats(),
AnySecondaryIndex::Vector(index) => index.stats(),
AnySecondaryIndex::Json(index) => index.stats(),
}
}
pub fn reset_stats(&mut self) {
match self {
AnySecondaryIndex::SortedArray(index) => index.reset_stats(),
AnySecondaryIndex::BTree(index) => index.reset_stats(),
AnySecondaryIndex::TTree(index) => index.reset_stats(),
AnySecondaryIndex::Vector(index) => index.reset_stats(),
AnySecondaryIndex::Json(index) => index.reset_stats(),
}
}
}
pub struct SecondaryIndex {
def: alloc::sync::Arc<TableDef>,
items: NonNull<SecondaryIndexItem>,
item_count: usize,
max_items: usize,
stats: IndexStats,
lock: u32,
}
impl SecondaryIndex {
pub unsafe fn new(
def: alloc::sync::Arc<TableDef>,
items_start: *mut SecondaryIndexItem,
max_items: usize,
) -> Self {
let items = NonNull::new_unchecked(items_start);
SecondaryIndex {
def,
items,
item_count: 0,
max_items,
stats: IndexStats {
access_count: 0,
hit_count: 0,
size: max_items * core::mem::size_of::<SecondaryIndexItem>(),
item_count: 0,
},
lock: 0,
}
}
pub const fn calculate_memory_size(max_items: usize) -> usize {
max_items * core::mem::size_of::<SecondaryIndexItem>()
}
fn compare_items(
&self,
item1: &SecondaryIndexItem,
item2: &SecondaryIndexItem,
) -> core::cmp::Ordering {
if item1.key_size != item2.key_size {
return item1.key_size.cmp(&item2.key_size);
}
if let Some(secondary_index) = &self.def.secondary_index {
if let Some(field_index) = secondary_index.first() {
if *field_index < self.def.fields.len() {
let field = &self.def.fields[*field_index];
if field.data_type == crate::types::DataType::VarChar
|| field.data_type == crate::types::DataType::Char
|| field.data_type == crate::types::DataType::Text
{
if let (Some(str1), Some(str2)) = (
crate::utf8::get_global_utf8_processor()
.to_string(&item1.key_data[..item1.key_size as usize]),
crate::utf8::get_global_utf8_processor()
.to_string(&item2.key_data[..item2.key_size as usize]),
) {
let cmp = str1.cmp(str2);
if cmp != core::cmp::Ordering::Equal {
return cmp;
}
}
}
}
}
}
let key_size = item1.key_size as usize;
for i in 0..key_size {
if item1.key_data[i] != item2.key_data[i] {
return item1.key_data[i].cmp(&item2.key_data[i]);
}
}
item1.record_id.cmp(&item2.record_id)
}
fn binary_search(&self, key: *const u8, key_size: usize) -> Result<usize> {
if self.item_count == 0 {
return Err(RemDbError::RecordNotFound);
}
let mut low = 0;
let mut high = self.item_count - 1;
while low <= high {
let mid = (low + high) / 2;
let mid_item = unsafe { &*self.items.as_ptr().add(mid) };
let cmp = if mid_item.key_size != key_size as u8 {
mid_item.key_size.cmp(&(key_size as u8))
} else {
let mut equal = true;
for i in 0..key_size {
if mid_item.key_data[i] != unsafe { *key.add(i) } {
equal = false;
break;
}
}
if equal {
core::cmp::Ordering::Equal
} else {
let mut ordering = core::cmp::Ordering::Equal;
for i in 0..key_size {
let b1 = mid_item.key_data[i];
let b2 = unsafe { *key.add(i) };
if b1 != b2 {
ordering = b1.cmp(&b2);
break;
}
}
ordering
}
};
match cmp {
core::cmp::Ordering::Equal => return Ok(mid),
core::cmp::Ordering::Less => low = mid + 1,
core::cmp::Ordering::Greater => {
if mid == 0 {
return Err(RemDbError::RecordNotFound);
}
high = mid - 1;
}
}
}
Err(RemDbError::RecordNotFound)
}
pub unsafe fn insert(&mut self, key: *const u8, key_size: usize, record_id: u16) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_inserts()
}
crate::platform::spin_lock(&mut self.lock);
if self.item_count >= self.max_items {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::OutOfMemory);
}
if key_size > 64 {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::UnsupportedOperation);
}
let new_item = SecondaryIndexItem {
key_size: key_size as u8,
record_id,
key_data: [0u8; 128],
};
memcpy(new_item.key_data.as_ptr() as *mut u8, key, key_size);
let mut insert_pos = self.item_count;
if self.item_count > 0 {
let mut low = 0;
let mut high = self.item_count - 1;
while low <= high {
let mid = (low + high) / 2;
let item = &*self.items.as_ptr().add(mid);
match self.compare_items(&new_item, item) {
core::cmp::Ordering::Less => {
insert_pos = mid;
high = mid - 1;
}
core::cmp::Ordering::Greater => {
low = mid + 1;
}
core::cmp::Ordering::Equal => {
insert_pos = mid + 1;
low = mid + 1;
}
}
}
}
if insert_pos < self.item_count {
for i in (insert_pos..self.item_count).rev() {
let src = self.items.as_ptr().add(i);
let dest = self.items.as_ptr().add(i + 1);
*dest = *src;
}
}
let insert_ptr = self.items.as_ptr().add(insert_pos);
*insert_ptr = new_item;
self.item_count += 1;
self.stats.item_count = self.item_count;
crate::platform::spin_unlock(&mut self.lock);
Ok(())
}
pub unsafe fn find(&mut self, key: *const u8, key_size: usize) -> Result<u16> {
self.stats.access_count += 1;
match self.binary_search(key, key_size) {
Ok(index) => {
self.stats.hit_count += 1;
Ok((*self.items.as_ptr().add(index)).record_id)
}
Err(e) => Err(e),
}
}
pub unsafe fn delete(&mut self, key: *const u8, key_size: usize) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_deletes()
}
crate::platform::spin_lock(&mut self.lock);
let result = match self.binary_search(key, key_size) {
Ok(index) => {
if index < self.item_count - 1 {
let dest_ptr = self.items.as_ptr().add(index);
let src_ptr = self.items.as_ptr().add(index + 1);
let move_size =
(self.item_count - index - 1) * core::mem::size_of::<SecondaryIndexItem>();
memcpy(dest_ptr as *mut u8, src_ptr as *const u8, move_size);
}
let last_ptr = self.items.as_ptr().add(self.item_count - 1);
memset(
last_ptr as *mut u8,
0,
core::mem::size_of::<SecondaryIndexItem>(),
);
self.item_count -= 1;
self.stats.item_count = self.item_count;
Ok(())
}
Err(e) => Err(e),
};
crate::platform::spin_unlock(&mut self.lock);
result
}
pub unsafe fn find_range(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
) -> Result<u16> {
self.stats.access_count += 1;
let mut start_pos = 0;
let mut low = 0;
let mut high = self.item_count - 1;
let start_item = SecondaryIndexItem {
key_size: start_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(
start_item.key_data.as_ptr() as *mut u8,
start_key,
start_key_size,
);
while low <= high {
let mid = (low + high) / 2;
let item = &*self.items.as_ptr().add(mid);
if self.compare_items(item, &start_item) == core::cmp::Ordering::Less {
start_pos = mid + 1;
low = mid + 1;
} else {
high = mid - 1;
}
}
for i in start_pos..self.item_count {
let item = &*self.items.as_ptr().add(i);
let mut le_end = false;
let key_size = item.key_size as usize;
if key_size > end_key_size {
break;
}
let min_size = core::cmp::min(key_size, end_key_size);
let mut all_equal = true;
for j in 0..min_size {
if item.key_data[j] < *end_key.add(j) {
le_end = true;
break;
} else if item.key_data[j] > *end_key.add(j) {
le_end = false;
all_equal = false;
break;
}
}
if all_equal && key_size == end_key_size {
le_end = true;
}
if le_end {
self.stats.hit_count += 1;
return Ok(item.record_id);
} else {
break;
}
}
Err(RemDbError::RecordNotFound)
}
pub unsafe fn find_range_all(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
out_record_ids: *mut u16,
max_records: usize,
) -> Result<usize> {
self.stats.access_count += 1;
if out_record_ids.is_null() {
return Err(RemDbError::UnsupportedOperation);
}
if self.item_count == 0 {
return Ok(0);
}
let mut start_pos = 0;
let mut low = 0;
let mut high = self.item_count - 1;
let start_item = SecondaryIndexItem {
key_size: start_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(
start_item.key_data.as_ptr() as *mut u8,
start_key,
start_key_size,
);
while low <= high {
let mid = (low + high) / 2;
let item = &*self.items.as_ptr().add(mid);
if self.compare_items(item, &start_item) == core::cmp::Ordering::Less {
start_pos = mid + 1;
low = mid + 1;
} else {
high = mid - 1;
}
}
let mut match_count = 0;
for i in start_pos..self.item_count {
if match_count >= max_records {
break;
}
let item = &*self.items.as_ptr().add(i);
let mut le_end = false;
let key_size = item.key_size as usize;
if key_size > end_key_size {
break;
}
let min_size = core::cmp::min(key_size, end_key_size);
let mut all_equal = true;
for j in 0..min_size {
if item.key_data[j] < *end_key.add(j) {
le_end = true;
break;
} else if item.key_data[j] > *end_key.add(j) {
le_end = false;
all_equal = false;
break;
}
}
if all_equal && key_size == end_key_size {
le_end = true;
}
if le_end {
*out_record_ids.add(match_count) = item.record_id;
match_count += 1;
} else {
break;
}
}
if match_count > 0 {
self.stats.hit_count += match_count;
}
Ok(match_count)
}
pub fn stats(&self) -> &IndexStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats.access_count = 0;
self.stats.hit_count = 0;
}
pub fn item_count(&self) -> usize {
self.item_count
}
pub fn max_items(&self) -> usize {
self.max_items
}
}
impl Drop for SecondaryIndex {
fn drop(&mut self) {
unsafe {
crate::memory::allocator::free(self.items.cast::<u8>());
}
}
}
impl BTreeIndex {
pub unsafe fn new(
def: alloc::sync::Arc<TableDef>,
nodes_start: *mut BTreeNode,
max_nodes: usize,
) -> Self {
let nodes = NonNull::new_unchecked(nodes_start);
let mut free_nodes: Option<NonNull<BTreeNode>> = None;
for i in (0..max_nodes).rev() {
let node_ptr = nodes.as_ptr().add(i);
let node_mut = &mut *node_ptr;
node_mut.is_leaf = true;
node_mut.key_count = 0;
for j in 0..BTREE_ORDER {
node_mut.keys[j].key_size = 0;
node_mut.keys[j].record_id = 0;
memset(node_mut.keys[j].key_data.as_mut_ptr(), 0, 64);
}
for j in 0..(BTREE_ORDER + 1) {
node_mut.children[j] = None;
}
let next_ptr = free_nodes
.map(|p: NonNull<BTreeNode>| p.as_ptr() as u64)
.unwrap_or(0);
memcpy(
node_mut.keys[0].key_data.as_mut_ptr(),
&next_ptr as *const u64 as *const u8,
core::mem::size_of::<u64>(),
);
free_nodes = Some(NonNull::new_unchecked(node_ptr));
}
BTreeIndex {
def,
root: None,
nodes,
free_nodes,
max_nodes,
stats: IndexStats {
access_count: 0,
hit_count: 0,
size: max_nodes * core::mem::size_of::<BTreeNode>(),
item_count: 0,
},
lock: 0,
}
}
pub const fn calculate_memory_size(max_nodes: usize) -> usize {
max_nodes * core::mem::size_of::<BTreeNode>()
}
unsafe fn allocate_node(&mut self) -> Result<NonNull<BTreeNode>> {
let node_ptr = self.free_nodes.ok_or(RemDbError::OutOfMemory)?;
let node_mut = &mut *node_ptr.as_ptr();
let mut next_ptr = 0u64;
memcpy(
&mut next_ptr as *mut u64 as *mut u8,
node_mut.keys[0].key_data.as_ptr(),
core::mem::size_of::<u64>(),
);
self.free_nodes = if next_ptr == 0 {
None
} else {
NonNull::new(next_ptr as *mut BTreeNode)
};
node_mut.is_leaf = true;
node_mut.key_count = 0;
for j in 0..BTREE_ORDER {
node_mut.keys[j].key_size = 0;
node_mut.keys[j].record_id = 0;
memset(node_mut.keys[j].key_data.as_mut_ptr(), 0, 64);
}
for j in 0..(BTREE_ORDER + 1) {
node_mut.children[j] = None;
}
Ok(node_ptr)
}
unsafe fn free_node(&mut self, node_ptr: NonNull<BTreeNode>) {
let node_mut = &mut *node_ptr.as_ptr();
let next_ptr = self.free_nodes.map(|p| p.as_ptr() as u64).unwrap_or(0);
memcpy(
node_mut.keys[0].key_data.as_mut_ptr(),
&next_ptr as *const u64 as *const u8,
core::mem::size_of::<u64>(),
);
self.free_nodes = Some(node_ptr);
}
fn compare_items(
&self,
item1: &SecondaryIndexItem,
item2: &SecondaryIndexItem,
) -> core::cmp::Ordering {
if item1.key_size != item2.key_size {
return item1.key_size.cmp(&item2.key_size);
}
if let Some(secondary_index) = &self.def.secondary_index {
if let Some(field_index) = secondary_index.first() {
if *field_index < self.def.fields.len() {
let field = &self.def.fields[*field_index];
if field.data_type == crate::types::DataType::VarChar
|| field.data_type == crate::types::DataType::Char
|| field.data_type == crate::types::DataType::Text
{
if let (Some(str1), Some(str2)) = (
crate::utf8::get_global_utf8_processor()
.to_string(&item1.key_data[..item1.key_size as usize]),
crate::utf8::get_global_utf8_processor()
.to_string(&item2.key_data[..item2.key_size as usize]),
) {
let cmp = str1.cmp(str2);
if cmp != core::cmp::Ordering::Equal {
return cmp;
}
}
}
}
}
}
let key_size = item1.key_size as usize;
for i in 0..key_size {
if item1.key_data[i] != item2.key_data[i] {
return item1.key_data[i].cmp(&item2.key_data[i]);
}
}
item1.record_id.cmp(&item2.record_id)
}
fn find_key_position(&self, node: &BTreeNode, key: &SecondaryIndexItem) -> usize {
let mut pos = 0;
while pos < node.key_count as usize
&& self.compare_items(&node.keys[pos], key) == core::cmp::Ordering::Less
{
pos += 1;
}
pos
}
unsafe fn split_child(
&mut self,
mut parent: NonNull<BTreeNode>,
child_idx: usize,
mut child: NonNull<BTreeNode>,
) -> Result<()> {
let parent_mut = parent.as_mut();
let child_mut = child.as_mut();
let mut new_node = self.allocate_node()?;
let new_node_mut = new_node.as_mut();
new_node_mut.is_leaf = child_mut.is_leaf;
new_node_mut.key_count = (BTREE_ORDER / 2) as u8;
for i in 0..(BTREE_ORDER / 2) {
new_node_mut.keys[i] = child_mut.keys[i + (BTREE_ORDER / 2) + 1];
}
if !child_mut.is_leaf {
for i in 0..(BTREE_ORDER / 2 + 1) {
new_node_mut.children[i] = child_mut.children[i + (BTREE_ORDER / 2) + 1];
}
}
child_mut.key_count = (BTREE_ORDER / 2) as u8;
for i in (child_idx + 1..=parent_mut.key_count as usize).rev() {
parent_mut.keys[i] = parent_mut.keys[i - 1];
parent_mut.children[i + 1] = parent_mut.children[i];
}
parent_mut.keys[child_idx] = child_mut.keys[BTREE_ORDER / 2];
parent_mut.children[child_idx + 1] = Some(new_node);
parent_mut.key_count += 1;
Ok(())
}
unsafe fn insert_non_full(
&mut self,
mut node: NonNull<BTreeNode>,
key: SecondaryIndexItem,
) -> Result<()> {
let node_mut = node.as_mut();
let mut pos = self.find_key_position(node_mut, &key);
if node_mut.is_leaf {
for i in (pos..node_mut.key_count as usize).rev() {
node_mut.keys[i + 1] = node_mut.keys[i];
}
node_mut.keys[pos] = key;
node_mut.key_count += 1;
self.stats.item_count += 1;
} else {
let child =
node_mut.children[pos].ok_or(RemDbError::InvalidData("Child node not found"))?;
if child.as_ref().key_count == BTREE_ORDER as u8 {
self.split_child(node, pos, child)?;
if self.compare_items(&node_mut.keys[pos], &key) == core::cmp::Ordering::Less {
pos += 1;
}
}
self.insert_non_full(
node_mut.children[pos]
.ok_or(RemDbError::InvalidData("Child node not found after split"))?,
key,
)?;
}
Ok(())
}
pub unsafe fn insert(&mut self, key: *const u8, key_size: usize, record_id: u16) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_inserts()
}
crate::platform::spin_lock(&mut self.lock);
if key_size > 64 {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::UnsupportedOperation);
}
let mut new_item = SecondaryIndexItem {
key_size: key_size as u8,
record_id,
key_data: [0u8; 128],
};
memcpy(new_item.key_data.as_mut_ptr(), key, key_size);
if self.root.is_none() {
let mut root_node = self.allocate_node()?;
let root_mut = root_node.as_mut();
root_mut.keys[0] = new_item;
root_mut.key_count = 1;
self.root = Some(root_node);
} else {
let root = self
.root
.ok_or(RemDbError::InvalidData("Root node missing"))?;
if root.as_ref().key_count == BTREE_ORDER as u8 {
let mut new_root = self.allocate_node()?;
let new_root_mut = new_root.as_mut();
new_root_mut.is_leaf = false;
new_root_mut.key_count = 0;
new_root_mut.children[0] = self.root;
let _ = self.split_child(new_root, 0, root);
let _ = self.insert_non_full(new_root, new_item);
self.root = Some(new_root);
} else {
let _ = self.insert_non_full(root, new_item);
}
}
crate::platform::spin_unlock(&mut self.lock);
Ok(())
}
pub unsafe fn find(&mut self, key: *const u8, key_size: usize) -> Result<u16> {
self.stats.access_count += 1;
if key_size > 64 {
return Err(RemDbError::UnsupportedOperation);
}
let mut search_item = SecondaryIndexItem {
key_size: key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(search_item.key_data.as_mut_ptr(), key, key_size);
let mut current = self.root;
while let Some(node) = current {
let node_ref = node.as_ref();
let mut pos = 0;
while pos < node_ref.key_count as usize
&& self.compare_items(&node_ref.keys[pos], &search_item)
== core::cmp::Ordering::Less
{
pos += 1;
}
if pos < node_ref.key_count as usize {
let cmp = self.compare_items(&node_ref.keys[pos], &search_item);
if cmp == core::cmp::Ordering::Equal {
self.stats.hit_count += 1;
return Ok(node_ref.keys[pos].record_id);
}
}
if node_ref.is_leaf {
break;
}
current = node_ref.children[pos];
}
Err(RemDbError::RecordNotFound)
}
pub unsafe fn find_range(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
) -> Result<u16> {
self.stats.access_count += 1;
let mut start_item = SecondaryIndexItem {
key_size: start_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(start_item.key_data.as_mut_ptr(), start_key, start_key_size);
let mut end_item = SecondaryIndexItem {
key_size: end_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(end_item.key_data.as_mut_ptr(), end_key, end_key_size);
let mut current = self.root;
let mut stack = [None; 64]; let mut stack_size = 0;
while let Some(node) = current {
stack[stack_size] = Some(node);
stack_size += 1;
let node_ref = node.as_ref();
if node_ref.is_leaf {
break;
}
let mut pos = 0;
while pos < node_ref.key_count as usize
&& self.compare_items(&node_ref.keys[pos], &start_item) == core::cmp::Ordering::Less
{
pos += 1;
}
current = node_ref.children[pos];
}
while stack_size > 0 {
stack_size -= 1;
let node = stack[stack_size].ok_or(RemDbError::InvalidData("Stack underflow"))?;
let node_ref = node.as_ref();
let mut start_pos = 0;
while start_pos < node_ref.key_count as usize
&& self.compare_items(&node_ref.keys[start_pos], &start_item)
== core::cmp::Ordering::Less
{
start_pos += 1;
}
for i in start_pos..node_ref.key_count as usize {
let key = &node_ref.keys[i];
if self.compare_items(key, &end_item) == core::cmp::Ordering::Greater {
continue; }
self.stats.hit_count += 1;
return Ok(key.record_id);
}
if !node_ref.is_leaf {
let mut child = node_ref.children[node_ref.key_count as usize];
while let Some(child_node) = child {
let child_ref = child_node.as_ref();
for i in 0..child_ref.key_count as usize {
let key = &child_ref.keys[i];
if self.compare_items(key, &end_item) == core::cmp::Ordering::Greater {
break; }
self.stats.hit_count += 1;
return Ok(key.record_id);
}
if child_ref.is_leaf {
break;
}
child = child_ref.children[0];
}
}
}
Err(RemDbError::RecordNotFound)
}
pub unsafe fn find_range_all(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
out_record_ids: *mut u16,
max_records: usize,
) -> Result<usize> {
self.stats.access_count += 1;
if out_record_ids.is_null() {
return Err(RemDbError::UnsupportedOperation);
}
let mut start_item = SecondaryIndexItem {
key_size: start_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(start_item.key_data.as_mut_ptr(), start_key, start_key_size);
let mut end_item = SecondaryIndexItem {
key_size: end_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(end_item.key_data.as_mut_ptr(), end_key, end_key_size);
let mut match_count = 0;
let mut stack = [None; 64]; let mut stack_size = 0;
let mut current = self.root;
while let Some(node) = current {
stack[stack_size] = Some(node);
stack_size += 1;
let node_ref = node.as_ref();
if node_ref.is_leaf {
break;
}
let mut pos = 0;
while pos < node_ref.key_count as usize
&& self.compare_items(&node_ref.keys[pos], &start_item) == core::cmp::Ordering::Less
{
pos += 1;
}
current = node_ref.children[pos];
}
while stack_size > 0 && match_count < max_records {
stack_size -= 1;
let node = stack[stack_size].ok_or(RemDbError::InvalidData("Stack underflow"))?;
let node_ref = node.as_ref();
let mut start_pos = 0;
while start_pos < node_ref.key_count as usize
&& self.compare_items(&node_ref.keys[start_pos], &start_item)
== core::cmp::Ordering::Less
{
start_pos += 1;
}
for i in start_pos..node_ref.key_count as usize {
if match_count >= max_records {
break;
}
let key = &node_ref.keys[i];
if self.compare_items(key, &end_item) == core::cmp::Ordering::Greater {
continue; }
*out_record_ids.add(match_count) = key.record_id;
match_count += 1;
}
if !node_ref.is_leaf && match_count < max_records {
let mut child = node_ref.children[node_ref.key_count as usize];
while let Some(child_node) = child {
let child_ref = child_node.as_ref();
for i in 0..child_ref.key_count as usize {
if match_count >= max_records {
break;
}
let key = &child_ref.keys[i];
if self.compare_items(key, &end_item) == core::cmp::Ordering::Greater {
break; }
*out_record_ids.add(match_count) = key.record_id;
match_count += 1;
}
if child_ref.is_leaf || match_count >= max_records {
break;
}
child = child_ref.children[0];
}
}
}
if match_count > 0 {
self.stats.hit_count += match_count;
}
Ok(match_count)
}
pub unsafe fn delete(&mut self, _key: *const u8, _key_size: usize) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_deletes()
}
crate::platform::spin_lock(&mut self.lock);
crate::platform::spin_unlock(&mut self.lock);
Err(RemDbError::UnsupportedOperation)
}
pub fn stats(&self) -> &IndexStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats.access_count = 0;
self.stats.hit_count = 0;
}
}
impl Drop for BTreeIndex {
fn drop(&mut self) {
unsafe {
crate::memory::allocator::free(self.nodes.cast::<u8>());
}
}
}
impl TTreeIndex {
pub unsafe fn new(
def: alloc::sync::Arc<TableDef>,
nodes_start: *mut TTreeNode,
max_nodes: usize,
) -> Self {
let nodes = NonNull::new_unchecked(nodes_start);
let mut free_nodes: Option<NonNull<TTreeNode>> = None;
for i in (0..max_nodes).rev() {
let node_ptr = nodes.as_ptr().add(i);
let node_mut = &mut *node_ptr;
node_mut.key_count = 0;
for j in 0..TTREE_ORDER {
node_mut.keys[j].key_size = 0;
node_mut.keys[j].record_id = 0;
memset(node_mut.keys[j].key_data.as_mut_ptr(), 0, 64);
}
node_mut.left = None;
node_mut.middle = None;
node_mut.right = None;
let next_ptr = free_nodes
.map(|p: NonNull<TTreeNode>| p.as_ptr() as u64)
.unwrap_or(0);
memcpy(
node_mut.keys[0].key_data.as_mut_ptr(),
&next_ptr as *const u64 as *const u8,
core::mem::size_of::<u64>(),
);
free_nodes = Some(NonNull::new_unchecked(node_ptr));
}
TTreeIndex {
def,
root: None,
nodes,
free_nodes,
max_nodes,
stats: IndexStats {
access_count: 0,
hit_count: 0,
size: max_nodes * core::mem::size_of::<TTreeNode>(),
item_count: 0,
},
lock: 0,
}
}
pub const fn calculate_memory_size(max_nodes: usize) -> usize {
max_nodes * core::mem::size_of::<TTreeNode>()
}
unsafe fn allocate_node(&mut self) -> Result<NonNull<TTreeNode>> {
let node_ptr = self.free_nodes.ok_or(RemDbError::OutOfMemory)?;
let node_mut = &mut *node_ptr.as_ptr();
let mut next_ptr = 0u64;
memcpy(
&mut next_ptr as *mut u64 as *mut u8,
node_mut.keys[0].key_data.as_ptr(),
core::mem::size_of::<u64>(),
);
self.free_nodes = if next_ptr == 0 {
None
} else {
NonNull::new(next_ptr as *mut TTreeNode)
};
node_mut.key_count = 0;
for j in 0..TTREE_ORDER {
node_mut.keys[j].key_size = 0;
node_mut.keys[j].record_id = 0;
memset(node_mut.keys[j].key_data.as_mut_ptr(), 0, 64);
}
node_mut.left = None;
node_mut.middle = None;
node_mut.right = None;
Ok(node_ptr)
}
unsafe fn free_node(&mut self, node_ptr: NonNull<TTreeNode>) {
let node_mut = &mut *node_ptr.as_ptr();
let next_ptr = self.free_nodes.map(|p| p.as_ptr() as u64).unwrap_or(0);
memcpy(
node_mut.keys[0].key_data.as_mut_ptr(),
&next_ptr as *const u64 as *const u8,
core::mem::size_of::<u64>(),
);
self.free_nodes = Some(node_ptr);
}
fn compare_items(
&self,
item1: &SecondaryIndexItem,
item2: &SecondaryIndexItem,
) -> core::cmp::Ordering {
if item1.key_size != item2.key_size {
return item1.key_size.cmp(&item2.key_size);
}
let key_size = item1.key_size as usize;
for i in 0..key_size {
if item1.key_data[i] != item2.key_data[i] {
return item1.key_data[i].cmp(&item2.key_data[i]);
}
}
item1.record_id.cmp(&item2.record_id)
}
fn find_key_position(
&self,
node: &TTreeNode,
key: &SecondaryIndexItem,
) -> (usize, core::cmp::Ordering) {
let mut pos = 0;
while pos < node.key_count as usize {
let cmp = self.compare_items(&node.keys[pos], key);
match cmp {
core::cmp::Ordering::Less => pos += 1,
_ => return (pos, cmp),
}
}
(pos, core::cmp::Ordering::Less)
}
unsafe fn insert_into_node(&mut self, mut node: NonNull<TTreeNode>, key: SecondaryIndexItem) {
let node_mut = node.as_mut();
let (pos, _) = self.find_key_position(node_mut, &key);
for i in (pos..node_mut.key_count as usize).rev() {
node_mut.keys[i + 1] = node_mut.keys[i];
}
node_mut.keys[pos] = key;
node_mut.key_count += 1;
self.stats.item_count += 1;
}
pub unsafe fn insert(&mut self, key: *const u8, key_size: usize, record_id: u16) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_inserts()
}
crate::platform::spin_lock(&mut self.lock);
if key_size > 64 {
crate::platform::spin_unlock(&mut self.lock);
return Err(RemDbError::UnsupportedOperation);
}
let mut new_item = SecondaryIndexItem {
key_size: key_size as u8,
record_id,
key_data: [0u8; 128],
};
memcpy(new_item.key_data.as_mut_ptr(), key, key_size);
if self.root.is_none() {
let mut root_node = self.allocate_node()?;
let root_mut = root_node.as_mut();
root_mut.keys[0] = new_item;
root_mut.key_count = 1;
self.root = Some(root_node);
} else {
let mut root = self
.root
.ok_or(RemDbError::InvalidData("Root node missing"))?;
if root.as_ref().key_count == TTREE_ORDER as u8 {
let mut new_root = self.allocate_node()?;
let new_root_mut = new_root.as_mut();
let mut keys = [SecondaryIndexItem::default(); TTREE_ORDER + 1];
for i in 0..TTREE_ORDER {
keys[i] = root.as_ref().keys[i];
}
let mut inserted = false;
for i in 0..TTREE_ORDER {
if self.compare_items(&keys[i], &new_item) == core::cmp::Ordering::Greater {
for j in (i..TTREE_ORDER).rev() {
keys[j + 1] = keys[j];
}
keys[i] = new_item;
inserted = true;
break;
}
}
if !inserted {
keys[TTREE_ORDER] = new_item;
}
let mut right_node = self.allocate_node()?;
let right_mut = right_node.as_mut();
let mid = TTREE_ORDER.div_ceil(2);
let root_mut = root.as_mut();
root_mut.key_count = mid as u8;
for i in 0..mid {
root_mut.keys[i] = keys[i];
}
right_mut.key_count = ((TTREE_ORDER + 1) - mid) as u8;
for i in 0..right_mut.key_count as usize {
right_mut.keys[i] = keys[mid + i];
}
new_root_mut.keys[0] = keys[mid - 1];
new_root_mut.key_count = 1;
new_root_mut.left = Some(root);
new_root_mut.right = Some(right_node);
self.root = Some(new_root);
} else {
let _ = self.insert_recursive_with_depth(root, new_item, 0);
}
}
crate::platform::spin_unlock(&mut self.lock);
Ok(())
}
unsafe fn insert_recursive_with_depth(
&mut self,
mut node: NonNull<TTreeNode>,
key: SecondaryIndexItem,
depth: usize,
) -> Result<()> {
const MAX_RECURSION_DEPTH: usize = 100;
if depth > MAX_RECURSION_DEPTH {
#[cfg(feature = "log")]
crate::log::error!("T-Tree recursion depth exceeded, potential infinite recursion");
return Ok(());
}
let node_mut = node.as_mut();
let (pos, cmp) = self.find_key_position(node_mut, &key);
if cmp == core::cmp::Ordering::Equal {
node_mut.keys[pos].record_id = key.record_id;
return Ok(());
}
let child = if pos == 0 {
&mut node_mut.left
} else if pos < node_mut.key_count as usize {
&mut node_mut.middle
} else {
&mut node_mut.right
};
if let Some(mut child_node) = *child {
if child_node.as_ref().key_count == TTREE_ORDER as u8 {
let mut keys = [SecondaryIndexItem::default(); TTREE_ORDER + 1];
for i in 0..TTREE_ORDER {
keys[i] = child_node.as_ref().keys[i];
}
let mut inserted = false;
for i in 0..TTREE_ORDER {
if self.compare_items(&keys[i], &key) == core::cmp::Ordering::Greater {
for j in (i..TTREE_ORDER).rev() {
keys[j + 1] = keys[j];
}
keys[i] = key;
inserted = true;
break;
}
}
if !inserted {
keys[TTREE_ORDER] = key;
}
let mut new_right = self.allocate_node()?;
let new_right_mut = new_right.as_mut();
let mid = TTREE_ORDER.div_ceil(2);
let child_mut = child_node.as_mut();
child_mut.key_count = mid as u8;
for i in 0..mid {
child_mut.keys[i] = keys[i];
}
new_right_mut.key_count = ((TTREE_ORDER + 1) - mid) as u8;
for i in 0..new_right_mut.key_count as usize {
new_right_mut.keys[i] = keys[mid + i];
}
let promoted_key = keys[mid - 1];
self.insert_into_node(node, promoted_key);
let (promoted_pos, _) = self.find_key_position(node_mut, &promoted_key);
if promoted_pos == 0 {
node_mut.left = Some(child_node);
node_mut.middle = Some(new_right);
} else if promoted_pos < node_mut.key_count as usize {
node_mut.middle = Some(child_node);
node_mut.right = Some(new_right);
} else {
node_mut.right = Some(new_right);
}
Ok(())
} else {
self.insert_recursive_with_depth(child_node, key, depth + 1)?;
Ok(())
}
} else {
if node_mut.key_count < TTREE_ORDER as u8 {
self.insert_into_node(node, key);
Ok(())
} else {
let mut new_node = self.allocate_node()?;
let new_node_mut = new_node.as_mut();
let mut keys = [SecondaryIndexItem::default(); TTREE_ORDER + 1];
for i in 0..TTREE_ORDER {
keys[i] = node_mut.keys[i];
}
let mut inserted = false;
for i in 0..TTREE_ORDER {
if self.compare_items(&keys[i], &key) == core::cmp::Ordering::Greater {
for j in (i..TTREE_ORDER).rev() {
keys[j + 1] = keys[j];
}
keys[i] = key;
inserted = true;
break;
}
}
if !inserted {
keys[TTREE_ORDER] = key;
}
let mid = TTREE_ORDER.div_ceil(2);
node_mut.key_count = mid as u8;
for i in 0..mid {
node_mut.keys[i] = keys[i];
}
new_node_mut.key_count = ((TTREE_ORDER + 1) - mid) as u8;
for i in 0..new_node_mut.key_count as usize {
new_node_mut.keys[i] = keys[mid + i];
}
if pos == 0 {
node_mut.left = Some(new_node);
} else if pos < node_mut.key_count as usize {
node_mut.middle = Some(new_node);
} else {
node_mut.right = Some(new_node);
}
Ok(())
}
}
}
pub unsafe fn find(&mut self, key: *const u8, key_size: usize) -> Result<u16> {
self.stats.access_count += 1;
if key_size > 64 {
return Err(RemDbError::UnsupportedOperation);
}
let mut search_item = SecondaryIndexItem {
key_size: key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(search_item.key_data.as_mut_ptr(), key, key_size);
let mut current = self.root;
while let Some(node) = current {
let node_ref = node.as_ref();
let (pos, cmp) = self.find_key_position(node_ref, &search_item);
if cmp == core::cmp::Ordering::Equal {
self.stats.hit_count += 1;
return Ok(node_ref.keys[pos].record_id);
}
current = if pos == 0 {
node_ref.left
} else if pos < node_ref.key_count as usize {
node_ref.middle
} else {
node_ref.right
};
}
Err(RemDbError::RecordNotFound)
}
pub unsafe fn find_range(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
) -> Result<u16> {
self.stats.access_count += 1;
let mut start_item = SecondaryIndexItem {
key_size: start_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(start_item.key_data.as_mut_ptr(), start_key, start_key_size);
let mut end_item = SecondaryIndexItem {
key_size: end_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(end_item.key_data.as_mut_ptr(), end_key, end_key_size);
let mut stack = [None; 64]; let mut stack_size = 0;
let mut current = self.root;
while let Some(node) = current {
stack[stack_size] = Some(node);
stack_size += 1;
current = node.as_ref().left;
}
while stack_size > 0 {
stack_size -= 1;
let node = stack[stack_size].ok_or(RemDbError::InvalidData("Stack underflow"))?;
let node_ref = node.as_ref();
for i in 0..node_ref.key_count as usize {
let key = &node_ref.keys[i];
if self.compare_items(key, &start_item) != core::cmp::Ordering::Less
&& self.compare_items(key, &end_item) != core::cmp::Ordering::Greater
{
self.stats.hit_count += 1;
return Ok(key.record_id);
}
if self.compare_items(key, &end_item) == core::cmp::Ordering::Greater {
break;
}
}
let mut child = node_ref.right;
while let Some(child_node) = child {
stack[stack_size] = Some(child_node);
stack_size += 1;
child = child_node.as_ref().left;
}
}
Err(RemDbError::RecordNotFound)
}
pub unsafe fn find_range_all(
&mut self,
start_key: *const u8,
start_key_size: usize,
end_key: *const u8,
end_key_size: usize,
out_record_ids: *mut u16,
max_records: usize,
) -> Result<usize> {
self.stats.access_count += 1;
if out_record_ids.is_null() {
return Err(RemDbError::UnsupportedOperation);
}
let mut start_item = SecondaryIndexItem {
key_size: start_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(start_item.key_data.as_mut_ptr(), start_key, start_key_size);
let mut end_item = SecondaryIndexItem {
key_size: end_key_size as u8,
record_id: 0,
key_data: [0u8; 128],
};
memcpy(end_item.key_data.as_mut_ptr(), end_key, end_key_size);
let mut match_count = 0;
let mut stack = [None; 64]; let mut stack_size = 0;
let mut current = self.root;
while let Some(node) = current {
stack[stack_size] = Some(node);
stack_size += 1;
current = node.as_ref().left;
}
while stack_size > 0 && match_count < max_records {
stack_size -= 1;
let node = stack[stack_size].ok_or(RemDbError::InvalidData("Stack underflow"))?;
let node_ref = node.as_ref();
for i in 0..node_ref.key_count as usize {
if match_count >= max_records {
break;
}
let key = &node_ref.keys[i];
if self.compare_items(key, &start_item) != core::cmp::Ordering::Less
&& self.compare_items(key, &end_item) != core::cmp::Ordering::Greater
{
*out_record_ids.add(match_count) = key.record_id;
match_count += 1;
}
if self.compare_items(key, &end_item) == core::cmp::Ordering::Greater {
break;
}
}
let mut child = node_ref.right;
while let Some(child_node) = child {
stack[stack_size] = Some(child_node);
stack_size += 1;
child = child_node.as_ref().left;
}
}
if match_count > 0 {
self.stats.hit_count += match_count;
}
Ok(match_count)
}
pub unsafe fn delete(&mut self, _key: *const u8, _key_size: usize) -> Result<()> {
if let Some(db) = crate::get_global_db() {
db.metrics.inc_index_deletes()
}
crate::platform::spin_lock(&mut self.lock);
crate::platform::spin_unlock(&mut self.lock);
Err(RemDbError::UnsupportedOperation)
}
pub fn stats(&self) -> &IndexStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats.access_count = 0;
self.stats.hit_count = 0;
}
}