1use crossbeam_epoch::{self as epoch, Atomic, Guard, Owned, Shared};
72use std::mem::MaybeUninit;
73use std::ptr;
74
75fn init_atomic_array_48() -> [Atomic<ArtNode>; 48] {
81 let mut arr: [MaybeUninit<Atomic<ArtNode>>; 48] = 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] = unsafe { MaybeUninit::uninit().assume_init() };
93 for elem in &mut arr {
94 elem.write(Atomic::null());
95 }
96 unsafe { std::mem::transmute(arr) }
97}
98use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
99use std::sync::Arc;
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() {
414 Some(child)
415 } else {
416 None
417 }
418 }
419
420 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
422 let n = self.header.num_children.load(Ordering::Acquire);
423 if n >= 48 {
424 return false;
425 }
426
427 unsafe {
428 let index_ptr = self.child_index.as_ptr() as *mut u8;
429 let children_ptr = self.children.as_ptr() as *mut Atomic<ArtNode>;
430
431 *index_ptr.add(key as usize) = (n + 1) as u8;
432 (*children_ptr.add(n)).store(child, Ordering::Release);
433 }
434
435 self.header.num_children.fetch_add(1, Ordering::Release);
436 true
437 }
438
439 pub fn is_full(&self) -> bool {
440 self.header.num_children.load(Ordering::Acquire) >= 48
441 }
442}
443
444impl Default for Node48 {
445 fn default() -> Self {
446 Self::new()
447 }
448}
449
450impl Node256 {
451 pub fn new() -> Self {
452 Self {
453 header: NodeHeader::new(),
454 children: init_atomic_array_256(),
455 }
456 }
457
458 pub fn with_prefix(prefix: &[u8]) -> Self {
459 Self {
460 header: NodeHeader::with_prefix(prefix),
461 children: init_atomic_array_256(),
462 }
463 }
464
465 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
467 let child = self.children[key as usize].load(Ordering::Acquire, guard);
468 if !child.is_null() {
469 Some(child)
470 } else {
471 None
472 }
473 }
474
475 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
477 let old = self.children[key as usize].swap(child, Ordering::AcqRel, guard);
478 if old.is_null() {
479 self.header.num_children.fetch_add(1, Ordering::Release);
480 }
481 true
482 }
483
484 pub fn is_full(&self) -> bool {
485 false }
487}
488
489impl Default for Node256 {
490 fn default() -> Self {
491 Self::new()
492 }
493}
494
495pub struct ConcurrentART {
508 root: Atomic<ArtNode>,
510 size: AtomicU64,
512 memory_usage: AtomicU64,
514}
515
516impl ConcurrentART {
517 pub fn new() -> Self {
519 Self {
520 root: Atomic::null(),
521 size: AtomicU64::new(0),
522 memory_usage: AtomicU64::new(0),
523 }
524 }
525
526 pub fn get(&self, key: &[u8]) -> Option<(Option<Vec<u8>>, u64)> {
530 let guard = &epoch::pin();
531 let mut node = self.root.load(Ordering::Acquire, guard);
532 let mut depth = 0;
533
534 while !node.is_null() {
535 let node_ref = unsafe { node.deref() };
536
537 match node_ref {
538 ArtNode::Leaf(leaf) => {
539 if leaf.key == key {
541 return Some((leaf.value.clone(), leaf.seqno));
542 }
543 return None;
544 }
545 ArtNode::Node4(n) => {
546 if !self.check_prefix(&n.header, key, depth) {
548 return None;
549 }
550 depth += n.header.prefix_len;
551
552 if depth >= key.len() {
553 return None;
554 }
555
556 match n.find_child(key[depth], guard) {
557 Some(child) => {
558 node = child;
559 depth += 1;
560 }
561 None => return None,
562 }
563 }
564 ArtNode::Node16(n) => {
565 if !self.check_prefix(&n.header, key, depth) {
566 return None;
567 }
568 depth += n.header.prefix_len;
569
570 if depth >= key.len() {
571 return None;
572 }
573
574 match n.find_child(key[depth], guard) {
575 Some(child) => {
576 node = child;
577 depth += 1;
578 }
579 None => return None,
580 }
581 }
582 ArtNode::Node48(n) => {
583 if !self.check_prefix(&n.header, key, depth) {
584 return None;
585 }
586 depth += n.header.prefix_len;
587
588 if depth >= key.len() {
589 return None;
590 }
591
592 match n.find_child(key[depth], guard) {
593 Some(child) => {
594 node = child;
595 depth += 1;
596 }
597 None => return None,
598 }
599 }
600 ArtNode::Node256(n) => {
601 if !self.check_prefix(&n.header, key, depth) {
602 return None;
603 }
604 depth += n.header.prefix_len;
605
606 if depth >= key.len() {
607 return None;
608 }
609
610 match n.find_child(key[depth], guard) {
611 Some(child) => {
612 node = child;
613 depth += 1;
614 }
615 None => return None,
616 }
617 }
618 }
619 }
620
621 None
622 }
623
624 fn check_prefix(&self, header: &NodeHeader, key: &[u8], depth: usize) -> bool {
626 let prefix_len = header.prefix_len.min(MAX_PREFIX_LEN);
627 if depth + prefix_len > key.len() {
628 return false;
629 }
630
631 for i in 0..prefix_len {
632 if header.prefix[i] != key[depth + i] {
633 return false;
634 }
635 }
636 true
637 }
638
639 pub fn insert(&self, key: Vec<u8>, value: Option<Vec<u8>>, seqno: u64) -> Option<(Option<Vec<u8>>, u64)> {
643 let guard = &epoch::pin();
644
645 loop {
646 let root = self.root.load(Ordering::Acquire, guard);
647
648 if root.is_null() {
649 let leaf = Box::new(LeafNode::new(key.clone(), value.clone(), seqno));
652 let leaf_node = Owned::new(ArtNode::Leaf(leaf));
653
654 match self.root.compare_exchange(
655 Shared::null(),
656 leaf_node,
657 Ordering::AcqRel,
658 Ordering::Relaxed,
659 guard,
660 ) {
661 Ok(_) => {
662 self.size.fetch_add(1, Ordering::Relaxed);
663 self.memory_usage.fetch_add(
664 key.len() as u64 + std::mem::size_of::<LeafNode>() as u64,
665 Ordering::Relaxed,
666 );
667 return None;
668 }
669 Err(_e) => {
670 continue;
672 }
673 }
674 }
675
676 match self.insert_recursive(root, &key, value.clone(), seqno, 0, guard) {
680 InsertResult::Success(old) => return old,
681 InsertResult::Retry => continue,
682 }
683 }
684 }
685
686 fn insert_recursive<'g>(
687 &self,
688 node: Shared<'g, ArtNode>,
689 key: &[u8],
690 value: Option<Vec<u8>>,
691 seqno: u64,
692 depth: usize,
693 guard: &'g Guard,
694 ) -> InsertResult {
695 if node.is_null() {
696 return InsertResult::Retry;
697 }
698
699 let node_ref = unsafe { node.deref() };
700
701 match node_ref {
702 ArtNode::Leaf(existing_leaf) => {
703 if existing_leaf.key == key {
704 return InsertResult::Success(Some((existing_leaf.value.clone(), existing_leaf.seqno)));
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 _ = new_node.add_child(existing_key[common_depth], existing_leaf_node, guard);
740 }
741
742 self.size.fetch_add(1, Ordering::Relaxed);
746 InsertResult::Success(None)
747 }
748 ArtNode::Node4(n) => {
749 if depth < key.len() {
750 if let Some(child) = n.find_child(key[depth], guard) {
751 return self.insert_recursive(child, key, value, seqno, depth + 1, guard);
752 }
753 if !n.is_full() {
755 let leaf = Box::new(LeafNode::new(key.to_vec(), value, seqno));
756 let leaf_node = Owned::new(ArtNode::Leaf(leaf));
757 if n.add_child(key[depth], leaf_node, guard) {
758 self.size.fetch_add(1, Ordering::Relaxed);
759 return InsertResult::Success(None);
760 }
761 }
762 }
763 InsertResult::Retry
764 }
765 _ => InsertResult::Retry,
766 }
767 }
768
769 pub fn len(&self) -> u64 {
771 self.size.load(Ordering::Relaxed)
772 }
773
774 pub fn is_empty(&self) -> bool {
776 self.len() == 0
777 }
778
779 pub fn memory_usage(&self) -> u64 {
781 self.memory_usage.load(Ordering::Relaxed)
782 }
783}
784
785impl Default for ConcurrentART {
786 fn default() -> Self {
787 Self::new()
788 }
789}
790
791enum InsertResult {
792 Success(Option<(Option<Vec<u8>>, u64)>),
793 Retry,
794}
795
796#[cfg(test)]
801mod tests {
802 use super::*;
803
804 #[test]
805 fn test_art_empty() {
806 let art = ConcurrentART::new();
807 assert!(art.is_empty());
808 assert_eq!(art.get(b"key"), None);
809 }
810
811 #[test]
812 fn test_art_insert_get() {
813 let art = ConcurrentART::new();
814
815 let old = art.insert(b"hello".to_vec(), Some(b"world".to_vec()), 1);
816 assert!(old.is_none());
817 assert_eq!(art.len(), 1);
818
819 let result = art.get(b"hello");
820 assert!(result.is_some());
821 let (value, seqno) = result.unwrap();
822 assert_eq!(value, Some(b"world".to_vec()));
823 assert_eq!(seqno, 1);
824 }
825
826 #[test]
827 fn test_art_multiple_keys() {
828 let art = ConcurrentART::new();
832
833 art.insert(b"key1".to_vec(), Some(b"value1".to_vec()), 1);
834 art.insert(b"key2".to_vec(), Some(b"value2".to_vec()), 2);
835 art.insert(b"key3".to_vec(), Some(b"value3".to_vec()), 3);
836
837 assert_eq!(art.len(), 3);
839 }
840
841 #[test]
842 fn test_node4_operations() {
843 let node = Node4::new();
844 let guard = &epoch::pin();
845
846 assert!(!node.is_full());
847 assert_eq!(node.find_child(b'a', guard), None);
848 }
849}