1use crossbeam_epoch::{self as epoch, Atomic, Guard, Owned, Shared};
69use std::mem::MaybeUninit;
70use std::ptr;
71
72fn init_atomic_array_48() -> [Atomic<ArtNode>; 48] {
78 let mut arr: [MaybeUninit<Atomic<ArtNode>>; 48] = unsafe { MaybeUninit::uninit().assume_init() };
80 for elem in &mut arr {
81 elem.write(Atomic::null());
82 }
83 unsafe { std::mem::transmute(arr) }
85}
86
87fn init_atomic_array_256() -> [Atomic<ArtNode>; 256] {
89 let mut arr: [MaybeUninit<Atomic<ArtNode>>; 256] = unsafe { MaybeUninit::uninit().assume_init() };
90 for elem in &mut arr {
91 elem.write(Atomic::null());
92 }
93 unsafe { std::mem::transmute(arr) }
94}
95use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
96use std::sync::Arc;
97
98const MAX_PREFIX_LEN: usize = 10;
104
105#[repr(C)]
107#[derive(Debug)]
108pub struct NodeHeader {
109 pub num_children: AtomicUsize,
111 pub prefix_len: usize,
113 pub prefix: [u8; MAX_PREFIX_LEN],
115}
116
117impl NodeHeader {
118 fn new() -> Self {
119 Self {
120 num_children: AtomicUsize::new(0),
121 prefix_len: 0,
122 prefix: [0; MAX_PREFIX_LEN],
123 }
124 }
125
126 fn with_prefix(prefix: &[u8]) -> Self {
127 let mut header = Self::new();
128 let len = prefix.len().min(MAX_PREFIX_LEN);
129 header.prefix[..len].copy_from_slice(&prefix[..len]);
130 header.prefix_len = prefix.len();
131 header
132 }
133}
134
135#[repr(C)]
138pub struct Node4 {
139 pub header: NodeHeader,
140 pub keys: [u8; 4],
142 pub children: [Atomic<ArtNode>; 4],
144}
145
146#[repr(C)]
148pub struct Node16 {
149 pub header: NodeHeader,
150 pub keys: [u8; 16],
152 pub children: [Atomic<ArtNode>; 16],
154}
155
156#[repr(C)]
159pub struct Node48 {
160 pub header: NodeHeader,
161 pub child_index: [u8; 256],
163 pub children: [Atomic<ArtNode>; 48],
165}
166
167#[repr(C)]
170pub struct Node256 {
171 pub header: NodeHeader,
172 pub children: [Atomic<ArtNode>; 256],
174}
175
176#[derive(Debug)]
178pub struct LeafNode {
179 pub key: Vec<u8>,
181 pub value: Option<Vec<u8>>,
183 pub seqno: u64,
185}
186
187impl LeafNode {
188 pub fn new(key: Vec<u8>, value: Option<Vec<u8>>, seqno: u64) -> Self {
189 Self { key, value, seqno }
190 }
191}
192
193pub enum ArtNode {
195 Node4(Box<Node4>),
196 Node16(Box<Node16>),
197 Node48(Box<Node48>),
198 Node256(Box<Node256>),
199 Leaf(Box<LeafNode>),
200}
201
202impl std::fmt::Debug for ArtNode {
203 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
204 match self {
205 ArtNode::Node4(_) => write!(f, "Node4"),
206 ArtNode::Node16(_) => write!(f, "Node16"),
207 ArtNode::Node48(_) => write!(f, "Node48"),
208 ArtNode::Node256(_) => write!(f, "Node256"),
209 ArtNode::Leaf(l) => write!(f, "Leaf({:?})", l.key),
210 }
211 }
212}
213
214impl Node4 {
219 pub fn new() -> Self {
220 Self {
221 header: NodeHeader::new(),
222 keys: [0; 4],
223 children: Default::default(),
224 }
225 }
226
227 pub fn with_prefix(prefix: &[u8]) -> Self {
228 Self {
229 header: NodeHeader::with_prefix(prefix),
230 keys: [0; 4],
231 children: Default::default(),
232 }
233 }
234
235 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
237 let n = self.header.num_children.load(Ordering::Acquire);
238 for i in 0..n {
239 if self.keys[i] == key {
240 let child = self.children[i].load(Ordering::Acquire, guard);
241 if !child.is_null() {
242 return Some(child);
243 }
244 }
245 }
246 None
247 }
248
249 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
251 let n = self.header.num_children.load(Ordering::Acquire);
252 if n >= 4 {
253 return false; }
255
256 let mut pos = n;
258 for i in 0..n {
259 if self.keys[i] > key {
260 pos = i;
261 break;
262 }
263 }
264
265 unsafe {
268 let keys_ptr = self.keys.as_ptr() as *mut u8;
269 let children_ptr = self.children.as_ptr() as *mut Atomic<ArtNode>;
270
271 for i in (pos..n).rev() {
272 *keys_ptr.add(i + 1) = *keys_ptr.add(i);
273 (*children_ptr.add(i + 1)).store(
274 (*children_ptr.add(i)).load(Ordering::Relaxed, guard),
275 Ordering::Relaxed,
276 );
277 }
278 *keys_ptr.add(pos) = key;
279 (*children_ptr.add(pos)).store(child, Ordering::Release);
280 }
281
282 self.header.num_children.fetch_add(1, Ordering::Release);
283 true
284 }
285
286 pub fn is_full(&self) -> bool {
288 self.header.num_children.load(Ordering::Acquire) >= 4
289 }
290}
291
292impl Default for Node4 {
293 fn default() -> Self {
294 Self::new()
295 }
296}
297
298impl Node16 {
299 pub fn new() -> Self {
300 Self {
301 header: NodeHeader::new(),
302 keys: [0; 16],
303 children: Default::default(),
304 }
305 }
306
307 pub fn with_prefix(prefix: &[u8]) -> Self {
308 Self {
309 header: NodeHeader::with_prefix(prefix),
310 keys: [0; 16],
311 children: Default::default(),
312 }
313 }
314
315 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
317 let n = self.header.num_children.load(Ordering::Acquire);
318
319 let mut lo = 0;
321 let mut hi = n;
322 while lo < hi {
323 let mid = lo + (hi - lo) / 2;
324 if self.keys[mid] < key {
325 lo = mid + 1;
326 } else {
327 hi = mid;
328 }
329 }
330
331 if lo < n && self.keys[lo] == key {
332 let child = self.children[lo].load(Ordering::Acquire, guard);
333 if !child.is_null() {
334 return Some(child);
335 }
336 }
337 None
338 }
339
340 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
342 let n = self.header.num_children.load(Ordering::Acquire);
343 if n >= 16 {
344 return false;
345 }
346
347 let mut pos = n;
349 for i in 0..n {
350 if self.keys[i] > key {
351 pos = i;
352 break;
353 }
354 }
355
356 unsafe {
357 let keys_ptr = self.keys.as_ptr() as *mut u8;
358 let children_ptr = self.children.as_ptr() as *mut Atomic<ArtNode>;
359
360 for i in (pos..n).rev() {
361 *keys_ptr.add(i + 1) = *keys_ptr.add(i);
362 (*children_ptr.add(i + 1)).store(
363 (*children_ptr.add(i)).load(Ordering::Relaxed, guard),
364 Ordering::Relaxed,
365 );
366 }
367 *keys_ptr.add(pos) = key;
368 (*children_ptr.add(pos)).store(child, Ordering::Release);
369 }
370
371 self.header.num_children.fetch_add(1, Ordering::Release);
372 true
373 }
374
375 pub fn is_full(&self) -> bool {
376 self.header.num_children.load(Ordering::Acquire) >= 16
377 }
378}
379
380impl Default for Node16 {
381 fn default() -> Self {
382 Self::new()
383 }
384}
385
386impl Node48 {
387 pub fn new() -> Self {
388 Self {
389 header: NodeHeader::new(),
390 child_index: [0; 256],
391 children: init_atomic_array_48(),
392 }
393 }
394
395 pub fn with_prefix(prefix: &[u8]) -> Self {
396 Self {
397 header: NodeHeader::with_prefix(prefix),
398 child_index: [0; 256],
399 children: init_atomic_array_48(),
400 }
401 }
402
403 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
405 let idx = self.child_index[key as usize];
406 if idx == 0 {
407 return None;
408 }
409 let child = self.children[(idx - 1) as usize].load(Ordering::Acquire, guard);
410 if !child.is_null() {
411 Some(child)
412 } else {
413 None
414 }
415 }
416
417 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
419 let n = self.header.num_children.load(Ordering::Acquire);
420 if n >= 48 {
421 return false;
422 }
423
424 unsafe {
425 let index_ptr = self.child_index.as_ptr() as *mut u8;
426 let children_ptr = self.children.as_ptr() as *mut Atomic<ArtNode>;
427
428 *index_ptr.add(key as usize) = (n + 1) as u8;
429 (*children_ptr.add(n)).store(child, Ordering::Release);
430 }
431
432 self.header.num_children.fetch_add(1, Ordering::Release);
433 true
434 }
435
436 pub fn is_full(&self) -> bool {
437 self.header.num_children.load(Ordering::Acquire) >= 48
438 }
439}
440
441impl Default for Node48 {
442 fn default() -> Self {
443 Self::new()
444 }
445}
446
447impl Node256 {
448 pub fn new() -> Self {
449 Self {
450 header: NodeHeader::new(),
451 children: init_atomic_array_256(),
452 }
453 }
454
455 pub fn with_prefix(prefix: &[u8]) -> Self {
456 Self {
457 header: NodeHeader::with_prefix(prefix),
458 children: init_atomic_array_256(),
459 }
460 }
461
462 pub fn find_child<'g>(&self, key: u8, guard: &'g Guard) -> Option<Shared<'g, ArtNode>> {
464 let child = self.children[key as usize].load(Ordering::Acquire, guard);
465 if !child.is_null() {
466 Some(child)
467 } else {
468 None
469 }
470 }
471
472 pub fn add_child(&self, key: u8, child: Owned<ArtNode>, guard: &Guard) -> bool {
474 let old = self.children[key as usize].swap(child, Ordering::AcqRel, guard);
475 if old.is_null() {
476 self.header.num_children.fetch_add(1, Ordering::Release);
477 }
478 true
479 }
480
481 pub fn is_full(&self) -> bool {
482 false }
484}
485
486impl Default for Node256 {
487 fn default() -> Self {
488 Self::new()
489 }
490}
491
492pub struct ConcurrentART {
505 root: Atomic<ArtNode>,
507 size: AtomicU64,
509 memory_usage: AtomicU64,
511}
512
513impl ConcurrentART {
514 pub fn new() -> Self {
516 Self {
517 root: Atomic::null(),
518 size: AtomicU64::new(0),
519 memory_usage: AtomicU64::new(0),
520 }
521 }
522
523 pub fn get(&self, key: &[u8]) -> Option<(Option<Vec<u8>>, u64)> {
527 let guard = &epoch::pin();
528 let mut node = self.root.load(Ordering::Acquire, guard);
529 let mut depth = 0;
530
531 while !node.is_null() {
532 let node_ref = unsafe { node.deref() };
533
534 match node_ref {
535 ArtNode::Leaf(leaf) => {
536 if leaf.key == key {
538 return Some((leaf.value.clone(), leaf.seqno));
539 }
540 return None;
541 }
542 ArtNode::Node4(n) => {
543 if !self.check_prefix(&n.header, key, depth) {
545 return None;
546 }
547 depth += n.header.prefix_len;
548
549 if depth >= key.len() {
550 return None;
551 }
552
553 match n.find_child(key[depth], guard) {
554 Some(child) => {
555 node = child;
556 depth += 1;
557 }
558 None => return None,
559 }
560 }
561 ArtNode::Node16(n) => {
562 if !self.check_prefix(&n.header, key, depth) {
563 return None;
564 }
565 depth += n.header.prefix_len;
566
567 if depth >= key.len() {
568 return None;
569 }
570
571 match n.find_child(key[depth], guard) {
572 Some(child) => {
573 node = child;
574 depth += 1;
575 }
576 None => return None,
577 }
578 }
579 ArtNode::Node48(n) => {
580 if !self.check_prefix(&n.header, key, depth) {
581 return None;
582 }
583 depth += n.header.prefix_len;
584
585 if depth >= key.len() {
586 return None;
587 }
588
589 match n.find_child(key[depth], guard) {
590 Some(child) => {
591 node = child;
592 depth += 1;
593 }
594 None => return None,
595 }
596 }
597 ArtNode::Node256(n) => {
598 if !self.check_prefix(&n.header, key, depth) {
599 return None;
600 }
601 depth += n.header.prefix_len;
602
603 if depth >= key.len() {
604 return None;
605 }
606
607 match n.find_child(key[depth], guard) {
608 Some(child) => {
609 node = child;
610 depth += 1;
611 }
612 None => return None,
613 }
614 }
615 }
616 }
617
618 None
619 }
620
621 fn check_prefix(&self, header: &NodeHeader, key: &[u8], depth: usize) -> bool {
623 let prefix_len = header.prefix_len.min(MAX_PREFIX_LEN);
624 if depth + prefix_len > key.len() {
625 return false;
626 }
627
628 for i in 0..prefix_len {
629 if header.prefix[i] != key[depth + i] {
630 return false;
631 }
632 }
633 true
634 }
635
636 pub fn insert(&self, key: Vec<u8>, value: Option<Vec<u8>>, seqno: u64) -> 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((existing_leaf.value.clone(), existing_leaf.seqno)));
704 }
705
706 let existing_key = &existing_leaf.key;
709 let mut common_depth = depth;
710 while common_depth < key.len().min(existing_key.len())
711 && key[common_depth] == existing_key[common_depth]
712 {
713 common_depth += 1;
714 }
715
716 let new_node = Node4::new();
718
719 let new_leaf = Box::new(LeafNode::new(key.to_vec(), value, seqno));
721 let new_leaf_node = Owned::new(ArtNode::Leaf(new_leaf));
722
723 if common_depth < key.len() {
725 let _ = new_node.add_child(key[common_depth], new_leaf_node, guard);
726 }
727
728 if common_depth < existing_key.len() {
730 let existing_leaf_clone = Box::new(LeafNode::new(
731 existing_key.clone(),
732 existing_leaf.value.clone(),
733 existing_leaf.seqno,
734 ));
735 let existing_leaf_node = Owned::new(ArtNode::Leaf(existing_leaf_clone));
736 let _ = new_node.add_child(existing_key[common_depth], existing_leaf_node, guard);
737 }
738
739 self.size.fetch_add(1, Ordering::Relaxed);
743 InsertResult::Success(None)
744 }
745 ArtNode::Node4(n) => {
746 if depth < key.len() {
747 if let Some(child) = n.find_child(key[depth], guard) {
748 return self.insert_recursive(child, key, value, seqno, depth + 1, guard);
749 }
750 if !n.is_full() {
752 let leaf = Box::new(LeafNode::new(key.to_vec(), value, seqno));
753 let leaf_node = Owned::new(ArtNode::Leaf(leaf));
754 if n.add_child(key[depth], leaf_node, guard) {
755 self.size.fetch_add(1, Ordering::Relaxed);
756 return InsertResult::Success(None);
757 }
758 }
759 }
760 InsertResult::Retry
761 }
762 _ => InsertResult::Retry,
763 }
764 }
765
766 pub fn len(&self) -> u64 {
768 self.size.load(Ordering::Relaxed)
769 }
770
771 pub fn is_empty(&self) -> bool {
773 self.len() == 0
774 }
775
776 pub fn memory_usage(&self) -> u64 {
778 self.memory_usage.load(Ordering::Relaxed)
779 }
780}
781
782impl Default for ConcurrentART {
783 fn default() -> Self {
784 Self::new()
785 }
786}
787
788enum InsertResult {
789 Success(Option<(Option<Vec<u8>>, u64)>),
790 Retry,
791}
792
793#[cfg(test)]
798mod tests {
799 use super::*;
800
801 #[test]
802 fn test_art_empty() {
803 let art = ConcurrentART::new();
804 assert!(art.is_empty());
805 assert_eq!(art.get(b"key"), None);
806 }
807
808 #[test]
809 fn test_art_insert_get() {
810 let art = ConcurrentART::new();
811
812 let old = art.insert(b"hello".to_vec(), Some(b"world".to_vec()), 1);
813 assert!(old.is_none());
814 assert_eq!(art.len(), 1);
815
816 let result = art.get(b"hello");
817 assert!(result.is_some());
818 let (value, seqno) = result.unwrap();
819 assert_eq!(value, Some(b"world".to_vec()));
820 assert_eq!(seqno, 1);
821 }
822
823 #[test]
824 fn test_art_multiple_keys() {
825 let art = ConcurrentART::new();
829
830 art.insert(b"key1".to_vec(), Some(b"value1".to_vec()), 1);
831 art.insert(b"key2".to_vec(), Some(b"value2".to_vec()), 2);
832 art.insert(b"key3".to_vec(), Some(b"value3".to_vec()), 3);
833
834 assert_eq!(art.len(), 3);
836 }
837
838 #[test]
839 fn test_node4_operations() {
840 let node = Node4::new();
841 let guard = &epoch::pin();
842
843 assert!(!node.is_full());
844 assert_eq!(node.find_child(b'a', guard), None);
845 }
846}