1pub const MAX_KEY_LENGTH: usize = 63;
47
48#[repr(C, align(64))]
55#[derive(Clone)]
56pub struct KeyBuffer {
57 len: u8,
58 data: [u8; MAX_KEY_LENGTH],
59}
60
61impl KeyBuffer {
62 #[inline(always)]
64 pub const fn new() -> Self {
65 Self {
66 len: 0,
67 data: [0; MAX_KEY_LENGTH],
68 }
69 }
70
71 #[inline]
80 pub fn format_row_key(table: &str, row_id: u64) -> Self {
81 let mut buf = Self::new();
82
83 let table_bytes = table.as_bytes();
85 let table_len = table_bytes.len().min(MAX_KEY_LENGTH - 20); buf.data[..table_len].copy_from_slice(&table_bytes[..table_len]);
87 buf.len = table_len as u8;
88
89 if buf.len < MAX_KEY_LENGTH as u8 {
91 buf.data[buf.len as usize] = b'/';
92 buf.len += 1;
93 }
94
95 buf.write_u64(row_id);
97
98 buf
99 }
100
101 #[inline]
103 pub fn format_column_key(table: &str, row_id: u64, column: &str) -> Self {
104 let mut buf = Self::format_row_key(table, row_id);
105
106 if buf.len < MAX_KEY_LENGTH as u8 {
108 buf.data[buf.len as usize] = b'/';
109 buf.len += 1;
110 }
111
112 let col_bytes = column.as_bytes();
114 let available = MAX_KEY_LENGTH - buf.len as usize;
115 let col_len = col_bytes.len().min(available);
116 buf.data[buf.len as usize..buf.len as usize + col_len]
117 .copy_from_slice(&col_bytes[..col_len]);
118 buf.len += col_len as u8;
119
120 buf
121 }
122
123 #[inline]
125 fn write_u64(&mut self, mut value: u64) {
126 if value == 0 {
127 if self.len < MAX_KEY_LENGTH as u8 {
128 self.data[self.len as usize] = b'0';
129 self.len += 1;
130 }
131 return;
132 }
133
134 let mut temp = value;
136 let mut digit_count = 0u8;
137 while temp > 0 {
138 digit_count += 1;
139 temp /= 10;
140 }
141
142 if self.len as usize + digit_count as usize > MAX_KEY_LENGTH {
144 return; }
146
147 let start = self.len as usize;
149 let end = start + digit_count as usize;
150 self.len += digit_count;
151
152 let mut pos = end;
153 while value > 0 {
154 pos -= 1;
155 self.data[pos] = b'0' + (value % 10) as u8;
156 value /= 10;
157 }
158 }
159
160 #[inline]
162 pub fn append(&mut self, bytes: &[u8]) {
163 let available = MAX_KEY_LENGTH - self.len as usize;
164 let copy_len = bytes.len().min(available);
165 self.data[self.len as usize..self.len as usize + copy_len]
166 .copy_from_slice(&bytes[..copy_len]);
167 self.len += copy_len as u8;
168 }
169
170 #[inline]
172 pub fn push(&mut self, byte: u8) {
173 if self.len < MAX_KEY_LENGTH as u8 {
174 self.data[self.len as usize] = byte;
175 self.len += 1;
176 }
177 }
178
179 #[inline(always)]
181 pub fn as_bytes(&self) -> &[u8] {
182 &self.data[..self.len as usize]
183 }
184
185 #[inline(always)]
187 pub fn len(&self) -> usize {
188 self.len as usize
189 }
190
191 #[inline(always)]
193 pub fn is_empty(&self) -> bool {
194 self.len == 0
195 }
196
197 #[inline(always)]
199 pub fn clear(&mut self) {
200 self.len = 0;
201 }
202
203 #[inline(always)]
205 pub fn remaining(&self) -> usize {
206 MAX_KEY_LENGTH - self.len as usize
207 }
208}
209
210impl Default for KeyBuffer {
211 fn default() -> Self {
212 Self::new()
213 }
214}
215
216impl std::fmt::Debug for KeyBuffer {
217 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
218 match std::str::from_utf8(self.as_bytes()) {
219 Ok(s) => write!(f, "KeyBuffer({:?})", s),
220 Err(_) => write!(f, "KeyBuffer({:?})", self.as_bytes()),
221 }
222 }
223}
224
225impl AsRef<[u8]> for KeyBuffer {
226 fn as_ref(&self) -> &[u8] {
227 self.as_bytes()
228 }
229}
230
231#[repr(C, align(64))]
236pub struct InternedTablePrefix {
237 prefix: [u8; 32],
239 prefix_len: u8,
241 hash: u64,
243}
244
245impl InternedTablePrefix {
246 pub fn new(table: &str) -> Self {
251 let mut prefix = [0u8; 32];
252 let bytes = table.as_bytes();
253 let len = bytes.len().min(30); prefix[..len].copy_from_slice(&bytes[..len]);
255 prefix[len] = b'/';
256
257 let hash = Self::compute_hash(&prefix[..len + 1]);
259
260 Self {
261 prefix,
262 prefix_len: (len + 1) as u8,
263 hash,
264 }
265 }
266
267 #[inline]
269 fn compute_hash(bytes: &[u8]) -> u64 {
270 const K: u64 = 0x517cc1b727220a95;
271 let mut hash = 0u64;
272 for &byte in bytes {
273 hash = (hash.rotate_left(5) ^ (byte as u64)).wrapping_mul(K);
274 }
275 hash
276 }
277
278 #[inline]
286 pub fn make_row_key(&self, row_id: u64) -> KeyBuffer {
287 let mut buf = KeyBuffer::new();
288
289 buf.data[..self.prefix_len as usize]
291 .copy_from_slice(&self.prefix[..self.prefix_len as usize]);
292 buf.len = self.prefix_len;
293
294 buf.write_u64(row_id);
296
297 buf
298 }
299
300 #[inline]
309 pub fn make_column_key(&self, row_id: u64, column: &str) -> KeyBuffer {
310 let mut buf = self.make_row_key(row_id);
311 buf.push(b'/');
312 buf.append(column.as_bytes());
313 buf
314 }
315
316 #[inline]
318 pub fn prefix(&self) -> &[u8] {
319 &self.prefix[..self.prefix_len as usize]
320 }
321
322 #[inline]
324 pub fn hash(&self) -> u64 {
325 self.hash
326 }
327
328 #[inline]
330 pub fn same_table(&self, other: &Self) -> bool {
331 self.hash == other.hash
332 && self.prefix_len == other.prefix_len
333 && self.prefix[..self.prefix_len as usize] == other.prefix[..other.prefix_len as usize]
334 }
335}
336
337impl std::fmt::Debug for InternedTablePrefix {
338 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
339 match std::str::from_utf8(self.prefix()) {
340 Ok(s) => write!(f, "InternedTablePrefix({:?})", s),
341 Err(_) => write!(f, "InternedTablePrefix({:?})", self.prefix()),
342 }
343 }
344}
345
346pub struct BatchKeyGenerator {
351 prefix: InternedTablePrefix,
352 buffer: KeyBuffer,
354}
355
356impl BatchKeyGenerator {
357 pub fn new(table: &str) -> Self {
359 Self {
360 prefix: InternedTablePrefix::new(table),
361 buffer: KeyBuffer::new(),
362 }
363 }
364
365 #[inline]
367 pub fn row_key(&mut self, row_id: u64) -> &[u8] {
368 self.buffer = self.prefix.make_row_key(row_id);
369 self.buffer.as_bytes()
370 }
371
372 #[inline]
374 pub fn column_key(&mut self, row_id: u64, column: &str) -> &[u8] {
375 self.buffer = self.prefix.make_column_key(row_id, column);
376 self.buffer.as_bytes()
377 }
378
379 #[inline]
381 pub fn prefix(&self) -> &InternedTablePrefix {
382 &self.prefix
383 }
384}
385
386#[cfg(test)]
387mod tests {
388 use super::*;
389
390 #[test]
391 fn test_key_buffer_row_key() {
392 let key = KeyBuffer::format_row_key("users", 12345);
393 assert_eq!(key.as_bytes(), b"users/12345");
394 }
395
396 #[test]
397 fn test_key_buffer_column_key() {
398 let key = KeyBuffer::format_column_key("users", 42, "name");
399 assert_eq!(key.as_bytes(), b"users/42/name");
400 }
401
402 #[test]
403 fn test_key_buffer_zero_id() {
404 let key = KeyBuffer::format_row_key("table", 0);
405 assert_eq!(key.as_bytes(), b"table/0");
406 }
407
408 #[test]
409 fn test_key_buffer_large_id() {
410 let key = KeyBuffer::format_row_key("t", u64::MAX);
411 let expected = format!("t/{}", u64::MAX);
412 assert_eq!(key.as_bytes(), expected.as_bytes());
413 }
414
415 #[test]
416 fn test_interned_prefix() {
417 let prefix = InternedTablePrefix::new("orders");
418 let key = prefix.make_row_key(999);
419 assert_eq!(key.as_bytes(), b"orders/999");
420 }
421
422 #[test]
423 fn test_interned_column_key() {
424 let prefix = InternedTablePrefix::new("products");
425 let key = prefix.make_column_key(100, "price");
426 assert_eq!(key.as_bytes(), b"products/100/price");
427 }
428
429 #[test]
430 fn test_batch_generator() {
431 let mut generator = BatchKeyGenerator::new("items");
432
433 assert_eq!(generator.row_key(1), b"items/1");
434 assert_eq!(generator.row_key(2), b"items/2");
435 assert_eq!(generator.column_key(3, "qty"), b"items/3/qty");
436 }
437
438 #[test]
439 fn test_same_table_check() {
440 let p1 = InternedTablePrefix::new("users");
441 let p2 = InternedTablePrefix::new("users");
442 let p3 = InternedTablePrefix::new("orders");
443
444 assert!(p1.same_table(&p2));
445 assert!(!p1.same_table(&p3));
446 }
447
448 #[test]
449 fn test_cache_line_alignment() {
450 assert_eq!(std::mem::align_of::<KeyBuffer>(), 64);
452 assert_eq!(std::mem::align_of::<InternedTablePrefix>(), 64);
453 }
454
455 #[test]
456 fn test_no_heap_allocation() {
457 let key = KeyBuffer::format_column_key("users", 12345678901234567890, "email_address");
459 assert!(key.len() <= MAX_KEY_LENGTH);
460 }
463
464 #[test]
465 fn test_arena_basic() {
466 KeyArena::with(|arena| {
467 let key1 = arena.alloc_key(b"test_key_1");
468 let key2 = arena.alloc_key(b"test_key_2");
469
470 assert_eq!(key1.as_bytes(), b"test_key_1");
471 assert_eq!(key2.as_bytes(), b"test_key_2");
472 });
473 }
474
475 #[test]
476 fn test_arena_reset() {
477 KeyArena::with(|arena| {
478 for i in 0..100 {
480 let key = format!("key_{}", i);
481 arena.alloc_key(key.as_bytes());
482 }
483
484 let used_before = arena.bytes_used();
485 assert!(used_before > 0);
486
487 arena.reset();
489
490 assert_eq!(arena.bytes_used(), 0);
491
492 let key = arena.alloc_key(b"after_reset");
494 assert_eq!(key.as_bytes(), b"after_reset");
495 });
496 }
497
498 #[test]
499 fn test_arena_large_allocation() {
500 KeyArena::with(|arena| {
501 let large_key = vec![b'x'; 1024];
503 let key = arena.alloc_key(&large_key);
504 assert_eq!(key.as_bytes(), large_key.as_slice());
505 });
506 }
507
508 #[test]
511 fn test_arena_key_handle_inline() {
512 let key = ArenaKeyHandle::new(b"short_key");
514 assert!(key.is_inline());
515 assert_eq!(key.as_bytes(), b"short_key");
516 assert_eq!(key.len(), 9);
517 }
518
519 #[test]
520 fn test_arena_key_handle_heap() {
521 let large = vec![b'x'; 50];
523 let key = ArenaKeyHandle::new(&large);
524 assert!(!key.is_inline());
525 assert_eq!(key.as_bytes(), large.as_slice());
526 assert_eq!(key.len(), 50);
527 }
528
529 #[test]
530 fn test_arena_key_handle_hash() {
531 let key1 = ArenaKeyHandle::new(b"test_key");
533 let key2 = ArenaKeyHandle::new(b"test_key");
534 assert_eq!(key1.hash(), key2.hash());
535
536 let key3 = ArenaKeyHandle::new(b"other_key");
538 assert_ne!(key1.hash(), key3.hash());
539 }
540
541 #[test]
542 fn test_arena_key_handle_equality() {
543 let key1 = ArenaKeyHandle::new(b"test");
544 let key2 = ArenaKeyHandle::new(b"test");
545 let key3 = ArenaKeyHandle::new(b"other");
546
547 assert_eq!(key1, key2);
548 assert_ne!(key1, key3);
549 }
550
551 #[test]
552 fn test_arena_key_handle_ordering() {
553 let key1 = ArenaKeyHandle::new(b"aaa");
554 let key2 = ArenaKeyHandle::new(b"bbb");
555 let key3 = ArenaKeyHandle::new(b"aaa");
556
557 assert!(key1 < key2);
558 assert!(key2 > key1);
559 assert_eq!(key1.cmp(&key3), std::cmp::Ordering::Equal);
560 }
561
562 #[test]
563 fn test_arena_key_handle_in_hashmap() {
564 use std::collections::HashMap;
565
566 let mut map = HashMap::new();
567 map.insert(ArenaKeyHandle::new(b"key1"), "value1");
568 map.insert(ArenaKeyHandle::new(b"key2"), "value2");
569
570 assert_eq!(map.get(&ArenaKeyHandle::new(b"key1")), Some(&"value1"));
572 assert_eq!(map.get(&ArenaKeyHandle::new(b"key2")), Some(&"value2"));
573 assert_eq!(map.get(&ArenaKeyHandle::new(b"key3")), None);
574 }
575
576 #[test]
577 fn test_arena_key_handle_from_arena_key() {
578 KeyArena::with(|arena| {
579 let arena_key = arena.alloc_key(b"test_data");
580 let handle = ArenaKeyHandle::from_arena_key(&arena_key);
581
582 assert_eq!(handle.as_bytes(), b"test_data");
583 });
584 }
585
586 #[test]
587 fn test_arena_key_handle_clone() {
588 let key = ArenaKeyHandle::new(b"clone_me");
589 let cloned = key.clone();
590
591 assert_eq!(key, cloned);
592 assert_eq!(key.hash(), cloned.hash());
593 }
594}
595
596use std::cell::{Cell, UnsafeCell};
601use std::marker::PhantomData;
602
603const ARENA_CHUNK_SIZE: usize = 64 * 1024;
605
606pub struct KeyArena {
628 chunks: UnsafeCell<Vec<ArenaChunk>>,
630 current_chunk: Cell<usize>,
632 offset: Cell<usize>,
634 total_allocated: Cell<usize>,
636}
637
638struct ArenaChunk {
640 data: Vec<u8>,
641 capacity: usize,
642}
643
644impl ArenaChunk {
645 fn new(capacity: usize) -> Self {
646 Self {
647 data: vec![0u8; capacity],
648 capacity,
649 }
650 }
651}
652
653impl KeyArena {
654 pub fn new() -> Self {
656 Self::with_chunk_size(ARENA_CHUNK_SIZE)
657 }
658
659 pub fn with_chunk_size(chunk_size: usize) -> Self {
661 let initial_chunk = ArenaChunk::new(chunk_size);
662 Self {
663 chunks: UnsafeCell::new(vec![initial_chunk]),
664 current_chunk: Cell::new(0),
665 offset: Cell::new(0),
666 total_allocated: Cell::new(0),
667 }
668 }
669
670 #[inline]
675 pub fn with<F, R>(f: F) -> R
676 where
677 F: FnOnce(&KeyArena) -> R,
678 {
679 thread_local! {
680 static ARENA: KeyArena = KeyArena::new();
681 }
682
683 ARENA.with(f)
684 }
685
686 #[inline]
696 pub fn alloc_key<'a>(&'a self, data: &[u8]) -> ArenaKey<'a> {
697 let len = data.len();
698 let offset = self.offset.get();
699
700 let chunks = unsafe { &mut *self.chunks.get() };
702 let current_idx = self.current_chunk.get();
703 let current = &mut chunks[current_idx];
704
705 if offset + len <= current.capacity {
706 current.data[offset..offset + len].copy_from_slice(data);
708 self.offset.set(offset + len);
709 self.total_allocated.set(self.total_allocated.get() + len);
710
711 return ArenaKey {
712 ptr: current.data[offset..offset + len].as_ptr(),
713 len,
714 _marker: PhantomData,
715 };
716 }
717
718 self.alloc_slow(data)
720 }
721
722 #[cold]
724 fn alloc_slow<'a>(&'a self, data: &[u8]) -> ArenaKey<'a> {
725 let len = data.len();
726 let chunks = unsafe { &mut *self.chunks.get() };
727
728 let next_idx = self.current_chunk.get() + 1;
730
731 if next_idx < chunks.len() {
732 self.current_chunk.set(next_idx);
734 self.offset.set(0);
735 } else {
736 let chunk_size = std::cmp::max(ARENA_CHUNK_SIZE, len);
738 chunks.push(ArenaChunk::new(chunk_size));
739 self.current_chunk.set(next_idx);
740 self.offset.set(0);
741 }
742
743 let current = &mut chunks[self.current_chunk.get()];
745 let offset = self.offset.get();
746
747 current.data[offset..offset + len].copy_from_slice(data);
748 self.offset.set(offset + len);
749 self.total_allocated.set(self.total_allocated.get() + len);
750
751 ArenaKey {
752 ptr: current.data[offset..offset + len].as_ptr(),
753 len,
754 _marker: PhantomData,
755 }
756 }
757
758 #[inline]
763 pub fn reset(&self) {
764 self.current_chunk.set(0);
765 self.offset.set(0);
766 self.total_allocated.set(0);
767 }
768
769 #[inline]
771 pub fn bytes_used(&self) -> usize {
772 self.total_allocated.get()
773 }
774
775 pub fn chunk_count(&self) -> usize {
777 let chunks = unsafe { &*self.chunks.get() };
778 chunks.len()
779 }
780}
781
782impl Default for KeyArena {
783 fn default() -> Self {
784 Self::new()
785 }
786}
787
788#[derive(Clone, Copy)]
796pub struct ArenaKey<'a> {
797 ptr: *const u8,
798 len: usize,
799 _marker: PhantomData<&'a ()>,
800}
801
802impl<'a> ArenaKey<'a> {
803 #[inline]
805 pub fn as_bytes(&self) -> &'a [u8] {
806 unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
807 }
808
809 #[inline]
811 pub fn len(&self) -> usize {
812 self.len
813 }
814
815 #[inline]
817 pub fn is_empty(&self) -> bool {
818 self.len == 0
819 }
820}
821
822impl<'a> AsRef<[u8]> for ArenaKey<'a> {
823 fn as_ref(&self) -> &[u8] {
824 self.as_bytes()
825 }
826}
827
828impl<'a> std::fmt::Debug for ArenaKey<'a> {
829 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
830 match std::str::from_utf8(self.as_bytes()) {
831 Ok(s) => write!(f, "ArenaKey({:?})", s),
832 Err(_) => write!(f, "ArenaKey({:?})", self.as_bytes()),
833 }
834 }
835}
836
837impl<'a> PartialEq for ArenaKey<'a> {
838 fn eq(&self, other: &Self) -> bool {
839 self.as_bytes() == other.as_bytes()
840 }
841}
842
843impl<'a> Eq for ArenaKey<'a> {}
844
845impl<'a> std::hash::Hash for ArenaKey<'a> {
846 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
847 self.as_bytes().hash(state);
848 }
849}
850
851#[derive(Clone)]
872pub struct ArenaKeyHandle {
873 hash: u64,
875 data: KeyData,
877}
878
879#[derive(Clone)]
881enum KeyData {
882 Inline { len: u8, bytes: [u8; 24] },
884 Heap(Vec<u8>),
886}
887
888impl ArenaKeyHandle {
889 const INLINE_MAX: usize = 24;
891
892 #[inline]
894 pub fn new(data: &[u8]) -> Self {
895 let hash = Self::compute_hash(data);
896 let data = if data.len() <= Self::INLINE_MAX {
897 let mut bytes = [0u8; 24];
898 bytes[..data.len()].copy_from_slice(data);
899 KeyData::Inline {
900 len: data.len() as u8,
901 bytes,
902 }
903 } else {
904 KeyData::Heap(data.to_vec())
905 };
906 Self { hash, data }
907 }
908
909 #[inline]
911 pub fn from_arena_key(key: &ArenaKey<'_>) -> Self {
912 Self::new(key.as_bytes())
913 }
914
915 #[inline]
917 pub fn as_bytes(&self) -> &[u8] {
918 match &self.data {
919 KeyData::Inline { len, bytes } => &bytes[..*len as usize],
920 KeyData::Heap(v) => v.as_slice(),
921 }
922 }
923
924 #[inline]
926 pub fn hash(&self) -> u64 {
927 self.hash
928 }
929
930 #[inline]
932 pub fn len(&self) -> usize {
933 match &self.data {
934 KeyData::Inline { len, .. } => *len as usize,
935 KeyData::Heap(v) => v.len(),
936 }
937 }
938
939 #[inline]
941 pub fn is_empty(&self) -> bool {
942 self.len() == 0
943 }
944
945 #[inline]
947 pub fn is_inline(&self) -> bool {
948 matches!(self.data, KeyData::Inline { .. })
949 }
950
951 #[inline]
953 fn compute_hash(data: &[u8]) -> u64 {
954 const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
955 const FNV_PRIME: u64 = 0x00000100000001B3;
956
957 let mut h = FNV_OFFSET_BASIS;
958 for &b in data {
959 h ^= b as u64;
960 h = h.wrapping_mul(FNV_PRIME);
961 }
962 h
963 }
964}
965
966impl PartialEq for ArenaKeyHandle {
967 #[inline]
968 fn eq(&self, other: &Self) -> bool {
969 self.hash == other.hash && self.len() == other.len() && self.as_bytes() == other.as_bytes()
971 }
972}
973
974impl Eq for ArenaKeyHandle {}
975
976impl std::hash::Hash for ArenaKeyHandle {
977 #[inline]
978 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
979 state.write_u64(self.hash);
981 }
982}
983
984impl PartialOrd for ArenaKeyHandle {
985 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
986 Some(self.cmp(other))
987 }
988}
989
990impl Ord for ArenaKeyHandle {
991 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
992 self.as_bytes().cmp(other.as_bytes())
993 }
994}
995
996impl AsRef<[u8]> for ArenaKeyHandle {
997 fn as_ref(&self) -> &[u8] {
998 self.as_bytes()
999 }
1000}
1001
1002impl std::fmt::Debug for ArenaKeyHandle {
1003 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1004 match std::str::from_utf8(self.as_bytes()) {
1005 Ok(s) => write!(f, "ArenaKeyHandle({:?})", s),
1006 Err(_) => write!(f, "ArenaKeyHandle({:?})", self.as_bytes()),
1007 }
1008 }
1009}
1010
1011impl From<&[u8]> for ArenaKeyHandle {
1012 fn from(data: &[u8]) -> Self {
1013 Self::new(data)
1014 }
1015}
1016
1017impl From<Vec<u8>> for ArenaKeyHandle {
1018 fn from(data: Vec<u8>) -> Self {
1019 Self::new(&data)
1020 }
1021}
1022
1023impl From<&str> for ArenaKeyHandle {
1024 fn from(s: &str) -> Self {
1025 Self::new(s.as_bytes())
1026 }
1027}