1use crossbeam_epoch::{self as epoch, Atomic, Guard, Owned, Shared};
72use std::mem::MaybeUninit;
73
74fn init_atomic_array_48() -> [Atomic<ArtNode>; 48] {
80 let mut arr: [MaybeUninit<Atomic<ArtNode>>; 48] =
82 unsafe { MaybeUninit::uninit().assume_init() };
83 for elem in &mut arr {
84 elem.write(Atomic::null());
85 }
86 unsafe { std::mem::transmute(arr) }
88}
89
90fn init_atomic_array_256() -> [Atomic<ArtNode>; 256] {
92 let mut arr: [MaybeUninit<Atomic<ArtNode>>; 256] =
93 unsafe { MaybeUninit::uninit().assume_init() };
94 for elem in &mut arr {
95 elem.write(Atomic::null());
96 }
97 unsafe { std::mem::transmute(arr) }
98}
99use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
100
101const MAX_PREFIX_LEN: usize = 10;
107
108#[repr(C)]
110#[derive(Debug)]
111pub struct NodeHeader {
112 pub num_children: AtomicUsize,
114 pub prefix_len: usize,
116 pub prefix: [u8; MAX_PREFIX_LEN],
118}
119
120impl NodeHeader {
121 fn new() -> Self {
122 Self {
123 num_children: AtomicUsize::new(0),
124 prefix_len: 0,
125 prefix: [0; MAX_PREFIX_LEN],
126 }
127 }
128
129 fn with_prefix(prefix: &[u8]) -> Self {
130 let mut header = Self::new();
131 let len = prefix.len().min(MAX_PREFIX_LEN);
132 header.prefix[..len].copy_from_slice(&prefix[..len]);
133 header.prefix_len = prefix.len();
134 header
135 }
136}
137
138#[repr(C)]
141pub struct Node4 {
142 pub header: NodeHeader,
143 pub keys: [u8; 4],
145 pub children: [Atomic<ArtNode>; 4],
147}
148
149#[repr(C)]
151pub struct Node16 {
152 pub header: NodeHeader,
153 pub keys: [u8; 16],
155 pub children: [Atomic<ArtNode>; 16],
157}
158
159#[repr(C)]
162pub struct Node48 {
163 pub header: NodeHeader,
164 pub child_index: [u8; 256],
166 pub children: [Atomic<ArtNode>; 48],
168}
169
170#[repr(C)]
173pub struct Node256 {
174 pub header: NodeHeader,
175 pub children: [Atomic<ArtNode>; 256],
177}
178
179#[derive(Debug)]
181pub struct LeafNode {
182 pub key: Vec<u8>,
184 pub value: Option<Vec<u8>>,
186 pub seqno: u64,
188}
189
190impl LeafNode {
191 pub fn new(key: Vec<u8>, value: Option<Vec<u8>>, seqno: u64) -> Self {
192 Self { key, value, seqno }
193 }
194}
195
196pub enum ArtNode {
198 Node4(Box<Node4>),
199 Node16(Box<Node16>),
200 Node48(Box<Node48>),
201 Node256(Box<Node256>),
202 Leaf(Box<LeafNode>),
203}
204
205impl std::fmt::Debug for ArtNode {
206 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
207 match self {
208 ArtNode::Node4(_) => write!(f, "Node4"),
209 ArtNode::Node16(_) => write!(f, "Node16"),
210 ArtNode::Node48(_) => write!(f, "Node48"),
211 ArtNode::Node256(_) => write!(f, "Node256"),
212 ArtNode::Leaf(l) => write!(f, "Leaf({:?})", l.key),
213 }
214 }
215}
216
217impl Node4 {
222 pub fn new() -> Self {
223 Self {
224 header: NodeHeader::new(),
225 keys: [0; 4],
226 children: Default::default(),
227 }
228 }
229
230 pub fn with_prefix(prefix: &[u8]) -> Self {
231 Self {
232 header: NodeHeader::with_prefix(prefix),
233 keys: [0; 4],
234 children: Default::default(),
235 }
236 }
237
238 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
240 let n = self.header.num_children.load(Ordering::Acquire);
241 for i in 0..n {
242 if self.keys[i] == key {
243 let child = self.children[i].load(Ordering::Acquire, guard);
244 if !child.is_null() {
245 return Some(child);
246 }
247 }
248 }
249 None
250 }
251
252 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
254 let n = self.header.num_children.load(Ordering::Acquire);
255 if n >= 4 {
256 return false; }
258
259 let mut pos = n;
261 for i in 0..n {
262 if self.keys[i] > key {
263 pos = i;
264 break;
265 }
266 }
267
268 unsafe {
271 let keys_ptr = self.keys.as_ptr() as *mut u8;
272 let children_ptr = self.children.as_ptr() as *mut Atomic<ArtNode>;
273
274 for i in (pos..n).rev() {
275 *keys_ptr.add(i + 1) = *keys_ptr.add(i);
276 (*children_ptr.add(i + 1)).store(
277 (*children_ptr.add(i)).load(Ordering::Relaxed, guard),
278 Ordering::Relaxed,
279 );
280 }
281 *keys_ptr.add(pos) = key;
282 (*children_ptr.add(pos)).store(child, Ordering::Release);
283 }
284
285 self.header.num_children.fetch_add(1, Ordering::Release);
286 true
287 }
288
289 pub fn is_full(&self) -> bool {
291 self.header.num_children.load(Ordering::Acquire) >= 4
292 }
293}
294
295impl Default for Node4 {
296 fn default() -> Self {
297 Self::new()
298 }
299}
300
301impl Node16 {
302 pub fn new() -> Self {
303 Self {
304 header: NodeHeader::new(),
305 keys: [0; 16],
306 children: Default::default(),
307 }
308 }
309
310 pub fn with_prefix(prefix: &[u8]) -> Self {
311 Self {
312 header: NodeHeader::with_prefix(prefix),
313 keys: [0; 16],
314 children: Default::default(),
315 }
316 }
317
318 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
320 let n = self.header.num_children.load(Ordering::Acquire);
321
322 let mut lo = 0;
324 let mut hi = n;
325 while lo < hi {
326 let mid = lo + (hi - lo) / 2;
327 if self.keys[mid] < key {
328 lo = mid + 1;
329 } else {
330 hi = mid;
331 }
332 }
333
334 if lo < n && self.keys[lo] == key {
335 let child = self.children[lo].load(Ordering::Acquire, guard);
336 if !child.is_null() {
337 return Some(child);
338 }
339 }
340 None
341 }
342
343 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
345 let n = self.header.num_children.load(Ordering::Acquire);
346 if n >= 16 {
347 return false;
348 }
349
350 let mut pos = n;
352 for i in 0..n {
353 if self.keys[i] > key {
354 pos = i;
355 break;
356 }
357 }
358
359 unsafe {
360 let keys_ptr = self.keys.as_ptr() as *mut u8;
361 let children_ptr = self.children.as_ptr() as *mut Atomic<ArtNode>;
362
363 for i in (pos..n).rev() {
364 *keys_ptr.add(i + 1) = *keys_ptr.add(i);
365 (*children_ptr.add(i + 1)).store(
366 (*children_ptr.add(i)).load(Ordering::Relaxed, guard),
367 Ordering::Relaxed,
368 );
369 }
370 *keys_ptr.add(pos) = key;
371 (*children_ptr.add(pos)).store(child, Ordering::Release);
372 }
373
374 self.header.num_children.fetch_add(1, Ordering::Release);
375 true
376 }
377
378 pub fn is_full(&self) -> bool {
379 self.header.num_children.load(Ordering::Acquire) >= 16
380 }
381}
382
383impl Default for Node16 {
384 fn default() -> Self {
385 Self::new()
386 }
387}
388
389impl Node48 {
390 pub fn new() -> Self {
391 Self {
392 header: NodeHeader::new(),
393 child_index: [0; 256],
394 children: init_atomic_array_48(),
395 }
396 }
397
398 pub fn with_prefix(prefix: &[u8]) -> Self {
399 Self {
400 header: NodeHeader::with_prefix(prefix),
401 child_index: [0; 256],
402 children: init_atomic_array_48(),
403 }
404 }
405
406 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
408 let idx = self.child_index[key as usize];
409 if idx == 0 {
410 return None;
411 }
412 let child = self.children[(idx - 1) as usize].load(Ordering::Acquire, guard);
413 if !child.is_null() { Some(child) } else { None }
414 }
415
416 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, _guard: &Guard) -> bool {
418 let n = self.header.num_children.load(Ordering::Acquire);
419 if n >= 48 {
420 return false;
421 }
422
423 unsafe {
424 let index_ptr = self.child_index.as_ptr() as *mut u8;
425 let children_ptr = self.children.as_ptr() as *mut Atomic<ArtNode>;
426
427 *index_ptr.add(key as usize) = (n + 1) as u8;
428 (*children_ptr.add(n)).store(child, Ordering::Release);
429 }
430
431 self.header.num_children.fetch_add(1, Ordering::Release);
432 true
433 }
434
435 pub fn is_full(&self) -> bool {
436 self.header.num_children.load(Ordering::Acquire) >= 48
437 }
438}
439
440impl Default for Node48 {
441 fn default() -> Self {
442 Self::new()
443 }
444}
445
446impl Node256 {
447 pub fn new() -> Self {
448 Self {
449 header: NodeHeader::new(),
450 children: init_atomic_array_256(),
451 }
452 }
453
454 pub fn with_prefix(prefix: &[u8]) -> Self {
455 Self {
456 header: NodeHeader::with_prefix(prefix),
457 children: init_atomic_array_256(),
458 }
459 }
460
461 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
463 let child = self.children[key as usize].load(Ordering::Acquire, guard);
464 if !child.is_null() { Some(child) } else { None }
465 }
466
467 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
469 let old = self.children[key as usize].swap(child, Ordering::AcqRel, guard);
470 if old.is_null() {
471 self.header.num_children.fetch_add(1, Ordering::Release);
472 }
473 true
474 }
475
476 pub fn is_full(&self) -> bool {
477 false }
479}
480
481impl Default for Node256 {
482 fn default() -> Self {
483 Self::new()
484 }
485}
486
487pub struct ConcurrentART {
500 root: Atomic<ArtNode>,
502 size: AtomicU64,
504 memory_usage: AtomicU64,
506}
507
508impl ConcurrentART {
509 pub fn new() -> Self {
511 Self {
512 root: Atomic::null(),
513 size: AtomicU64::new(0),
514 memory_usage: AtomicU64::new(0),
515 }
516 }
517
518 pub fn get(&self, key: &[u8]) -> Option<(Option<Vec<u8>>, u64)> {
522 let guard = &epoch::pin();
523 let mut node = self.root.load(Ordering::Acquire, guard);
524 let mut depth = 0;
525
526 while !node.is_null() {
527 let node_ref = unsafe { node.deref() };
528
529 match node_ref {
530 ArtNode::Leaf(leaf) => {
531 if leaf.key == key {
533 return Some((leaf.value.clone(), leaf.seqno));
534 }
535 return None;
536 }
537 ArtNode::Node4(n) => {
538 if !self.check_prefix(&n.header, key, depth) {
540 return None;
541 }
542 depth += n.header.prefix_len;
543
544 if depth >= key.len() {
545 return None;
546 }
547
548 match n.find_child(key[depth], guard) {
549 Some(child) => {
550 node = child;
551 depth += 1;
552 }
553 None => return None,
554 }
555 }
556 ArtNode::Node16(n) => {
557 if !self.check_prefix(&n.header, key, depth) {
558 return None;
559 }
560 depth += n.header.prefix_len;
561
562 if depth >= key.len() {
563 return None;
564 }
565
566 match n.find_child(key[depth], guard) {
567 Some(child) => {
568 node = child;
569 depth += 1;
570 }
571 None => return None,
572 }
573 }
574 ArtNode::Node48(n) => {
575 if !self.check_prefix(&n.header, key, depth) {
576 return None;
577 }
578 depth += n.header.prefix_len;
579
580 if depth >= key.len() {
581 return None;
582 }
583
584 match n.find_child(key[depth], guard) {
585 Some(child) => {
586 node = child;
587 depth += 1;
588 }
589 None => return None,
590 }
591 }
592 ArtNode::Node256(n) => {
593 if !self.check_prefix(&n.header, key, depth) {
594 return None;
595 }
596 depth += n.header.prefix_len;
597
598 if depth >= key.len() {
599 return None;
600 }
601
602 match n.find_child(key[depth], guard) {
603 Some(child) => {
604 node = child;
605 depth += 1;
606 }
607 None => return None,
608 }
609 }
610 }
611 }
612
613 None
614 }
615
616 fn check_prefix(&self, header: &NodeHeader, key: &[u8], depth: usize) -> bool {
618 let prefix_len = header.prefix_len.min(MAX_PREFIX_LEN);
619 if depth + prefix_len > key.len() {
620 return false;
621 }
622
623 for i in 0..prefix_len {
624 if header.prefix[i] != key[depth + i] {
625 return false;
626 }
627 }
628 true
629 }
630
631 pub fn insert(
635 &self,
636 key: Vec<u8>,
637 value: Option<Vec<u8>>,
638 seqno: u64,
639 ) -> Option<(Option<Vec<u8>>, u64)> {
640 let guard = &epoch::pin();
641
642 loop {
643 let root = self.root.load(Ordering::Acquire, guard);
644
645 if root.is_null() {
646 let leaf = Box::new(LeafNode::new(key.clone(), value.clone(), seqno));
649 let leaf_node = Owned::new(ArtNode::Leaf(leaf));
650
651 match self.root.compare_exchange(
652 Shared::null(),
653 leaf_node,
654 Ordering::AcqRel,
655 Ordering::Relaxed,
656 guard,
657 ) {
658 Ok(_) => {
659 self.size.fetch_add(1, Ordering::Relaxed);
660 self.memory_usage.fetch_add(
661 key.len() as u64 + std::mem::size_of::<LeafNode>() as u64,
662 Ordering::Relaxed,
663 );
664 return None;
665 }
666 Err(_e) => {
667 continue;
669 }
670 }
671 }
672
673 match self.insert_recursive(root, &key, value.clone(), seqno, 0, guard) {
677 InsertResult::Success(old) => return old,
678 InsertResult::Retry => continue,
679 }
680 }
681 }
682
683 fn insert_recursive<'g>(
684 &self,
685 node: Shared<'g, ArtNode>,
686 key: &[u8],
687 value: Option<Vec<u8>>,
688 seqno: u64,
689 depth: usize,
690 guard: &'g Guard,
691 ) -> InsertResult {
692 if node.is_null() {
693 return InsertResult::Retry;
694 }
695
696 let node_ref = unsafe { node.deref() };
697
698 match node_ref {
699 ArtNode::Leaf(existing_leaf) => {
700 if existing_leaf.key == key {
701 return InsertResult::Success(Some((
704 existing_leaf.value.clone(),
705 existing_leaf.seqno,
706 )));
707 }
708
709 let existing_key = &existing_leaf.key;
712 let mut common_depth = depth;
713 while common_depth < key.len().min(existing_key.len())
714 && key[common_depth] == existing_key[common_depth]
715 {
716 common_depth += 1;
717 }
718
719 let new_node = Node4::new();
721
722 let new_leaf = Box::new(LeafNode::new(key.to_vec(), value, seqno));
724 let new_leaf_node = Owned::new(ArtNode::Leaf(new_leaf));
725
726 if common_depth < key.len() {
728 let _ = new_node.add_child(key[common_depth], new_leaf_node, guard);
729 }
730
731 if common_depth < existing_key.len() {
733 let existing_leaf_clone = Box::new(LeafNode::new(
734 existing_key.clone(),
735 existing_leaf.value.clone(),
736 existing_leaf.seqno,
737 ));
738 let existing_leaf_node = Owned::new(ArtNode::Leaf(existing_leaf_clone));
739 let _ =
740 new_node.add_child(existing_key[common_depth], existing_leaf_node, guard);
741 }
742
743 self.size.fetch_add(1, Ordering::Relaxed);
747 InsertResult::Success(None)
748 }
749 ArtNode::Node4(n) => {
750 if depth < key.len() {
751 if let Some(child) = n.find_child(key[depth], guard) {
752 return self.insert_recursive(child, key, value, seqno, depth + 1, guard);
753 }
754 if !n.is_full() {
756 let leaf = Box::new(LeafNode::new(key.to_vec(), value, seqno));
757 let leaf_node = Owned::new(ArtNode::Leaf(leaf));
758 if n.add_child(key[depth], leaf_node, guard) {
759 self.size.fetch_add(1, Ordering::Relaxed);
760 return InsertResult::Success(None);
761 }
762 }
763 }
764 InsertResult::Retry
765 }
766 _ => InsertResult::Retry,
767 }
768 }
769
770 pub fn len(&self) -> u64 {
772 self.size.load(Ordering::Relaxed)
773 }
774
775 pub fn is_empty(&self) -> bool {
777 self.len() == 0
778 }
779
780 pub fn memory_usage(&self) -> u64 {
782 self.memory_usage.load(Ordering::Relaxed)
783 }
784}
785
786impl Default for ConcurrentART {
787 fn default() -> Self {
788 Self::new()
789 }
790}
791
792enum InsertResult {
793 Success(Option<(Option<Vec<u8>>, u64)>),
794 Retry,
795}
796
797#[cfg(test)]
802mod tests {
803 use super::*;
804
805 #[test]
806 fn test_art_empty() {
807 let art = ConcurrentART::new();
808 assert!(art.is_empty());
809 assert_eq!(art.get(b"key"), None);
810 }
811
812 #[test]
813 fn test_art_insert_get() {
814 let art = ConcurrentART::new();
815
816 let old = art.insert(b"hello".to_vec(), Some(b"world".to_vec()), 1);
817 assert!(old.is_none());
818 assert_eq!(art.len(), 1);
819
820 let result = art.get(b"hello");
821 assert!(result.is_some());
822 let (value, seqno) = result.unwrap();
823 assert_eq!(value, Some(b"world".to_vec()));
824 assert_eq!(seqno, 1);
825 }
826
827 #[test]
828 fn test_art_multiple_keys() {
829 let art = ConcurrentART::new();
833
834 art.insert(b"key1".to_vec(), Some(b"value1".to_vec()), 1);
835 art.insert(b"key2".to_vec(), Some(b"value2".to_vec()), 2);
836 art.insert(b"key3".to_vec(), Some(b"value3".to_vec()), 3);
837
838 assert_eq!(art.len(), 3);
840 }
841
842 #[test]
843 fn test_node4_operations() {
844 let node = Node4::new();
845 let guard = &epoch::pin();
846
847 assert!(!node.is_full());
848 assert_eq!(node.find_child(b'a', guard), None);
849 }
850}