1pub const MAX_KEY_LENGTH: usize = 63;
50
51#[repr(C, align(64))]
58#[derive(Clone)]
59pub struct KeyBuffer {
60 len: u8,
61 data: [u8; MAX_KEY_LENGTH],
62}
63
64impl KeyBuffer {
65 #[inline(always)]
67 pub const fn new() -> Self {
68 Self {
69 len: 0,
70 data: [0; MAX_KEY_LENGTH],
71 }
72 }
73
74 #[inline]
83 pub fn format_row_key(table: &str, row_id: u64) -> Self {
84 let mut buf = Self::new();
85
86 let table_bytes = table.as_bytes();
88 let table_len = table_bytes.len().min(MAX_KEY_LENGTH - 20); buf.data[..table_len].copy_from_slice(&table_bytes[..table_len]);
90 buf.len = table_len as u8;
91
92 if buf.len < MAX_KEY_LENGTH as u8 {
94 buf.data[buf.len as usize] = b'/';
95 buf.len += 1;
96 }
97
98 buf.write_u64(row_id);
100
101 buf
102 }
103
104 #[inline]
106 pub fn format_column_key(table: &str, row_id: u64, column: &str) -> Self {
107 let mut buf = Self::format_row_key(table, row_id);
108
109 if buf.len < MAX_KEY_LENGTH as u8 {
111 buf.data[buf.len as usize] = b'/';
112 buf.len += 1;
113 }
114
115 let col_bytes = column.as_bytes();
117 let available = MAX_KEY_LENGTH - buf.len as usize;
118 let col_len = col_bytes.len().min(available);
119 buf.data[buf.len as usize..buf.len as usize + col_len]
120 .copy_from_slice(&col_bytes[..col_len]);
121 buf.len += col_len as u8;
122
123 buf
124 }
125
126 #[inline]
128 fn write_u64(&mut self, mut value: u64) {
129 if value == 0 {
130 if self.len < MAX_KEY_LENGTH as u8 {
131 self.data[self.len as usize] = b'0';
132 self.len += 1;
133 }
134 return;
135 }
136
137 let mut temp = value;
139 let mut digit_count = 0u8;
140 while temp > 0 {
141 digit_count += 1;
142 temp /= 10;
143 }
144
145 if self.len as usize + digit_count as usize > MAX_KEY_LENGTH {
147 return; }
149
150 let start = self.len as usize;
152 let end = start + digit_count as usize;
153 self.len += digit_count;
154
155 let mut pos = end;
156 while value > 0 {
157 pos -= 1;
158 self.data[pos] = b'0' + (value % 10) as u8;
159 value /= 10;
160 }
161 }
162
163 #[inline]
165 pub fn append(&mut self, bytes: &[u8]) {
166 let available = MAX_KEY_LENGTH - self.len as usize;
167 let copy_len = bytes.len().min(available);
168 self.data[self.len as usize..self.len as usize + copy_len]
169 .copy_from_slice(&bytes[..copy_len]);
170 self.len += copy_len as u8;
171 }
172
173 #[inline]
175 pub fn push(&mut self, byte: u8) {
176 if self.len < MAX_KEY_LENGTH as u8 {
177 self.data[self.len as usize] = byte;
178 self.len += 1;
179 }
180 }
181
182 #[inline(always)]
184 pub fn as_bytes(&self) -> &[u8] {
185 &self.data[..self.len as usize]
186 }
187
188 #[inline(always)]
190 pub fn len(&self) -> usize {
191 self.len as usize
192 }
193
194 #[inline(always)]
196 pub fn is_empty(&self) -> bool {
197 self.len == 0
198 }
199
200 #[inline(always)]
202 pub fn clear(&mut self) {
203 self.len = 0;
204 }
205
206 #[inline(always)]
208 pub fn remaining(&self) -> usize {
209 MAX_KEY_LENGTH - self.len as usize
210 }
211}
212
213impl Default for KeyBuffer {
214 fn default() -> Self {
215 Self::new()
216 }
217}
218
219impl std::fmt::Debug for KeyBuffer {
220 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221 match std::str::from_utf8(self.as_bytes()) {
222 Ok(s) => write!(f, "KeyBuffer({:?})", s),
223 Err(_) => write!(f, "KeyBuffer({:?})", self.as_bytes()),
224 }
225 }
226}
227
228impl AsRef<[u8]> for KeyBuffer {
229 fn as_ref(&self) -> &[u8] {
230 self.as_bytes()
231 }
232}
233
234#[repr(C, align(64))]
239pub struct InternedTablePrefix {
240 prefix: [u8; 32],
242 prefix_len: u8,
244 hash: u64,
246}
247
248impl InternedTablePrefix {
249 pub fn new(table: &str) -> Self {
254 let mut prefix = [0u8; 32];
255 let bytes = table.as_bytes();
256 let len = bytes.len().min(30); prefix[..len].copy_from_slice(&bytes[..len]);
258 prefix[len] = b'/';
259
260 let hash = Self::compute_hash(&prefix[..len + 1]);
262
263 Self {
264 prefix,
265 prefix_len: (len + 1) as u8,
266 hash,
267 }
268 }
269
270 #[inline]
272 fn compute_hash(bytes: &[u8]) -> u64 {
273 const K: u64 = 0x517cc1b727220a95;
274 let mut hash = 0u64;
275 for &byte in bytes {
276 hash = (hash.rotate_left(5) ^ (byte as u64)).wrapping_mul(K);
277 }
278 hash
279 }
280
281 #[inline]
289 pub fn make_row_key(&self, row_id: u64) -> KeyBuffer {
290 let mut buf = KeyBuffer::new();
291
292 buf.data[..self.prefix_len as usize]
294 .copy_from_slice(&self.prefix[..self.prefix_len as usize]);
295 buf.len = self.prefix_len;
296
297 buf.write_u64(row_id);
299
300 buf
301 }
302
303 #[inline]
312 pub fn make_column_key(&self, row_id: u64, column: &str) -> KeyBuffer {
313 let mut buf = self.make_row_key(row_id);
314 buf.push(b'/');
315 buf.append(column.as_bytes());
316 buf
317 }
318
319 #[inline]
321 pub fn prefix(&self) -> &[u8] {
322 &self.prefix[..self.prefix_len as usize]
323 }
324
325 #[inline]
327 pub fn hash(&self) -> u64 {
328 self.hash
329 }
330
331 #[inline]
333 pub fn same_table(&self, other: &Self) -> bool {
334 self.hash == other.hash
335 && self.prefix_len == other.prefix_len
336 && self.prefix[..self.prefix_len as usize] == other.prefix[..other.prefix_len as usize]
337 }
338}
339
340impl std::fmt::Debug for InternedTablePrefix {
341 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
342 match std::str::from_utf8(self.prefix()) {
343 Ok(s) => write!(f, "InternedTablePrefix({:?})", s),
344 Err(_) => write!(f, "InternedTablePrefix({:?})", self.prefix()),
345 }
346 }
347}
348
349pub struct BatchKeyGenerator {
354 prefix: InternedTablePrefix,
355 buffer: KeyBuffer,
357}
358
359impl BatchKeyGenerator {
360 pub fn new(table: &str) -> Self {
362 Self {
363 prefix: InternedTablePrefix::new(table),
364 buffer: KeyBuffer::new(),
365 }
366 }
367
368 #[inline]
370 pub fn row_key(&mut self, row_id: u64) -> &[u8] {
371 self.buffer = self.prefix.make_row_key(row_id);
372 self.buffer.as_bytes()
373 }
374
375 #[inline]
377 pub fn column_key(&mut self, row_id: u64, column: &str) -> &[u8] {
378 self.buffer = self.prefix.make_column_key(row_id, column);
379 self.buffer.as_bytes()
380 }
381
382 #[inline]
384 pub fn prefix(&self) -> &InternedTablePrefix {
385 &self.prefix
386 }
387}
388
389#[cfg(test)]
390mod tests {
391 use super::*;
392
393 #[test]
394 fn test_key_buffer_row_key() {
395 let key = KeyBuffer::format_row_key("users", 12345);
396 assert_eq!(key.as_bytes(), b"users/12345");
397 }
398
399 #[test]
400 fn test_key_buffer_column_key() {
401 let key = KeyBuffer::format_column_key("users", 42, "name");
402 assert_eq!(key.as_bytes(), b"users/42/name");
403 }
404
405 #[test]
406 fn test_key_buffer_zero_id() {
407 let key = KeyBuffer::format_row_key("table", 0);
408 assert_eq!(key.as_bytes(), b"table/0");
409 }
410
411 #[test]
412 fn test_key_buffer_large_id() {
413 let key = KeyBuffer::format_row_key("t", u64::MAX);
414 let expected = format!("t/{}", u64::MAX);
415 assert_eq!(key.as_bytes(), expected.as_bytes());
416 }
417
418 #[test]
419 fn test_interned_prefix() {
420 let prefix = InternedTablePrefix::new("orders");
421 let key = prefix.make_row_key(999);
422 assert_eq!(key.as_bytes(), b"orders/999");
423 }
424
425 #[test]
426 fn test_interned_column_key() {
427 let prefix = InternedTablePrefix::new("products");
428 let key = prefix.make_column_key(100, "price");
429 assert_eq!(key.as_bytes(), b"products/100/price");
430 }
431
432 #[test]
433 fn test_batch_generator() {
434 let mut generator = BatchKeyGenerator::new("items");
435
436 assert_eq!(generator.row_key(1), b"items/1");
437 assert_eq!(generator.row_key(2), b"items/2");
438 assert_eq!(generator.column_key(3, "qty"), b"items/3/qty");
439 }
440
441 #[test]
442 fn test_same_table_check() {
443 let p1 = InternedTablePrefix::new("users");
444 let p2 = InternedTablePrefix::new("users");
445 let p3 = InternedTablePrefix::new("orders");
446
447 assert!(p1.same_table(&p2));
448 assert!(!p1.same_table(&p3));
449 }
450
451 #[test]
452 fn test_cache_line_alignment() {
453 assert_eq!(std::mem::align_of::<KeyBuffer>(), 64);
455 assert_eq!(std::mem::align_of::<InternedTablePrefix>(), 64);
456 }
457
458 #[test]
459 fn test_no_heap_allocation() {
460 let key = KeyBuffer::format_column_key("users", 12345678901234567890, "email_address");
462 assert!(key.len() <= MAX_KEY_LENGTH);
463 }
466
467 #[test]
468 fn test_arena_basic() {
469 KeyArena::with(|arena| {
470 let key1 = arena.alloc_key(b"test_key_1");
471 let key2 = arena.alloc_key(b"test_key_2");
472
473 assert_eq!(key1.as_bytes(), b"test_key_1");
474 assert_eq!(key2.as_bytes(), b"test_key_2");
475 });
476 }
477
478 #[test]
479 fn test_arena_reset() {
480 KeyArena::with(|arena| {
481 for i in 0..100 {
483 let key = format!("key_{}", i);
484 arena.alloc_key(key.as_bytes());
485 }
486
487 let used_before = arena.bytes_used();
488 assert!(used_before > 0);
489
490 arena.reset();
492
493 assert_eq!(arena.bytes_used(), 0);
494
495 let key = arena.alloc_key(b"after_reset");
497 assert_eq!(key.as_bytes(), b"after_reset");
498 });
499 }
500
501 #[test]
502 fn test_arena_large_allocation() {
503 KeyArena::with(|arena| {
504 let large_key = vec![b'x'; 1024];
506 let key = arena.alloc_key(&large_key);
507 assert_eq!(key.as_bytes(), large_key.as_slice());
508 });
509 }
510
511 #[test]
514 fn test_arena_key_handle_inline() {
515 let key = ArenaKeyHandle::new(b"short_key");
517 assert!(key.is_inline());
518 assert_eq!(key.as_bytes(), b"short_key");
519 assert_eq!(key.len(), 9);
520 }
521
522 #[test]
523 fn test_arena_key_handle_heap() {
524 let large = vec![b'x'; 50];
526 let key = ArenaKeyHandle::new(&large);
527 assert!(!key.is_inline());
528 assert_eq!(key.as_bytes(), large.as_slice());
529 assert_eq!(key.len(), 50);
530 }
531
532 #[test]
533 fn test_arena_key_handle_hash() {
534 let key1 = ArenaKeyHandle::new(b"test_key");
536 let key2 = ArenaKeyHandle::new(b"test_key");
537 assert_eq!(key1.hash(), key2.hash());
538
539 let key3 = ArenaKeyHandle::new(b"other_key");
541 assert_ne!(key1.hash(), key3.hash());
542 }
543
544 #[test]
545 fn test_arena_key_handle_equality() {
546 let key1 = ArenaKeyHandle::new(b"test");
547 let key2 = ArenaKeyHandle::new(b"test");
548 let key3 = ArenaKeyHandle::new(b"other");
549
550 assert_eq!(key1, key2);
551 assert_ne!(key1, key3);
552 }
553
554 #[test]
555 fn test_arena_key_handle_ordering() {
556 let key1 = ArenaKeyHandle::new(b"aaa");
557 let key2 = ArenaKeyHandle::new(b"bbb");
558 let key3 = ArenaKeyHandle::new(b"aaa");
559
560 assert!(key1 < key2);
561 assert!(key2 > key1);
562 assert_eq!(key1.cmp(&key3), std::cmp::Ordering::Equal);
563 }
564
565 #[test]
566 fn test_arena_key_handle_in_hashmap() {
567 use std::collections::HashMap;
568
569 let mut map = HashMap::new();
570 map.insert(ArenaKeyHandle::new(b"key1"), "value1");
571 map.insert(ArenaKeyHandle::new(b"key2"), "value2");
572
573 assert_eq!(map.get(&ArenaKeyHandle::new(b"key1")), Some(&"value1"));
575 assert_eq!(map.get(&ArenaKeyHandle::new(b"key2")), Some(&"value2"));
576 assert_eq!(map.get(&ArenaKeyHandle::new(b"key3")), None);
577 }
578
579 #[test]
580 fn test_arena_key_handle_from_arena_key() {
581 KeyArena::with(|arena| {
582 let arena_key = arena.alloc_key(b"test_data");
583 let handle = ArenaKeyHandle::from_arena_key(&arena_key);
584
585 assert_eq!(handle.as_bytes(), b"test_data");
586 });
587 }
588
589 #[test]
590 fn test_arena_key_handle_clone() {
591 let key = ArenaKeyHandle::new(b"clone_me");
592 let cloned = key.clone();
593
594 assert_eq!(key, cloned);
595 assert_eq!(key.hash(), cloned.hash());
596 }
597}
598
599use std::cell::{Cell, UnsafeCell};
604use std::marker::PhantomData;
605
606const ARENA_CHUNK_SIZE: usize = 64 * 1024;
608
609pub struct KeyArena {
631 chunks: UnsafeCell<Vec<ArenaChunk>>,
633 current_chunk: Cell<usize>,
635 offset: Cell<usize>,
637 total_allocated: Cell<usize>,
639}
640
641struct ArenaChunk {
643 data: Vec<u8>,
644 capacity: usize,
645}
646
647impl ArenaChunk {
648 fn new(capacity: usize) -> Self {
649 Self {
650 data: vec![0u8; capacity],
651 capacity,
652 }
653 }
654}
655
656impl KeyArena {
657 pub fn new() -> Self {
659 Self::with_chunk_size(ARENA_CHUNK_SIZE)
660 }
661
662 pub fn with_chunk_size(chunk_size: usize) -> Self {
664 let initial_chunk = ArenaChunk::new(chunk_size);
665 Self {
666 chunks: UnsafeCell::new(vec![initial_chunk]),
667 current_chunk: Cell::new(0),
668 offset: Cell::new(0),
669 total_allocated: Cell::new(0),
670 }
671 }
672
673 #[inline]
678 pub fn with<F, R>(f: F) -> R
679 where
680 F: FnOnce(&KeyArena) -> R,
681 {
682 thread_local! {
683 static ARENA: KeyArena = KeyArena::new();
684 }
685
686 ARENA.with(f)
687 }
688
689 #[inline]
699 pub fn alloc_key<'a>(&'a self, data: &[u8]) -> ArenaKey<'a> {
700 let len = data.len();
701 let offset = self.offset.get();
702
703 let chunks = unsafe { &mut *self.chunks.get() };
705 let current_idx = self.current_chunk.get();
706 let current = &mut chunks[current_idx];
707
708 if offset + len <= current.capacity {
709 current.data[offset..offset + len].copy_from_slice(data);
711 self.offset.set(offset + len);
712 self.total_allocated.set(self.total_allocated.get() + len);
713
714 return ArenaKey {
715 ptr: current.data[offset..offset + len].as_ptr(),
716 len,
717 _marker: PhantomData,
718 };
719 }
720
721 self.alloc_slow(data)
723 }
724
725 #[cold]
727 fn alloc_slow<'a>(&'a self, data: &[u8]) -> ArenaKey<'a> {
728 let len = data.len();
729 let chunks = unsafe { &mut *self.chunks.get() };
730
731 let next_idx = self.current_chunk.get() + 1;
733
734 if next_idx < chunks.len() {
735 self.current_chunk.set(next_idx);
737 self.offset.set(0);
738 } else {
739 let chunk_size = std::cmp::max(ARENA_CHUNK_SIZE, len);
741 chunks.push(ArenaChunk::new(chunk_size));
742 self.current_chunk.set(next_idx);
743 self.offset.set(0);
744 }
745
746 let current = &mut chunks[self.current_chunk.get()];
748 let offset = self.offset.get();
749
750 current.data[offset..offset + len].copy_from_slice(data);
751 self.offset.set(offset + len);
752 self.total_allocated.set(self.total_allocated.get() + len);
753
754 ArenaKey {
755 ptr: current.data[offset..offset + len].as_ptr(),
756 len,
757 _marker: PhantomData,
758 }
759 }
760
761 #[inline]
766 pub fn reset(&self) {
767 self.current_chunk.set(0);
768 self.offset.set(0);
769 self.total_allocated.set(0);
770 }
771
772 #[inline]
774 pub fn bytes_used(&self) -> usize {
775 self.total_allocated.get()
776 }
777
778 pub fn chunk_count(&self) -> usize {
780 let chunks = unsafe { &*self.chunks.get() };
781 chunks.len()
782 }
783}
784
785impl Default for KeyArena {
786 fn default() -> Self {
787 Self::new()
788 }
789}
790
791#[derive(Clone, Copy)]
799pub struct ArenaKey<'a> {
800 ptr: *const u8,
801 len: usize,
802 _marker: PhantomData<&'a ()>,
803}
804
805impl<'a> ArenaKey<'a> {
806 #[inline]
808 pub fn as_bytes(&self) -> &'a [u8] {
809 unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
810 }
811
812 #[inline]
814 pub fn len(&self) -> usize {
815 self.len
816 }
817
818 #[inline]
820 pub fn is_empty(&self) -> bool {
821 self.len == 0
822 }
823}
824
825impl<'a> AsRef<[u8]> for ArenaKey<'a> {
826 fn as_ref(&self) -> &[u8] {
827 self.as_bytes()
828 }
829}
830
831impl<'a> std::fmt::Debug for ArenaKey<'a> {
832 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
833 match std::str::from_utf8(self.as_bytes()) {
834 Ok(s) => write!(f, "ArenaKey({:?})", s),
835 Err(_) => write!(f, "ArenaKey({:?})", self.as_bytes()),
836 }
837 }
838}
839
840impl<'a> PartialEq for ArenaKey<'a> {
841 fn eq(&self, other: &Self) -> bool {
842 self.as_bytes() == other.as_bytes()
843 }
844}
845
846impl<'a> Eq for ArenaKey<'a> {}
847
848impl<'a> std::hash::Hash for ArenaKey<'a> {
849 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
850 self.as_bytes().hash(state);
851 }
852}
853
854#[derive(Clone)]
875pub struct ArenaKeyHandle {
876 hash: u64,
878 data: KeyData,
880}
881
882#[derive(Clone)]
884enum KeyData {
885 Inline { len: u8, bytes: [u8; 24] },
887 Heap(Vec<u8>),
889}
890
891impl ArenaKeyHandle {
892 const INLINE_MAX: usize = 24;
894
895 #[inline]
897 pub fn new(data: &[u8]) -> Self {
898 let hash = Self::compute_hash(data);
899 let data = if data.len() <= Self::INLINE_MAX {
900 let mut bytes = [0u8; 24];
901 bytes[..data.len()].copy_from_slice(data);
902 KeyData::Inline {
903 len: data.len() as u8,
904 bytes,
905 }
906 } else {
907 KeyData::Heap(data.to_vec())
908 };
909 Self { hash, data }
910 }
911
912 #[inline]
914 pub fn from_arena_key(key: &ArenaKey<'_>) -> Self {
915 Self::new(key.as_bytes())
916 }
917
918 #[inline]
920 pub fn as_bytes(&self) -> &[u8] {
921 match &self.data {
922 KeyData::Inline { len, bytes } => &bytes[..*len as usize],
923 KeyData::Heap(v) => v.as_slice(),
924 }
925 }
926
927 #[inline]
929 pub fn hash(&self) -> u64 {
930 self.hash
931 }
932
933 #[inline]
935 pub fn len(&self) -> usize {
936 match &self.data {
937 KeyData::Inline { len, .. } => *len as usize,
938 KeyData::Heap(v) => v.len(),
939 }
940 }
941
942 #[inline]
944 pub fn is_empty(&self) -> bool {
945 self.len() == 0
946 }
947
948 #[inline]
950 pub fn is_inline(&self) -> bool {
951 matches!(self.data, KeyData::Inline { .. })
952 }
953
954 #[inline]
956 fn compute_hash(data: &[u8]) -> u64 {
957 const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
958 const FNV_PRIME: u64 = 0x00000100000001B3;
959
960 let mut h = FNV_OFFSET_BASIS;
961 for &b in data {
962 h ^= b as u64;
963 h = h.wrapping_mul(FNV_PRIME);
964 }
965 h
966 }
967}
968
969impl PartialEq for ArenaKeyHandle {
970 #[inline]
971 fn eq(&self, other: &Self) -> bool {
972 self.hash == other.hash && self.len() == other.len() && self.as_bytes() == other.as_bytes()
974 }
975}
976
977impl Eq for ArenaKeyHandle {}
978
979impl std::hash::Hash for ArenaKeyHandle {
980 #[inline]
981 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
982 state.write_u64(self.hash);
984 }
985}
986
987impl PartialOrd for ArenaKeyHandle {
988 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
989 Some(self.cmp(other))
990 }
991}
992
993impl Ord for ArenaKeyHandle {
994 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
995 self.as_bytes().cmp(other.as_bytes())
996 }
997}
998
999impl AsRef<[u8]> for ArenaKeyHandle {
1000 fn as_ref(&self) -> &[u8] {
1001 self.as_bytes()
1002 }
1003}
1004
1005impl std::fmt::Debug for ArenaKeyHandle {
1006 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1007 match std::str::from_utf8(self.as_bytes()) {
1008 Ok(s) => write!(f, "ArenaKeyHandle({:?})", s),
1009 Err(_) => write!(f, "ArenaKeyHandle({:?})", self.as_bytes()),
1010 }
1011 }
1012}
1013
1014impl From<&[u8]> for ArenaKeyHandle {
1015 fn from(data: &[u8]) -> Self {
1016 Self::new(data)
1017 }
1018}
1019
1020impl From<Vec<u8>> for ArenaKeyHandle {
1021 fn from(data: Vec<u8>) -> Self {
1022 Self::new(&data)
1023 }
1024}
1025
1026impl From<&str> for ArenaKeyHandle {
1027 fn from(s: &str) -> Self {
1028 Self::new(s.as_bytes())
1029 }
1030}