pub const MAX_KEY_LENGTH: usize = 63;
#[repr(C, align(64))]
#[derive(Clone)]
pub struct KeyBuffer {
len: u8,
data: [u8; MAX_KEY_LENGTH],
}
impl KeyBuffer {
#[inline(always)]
pub const fn new() -> Self {
Self {
len: 0,
data: [0; MAX_KEY_LENGTH],
}
}
#[inline]
pub fn format_row_key(table: &str, row_id: u64) -> Self {
let mut buf = Self::new();
let table_bytes = table.as_bytes();
let table_len = table_bytes.len().min(MAX_KEY_LENGTH - 20); buf.data[..table_len].copy_from_slice(&table_bytes[..table_len]);
buf.len = table_len as u8;
if buf.len < MAX_KEY_LENGTH as u8 {
buf.data[buf.len as usize] = b'/';
buf.len += 1;
}
buf.write_u64(row_id);
buf
}
#[inline]
pub fn format_column_key(table: &str, row_id: u64, column: &str) -> Self {
let mut buf = Self::format_row_key(table, row_id);
if buf.len < MAX_KEY_LENGTH as u8 {
buf.data[buf.len as usize] = b'/';
buf.len += 1;
}
let col_bytes = column.as_bytes();
let available = MAX_KEY_LENGTH - buf.len as usize;
let col_len = col_bytes.len().min(available);
buf.data[buf.len as usize..buf.len as usize + col_len]
.copy_from_slice(&col_bytes[..col_len]);
buf.len += col_len as u8;
buf
}
#[inline]
fn write_u64(&mut self, mut value: u64) {
if value == 0 {
if self.len < MAX_KEY_LENGTH as u8 {
self.data[self.len as usize] = b'0';
self.len += 1;
}
return;
}
let mut temp = value;
let mut digit_count = 0u8;
while temp > 0 {
digit_count += 1;
temp /= 10;
}
if self.len as usize + digit_count as usize > MAX_KEY_LENGTH {
return; }
let start = self.len as usize;
let end = start + digit_count as usize;
self.len += digit_count;
let mut pos = end;
while value > 0 {
pos -= 1;
self.data[pos] = b'0' + (value % 10) as u8;
value /= 10;
}
}
#[inline]
pub fn append(&mut self, bytes: &[u8]) {
let available = MAX_KEY_LENGTH - self.len as usize;
let copy_len = bytes.len().min(available);
self.data[self.len as usize..self.len as usize + copy_len]
.copy_from_slice(&bytes[..copy_len]);
self.len += copy_len as u8;
}
#[inline]
pub fn push(&mut self, byte: u8) {
if self.len < MAX_KEY_LENGTH as u8 {
self.data[self.len as usize] = byte;
self.len += 1;
}
}
#[inline(always)]
pub fn as_bytes(&self) -> &[u8] {
&self.data[..self.len as usize]
}
#[inline(always)]
pub fn len(&self) -> usize {
self.len as usize
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline(always)]
pub fn clear(&mut self) {
self.len = 0;
}
#[inline(always)]
pub fn remaining(&self) -> usize {
MAX_KEY_LENGTH - self.len as usize
}
}
impl Default for KeyBuffer {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for KeyBuffer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match std::str::from_utf8(self.as_bytes()) {
Ok(s) => write!(f, "KeyBuffer({:?})", s),
Err(_) => write!(f, "KeyBuffer({:?})", self.as_bytes()),
}
}
}
impl AsRef<[u8]> for KeyBuffer {
fn as_ref(&self) -> &[u8] {
self.as_bytes()
}
}
#[repr(C, align(64))]
pub struct InternedTablePrefix {
prefix: [u8; 32],
prefix_len: u8,
hash: u64,
}
impl InternedTablePrefix {
pub fn new(table: &str) -> Self {
let mut prefix = [0u8; 32];
let bytes = table.as_bytes();
let len = bytes.len().min(30); prefix[..len].copy_from_slice(&bytes[..len]);
prefix[len] = b'/';
let hash = Self::compute_hash(&prefix[..len + 1]);
Self {
prefix,
prefix_len: (len + 1) as u8,
hash,
}
}
#[inline]
fn compute_hash(bytes: &[u8]) -> u64 {
const K: u64 = 0x517cc1b727220a95;
let mut hash = 0u64;
for &byte in bytes {
hash = (hash.rotate_left(5) ^ (byte as u64)).wrapping_mul(K);
}
hash
}
#[inline]
pub fn make_row_key(&self, row_id: u64) -> KeyBuffer {
let mut buf = KeyBuffer::new();
buf.data[..self.prefix_len as usize]
.copy_from_slice(&self.prefix[..self.prefix_len as usize]);
buf.len = self.prefix_len;
buf.write_u64(row_id);
buf
}
#[inline]
pub fn make_column_key(&self, row_id: u64, column: &str) -> KeyBuffer {
let mut buf = self.make_row_key(row_id);
buf.push(b'/');
buf.append(column.as_bytes());
buf
}
#[inline]
pub fn prefix(&self) -> &[u8] {
&self.prefix[..self.prefix_len as usize]
}
#[inline]
pub fn hash(&self) -> u64 {
self.hash
}
#[inline]
pub fn same_table(&self, other: &Self) -> bool {
self.hash == other.hash
&& self.prefix_len == other.prefix_len
&& self.prefix[..self.prefix_len as usize] == other.prefix[..other.prefix_len as usize]
}
}
impl std::fmt::Debug for InternedTablePrefix {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match std::str::from_utf8(self.prefix()) {
Ok(s) => write!(f, "InternedTablePrefix({:?})", s),
Err(_) => write!(f, "InternedTablePrefix({:?})", self.prefix()),
}
}
}
pub struct BatchKeyGenerator {
prefix: InternedTablePrefix,
buffer: KeyBuffer,
}
impl BatchKeyGenerator {
pub fn new(table: &str) -> Self {
Self {
prefix: InternedTablePrefix::new(table),
buffer: KeyBuffer::new(),
}
}
#[inline]
pub fn row_key(&mut self, row_id: u64) -> &[u8] {
self.buffer = self.prefix.make_row_key(row_id);
self.buffer.as_bytes()
}
#[inline]
pub fn column_key(&mut self, row_id: u64, column: &str) -> &[u8] {
self.buffer = self.prefix.make_column_key(row_id, column);
self.buffer.as_bytes()
}
#[inline]
pub fn prefix(&self) -> &InternedTablePrefix {
&self.prefix
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_key_buffer_row_key() {
let key = KeyBuffer::format_row_key("users", 12345);
assert_eq!(key.as_bytes(), b"users/12345");
}
#[test]
fn test_key_buffer_column_key() {
let key = KeyBuffer::format_column_key("users", 42, "name");
assert_eq!(key.as_bytes(), b"users/42/name");
}
#[test]
fn test_key_buffer_zero_id() {
let key = KeyBuffer::format_row_key("table", 0);
assert_eq!(key.as_bytes(), b"table/0");
}
#[test]
fn test_key_buffer_large_id() {
let key = KeyBuffer::format_row_key("t", u64::MAX);
let expected = format!("t/{}", u64::MAX);
assert_eq!(key.as_bytes(), expected.as_bytes());
}
#[test]
fn test_interned_prefix() {
let prefix = InternedTablePrefix::new("orders");
let key = prefix.make_row_key(999);
assert_eq!(key.as_bytes(), b"orders/999");
}
#[test]
fn test_interned_column_key() {
let prefix = InternedTablePrefix::new("products");
let key = prefix.make_column_key(100, "price");
assert_eq!(key.as_bytes(), b"products/100/price");
}
#[test]
fn test_batch_generator() {
let mut generator = BatchKeyGenerator::new("items");
assert_eq!(generator.row_key(1), b"items/1");
assert_eq!(generator.row_key(2), b"items/2");
assert_eq!(generator.column_key(3, "qty"), b"items/3/qty");
}
#[test]
fn test_same_table_check() {
let p1 = InternedTablePrefix::new("users");
let p2 = InternedTablePrefix::new("users");
let p3 = InternedTablePrefix::new("orders");
assert!(p1.same_table(&p2));
assert!(!p1.same_table(&p3));
}
#[test]
fn test_cache_line_alignment() {
assert_eq!(std::mem::align_of::<KeyBuffer>(), 64);
assert_eq!(std::mem::align_of::<InternedTablePrefix>(), 64);
}
#[test]
fn test_no_heap_allocation() {
let key = KeyBuffer::format_column_key("users", 12345678901234567890, "email_address");
assert!(key.len() <= MAX_KEY_LENGTH);
}
#[test]
fn test_arena_basic() {
KeyArena::with(|arena| {
let key1 = arena.alloc_key(b"test_key_1");
let key2 = arena.alloc_key(b"test_key_2");
assert_eq!(key1.as_bytes(), b"test_key_1");
assert_eq!(key2.as_bytes(), b"test_key_2");
});
}
#[test]
fn test_arena_reset() {
KeyArena::with(|arena| {
for i in 0..100 {
let key = format!("key_{}", i);
arena.alloc_key(key.as_bytes());
}
let used_before = arena.bytes_used();
assert!(used_before > 0);
arena.reset();
assert_eq!(arena.bytes_used(), 0);
let key = arena.alloc_key(b"after_reset");
assert_eq!(key.as_bytes(), b"after_reset");
});
}
#[test]
fn test_arena_large_allocation() {
KeyArena::with(|arena| {
let large_key = vec![b'x'; 1024];
let key = arena.alloc_key(&large_key);
assert_eq!(key.as_bytes(), large_key.as_slice());
});
}
#[test]
fn test_arena_key_handle_inline() {
let key = ArenaKeyHandle::new(b"short_key");
assert!(key.is_inline());
assert_eq!(key.as_bytes(), b"short_key");
assert_eq!(key.len(), 9);
}
#[test]
fn test_arena_key_handle_heap() {
let large = vec![b'x'; 50];
let key = ArenaKeyHandle::new(&large);
assert!(!key.is_inline());
assert_eq!(key.as_bytes(), large.as_slice());
assert_eq!(key.len(), 50);
}
#[test]
fn test_arena_key_handle_hash() {
let key1 = ArenaKeyHandle::new(b"test_key");
let key2 = ArenaKeyHandle::new(b"test_key");
assert_eq!(key1.hash(), key2.hash());
let key3 = ArenaKeyHandle::new(b"other_key");
assert_ne!(key1.hash(), key3.hash());
}
#[test]
fn test_arena_key_handle_equality() {
let key1 = ArenaKeyHandle::new(b"test");
let key2 = ArenaKeyHandle::new(b"test");
let key3 = ArenaKeyHandle::new(b"other");
assert_eq!(key1, key2);
assert_ne!(key1, key3);
}
#[test]
fn test_arena_key_handle_ordering() {
let key1 = ArenaKeyHandle::new(b"aaa");
let key2 = ArenaKeyHandle::new(b"bbb");
let key3 = ArenaKeyHandle::new(b"aaa");
assert!(key1 < key2);
assert!(key2 > key1);
assert_eq!(key1.cmp(&key3), std::cmp::Ordering::Equal);
}
#[test]
fn test_arena_key_handle_in_hashmap() {
use std::collections::HashMap;
let mut map = HashMap::new();
map.insert(ArenaKeyHandle::new(b"key1"), "value1");
map.insert(ArenaKeyHandle::new(b"key2"), "value2");
assert_eq!(map.get(&ArenaKeyHandle::new(b"key1")), Some(&"value1"));
assert_eq!(map.get(&ArenaKeyHandle::new(b"key2")), Some(&"value2"));
assert_eq!(map.get(&ArenaKeyHandle::new(b"key3")), None);
}
#[test]
fn test_arena_key_handle_from_arena_key() {
KeyArena::with(|arena| {
let arena_key = arena.alloc_key(b"test_data");
let handle = ArenaKeyHandle::from_arena_key(&arena_key);
assert_eq!(handle.as_bytes(), b"test_data");
});
}
#[test]
fn test_arena_key_handle_clone() {
let key = ArenaKeyHandle::new(b"clone_me");
let cloned = key.clone();
assert_eq!(key, cloned);
assert_eq!(key.hash(), cloned.hash());
}
}
use std::cell::{Cell, UnsafeCell};
use std::marker::PhantomData;
const ARENA_CHUNK_SIZE: usize = 64 * 1024;
pub struct KeyArena {
chunks: UnsafeCell<Vec<ArenaChunk>>,
current_chunk: Cell<usize>,
offset: Cell<usize>,
total_allocated: Cell<usize>,
}
struct ArenaChunk {
data: Vec<u8>,
capacity: usize,
}
impl ArenaChunk {
fn new(capacity: usize) -> Self {
Self {
data: vec![0u8; capacity],
capacity,
}
}
}
impl KeyArena {
pub fn new() -> Self {
Self::with_chunk_size(ARENA_CHUNK_SIZE)
}
pub fn with_chunk_size(chunk_size: usize) -> Self {
let initial_chunk = ArenaChunk::new(chunk_size);
Self {
chunks: UnsafeCell::new(vec![initial_chunk]),
current_chunk: Cell::new(0),
offset: Cell::new(0),
total_allocated: Cell::new(0),
}
}
#[inline]
pub fn with<F, R>(f: F) -> R
where
F: FnOnce(&KeyArena) -> R,
{
thread_local! {
static ARENA: KeyArena = KeyArena::new();
}
ARENA.with(f)
}
#[inline]
pub fn alloc_key<'a>(&'a self, data: &[u8]) -> ArenaKey<'a> {
let len = data.len();
let offset = self.offset.get();
let chunks = unsafe { &mut *self.chunks.get() };
let current_idx = self.current_chunk.get();
let current = &mut chunks[current_idx];
if offset + len <= current.capacity {
current.data[offset..offset + len].copy_from_slice(data);
self.offset.set(offset + len);
self.total_allocated.set(self.total_allocated.get() + len);
return ArenaKey {
ptr: current.data[offset..offset + len].as_ptr(),
len,
_marker: PhantomData,
};
}
self.alloc_slow(data)
}
#[cold]
fn alloc_slow<'a>(&'a self, data: &[u8]) -> ArenaKey<'a> {
let len = data.len();
let chunks = unsafe { &mut *self.chunks.get() };
let next_idx = self.current_chunk.get() + 1;
if next_idx < chunks.len() {
self.current_chunk.set(next_idx);
self.offset.set(0);
} else {
let chunk_size = std::cmp::max(ARENA_CHUNK_SIZE, len);
chunks.push(ArenaChunk::new(chunk_size));
self.current_chunk.set(next_idx);
self.offset.set(0);
}
let current = &mut chunks[self.current_chunk.get()];
let offset = self.offset.get();
current.data[offset..offset + len].copy_from_slice(data);
self.offset.set(offset + len);
self.total_allocated.set(self.total_allocated.get() + len);
ArenaKey {
ptr: current.data[offset..offset + len].as_ptr(),
len,
_marker: PhantomData,
}
}
#[inline]
pub fn reset(&self) {
self.current_chunk.set(0);
self.offset.set(0);
self.total_allocated.set(0);
}
#[inline]
pub fn bytes_used(&self) -> usize {
self.total_allocated.get()
}
pub fn chunk_count(&self) -> usize {
let chunks = unsafe { &*self.chunks.get() };
chunks.len()
}
}
impl Default for KeyArena {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Copy)]
pub struct ArenaKey<'a> {
ptr: *const u8,
len: usize,
_marker: PhantomData<&'a ()>,
}
impl<'a> ArenaKey<'a> {
#[inline]
pub fn as_bytes(&self) -> &'a [u8] {
unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
}
impl<'a> AsRef<[u8]> for ArenaKey<'a> {
fn as_ref(&self) -> &[u8] {
self.as_bytes()
}
}
impl<'a> std::fmt::Debug for ArenaKey<'a> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match std::str::from_utf8(self.as_bytes()) {
Ok(s) => write!(f, "ArenaKey({:?})", s),
Err(_) => write!(f, "ArenaKey({:?})", self.as_bytes()),
}
}
}
impl<'a> PartialEq for ArenaKey<'a> {
fn eq(&self, other: &Self) -> bool {
self.as_bytes() == other.as_bytes()
}
}
impl<'a> Eq for ArenaKey<'a> {}
impl<'a> std::hash::Hash for ArenaKey<'a> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.as_bytes().hash(state);
}
}
#[derive(Clone)]
pub struct ArenaKeyHandle {
hash: u64,
data: KeyData,
}
#[derive(Clone)]
enum KeyData {
Inline { len: u8, bytes: [u8; 24] },
Heap(Vec<u8>),
}
impl ArenaKeyHandle {
const INLINE_MAX: usize = 24;
#[inline]
pub fn new(data: &[u8]) -> Self {
let hash = Self::compute_hash(data);
let data = if data.len() <= Self::INLINE_MAX {
let mut bytes = [0u8; 24];
bytes[..data.len()].copy_from_slice(data);
KeyData::Inline {
len: data.len() as u8,
bytes,
}
} else {
KeyData::Heap(data.to_vec())
};
Self { hash, data }
}
#[inline]
pub fn from_arena_key(key: &ArenaKey<'_>) -> Self {
Self::new(key.as_bytes())
}
#[inline]
pub fn as_bytes(&self) -> &[u8] {
match &self.data {
KeyData::Inline { len, bytes } => &bytes[..*len as usize],
KeyData::Heap(v) => v.as_slice(),
}
}
#[inline]
pub fn hash(&self) -> u64 {
self.hash
}
#[inline]
pub fn len(&self) -> usize {
match &self.data {
KeyData::Inline { len, .. } => *len as usize,
KeyData::Heap(v) => v.len(),
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[inline]
pub fn is_inline(&self) -> bool {
matches!(self.data, KeyData::Inline { .. })
}
#[inline]
fn compute_hash(data: &[u8]) -> u64 {
const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
const FNV_PRIME: u64 = 0x00000100000001B3;
let mut h = FNV_OFFSET_BASIS;
for &b in data {
h ^= b as u64;
h = h.wrapping_mul(FNV_PRIME);
}
h
}
}
impl PartialEq for ArenaKeyHandle {
#[inline]
fn eq(&self, other: &Self) -> bool {
self.hash == other.hash && self.len() == other.len() && self.as_bytes() == other.as_bytes()
}
}
impl Eq for ArenaKeyHandle {}
impl std::hash::Hash for ArenaKeyHandle {
#[inline]
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
state.write_u64(self.hash);
}
}
impl PartialOrd for ArenaKeyHandle {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for ArenaKeyHandle {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.as_bytes().cmp(other.as_bytes())
}
}
impl AsRef<[u8]> for ArenaKeyHandle {
fn as_ref(&self) -> &[u8] {
self.as_bytes()
}
}
impl std::fmt::Debug for ArenaKeyHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match std::str::from_utf8(self.as_bytes()) {
Ok(s) => write!(f, "ArenaKeyHandle({:?})", s),
Err(_) => write!(f, "ArenaKeyHandle({:?})", self.as_bytes()),
}
}
}
impl From<&[u8]> for ArenaKeyHandle {
fn from(data: &[u8]) -> Self {
Self::new(data)
}
}
impl From<Vec<u8>> for ArenaKeyHandle {
fn from(data: Vec<u8>) -> Self {
Self::new(&data)
}
}
impl From<&str> for ArenaKeyHandle {
fn from(s: &str) -> Self {
Self::new(s.as_bytes())
}
}