art_tree/
lib.rs

1//! # Adaptive Radix Tree
2//! The radix tree based on ([The Adaptive Radix Tree:
3//! ARTful Indexing for Main-Memory Databases](https://15721.courses.cs.cmu.edu/spring2016/papers/leis-icde2013.pdf)).
4//!
5//! # Examples
6//! ```
7//! use art_tree::ByteString;
8//! use art_tree::KeyBuilder;
9//! use art_tree::Art;
10//!
11//! let mut art = Art::<u16, u16>::new();
12//! for i in 0..u8::MAX as u16 {
13//!     assert!(art.insert(i, i), "{}", i);
14//!     assert!(matches!(art.get(&i), Some(val) if val == &i));
15//! }
16//! for i in 0..u8::MAX as u16 {
17//!     assert!(matches!(art.remove(&i), Some(val) if val == i));
18//! }
19//!
20//! let mut art = Art::<ByteString, u16>::new();
21//! for i in 0..u8::MAX as u16 {
22//!     let key = KeyBuilder::new().append(i).append(ByteString::new("abc".to_string().as_bytes())).build();
23//!     art.upsert(key.clone(), i + 1);
24//!     assert!(matches!(art.get(&key), Some(val) if val == &(i + 1)));
25//! }
26//!
27//! let from_key = KeyBuilder::new().append(16u16).append(ByteString::new("abc".to_string().as_bytes())).build();
28//! let until_key = KeyBuilder::new().append(20u16).append(ByteString::new("abc".to_string().as_bytes())).build();
29//! assert_eq!(art.range(from_key..=until_key).count(), 5);
30//! assert_eq!(art.iter().count(), u8::MAX as usize);
31//! ```
32
33mod keys;
34mod node;
35mod scanner;
36
37use crate::scanner::Scanner;
38pub use keys::ByteString;
39pub use keys::*;
40use node::*;
41use std::cmp::Ordering;
42use std::marker::PhantomData;
43use std::ops::RangeBounds;
44use std::option::Option::Some;
45use std::rc::Rc;
46use std::{cmp, mem, ptr};
47
48/// Adaptive Radix Tree.  
49///
50/// Radix tree is ordered according to key. Radix tree requires that key to be representable as
51/// comparable by sequence, e.g. key should implement [Key] trait which used to convert it to
52/// byte sequence.
53///
54/// This crate provides [Key] implementations for most commonly used data types:
55/// - unsigned integers(u8, u16, u32, u64, u128)  
56/// - signed integers(i8, i16, i32, i64, i128)
57/// - usize
58/// - floating point numbers through [Float32]/[Float64] types
59/// - [ByteString] for raw byte sequences. It can be used for ASCII strings(UTF-8 strings
60/// not supported now, they require additional library to convert into comparable byte sequence).
61pub struct Art<K, V> {
62    root: Option<TypedNode<K, V>>,
63    // to make type !Send and !Sync
64    _phantom: PhantomData<Rc<K>>,
65}
66
67impl<K: Key, V> Default for Art<K, V> {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73impl<K: Key, V> Art<K, V> {
74    /// Create empty [ART] tree.
75    pub fn new() -> Self {
76        Self {
77            root: None,
78            _phantom: PhantomData {},
79        }
80    }
81
82    /// Insert key-value pair into tree.  
83    /// Return `true` if key-value successfully inserted into tree, otherwise `false` if tree
84    /// already contains same key.
85    pub fn insert(&mut self, key: K, value: V) -> bool {
86        self.insert_internal(key, value, false)
87    }
88
89    /// Insert key-value pair into tree.
90    /// If key already exists in tree, existing value will be replaced, otherwise inserts new KV
91    /// into tree.
92    pub fn upsert(&mut self, key: K, value: V) {
93        self.insert_internal(key, value, true);
94    }
95
96    fn insert_internal(&mut self, key: K, value: V, upsert: bool) -> bool {
97        let key_bytes = key.to_bytes();
98        assert!(
99            !key_bytes.is_empty(),
100            "Key must have non empty byte string representation"
101        );
102
103        if self.root.is_none() {
104            self.root.replace(TypedNode::Leaf(Leaf { key, value }));
105            true
106        } else {
107            let mut node = self.root.as_mut().unwrap();
108            let mut key = key;
109            let mut offset = 0;
110            let mut val = value;
111            loop {
112                let node_ptr = node as *mut TypedNode<K, V>;
113                let res = match node {
114                    TypedNode::Leaf(_) => Ok(Self::replace_leaf(
115                        node, key, val, &key_bytes, offset, upsert,
116                    )),
117                    TypedNode::Combined(interim, leaf) => {
118                        match leaf.key.to_bytes().cmp(&key_bytes) {
119                            Ordering::Equal => {
120                                if upsert {
121                                    leaf.value = val;
122                                    Ok(true)
123                                } else {
124                                    Ok(false)
125                                }
126                            }
127                            Ordering::Greater => {
128                                // new key is 'less' than any key in this level
129                                Self::replace_combined(unsafe { &mut *node_ptr }, key, val);
130                                Ok(true)
131                            }
132                            _ => Err(InsertOp {
133                                node: interim.as_mut(),
134                                key_byte_offset: offset,
135                                key,
136                                value: val,
137                            }),
138                        }
139                    }
140                    TypedNode::Interim(_) => {
141                        Self::interim_insert(node, key, val, &key_bytes, offset)
142                    }
143                };
144                match res {
145                    Ok(is_inserted) => return is_inserted,
146                    Err(op) => {
147                        node = op.node;
148                        offset = op.key_byte_offset;
149                        key = op.key;
150                        val = op.value;
151                    }
152                }
153            }
154        }
155    }
156
157    /// Remove value associated with key.  
158    /// Returns `Some(V)` if key found in tree, otherwise `None`.
159    pub fn remove(&mut self, key: &K) -> Option<V> {
160        if let Some(root) = &mut self.root {
161            let key_bytes_vec = key.to_bytes();
162            let mut key_bytes = key_bytes_vec.as_slice();
163            let key_ro_buffer = key_bytes;
164            let mut parent_link = 0;
165            let mut parent: Option<&mut BoxedNode<TypedNode<K, V>>> = None;
166            let mut node_ptr = root as *mut TypedNode<K, V>;
167            loop {
168                let x: &mut TypedNode<K, V> = unsafe { &mut *node_ptr };
169                match x {
170                    TypedNode::Leaf(leaf) => {
171                        // TODO: merge nodes if parent contains only link to child node
172                        return if key_ro_buffer == leaf.key.to_bytes() {
173                            if let Some(p) = parent {
174                                if p.should_shrink() {
175                                    unsafe {
176                                        let new_node = ptr::read(p).shrink();
177                                        ptr::write(p, new_node);
178                                    };
179                                }
180                                Some(p.remove(parent_link).unwrap().take_leaf().value)
181                            } else {
182                                Some(
183                                    mem::replace(&mut self.root, None)
184                                        .unwrap()
185                                        .take_leaf()
186                                        .value,
187                                )
188                            }
189                        } else {
190                            None
191                        };
192                    }
193                    TypedNode::Interim(interim) => {
194                        if let Some((next_node, rem_key_bytes, key)) =
195                            Self::find_in_interim_mut(interim, key_bytes)
196                        {
197                            node_ptr = next_node as *mut TypedNode<K, V>;
198                            parent = Some(interim);
199                            parent_link = key;
200                            key_bytes = rem_key_bytes;
201                        } else {
202                            return None;
203                        }
204                    }
205                    TypedNode::Combined(interim, leaf) => {
206                        if key_ro_buffer == leaf.key.to_bytes() {
207                            let leaf = unsafe { ptr::read(leaf) };
208                            unsafe { ptr::write(node_ptr, *ptr::read(interim)) };
209                            return Some(leaf.value);
210                        } else {
211                            node_ptr = interim.as_mut() as *mut TypedNode<K, V>;
212                        }
213                    }
214                }
215            }
216        } else {
217            None
218        }
219    }
220
221    /// Get value associated with key.  
222    /// Returns `Some(V)` if key found in tree, otherwise `None`.
223    pub fn get(&self, key: &K) -> Option<&V> {
224        let key_vec = key.to_bytes();
225        assert!(
226            !key_vec.is_empty(),
227            "Key must have non empty byte string representation"
228        );
229
230        let mut node = self.root.as_ref();
231        let mut key_bytes = key_vec.as_slice();
232        let key_ro_buffer = key_bytes;
233        while let Some(typed_node) = node {
234            match typed_node {
235                TypedNode::Leaf(leaf) => {
236                    return if leaf.key.to_bytes() == key_ro_buffer {
237                        Some(&leaf.value)
238                    } else {
239                        None
240                    }
241                }
242                TypedNode::Interim(interim) => {
243                    if let Some((next_node, rem_key_bytes, _)) =
244                        Self::find_in_interim(interim, key_bytes)
245                    {
246                        node = Some(next_node);
247                        key_bytes = rem_key_bytes;
248                    } else {
249                        node = None;
250                    }
251                }
252                TypedNode::Combined(interim, leaf) => {
253                    if key_ro_buffer == leaf.key.to_bytes() {
254                        return Some(&leaf.value);
255                    } else {
256                        node = Some(interim);
257                    }
258                }
259            }
260        }
261        None
262    }
263
264    /// Execute tree range scan.  
265    pub fn range(&self, range: impl RangeBounds<K>) -> impl DoubleEndedIterator<Item = (&K, &V)>
266    where
267        K: Ord,
268    {
269        if let Some(root) = self.root.as_ref() {
270            Scanner::new(root, range)
271        } else {
272            Scanner::empty(range)
273        }
274    }
275
276    /// Returns tree iterator.
277    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (&K, &V)>
278    where
279        K: Ord,
280    {
281        self.range(..)
282    }
283
284    fn find_in_interim<'n, 'k>(
285        interim: &'n BoxedNode<TypedNode<K, V>>,
286        key_bytes: &'k [u8],
287    ) -> Option<(&'n TypedNode<K, V>, &'k [u8], u8)> {
288        let node = unsafe {
289            #[allow(clippy::cast_ref_to_mut)]
290            &mut *(interim as *const BoxedNode<TypedNode<K, V>> as *mut BoxedNode<TypedNode<K, V>>)
291        };
292        Self::find_in_interim_mut(node, key_bytes)
293            .map(|(node, bytes, key)| (unsafe { &*(node as *const TypedNode<K, V>) }, bytes, key))
294    }
295
296    fn find_in_interim_mut<'n, 'k>(
297        interim: &'n mut BoxedNode<TypedNode<K, V>>,
298        key_bytes: &'k [u8],
299    ) -> Option<(&'n mut TypedNode<K, V>, &'k [u8], u8)> {
300        let prefix = interim.prefix().to_vec();
301        if key_bytes.len() == prefix.len() || common_prefix_len(&prefix, key_bytes) != prefix.len()
302        {
303            // prefix of node exactly the same as key => no matches to key
304            // because all keys inside interim node longer at least by 1 byte
305            // or
306            // node has prefix which is not prefix of search key.
307            None
308        } else {
309            // we have a prefix match, now try to find next byte of key which follows immediately
310            // after prefix.
311            interim.get_mut(key_bytes[prefix.len()]).map(|node| {
312                let key_in_parent = key_bytes[prefix.len()];
313                let key_bytes = if key_bytes.len() > prefix.len() + 1 {
314                    &key_bytes[prefix.len() + 1..]
315                } else {
316                    &[]
317                };
318                (node, key_bytes, key_in_parent)
319            })
320        }
321    }
322
323    fn replace_combined(node: &mut TypedNode<K, V>, key: K, value: V) {
324        let existing_node = unsafe { ptr::read(node) };
325        let new_node = TypedNode::Combined(Box::new(existing_node), Leaf::new(key, value));
326        unsafe { ptr::write(node, new_node) };
327    }
328
329    fn replace_leaf(
330        node: &mut TypedNode<K, V>,
331        key: K,
332        value: V,
333        key_bytes: &[u8],
334        key_start_offset: usize,
335        upsert: bool,
336    ) -> bool {
337        let leaf = node.as_leaf_mut();
338        if key_bytes == leaf.key.to_bytes() {
339            return if upsert {
340                leaf.value = value;
341                true
342            } else {
343                false
344            };
345        }
346
347        let leaf_key_bytes = leaf.key.to_bytes();
348        let leaf_key = &leaf_key_bytes[key_start_offset..];
349        let key_bytes = &key_bytes[key_start_offset..];
350
351        let prefix_size = common_prefix_len(leaf_key, key_bytes);
352
353        let prefix = &leaf_key[..prefix_size];
354        let leaf_key = &leaf_key[prefix_size..];
355        let key_bytes = &key_bytes[prefix_size..];
356
357        let new_interim = if leaf_key.is_empty() {
358            // existing leaf key is shorter than new key.
359            let mut new_interim = FlatNode::new(prefix);
360            // safely move out value from node holder because
361            // later we will override it without drop
362            let err = new_interim.insert(key_bytes[0], TypedNode::Leaf(Leaf::new(key, value)));
363            debug_assert!(err.is_none());
364
365            let existing_leaf = unsafe { ptr::read(leaf) };
366            TypedNode::Combined(
367                Box::new(TypedNode::Interim(BoxedNode::Size4(Box::new(new_interim)))),
368                existing_leaf,
369            )
370        } else if key_bytes.is_empty() {
371            // no more bytes left in key of new KV => create combined node which will
372            // point to new key and existing leaf will be moved into new interim node.
373            // in this case, key of existing leaf is always longer(length in bytes) than new
374            // key(if leaf key has the same length as new key, then keys are equal).
375            let mut new_interim = FlatNode::new(prefix);
376            // safely move out value from node holder because
377            // later we will override it without drop
378            let existing_leaf = unsafe { ptr::read(leaf) };
379            let err = new_interim.insert(leaf_key[0], TypedNode::Leaf(existing_leaf));
380            debug_assert!(err.is_none());
381
382            TypedNode::Combined(
383                Box::new(TypedNode::Interim(BoxedNode::Size4(Box::new(new_interim)))),
384                Leaf::new(key, value),
385            )
386        } else {
387            // safely move out value from node holder because
388            // later we will override it without drop
389            let leaf = unsafe { ptr::read(leaf) };
390            let mut new_interim = FlatNode::new(prefix);
391            let err = new_interim.insert(key_bytes[0], TypedNode::Leaf(Leaf::new(key, value)));
392            debug_assert!(err.is_none());
393            let err = new_interim.insert(leaf_key[0], TypedNode::Leaf(leaf));
394            debug_assert!(err.is_none());
395            TypedNode::Interim(BoxedNode::Size4(Box::new(new_interim)))
396        };
397        unsafe { ptr::write(node, new_interim) };
398        true
399    }
400
401    fn interim_insert<'n>(
402        node: &'n mut TypedNode<K, V>,
403        key: K,
404        value: V,
405        key_bytes: &[u8],
406        key_start_offset: usize,
407    ) -> Result<bool, InsertOp<'n, K, V>> {
408        let node_ptr = node as *mut TypedNode<K, V>;
409        let interim = node.as_interim_mut();
410        if key_bytes.len() <= key_start_offset {
411            // no more bytes in key to go deeper inside the tree => replace interim by combined node
412            // which will contains link to existing interim and leaf node with new KV.
413            unsafe {
414                let interim = ptr::read(interim);
415                ptr::write(
416                    node_ptr,
417                    TypedNode::Combined(
418                        Box::new(TypedNode::Interim(interim)),
419                        Leaf::new(key, value),
420                    ),
421                )
422            }
423            return Ok(true);
424        }
425
426        let key_bytes = &key_bytes[key_start_offset..];
427        let prefix = interim.prefix();
428        let prefix_size = common_prefix_len(prefix, key_bytes);
429
430        if prefix_size == key_bytes.len() {
431            // existing interim prefix fully equals to new key: replace existing interim by combined
432            // node which will contain new key and link to existing values of interim
433            unsafe {
434                let interim = ptr::read(interim);
435                ptr::write(
436                    node_ptr,
437                    TypedNode::Combined(
438                        Box::new(TypedNode::Interim(interim)),
439                        Leaf::new(key, value),
440                    ),
441                )
442            };
443            Ok(true)
444        } else if prefix_size < prefix.len() {
445            // Node prefix and key has common byte sequence. For instance, node prefix is
446            // "abc" and key is "abde", common sequence will be "ab". This sequence will be
447            // used as prefix for new interim node and this interim will point to new leaf(with passed
448            // KV) and previous interim(with prefix "abc").
449            let mut new_interim = FlatNode::new(&prefix[..prefix_size]);
450            let res = new_interim.insert(
451                key_bytes[prefix_size],
452                TypedNode::Leaf(Leaf::new(key, value)),
453            );
454            debug_assert!(res.is_none());
455            let mut interim = unsafe { ptr::read(interim) };
456            let interim_key = prefix[prefix_size];
457            // update prefix of existing interim to remaining part of old prefix.
458            // e.g, prefix was "abcd", prefix of new node is "ab".
459            // Updated prefix will be "d" because "c" used as pointer inside new interim.
460            if prefix_size + 1 < prefix.len() {
461                interim.set_prefix(&prefix[prefix_size + 1..]);
462            } else {
463                interim.set_prefix(&[]);
464            }
465            let res = new_interim.insert(interim_key, TypedNode::Interim(interim));
466            debug_assert!(res.is_none());
467            unsafe {
468                ptr::write(
469                    node_ptr,
470                    TypedNode::Interim(BoxedNode::Size4(Box::new(new_interim))),
471                )
472            };
473            Ok(true)
474        } else {
475            let interim_ptr = unsafe { &mut *(interim as *mut BoxedNode<TypedNode<K, V>>) };
476            if let Some(next_node) = interim.get_mut(key_bytes[prefix_size]) {
477                // try to insert on the next level of tree
478                Err(InsertOp {
479                    node: next_node,
480                    key_byte_offset: key_start_offset + prefix_size + 1,
481                    key,
482                    value,
483                })
484            } else {
485                // we find interim node which should contain new KV
486                let leaf = TypedNode::Leaf(Leaf::new(key, value));
487                match interim_ptr.insert(key_bytes[prefix_size], leaf) {
488                    Some(InsertError::Overflow(val)) => {
489                        let interim = unsafe { ptr::read(interim_ptr) };
490                        let mut new_interim = interim.expand();
491                        let err = new_interim.insert(key_bytes[prefix_size], val);
492                        debug_assert!(
493                            err.is_none(),
494                            "Insert failed after node expand (unexpected duplicate key)"
495                        );
496                        unsafe { ptr::write(node_ptr, TypedNode::Interim(new_interim)) }
497                        Ok(true)
498                    }
499                    Some(InsertError::DuplicateKey) => panic!("Should not happen"),
500                    None => Ok(true),
501                }
502            }
503        }
504    }
505}
506
507fn common_prefix_len(vec1: &[u8], vec2: &[u8]) -> usize {
508    let mut len = 0;
509    for i in 0..cmp::min(vec1.len(), vec2.len()) {
510        if vec1[i] != vec2[i] {
511            break;
512        }
513        len += 1;
514    }
515    len
516}
517
518struct InsertOp<'n, K, V> {
519    node: &'n mut TypedNode<K, V>,
520    // offset of byte in key which should be used to insert KV pair inside `node`
521    key_byte_offset: usize,
522    key: K,
523    value: V,
524}
525
526#[cfg(test)]
527mod tests {
528    use crate::{Art, ByteString, Float32, Float64, KeyBuilder};
529    use rand::prelude::IteratorRandom;
530    use rand::seq::SliceRandom;
531    use rand::{thread_rng, Rng};
532    use std::collections::HashSet;
533
534    #[test]
535    fn seq_insert_combined_key() {
536        let mut art = Art::new();
537        for i in 0..=u8::MAX {
538            for j in i8::MIN..=i8::MAX {
539                let key = KeyBuilder::new().append(i).append(j).build();
540                assert!(art.insert(key, i.to_string()), "{}", i);
541            }
542        }
543
544        for i in 0..=u8::MAX {
545            for j in i8::MIN..=i8::MAX {
546                let key = KeyBuilder::new().append(i).append(j).build();
547                assert!(matches!(art.get(&key), Some(val) if val == &i.to_string()));
548            }
549        }
550    }
551
552    #[test]
553    fn seq_insert_u8() {
554        let mut art = Art::new();
555        for i in 0..=u8::MAX {
556            assert!(art.insert(i, i.to_string()), "{}", i);
557        }
558
559        for i in 0..=u8::MAX {
560            assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
561        }
562    }
563
564    #[test]
565    fn seq_insert_i8() {
566        let mut art = Art::new();
567        for i in i8::MIN..=i8::MAX {
568            assert!(art.insert(i, i.to_string()), "{}", i);
569        }
570
571        for i in i8::MIN..=i8::MAX {
572            assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
573        }
574    }
575
576    #[test]
577    fn seq_insert_u16() {
578        let mut art = Art::new();
579        for i in 0..=u16::MAX {
580            assert!(art.insert(i, i.to_string()), "{}", i);
581        }
582
583        for i in 0..=u16::MAX {
584            assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
585        }
586    }
587
588    #[test]
589    fn seq_insert_i16() {
590        let mut art = Art::new();
591        for i in i16::MIN..=i16::MAX {
592            assert!(art.insert(i, i.to_string()), "{}", i);
593        }
594
595        for i in i16::MIN..=i16::MAX {
596            assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
597        }
598    }
599
600    #[test]
601    fn seq_insert_u32() {
602        let mut art = Art::new();
603        for shift in 0..2 {
604            let start = (u16::MAX as u32 + 1) << (shift * 8);
605            let end = start + 10000;
606            for i in start..=end {
607                assert!(art.insert(i, i.to_string()), "{}", i);
608            }
609            for i in start..=end {
610                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
611            }
612        }
613    }
614
615    #[test]
616    fn seq_insert_i32() {
617        let mut art = Art::new();
618        for shift in 0..2 {
619            let start = i32::MIN >> (shift * 8);
620            let end = start + 10000;
621            for i in start..=end {
622                assert!(art.insert(i, i.to_string()), "{}", i);
623            }
624            for i in start..=end {
625                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
626            }
627        }
628
629        assert!(art.insert(0, "0".to_string()), "{}", 0);
630
631        for shift in 0..2 {
632            let start = (i16::MAX as i32 + 1) << (shift * 8);
633            let end = start + 10000;
634            for i in start..=end {
635                assert!(art.insert(i, i.to_string()), "{}", i);
636            }
637            for i in start..=end {
638                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
639            }
640        }
641    }
642
643    #[test]
644    fn seq_insert_f32() {
645        let mut art: Art<Float32, String> = Art::new();
646        assert!(
647            art.insert(f32::MIN.into(), f32::MIN.to_string()),
648            "{}",
649            f32::MIN
650        );
651        assert!(matches!(art.get(&f32::MIN.into()), Some(val) if val == &f32::MIN.to_string()));
652        assert!(
653            art.insert(f32::MAX.into(), f32::MAX.to_string()),
654            "{}",
655            f32::MAX
656        );
657        assert!(matches!(art.get(&f32::MAX.into()), Some(val) if val == &f32::MAX.to_string()));
658        assert!(art.insert(0.0.into(), 0.to_string()), "{}", 0);
659        assert!(matches!(art.get(&0.0.into()), Some(val) if val == &0.to_string()));
660        assert!(
661            art.insert(Float32::from(-1.0000001), 0.to_string()),
662            "{}",
663            0
664        );
665        assert!(
666            matches!(art.get(&Float32::from(-1.0000001)), Some(val) if val == &0.to_string
667            ())
668        );
669
670        assert!(
671            art.insert(f32::NAN.into(), f32::NAN.to_string()),
672            "{}",
673            f32::NAN
674        );
675        assert!(matches!(art.get(&f32::NAN.into()), Some(val) if val == &f32::NAN.to_string()));
676
677        assert!(
678            art.insert(f32::NEG_INFINITY.into(), f32::NEG_INFINITY.to_string()),
679            "{}",
680            f32::NEG_INFINITY
681        );
682        assert!(
683            matches!(art.get(&f32::NEG_INFINITY.into()), Some(val) if val == &f32::NEG_INFINITY.to_string())
684        );
685
686        assert!(
687            art.insert(f32::INFINITY.into(), f32::INFINITY.to_string()),
688            "{}",
689            f32::INFINITY
690        );
691        assert!(
692            matches!(art.get(&f32::INFINITY.into()), Some(val) if val == &f32::INFINITY.to_string())
693        );
694    }
695
696    #[test]
697    fn seq_insert_f64() {
698        let mut art: Art<Float64, String> = Art::new();
699        assert!(
700            art.insert(f64::MIN.into(), f64::MIN.to_string()),
701            "{}",
702            f64::MIN
703        );
704        assert!(matches!(art.get(&f64::MIN.into()), Some(val) if val == &f64::MIN.to_string()));
705        assert!(
706            art.insert(f64::MAX.into(), f64::MAX.to_string()),
707            "{}",
708            f64::MAX
709        );
710        assert!(matches!(art.get(&f64::MAX.into()), Some(val) if val == &f64::MAX.to_string()));
711        assert!(art.insert(0.0.into(), 0.to_string()), "{}", 0);
712        assert!(matches!(art.get(&0.0.into()), Some(val) if val == &0.to_string()));
713        assert!(
714            art.insert(Float64::from(-1.00000012), 0.to_string()),
715            "{}",
716            0
717        );
718        assert!(
719            matches!(art.get(&Float64::from(-1.00000012)), Some(val) if val == &0.to_string
720            ())
721        );
722
723        assert!(
724            art.insert(f64::NAN.into(), f64::NAN.to_string()),
725            "{}",
726            f64::NAN
727        );
728        assert!(matches!(art.get(&f64::NAN.into()), Some(val) if val == &f64::NAN.to_string()));
729
730        assert!(
731            art.insert(f64::NEG_INFINITY.into(), f64::NEG_INFINITY.to_string()),
732            "{}",
733            f64::NEG_INFINITY
734        );
735        assert!(
736            matches!(art.get(&f64::NEG_INFINITY.into()), Some(val) if val == &f64::NEG_INFINITY.to_string())
737        );
738
739        assert!(
740            art.insert(f64::INFINITY.into(), f64::INFINITY.to_string()),
741            "{}",
742            f64::INFINITY
743        );
744        assert!(
745            matches!(art.get(&f64::INFINITY.into()), Some(val) if val == &f64::INFINITY.to_string())
746        );
747    }
748
749    #[test]
750    fn seq_insert_u64() {
751        let mut art = Art::new();
752        for shift in 0..4 {
753            let start = (u32::MAX as u64 + 1) << (shift * 8);
754            let end = start + 100_000;
755            for i in start..=end {
756                assert!(art.insert(i, i.to_string()), "{}", i);
757            }
758            for i in start..=end {
759                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
760            }
761        }
762    }
763
764    #[test]
765    fn seq_insert_i64() {
766        let mut art = Art::new();
767        for shift in 0..4 {
768            let start = i64::MIN >> (shift * 8);
769            let end = start + 10000;
770            for i in start..=end {
771                assert!(art.insert(i, i.to_string()), "{}", i);
772            }
773            for i in start..=end {
774                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
775            }
776        }
777
778        assert!(art.insert(0, "0".to_string()), "{}", 0);
779
780        for shift in 0..4 {
781            let start = (i32::MAX as i64 + 1) << (shift * 8);
782            let end = start + 10000;
783            for i in start..=end {
784                assert!(art.insert(i, i.to_string()), "{}", i);
785            }
786            for i in start..=end {
787                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
788            }
789        }
790    }
791
792    #[test]
793    fn seq_insert_u128() {
794        let mut art = Art::new();
795        for shift in 0..8 {
796            let start = (u64::MAX as u128 + 1) << (shift * 8);
797            let end = start + 10000;
798            for i in start..=end {
799                assert!(art.insert(i, i.to_string()), "{}", i);
800            }
801            for i in start..=end {
802                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
803            }
804        }
805    }
806
807    #[test]
808    fn seq_insert_i128() {
809        let mut art = Art::new();
810        for shift in 0..8 {
811            let start = i128::MIN >> (shift * 8);
812            let end = start + 10000;
813            for i in start..=end {
814                assert!(art.insert(i, i.to_string()), "{}", i);
815            }
816            for i in start..=end {
817                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
818            }
819        }
820
821        assert!(art.insert(0, "0".to_string()), "{}", 0);
822
823        for shift in 0..8 {
824            let start = (i64::MAX as i128 + 1) << (shift * 8);
825            let end = start + 10000;
826            for i in start..=end {
827                assert!(art.insert(i, i.to_string()), "{}", i);
828            }
829            for i in start..=end {
830                assert!(matches!(art.get(&i), Some(val) if val == &i.to_string()));
831            }
832        }
833    }
834
835    #[test]
836    fn seq_remove_combined_key() {
837        let mut art = Art::new();
838        for i in 0..=u8::MAX {
839            for j in i8::MIN..=i8::MAX {
840                let key = KeyBuilder::new().append(i).append(j).build();
841                assert!(art.insert(key, i.to_string()), "{}", i);
842            }
843        }
844
845        for i in 0..=u8::MAX {
846            for j in i8::MIN..=i8::MAX {
847                let key = KeyBuilder::new().append(i).append(j).build();
848                assert!(matches!(art.remove(&key), Some(val) if val == i.to_string()));
849                assert!(matches!(art.get(&key), None));
850            }
851        }
852    }
853
854    #[test]
855    fn seq_remove_u8() {
856        let mut art = Art::new();
857        for i in 0..=u8::MAX {
858            assert!(art.insert(i, i.to_string()), "{}", i);
859        }
860
861        for i in 0..=u8::MAX {
862            assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
863            assert!(matches!(art.get(&i), None));
864        }
865    }
866
867    #[test]
868    fn seq_remove_i8() {
869        let mut art = Art::new();
870        for i in i8::MIN..=i8::MAX {
871            assert!(art.insert(i, i.to_string()), "{}", i);
872        }
873
874        for i in i8::MIN..=i8::MAX {
875            assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
876            assert!(matches!(art.get(&i), None));
877        }
878    }
879
880    #[test]
881    fn seq_remove_u16() {
882        let mut art = Art::new();
883        for i in 0..=u16::MAX {
884            assert!(art.insert(i, i.to_string()), "{}", i);
885        }
886
887        for i in 0..=u16::MAX {
888            assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
889            assert!(matches!(art.get(&i), None));
890        }
891    }
892
893    #[test]
894    fn seq_remove_i16() {
895        let mut art = Art::new();
896        for i in i16::MIN..=i16::MAX {
897            assert!(art.insert(i, i.to_string()), "{}", i);
898        }
899
900        for i in i16::MIN..=i16::MAX {
901            assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
902            assert!(matches!(art.get(&i), None));
903        }
904    }
905
906    #[test]
907    fn seq_remove_u32() {
908        let mut art = Art::new();
909        for shift in 0..2 {
910            let start = (u16::MAX as u32 + 1) << (shift * 8);
911            let end = start + 10000;
912            for i in start..=end {
913                assert!(art.insert(i, i.to_string()), "{}", i);
914            }
915            for i in start..=end {
916                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
917                assert!(matches!(art.get(&i), None));
918            }
919        }
920    }
921
922    #[test]
923    fn seq_remove_i32() {
924        let mut art = Art::new();
925        for shift in 0..2 {
926            let start = i32::MIN >> (shift * 8);
927            let end = start + 10000;
928            for i in start..=end {
929                assert!(art.insert(i, i.to_string()), "{}", i);
930            }
931            for i in start..=end {
932                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
933                assert!(matches!(art.get(&i), None));
934            }
935        }
936
937        assert!(art.insert(0, "0".to_string()), "{}", 0);
938        assert!(matches!(art.remove(&0), Some(val) if val == 0.to_string()));
939
940        for shift in 0..2 {
941            let start = (i16::MAX as i32 + 1) << (shift * 8);
942            let end = start + 10000;
943            for i in start..=end {
944                assert!(art.insert(i, i.to_string()), "{}", i);
945            }
946            for i in start..=end {
947                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
948                assert!(matches!(art.get(&i), None));
949            }
950        }
951    }
952
953    #[test]
954    fn seq_remove_u64() {
955        let mut art = Art::new();
956        for shift in 0..4 {
957            let start = (u32::MAX as u64 + 1) << (shift * 8);
958            let end = start + 100_000;
959            for i in start..=end {
960                assert!(art.insert(i, i.to_string()), "{}", i);
961            }
962            for i in start..=end {
963                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
964                assert!(matches!(art.get(&i), None));
965            }
966        }
967    }
968
969    #[test]
970    fn seq_remove_i64() {
971        let mut art = Art::new();
972        for shift in 0..4 {
973            let start = i64::MIN >> (shift * 8);
974            let end = start + 10000;
975            for i in start..=end {
976                assert!(art.insert(i, i.to_string()), "{}", i);
977            }
978            for i in start..=end {
979                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
980                assert!(matches!(art.get(&i), None));
981            }
982        }
983
984        assert!(art.insert(0, "0".to_string()), "{}", 0);
985        assert!(matches!(art.remove(&0), Some(val) if val == 0.to_string()));
986
987        for shift in 0..4 {
988            let start = (i32::MAX as i64 + 1) << (shift * 8);
989            let end = start + 10000;
990            for i in start..=end {
991                assert!(art.insert(i, i.to_string()), "{}", i);
992            }
993            for i in start..=end {
994                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
995                assert!(matches!(art.get(&i), None));
996            }
997        }
998    }
999
1000    #[test]
1001    fn seq_remove_u128() {
1002        let mut art = Art::new();
1003        for shift in 0..8 {
1004            let start = (u64::MAX as u128 + 1) << (shift * 8);
1005            let end = start + 10000;
1006            for i in start..=end {
1007                assert!(art.insert(i, i.to_string()), "{}", i);
1008            }
1009            for i in start..=end {
1010                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
1011                assert!(matches!(art.get(&i), None));
1012            }
1013        }
1014    }
1015
1016    #[test]
1017    fn seq_remove_i128() {
1018        let mut art = Art::new();
1019        for shift in 0..8 {
1020            let start = i128::MIN >> (shift * 8);
1021            let end = start + 10000;
1022            for i in start..=end {
1023                assert!(art.insert(i, i.to_string()), "{}", i);
1024            }
1025            for i in start..=end {
1026                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
1027                assert!(matches!(art.get(&i), None));
1028            }
1029        }
1030
1031        assert!(art.insert(0, "0".to_string()), "{}", 0);
1032
1033        for shift in 0..8 {
1034            let start = (i64::MAX as i128 + 1) << (shift * 8);
1035            let end = start + 10000;
1036            for i in start..=end {
1037                assert!(art.insert(i, i.to_string()), "{}", i);
1038            }
1039            for i in start..=end {
1040                assert!(matches!(art.remove(&i), Some(val) if val == i.to_string()));
1041                assert!(matches!(art.get(&i), None));
1042            }
1043        }
1044    }
1045
1046    #[test]
1047    fn modifications_with_seq_keys_with_increasing_size() {
1048        let mut art = Art::new();
1049        for i in 0..=u8::MAX {
1050            let key = KeyBuilder::new().append(i).build();
1051            assert!(art.insert(key.clone(), i.to_string()), "{}", i);
1052            assert!(matches!(art.get(&key), Some(val) if val == &i.to_string()));
1053        }
1054        for i in 0..=u8::MAX {
1055            let key = KeyBuilder::new().append(i).build();
1056            assert!(matches!(art.get(&key), Some(val) if val == &i.to_string()));
1057        }
1058
1059        for i in u8::MAX as u16 + 1..=u16::MAX {
1060            let key = KeyBuilder::new().append(i).build();
1061            assert!(art.insert(key.clone(), i.to_string()), "{}", i);
1062            assert!(matches!(art.get(&key), Some(val) if val == &i.to_string()));
1063        }
1064        for i in u8::MAX as u16 + 1..=u16::MAX {
1065            let key = KeyBuilder::new().append(i).build();
1066            assert!(matches!(art.get(&key), Some(val) if val == &i.to_string()));
1067        }
1068
1069        for i in u16::MAX as u32 + 1..=(1 << 21) as u32 {
1070            let key = KeyBuilder::new().append(i).build();
1071            assert!(art.insert(key.clone(), i.to_string()), "{}", i);
1072            assert!(matches!(art.get(&key), Some(val) if val == &i.to_string()));
1073        }
1074        for i in u16::MAX as u32 + 1..=(1 << 21) as u32 {
1075            let key = KeyBuilder::new().append(i).build();
1076            assert!(matches!(art.get(&key), Some(val) if val == &i.to_string()));
1077        }
1078
1079        for i in 0..=u8::MAX {
1080            let key = KeyBuilder::new().append(i).build();
1081            assert!(matches!(art.remove(&key), Some(val) if val == i.to_string()));
1082        }
1083        for i in u8::MAX as u16 + 1..=u16::MAX {
1084            let key = KeyBuilder::new().append(i).build();
1085            assert!(matches!(art.remove(&key), Some(val) if val == i.to_string()));
1086        }
1087        for i in u16::MAX as u32 + 1..=(1 << 21) as u32 {
1088            let key = KeyBuilder::new().append(i).build();
1089            assert!(matches!(art.remove(&key), Some(val) if val == i.to_string()));
1090        }
1091    }
1092
1093    #[test]
1094    fn insert_with_long_prefix() {
1095        let mut art = Art::new();
1096        long_prefix_test(&mut art, |art, key| {
1097            assert!(
1098                art.insert(ByteString::new(key.as_bytes()), key.clone()),
1099                "{}",
1100                key
1101            );
1102            assert!(matches!(art.get(&ByteString::new(key.as_bytes())), Some(val) if val == &key));
1103        });
1104    }
1105
1106    #[test]
1107    fn mixed_upsert_and_delete() {
1108        let mut art = Art::new();
1109        let mut existing = HashSet::new();
1110        long_prefix_test(&mut art, |art, key| {
1111            if thread_rng().gen_bool(0.3) && !existing.is_empty() {
1112                let key: &String = existing.iter().choose(&mut thread_rng()).unwrap();
1113                let key = key.clone();
1114                art.remove(&ByteString::new(key.as_bytes())).unwrap();
1115                existing.remove(&key);
1116            } else {
1117                art.upsert(ByteString::new(key.as_bytes()), key.clone());
1118                existing.insert(key);
1119            }
1120        });
1121
1122        let res: Vec<&String> = art.iter().map(|(_, v)| v).collect();
1123        let mut expected: Vec<&String> = existing.iter().collect();
1124        expected.sort();
1125        assert_eq!(expected, res);
1126    }
1127
1128    #[test]
1129    fn upsert() {
1130        let mut art = Art::new();
1131        let mut existing = HashSet::new();
1132        long_prefix_test(&mut art, |art, key| {
1133            if existing.insert(key.clone()) {
1134                art.insert(ByteString::new(key.as_bytes()), key.clone());
1135            }
1136        });
1137
1138        for (i, key) in existing.iter().enumerate() {
1139            let new_val = i.to_string();
1140            art.upsert(ByteString::new(key.as_bytes()), new_val.clone());
1141            assert!(matches!(
1142                art.get(&ByteString::new(key.as_bytes())),
1143                Some(v) if v == &new_val
1144            ));
1145        }
1146    }
1147
1148    #[test]
1149    fn existed_elements_cannot_be_inserted() {
1150        let mut art = Art::new();
1151        let mut existing = HashSet::new();
1152        long_prefix_test(&mut art, |art, key| {
1153            if existing.insert(key.clone()) {
1154                assert!(
1155                    art.insert(ByteString::new(key.as_bytes()), key.clone()),
1156                    "{} not exist in tree, but insertion failed",
1157                    key
1158                );
1159            } else {
1160                assert!(
1161                    !art.insert(ByteString::new(key.as_bytes()), key.clone()),
1162                    "{} already exists but inserted again",
1163                    key
1164                );
1165            }
1166        });
1167    }
1168
1169    #[test]
1170    fn remove_with_long_prefix() {
1171        let mut art = Art::new();
1172        let mut existing = HashSet::new();
1173        long_prefix_test(&mut art, |art, key| {
1174            assert!(
1175                art.insert(ByteString::new(key.as_bytes()), key.clone()),
1176                "{}",
1177                key
1178            );
1179            existing.insert(key);
1180        });
1181
1182        for key in existing {
1183            assert!(
1184                matches!(art.remove(&ByteString::new(key.as_bytes())), Some(val) if val == key),
1185                "{}",
1186                key
1187            );
1188            assert!(matches!(art.get(&ByteString::new(key.as_bytes())), None));
1189        }
1190    }
1191
1192    fn long_prefix_test<F: FnMut(&mut Art<ByteString, String>, String)>(
1193        art: &mut Art<ByteString, String>,
1194        mut test_fn: F,
1195    ) {
1196        let mut existing = HashSet::new();
1197        let mut chars: Vec<char> = ('a'..='z').collect();
1198        chars.shuffle(&mut thread_rng());
1199        let chars = &chars[..thread_rng().gen_range(1..chars.len())];
1200        for i in 0..chars.len() {
1201            let level1_prefix = chars[i].to_string().repeat(thread_rng().gen_range(1..8));
1202            for i in 0..chars.len() {
1203                let level2_prefix = chars[i].to_string().repeat(thread_rng().gen_range(1..8));
1204                let key_prefix = level1_prefix.clone() + &level2_prefix;
1205                for _ in 0..=u8::MAX {
1206                    let suffix: String = (0..thread_rng().gen_range(0..8))
1207                        .map(|_| chars[thread_rng().gen_range(0..chars.len())])
1208                        .collect();
1209                    let string = key_prefix.clone() + &suffix;
1210                    if !existing.insert(string.clone()) {
1211                        continue;
1212                    }
1213                    test_fn(art, string);
1214                    if existing.len() >= 10_000 {
1215                        return;
1216                    }
1217                }
1218            }
1219        }
1220    }
1221}
art_tree/lib.rs

art_tree/
lib.rs
