Skip to main content

rax/
lib.rs

1#![cfg_attr(test, feature(test))]
2
3///  Representation of a radix tree as implemented in this file, that contains
4///  the strings "foo", "foobar" and "footer" after the insertion of each
5///  word. When the node represents a key inside the radix tree, we write it
6///  between [], otherwise it is written between ().
7///
8///  This is the vanilla representation:
9///
10/// ```text
11///               (f) ""
12///                 \
13///                 (o) "f"
14///                   \
15///                   (o) "fo"
16///                     \
17///                   [t   b] "foo"
18///                   /     \
19///          "foot" (e)     (a) "foob"
20///                 /         \
21///       "foote" (r)         (r) "fooba"
22///               /             \
23///     "footer" []             [] "foobar"
24/// ```
25///
26///  However, this implementation implements a very common optimization where
27///  successive nodes having a single child are "compressed" into the node
28///  itself as a string of characters, each representing a next-level child,
29///  and only the link to the node representing the last character node is
30///  provided inside the representation. So the above representation is turned
31///  into:
32///
33/// ```text
34///                   ["foo"] ""
35///                      |
36///                   [t   b] "foo"
37///                   /     \
38///         "foot" ("er")    ("ar") "foob"
39///                  /          \
40///        "footer" []          [] "foobar"
41/// ```
42///
43///  However this optimization makes the implementation a bit more complex.
44///  For instance if a key "first" is added in the above radix tree, a
45///  "node splitting" operation is needed, since the "foo" prefix is no longer
46///  composed of nodes having a single child one after the other. This is the
47///  above tree and the resulting node splitting after this event happens:
48///
49///
50/// ```text
51///                     (f) ""
52///                     /
53///                  (i o) "f"
54///                  /   \
55///     "firs"  ("rst")  (o) "fo"
56///               /        \
57///     "first" []       [t   b] "foo"
58///                      /     \
59///            "foot" ("er")    ("ar") "foob"
60///                     /          \
61///           "footer" []          [] "foobar"
62/// ```
63///
64///  Similarly after deletion, if a new chain of nodes having a single child
65///  is created (the chain must also not include nodes that represent keys),
66///  it must be compressed back into a single node.
67extern crate libc;
68extern crate nix;
69
70use std::{
71    error, fmt,
72    mem::{size_of, MaybeUninit},
73    ptr,
74};
75
76use nix::errno::Errno;
77
78pub const GREATER: &str = ">";
79pub const GREATER_EQUAL: &str = ">=";
80pub const LESSER: &str = "<";
81pub const LESSER_EQUAL: &str = "<=";
82pub const EQUAL: &str = "=";
83pub const BEGIN: &str = "^";
84pub const END: &str = "$";
85
86pub const RAX_NODE_MAX_SIZE: libc::c_int = (1 << 29) - 1;
87pub const RAX_STACK_STATIC_ITEMS: libc::c_int = 32;
88pub const RAX_ITER_STATIC_LEN: libc::c_int = 128;
89pub const RAX_ITER_JUST_SEEKED: libc::c_int = 1 << 0;
90pub const RAX_ITER_EOF: libc::c_int = 1 << 1;
91pub const RAX_ITER_SAFE: libc::c_int = 1 << 2;
92
93/// Return the existing Rax allocator.
94///
95/// # Safety
96///
97/// Must only be called when no other thread is modifying the allocator.
98pub unsafe fn allocator() -> (
99    extern "C" fn(size: libc::size_t) -> *mut u8,
100    extern "C" fn(ptr: *mut libc::c_void, size: libc::size_t) -> *mut u8,
101    extern "C" fn(ptr: *mut libc::c_void),
102) {
103    (rax_malloc, rax_realloc, rax_free)
104}
105
106/// Rax internally makes calls to "malloc", "realloc" and "free" for all of it's
107/// heap memory needs. These calls can be patched with the supplied hooks.
108/// Do not call this method after Rax has been used at all. This must
109/// be called before using or calling any other Rax API function.
110///
111/// # Safety
112///
113/// Must be called before any Rax allocation occurs. Not thread-safe.
114pub unsafe fn set_allocator(
115    malloc: extern "C" fn(size: libc::size_t) -> *mut u8,
116    realloc: extern "C" fn(ptr: *mut libc::c_void, size: libc::size_t) -> *mut u8,
117    free: extern "C" fn(ptr: *mut libc::c_void),
118) {
119    rax_malloc = malloc;
120    rax_realloc = realloc;
121    rax_free = free;
122}
123
124#[derive(Debug)]
125pub enum RaxError {
126    Generic(GenericError),
127    OutOfMemory(),
128}
129
130impl RaxError {
131    pub fn generic(message: &str) -> RaxError {
132        RaxError::Generic(GenericError::new(message))
133    }
134}
135
136/// Redis has a beautiful Radix Tree implementation in ANSI C.
137/// This brings it to Rust and creates a safe Map like wrapper
138/// for it. This is very similar in utility to a BTreeMap, but
139/// RAX is likely much faster and more efficient. Naive testing
140/// showed a 2x-4x improvement for all common operations. The only
141/// disadvantage to BTreeMap is that BTree's allow much more flexibility
142/// in regards to comparing keys. Radix trees are lexicographically only.
143/// Composite keys where the non-last member is variable length could
144/// be something BTrees could handle much easier.
145///
146/// Internal RAX Node Layout
147///
148/// uint32_t iskey:1;     /* Does this node contain a key? */
149/// uint32_t isnull:1;    /* Associated value is NULL (don't store it). */
150/// uint32_t iscompr:1;   /* Node is compressed. */
151/// uint32_t size:29;     /* Number of children, or compressed string len. */
152///
153/// +----+---+--------+--------+--------+--------+
154/// |HDR |xyz| x-ptr  | y-ptr  | z-ptr  |dataptr |
155/// +----+---+--------+--------+--------+--------+
156///
157/// As is evident above, there is no storage penalty for NULL values.
158///
159/// Keys are represented in compressed form therefore, there is no
160/// need to pump in Boxed keys or any sort of heap allocated chunk of
161/// memory. Stack or heap keys may be used from rust. Values can either
162/// be a sizeof<usize> size integer or it's a data pointer to a heap
163/// allocated / Boxed value.
164///
165/// Iterators were designed to be fast and attempt to only use stack
166/// allocated memory. RaxMap provides a model to take full advantage
167/// of stack allocated iterators through wrapping in a closure.
168///
169/// #Examples
170///
171/// ```
172/// use rax::RaxMap;
173/// let mut r = RaxMap::new();
174/// r.insert(1, Box::new("my heap allocation".to_string()));
175/// r.insert(2, Box::new("my other heap allocation".to_string()));
176///
177/// r.iter(|r, iter| {
178///     // Place iterator at the first entry.
179///     if !iter.seek_min() {
180///         // EOF
181///         return;
182///     }
183///
184///     // Can test EOF at any time.
185///     if iter.eof() {
186///         // EOF
187///         return;
188///     }
189///
190///     while iter.forward() {
191///         iter.key();
192///         iter.value();
193///     }
194///     // In reverse
195///     // Place iterator at the end.
196///     if !iter.seek_max() {
197///         // EOF
198///         return;
199///     }
200///     while iter.back() {
201///         iter.key();
202///         iter.value();
203///     }
204///
205///     // Seek
206///     if !iter.seek(">=", 2) {
207///         // EOF
208///     }
209///     while iter.forward() {
210///         iter.key();
211///         iter.value();
212///     }
213/// });
214/// ```
215pub struct RaxMap<K: RaxKey, V> {
216    pub rax: *mut rax,
217    phantom: std::marker::PhantomData<(K, V)>,
218}
219
220impl<K: RaxKey, V> Drop for RaxMap<K, V> {
221    fn drop(&mut self) {
222        // SAFETY: self.rax is non-null (guaranteed by try_new/new).
223        // raxFreeWithCallback frees all nodes and invokes the callback
224        // on each stored value pointer so we can drop the Box<V>.
225        unsafe {
226            raxFreeWithCallback(self.rax, RaxFreeWithCallbackWrapper::<V>);
227        }
228    }
229}
230
231// SAFETY: RaxMap owns a unique heap allocation. No shared mutable
232// state, therefore safe to transfer between threads.
233unsafe impl<K: RaxKey + Send, V: Send> Send for RaxMap<K, V> {}
234
235// SAFETY: &RaxMap only allows read-only access (find/get) which don't
236// mutate the underlying C tree, therefore safe to share across threads.
237unsafe impl<K: RaxKey + Sync, V: Sync> Sync for RaxMap<K, V> {}
238
239impl<K: RaxKey, V> Default for RaxMap<K, V> {
240    fn default() -> Self {
241        Self::new()
242    }
243}
244
245/// Implementation of RaxMap
246impl<K: RaxKey, V> RaxMap<K, V> {
247    /// Create a new RaxMap.
248    ///
249    /// # Panics
250    ///
251    /// Panics if `raxNew()` returns NULL (out of memory).
252    pub fn new() -> RaxMap<K, V> {
253        Self::try_new().expect("raxNew: out of memory")
254    }
255
256    /// Fallible constructor.
257    ///
258    /// Returns `Err(OutOfMemory)` when `raxNew()` returns NULL.
259    pub fn try_new() -> Result<RaxMap<K, V>, RaxError> {
260        // SAFETY: raxNew() allocates a new rax tree.
261        // We check for NULL before storing the pointer.
262        let ptr = unsafe { raxNew() };
263
264        if ptr.is_null() {
265            return Err(RaxError::OutOfMemory());
266        }
267
268        Ok(RaxMap {
269            rax: ptr,
270            phantom: std::marker::PhantomData,
271        })
272    }
273
274    /// The number of entries in the RAX
275    pub fn len(&self) -> u64 {
276        unsafe { raxSize(self.rax) }
277    }
278
279    /// The number of entries in the RAX
280    pub fn size(&self) -> u64 {
281        unsafe { raxSize(self.rax) }
282    }
283
284    /// Returns true if the RAX is empty.
285    pub fn is_empty(&self) -> bool {
286        self.len() == 0
287    }
288
289    /// Prints the Rax as ASCII art to stdout.
290    pub fn show(&self) {
291        unsafe { raxShow(self.rax) }
292    }
293
294    /// Insert or replace existing key with a NULL value.
295    pub fn insert_null(&mut self, key: K) -> Result<Option<Box<V>>, RaxError> {
296        unsafe {
297            // Allocate a pointer to catch the old value.
298            let old: &mut *mut u8 = &mut ptr::null_mut();
299
300            // Integer values require Big Endian to allow the Rax to fully optimize
301            // storing them since it will be able to compress the prefixes especially
302            // for 64/128bit numbers.
303            let k = key.encode();
304            let (ptr, len) = k.to_buf();
305
306            let r = raxInsert(
307                self.rax,
308                // Grab a raw pointer to the key. Keys are most likely allocated
309                // on the stack. The rax will keep it's own copy of the key so we
310                // don't want to keep in in the heap twice and it exists in the
311                // rax in it's compressed form.
312                ptr,
313                len,
314                std::ptr::null_mut(),
315                old,
316            );
317
318            if r == 0 && Errno::last() == Errno::ENOMEM {
319                Err(RaxError::OutOfMemory())
320            } else if old.is_null() {
321                Ok(None)
322            } else {
323                // Box the previous value since Rax is done with it and it's our
324                // responsibility now to drop it. Once this Box goes out of scope
325                // the value is dropped and memory reclaimed.
326                Ok(Some(Box::from_raw(*old as *mut V)))
327            }
328        }
329    }
330
331    /// Insert a new entry into the RAX if an existing one does not exist.
332    pub fn try_insert(&mut self, key: K, data: Box<V>) -> Result<Option<Box<V>>, RaxError> {
333        unsafe {
334            // Allocate a pointer to catch the old value.
335            let old: &mut *mut u8 = &mut ptr::null_mut();
336
337            // Leak the boxed value as we hand it over to Rax to keep track of.
338            // These must be heap allocated unless we want to store sizeof(usize) or
339            // less bytes, then the value can be the pointer.
340            let value: &mut V = Box::leak(data);
341
342            // Integer values require Big Endian to allow the Rax to fully optimize
343            // storing them since it will be able to compress the prefixes especially
344            // for 64/128bit numbers.
345            let k = key.encode();
346            let (ptr, len) = k.to_buf();
347
348            let r = raxTryInsert(
349                self.rax,
350                // Grab a raw pointer to the key. Keys are most likely allocated
351                // on the stack. The rax will keep it's own copy of the key so we
352                // don't want to keep in in the heap twice and it exists in the
353                // rax in it's compressed form.
354                ptr,
355                len,
356                value as *mut V as *mut u8,
357                old,
358            );
359
360            if r == 0 {
361                if Errno::last() == Errno::ENOMEM {
362                    Err(RaxError::OutOfMemory())
363                } else {
364                    Ok(Some(Box::from_raw(value as *mut V)))
365                }
366            } else if old.is_null() {
367                Ok(None)
368            } else {
369                // This shouldn't happen, but if it does let's be safe and
370                // not leak memory.
371                Ok(Some(Box::from_raw(*old as *mut V)))
372            }
373        }
374    }
375
376    /// Try to insert a raw pointer value into the RAX.
377    ///
378    /// # Safety
379    ///
380    /// `value` must be a valid pointer or null.
381    pub unsafe fn try_insert_ptr(
382        &mut self,
383        key: K,
384        value: *mut u8,
385    ) -> Result<Option<*mut u8>, RaxError> {
386        // Allocate a pointer to catch the old value.
387        let old: &mut *mut u8 = &mut ptr::null_mut();
388
389        // Integer values require Big Endian to allow the Rax to fully optimize
390        // storing them since it will be able to compress the prefixes especially
391        // for 64/128bit numbers.
392        let k = key.encode();
393        let (ptr, len) = k.to_buf();
394
395        let r = raxTryInsert(
396            self.rax,
397            // Grab a raw pointer to the key. Keys are most likely allocated
398            // on the stack. The rax will keep it's own copy of the key so we
399            // don't want to keep in in the heap twice and it exists in the
400            // rax in it's compressed form.
401            ptr, len, value, old,
402        );
403
404        if r == 0 {
405            if Errno::last() == Errno::ENOMEM {
406                Err(RaxError::OutOfMemory())
407            } else {
408                Ok(Some(value))
409            }
410        } else if old.is_null() {
411            Ok(None)
412        } else {
413            // This shouldn't happen, but if it does let's be safe and
414            // not leak memory.
415            Ok(Some(*old))
416        }
417    }
418
419    /// Insert a new entry into the RAX replacing and returning the existing
420    /// entry for the supplied key.
421    pub fn insert(&mut self, key: K, data: Box<V>) -> Result<Option<Box<V>>, RaxError> {
422        unsafe {
423            // Allocate a pointer to catch the old value.
424            let old: &mut *mut u8 = &mut ptr::null_mut();
425
426            // Leak the boxed value as we hand it over to Rax to keep track of.
427            // These must be heap allocated unless we want to store sizeof(usize) or
428            // less bytes, then the value can be the pointer.
429            let value: &mut V = Box::leak(data);
430
431            // Integer values require Big Endian to allow the Rax to fully optimize
432            // storing them since it will be able to compress the prefixes especially
433            // for 64/128bit numbers.
434            let k = key.encode();
435            let (ptr, len) = k.to_buf();
436
437            let r = raxInsert(
438                self.rax,
439                // Grab a raw pointer to the key. Keys are most likely allocated
440                // on the stack. The rax will keep it's own copy of the key so we
441                // don't want to keep in in the heap twice and it exists in the
442                // rax in it's compressed form.
443                ptr,
444                len,
445                value as *mut V as *mut u8,
446                old,
447            );
448
449            if r == 0 && Errno::last() == Errno::ENOMEM {
450                Err(RaxError::OutOfMemory())
451            } else if old.is_null() {
452                Ok(None)
453            } else {
454                // Box the previous value since Rax is done with it and it's our
455                // responsibility now to drop it. Once this Box goes out of scope
456                // the value is dropped and memory reclaimed.
457                Ok(Some(Box::from_raw(*old as *mut V)))
458            }
459        }
460    }
461
462    /// Insert a raw pointer value into the RAX.
463    ///
464    /// # Safety
465    ///
466    /// `value` must be a valid pointer or null.
467    pub unsafe fn insert_ptr(
468        &mut self,
469        key: K,
470        value: *mut u8,
471    ) -> Result<Option<*mut u8>, RaxError> {
472        // Allocate a pointer to catch the old value.
473        let old: &mut *mut u8 = &mut ptr::null_mut();
474
475        // Integer values require Big Endian to allow the Rax to fully optimize
476        // storing them since it will be able to compress the prefixes especially
477        // for 64/128bit numbers.
478        let k = key.encode();
479        let (ptr, len) = k.to_buf();
480
481        let r = raxInsert(
482            self.rax,
483            // Grab a raw pointer to the key. Keys are most likely allocated
484            // on the stack. The rax will keep it's own copy of the key so we
485            // don't want to keep in in the heap twice and it exists in the
486            // rax in it's compressed form.
487            ptr, len, value, old,
488        );
489
490        if r == 0 && Errno::last() == Errno::ENOMEM {
491            Err(RaxError::OutOfMemory())
492        } else if old.is_null() {
493            Ok(None)
494        } else {
495            // Box the previous value since Rax is done with it and it's our
496            // responsibility now to drop it. Once this Box goes out of scope
497            // the value is dropped and memory reclaimed.
498            Ok(Some(*old))
499        }
500    }
501
502    /// Remove an entry from the RAX and return the associated value.
503    pub fn remove(&mut self, key: K) -> (bool, Option<Box<V>>) {
504        unsafe {
505            let old: &mut *mut u8 = &mut ptr::null_mut();
506            let k = key.encode();
507            let (ptr, len) = k.to_buf();
508
509            let r = raxRemove(self.rax, ptr, len, old);
510
511            if old.is_null() {
512                (r == 1, None)
513            } else {
514                (r == 1, Some(Box::from_raw(*old as *mut V)))
515            }
516        }
517    }
518
519    /// Find a key and return whether it exists along with its value.
520    pub fn find_exists(&self, key: K) -> (bool, Option<&V>) {
521        unsafe {
522            let k = key.encode();
523            let (ptr, len) = k.to_buf();
524
525            let value = raxFind(self.rax, ptr, len);
526
527            if value.is_null() {
528                (true, None)
529            } else if value == raxNotFound {
530                (false, None)
531            } else {
532                // transmute to the value so we don't drop the actual value accidentally.
533                // While the key associated to the value is in the RAX then we cannot
534                // drop it.
535                (true, Some(&*(value as *const V)))
536            }
537        }
538    }
539
540    /// Same as get but added for semantics parity.
541    pub fn find(&self, key: K) -> Option<&V> {
542        unsafe {
543            let k = key.encode();
544            let (ptr, len) = k.to_buf();
545
546            let value = raxFind(self.rax, ptr, len);
547
548            if value.is_null() || value == raxNotFound {
549                None
550            } else {
551                // transmute to the value so we don't drop the actual value accidentally.
552                // While the key associated to the value is in the RAX then we cannot
553                // drop it.
554                Some(&*(value as *const V))
555            }
556        }
557    }
558
559    /// Get the value associated with the key.
560    pub fn get(&self, key: K) -> Option<&V> {
561        unsafe {
562            let k = key.encode();
563            let (ptr, len) = k.to_buf();
564
565            let value = raxFind(self.rax, ptr, len);
566
567            if value.is_null() || value == raxNotFound {
568                None
569            } else {
570                // transmute to the value so we don't drop the actual value accidentally.
571                // While the key associated to the value is in the RAX then we cannot
572                // drop it.
573                Some(&*(value as *const V))
574            }
575        }
576    }
577
578    /// Determines if the supplied key exists in the Rax.
579    pub fn exists(&self, key: K) -> bool {
580        unsafe {
581            let k = key.encode();
582            let (ptr, len) = k.to_buf();
583
584            let value = raxFind(self.rax, ptr, len);
585
586            !(value.is_null() || value == raxNotFound)
587        }
588    }
589
590    /// Seek to the minimum key and execute the closure.
591    ///
592    /// # Safety
593    ///
594    /// Mutating the map inside the closure is undefined behaviour.
595    pub fn seek_min<F>(&mut self, f: F)
596    where
597        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>),
598    {
599        // SAFETY: Iterator is initialized with raxStart before use.
600        unsafe {
601            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
602            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
603            let mut iter = iter.assume_init();
604            iter.seek_min();
605            f(self, &mut iter)
606        }
607    }
608
609    /// Seek to the minimum key and execute the closure, returning a result.
610    ///
611    /// # Safety
612    ///
613    /// Mutating the map inside the closure is undefined behaviour.
614    pub fn seek_min_result<R, F>(&mut self, f: F) -> Result<R, RaxError>
615    where
616        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>) -> Result<R, RaxError>,
617    {
618        // SAFETY: Iterator is initialized with raxStart before use.
619        unsafe {
620            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
621            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
622            let mut iter = iter.assume_init();
623            iter.seek_min();
624            f(self, &mut iter)
625        }
626    }
627
628    /// Seek to the maximum key and execute the closure.
629    ///
630    /// # Safety
631    ///
632    /// Mutating the map inside the closure is undefined behaviour.
633    pub fn seek_max<F>(&mut self, f: F)
634    where
635        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>),
636    {
637        // SAFETY: Iterator is initialized with raxStart before use.
638        unsafe {
639            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
640            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
641            let mut iter = iter.assume_init();
642            iter.seek_max();
643            f(self, &mut iter)
644        }
645    }
646
647    /// Seek to the maximum key and execute the closure, returning a result.
648    ///
649    /// # Safety
650    ///
651    /// Mutating the map inside the closure is undefined behaviour.
652    pub fn seek_max_result<R, F>(&mut self, f: F) -> Result<R, RaxError>
653    where
654        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>) -> Result<R, RaxError>,
655    {
656        // SAFETY: Iterator is initialized with raxStart before use.
657        unsafe {
658            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
659            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
660            let mut iter = iter.assume_init();
661            iter.seek_max();
662            f(self, &mut iter)
663        }
664    }
665
666    /// Seek to the given key using the specified operator and execute the closure.
667    ///
668    /// # Safety
669    ///
670    /// Mutating the map inside the closure is undefined behaviour.
671    pub fn seek<F>(&mut self, op: &str, key: K, f: F)
672    where
673        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>),
674    {
675        // SAFETY: Iterator is initialized with raxStart before use.
676        unsafe {
677            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
678            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
679            let mut iter = iter.assume_init();
680            iter.seek(op, key);
681            f(self, &mut iter)
682        }
683    }
684
685    /// Seek to the given key using the specified operator and execute the closure, returning a result.
686    ///
687    /// # Safety
688    ///
689    /// Mutating the map inside the closure is undefined behaviour.
690    pub fn seek_result<R, F>(&mut self, op: &str, key: K, f: F) -> Result<R, RaxError>
691    where
692        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>) -> Result<R, RaxError>,
693    {
694        // SAFETY: Iterator is initialized with raxStart before use.
695        unsafe {
696            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
697            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
698            let mut iter = iter.assume_init();
699            iter.seek(op, key);
700            f(self, &mut iter)
701        }
702    }
703
704    /// Create an iterator and execute the closure.
705    ///
706    /// # Safety
707    ///
708    /// Mutating the map inside the closure is undefined behaviour.
709    pub fn iter<F>(&mut self, f: F)
710    where
711        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>),
712    {
713        // SAFETY: Iterator is initialized with raxStart before use.
714        unsafe {
715            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
716            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
717            let mut iter = iter.assume_init();
718            f(self, &mut iter)
719        }
720    }
721
722    /// Create an iterator and execute the closure, returning a result.
723    ///
724    /// # Safety
725    ///
726    /// Mutating the map inside the closure is undefined behaviour.
727    pub fn iter_result<F, R>(&mut self, f: F) -> Result<R, RaxError>
728    where
729        F: Fn(&mut RaxMap<K, V>, &mut RaxIterator<K, V>) -> Result<R, RaxError>,
730    {
731        // SAFETY: Iterator is initialized with raxStart before use.
732        unsafe {
733            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
734            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
735            let mut iter = iter.assume_init();
736            f(self, &mut iter)
737        }
738    }
739}
740
741/// RaxMap but without the values. The "isnull" bit will be set for
742/// all entries.
743/// #Examples
744///
745/// ```
746/// use rax::RaxSet;
747/// let mut r = RaxSet::new();
748/// r.insert(1);
749/// r.insert(2);
750///
751/// r.iter(|r, iter| {
752///     // Place iterator at the first entry.
753///     if !iter.seek_min() {
754///         // EOF
755///         return;
756///     }
757///
758///     // Can test EOF at any time.
759///     if iter.eof() {
760///         // EOF
761///         return;
762///     }
763///
764///     while iter.forward() {
765///         iter.key();
766///     }
767///     // In reverse
768///     // Place iterator at the end.
769///     if !iter.seek_max() {
770///         // EOF
771///         return;
772///     }
773///     while iter.back() {
774///         iter.key();
775///     }
776///
777///     // Seek
778///     if !iter.seek(">=", 2) {
779///         // EOF
780///     }
781///     while iter.forward() {
782///         iter.key();
783///     }
784/// });
785/// ```
786pub struct RaxSet<K: RaxKey> {
787    rax: *mut rax,
788    _marker: std::marker::PhantomData<K>,
789}
790
791impl<K: RaxKey> Drop for RaxSet<K> {
792    fn drop(&mut self) {
793        // SAFETY: self.rax is non-null.
794        unsafe { raxFree(self.rax) }
795    }
796}
797
798// SAFETY: RaxSet owns a unique heap allocation. No shared mutable
799// state, therefore safe to transfer between threads.
800unsafe impl<K: RaxKey + Send> Send for RaxSet<K> {}
801
802// SAFETY: &RaxSet only allows read-only access (find/get) which don't
803// mutate the underlying C tree, therefore safe to share across threads.
804unsafe impl<K: RaxKey + Sync> Sync for RaxSet<K> {}
805
806impl<K: RaxKey> Default for RaxSet<K> {
807    fn default() -> Self {
808        Self::new()
809    }
810}
811
812/// Implementation of RaxSet.
813impl<K: RaxKey> RaxSet<K> {
814    /// Create a new RaxSet.
815    ///
816    /// # Panics
817    ///
818    /// Panics if `raxNew()` returns NULL (out of memory).
819    pub fn new() -> RaxSet<K> {
820        Self::try_new().expect("raxNew: out of memory")
821    }
822
823    /// Fallible constructor.
824    ///
825    /// Returns `Err(OutOfMemory)` when `raxNew()` returns NULL.
826    pub fn try_new() -> Result<RaxSet<K>, RaxError> {
827        // SAFETY: raxNew() allocates a new rax tree. We check for NULL
828        // before storing the pointer.
829        let ptr = unsafe { raxNew() };
830
831        if ptr.is_null() {
832            return Err(RaxError::OutOfMemory());
833        }
834
835        Ok(RaxSet {
836            rax: ptr,
837            _marker: std::marker::PhantomData,
838        })
839    }
840
841    /// The number of entries in the RAX
842    #[inline]
843    pub fn len(&self) -> u64 {
844        unsafe { raxSize(self.rax) }
845    }
846
847    /// The number of entries in the RAX
848    #[inline]
849    pub fn size(&self) -> u64 {
850        unsafe { raxSize(self.rax) }
851    }
852
853    /// Prints the Rax as ASCII art to stdout.
854    #[inline]
855    pub fn show(&self) {
856        unsafe { raxShow(self.rax) }
857    }
858
859    /// Returns true if the RAX is empty.
860    pub fn is_empty(&self) -> bool {
861        self.len() == 0
862    }
863
864    /// Insert a new entry into the RAX replacing and returning the existing
865    /// entry for the supplied key.
866    pub fn insert(&mut self, key: K) -> Result<bool, RaxError> {
867        unsafe {
868            // Integer values require Big Endian to allow the Rax to fully optimize
869            // storing them since it will be able to compress the prefixes especially
870            // for 64/128bit numbers.
871            let k = key.encode();
872            let (ptr, len) = k.to_buf();
873
874            let r = raxTryInsert(
875                self.rax,
876                // Grab a raw pointer to the key. Keys are most likely allocated
877                // on the stack. The rax will keep it's own copy of the key so we
878                // don't want to keep in in the heap twice and it exists in the
879                // rax in it's compressed form.
880                ptr,
881                len,
882                std::ptr::null_mut(),
883                std::ptr::null_mut(),
884            );
885
886            if r == 0 {
887                if Errno::last() == Errno::ENOMEM {
888                    Err(RaxError::OutOfMemory())
889                } else {
890                    Ok(false)
891                }
892            } else {
893                Ok(true)
894            }
895        }
896    }
897
898    pub fn remove(&mut self, key: K) -> bool {
899        unsafe {
900            let k = key.encode();
901            let (ptr, len) = k.to_buf();
902
903            let r = raxRemove(self.rax, ptr, len, &mut std::ptr::null_mut());
904
905            r == 1
906        }
907    }
908
909    /// Determines if the supplied key exists in the Rax.
910    pub fn exists(&self, key: K) -> bool {
911        unsafe {
912            let k = key.encode();
913            let (ptr, len) = k.to_buf();
914
915            let value = raxFind(self.rax, ptr, len);
916
917            value != raxNotFound
918        }
919    }
920
921    /// Seek to the minimum key and execute the closure.
922    ///
923    /// # Safety
924    ///
925    /// Mutating the set inside the closure is undefined behaviour.
926    pub fn seek_min<F>(&mut self, f: F)
927    where
928        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>),
929    {
930        // SAFETY: Iterator is initialized with raxStart before use.
931        unsafe {
932            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
933            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
934            let mut iter = iter.assume_init();
935            iter.seek_min();
936            f(self, &mut iter)
937        }
938    }
939
940    /// Seek to the minimum key and execute the closure, returning a result.
941    ///
942    /// # Safety
943    ///
944    /// Mutating the set inside the closure is undefined behaviour.
945    pub fn seek_min_result<R, F>(&mut self, f: F) -> Result<R, RaxError>
946    where
947        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>) -> Result<R, RaxError>,
948    {
949        // SAFETY: Iterator is initialized with raxStart before use.
950        unsafe {
951            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
952            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
953            let mut iter = iter.assume_init();
954            iter.seek_min();
955            f(self, &mut iter)
956        }
957    }
958
959    /// Seek to the maximum key and execute the closure.
960    ///
961    /// # Safety
962    ///
963    /// Mutating the set inside the closure is undefined behaviour.
964    pub fn seek_max<F>(&mut self, f: F)
965    where
966        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>),
967    {
968        // SAFETY: Iterator is initialized with raxStart before use.
969        unsafe {
970            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
971            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
972            let mut iter = iter.assume_init();
973            iter.seek_max();
974            f(self, &mut iter)
975        }
976    }
977
978    /// Seek to the maximum key and execute the closure, returning a result.
979    ///
980    /// # Safety
981    ///
982    /// Mutating the set inside the closure is undefined behaviour.
983    pub fn seek_max_result<R, F>(&mut self, f: F) -> Result<R, RaxError>
984    where
985        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>) -> Result<R, RaxError>,
986    {
987        // SAFETY: Iterator is initialized with raxStart before use.
988        unsafe {
989            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
990            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
991            let mut iter = iter.assume_init();
992            iter.seek_max();
993            f(self, &mut iter)
994        }
995    }
996
997    /// Seek to the given key using the specified operator and execute the closure.
998    ///
999    /// # Safety
1000    ///
1001    /// Mutating the set inside the closure is undefined behaviour.
1002    pub fn seek<F>(&mut self, op: &str, key: K, f: F)
1003    where
1004        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>),
1005    {
1006        // SAFETY: Iterator is initialized with raxStart before use.
1007        unsafe {
1008            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
1009            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
1010            let mut iter = iter.assume_init();
1011            iter.seek(op, key);
1012            f(self, &mut iter)
1013        }
1014    }
1015
1016    /// Seek to the given key using the specified operator and execute the closure, returning a result.
1017    ///
1018    /// # Safety
1019    ///
1020    /// Mutating the set inside the closure is undefined behaviour.
1021    pub fn seek_result<R, F>(&mut self, op: &str, key: K, f: F) -> Result<R, RaxError>
1022    where
1023        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>) -> Result<R, RaxError>,
1024    {
1025        // SAFETY: Iterator is initialized with raxStart before use.
1026        unsafe {
1027            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
1028            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
1029            let mut iter = iter.assume_init();
1030            iter.seek(op, key);
1031            f(self, &mut iter)
1032        }
1033    }
1034
1035    /// Create an iterator and execute the closure.
1036    ///
1037    /// # Safety
1038    ///
1039    /// Mutating the set inside the closure is undefined behaviour.
1040    pub fn iter<F>(&mut self, f: F)
1041    where
1042        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>),
1043    {
1044        // SAFETY: Iterator is initialized with raxStart before use.
1045        unsafe {
1046            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
1047            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
1048            let mut iter = iter.assume_init();
1049            f(self, &mut iter)
1050        }
1051    }
1052
1053    /// Create an iterator and execute the closure, returning a result.
1054    ///
1055    /// # Safety
1056    ///
1057    /// Mutating the set inside the closure is undefined behaviour.
1058    pub fn iter_result<F, R>(&mut self, f: F) -> Result<R, RaxError>
1059    where
1060        F: Fn(&mut RaxSet<K>, &mut RaxIterator<K, usize>) -> Result<R, RaxError>,
1061    {
1062        // SAFETY: Iterator is initialized with raxStart before use.
1063        unsafe {
1064            let mut iter = MaybeUninit::<RaxIterator<K, usize>>::uninit();
1065            raxStart(iter.as_mut_ptr() as *mut raxIterator, self.rax);
1066            let mut iter = iter.assume_init();
1067            f(self, &mut iter)
1068        }
1069    }
1070}
1071
1072// Same as RaxMap except values are not pointers to heap allocations.
1073// Instead the "data pointer" in the RAX is the value. This means we
1074// have sizeof<usize> worth of bytes to play with. Perhaps, in the future
1075// we could create data values of any size, but for now we have the size
1076// of pointers to work with or null which has no added size to a rax node.
1077//pub struct RaxIntMap<K: RaxKey> {
1078//    rax: *mut rax,
1079//    _marker: std::marker::PhantomData<K>,
1080//}
1081//
1082//impl<K: RaxKey> RaxIntMap<K> {
1083//    pub fn new() -> RaxIntMap<K> {
1084//        RaxIntMap {
1085//            rax: unsafe { raxNew() },
1086//            _marker: std::marker::PhantomData,
1087//        }
1088//    }
1089//
1090//    /// Insert a new entry into the RAX replacing and returning the existing
1091//    /// entry for the supplied key.
1092//    pub fn insert(&mut self, key: K, value: usize) -> Result<Option<usize>, RaxError> {
1093//        unsafe {
1094//            // Allocate a pointer to catch the old value.
1095//            let old: &mut *mut u8 = &mut ptr::null_mut();
1096//
1097//            // Integer values require Big Endian to allow the Rax to fully optimize
1098//            // storing them since it will be able to compress the prefixes especially
1099//            // for 64/128bit numbers.
1100//            let k = key.encode();
1101//            let (ptr, len) = k.to_buf();
1102//
1103//            let r = raxInsert(
1104//                self.rax,
1105//                // Grab a raw pointer to the key. Keys are most likely allocated
1106//                // on the stack. The rax will keep it's own copy of the key so we
1107//                // don't want to keep in in the heap twice and it exists in the
1108//                // rax in it's compressed form.
1109//                ptr,
1110//                len,
1111//                &value as *const _ as *mut u8,
1112//                old,
1113//            );
1114//
1115//            if r == 0 && Errno::last() == Errno::ENOMEM {
1116//                Err(RaxError::OutOfMemory())
1117//            } else if old.is_null() {
1118//                Ok(None)
1119//            } else {
1120//                Ok(Some(std::mem::transmute(*old)))
1121//            }
1122//        }
1123//    }
1124//
1125//    /// Insert a new entry into the RAX if an existing one does not exist.
1126//    pub fn try_insert(&mut self, key: K, data: usize) -> Result<Option<usize>, RaxError> {
1127//        unsafe {
1128//            // Allocate a pointer to catch the old value.
1129//            let old: &mut *mut u8 = &mut ptr::null_mut();
1130//
1131//            // Integer values require Big Endian to allow the Rax to fully optimize
1132//            // storing them since it will be able to compress the prefixes especially
1133//            // for 64/128bit numbers.
1134//            let k = key.encode();
1135//            let (ptr, len) = k.to_buf();
1136//
1137//            let r = raxTryInsert(
1138//                self.rax,
1139//                // Grab a raw pointer to the key. Keys are most likely allocated
1140//                // on the stack. The rax will keep it's own copy of the key so we
1141//                // don't want to keep in in the heap twice and it exists in the
1142//                // rax in it's compressed form.
1143//                ptr,
1144//                len,
1145//                &data as *const _ as *mut u8,
1146//                old,
1147//            );
1148//
1149//            if r == 0 {
1150//                if Errno::last() == Errno::ENOMEM {
1151//                    Err(RaxError::OutOfMemory())
1152//                } else if old.is_null() {
1153//                    Ok(None)
1154//                } else {
1155//                    Ok(Some(transmute(*old)))
1156//                }
1157//            } else if old.is_null() {
1158//                Ok(None)
1159//            } else {
1160//                Ok(Some(std::mem::transmute(*old)))
1161//            }
1162//        }
1163//    }
1164//}
1165
1166pub trait RaxKey<RHS = Self>: Clone + Default + std::fmt::Debug {
1167    type Output: RaxKey;
1168
1169    fn encode(self) -> Self::Output;
1170
1171    fn to_buf(&self) -> (*const u8, usize);
1172
1173    /// # Safety
1174    ///
1175    /// `ptr` must be valid for reads of `len` bytes.
1176    unsafe fn from_buf(ptr: *const u8, len: usize) -> RHS;
1177}
1178
1179impl RaxKey for f32 {
1180    type Output = u32;
1181
1182    #[inline]
1183    fn encode(self) -> Self::Output {
1184        // Encode as u32 Big Endian
1185        self.to_bits().to_be()
1186    }
1187
1188    #[inline]
1189    fn to_buf(&self) -> (*const u8, usize) {
1190        // This should never get called since we represent as a u32
1191        (
1192            self as *const _ as *const u8,
1193            std::mem::size_of::<Self::Output>(),
1194        )
1195    }
1196
1197    #[inline]
1198    unsafe fn from_buf(ptr: *const u8, len: usize) -> f32 {
1199        assert_eq!(len, std::mem::size_of::<f32>());
1200        if len != size_of::<Self::Output>() {
1201            return Self::default();
1202        }
1203        unsafe {
1204            // We used a BigEndian u32 to encode so let's reverse it
1205            f32::from_bits(u32::from_be(
1206                *(ptr as *mut [u8; std::mem::size_of::<Self::Output>()] as *mut u32),
1207            ))
1208        }
1209    }
1210}
1211
1212impl RaxKey for f64 {
1213    type Output = u64;
1214
1215    #[inline]
1216    fn encode(self) -> Self::Output {
1217        // Encode as u64 Big Endian
1218        self.to_bits().to_be()
1219    }
1220
1221    #[inline]
1222    fn to_buf(&self) -> (*const u8, usize) {
1223        // This should never get called since we represent as a u64
1224        (self as *const _ as *const u8, size_of::<Self::Output>())
1225    }
1226
1227    #[inline]
1228    unsafe fn from_buf(ptr: *const u8, len: usize) -> f64 {
1229        assert_eq!(len, std::mem::size_of::<f64>());
1230        if len != size_of::<Self::Output>() {
1231            return Self::default();
1232        }
1233        unsafe {
1234            // We used a BigEndian u64 to encode so let's reverse it
1235            f64::from_bits(u64::from_be(
1236                *(ptr as *mut [u8; size_of::<Self::Output>()] as *mut u64),
1237            ))
1238        }
1239    }
1240}
1241
1242impl RaxKey for isize {
1243    type Output = isize;
1244
1245    #[inline]
1246    fn encode(self) -> Self::Output {
1247        self.to_be()
1248    }
1249
1250    #[inline]
1251    fn to_buf(&self) -> (*const u8, usize) {
1252        (self as *const _ as *const u8, size_of::<Self::Output>())
1253    }
1254
1255    #[inline]
1256    unsafe fn from_buf(ptr: *const u8, len: usize) -> isize {
1257        assert_eq!(len, std::mem::size_of::<isize>());
1258        if len != size_of::<Self::Output>() {
1259            return Self::default();
1260        }
1261        unsafe { isize::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut isize)) }
1262    }
1263}
1264
1265impl RaxKey for usize {
1266    type Output = usize;
1267
1268    #[inline]
1269    fn encode(self) -> Self::Output {
1270        self.to_be()
1271    }
1272
1273    #[inline]
1274    fn to_buf(&self) -> (*const u8, usize) {
1275        (
1276            self as *const _ as *const u8,
1277            std::mem::size_of::<Self::Output>(),
1278        )
1279    }
1280
1281    #[inline]
1282    unsafe fn from_buf(ptr: *const u8, len: usize) -> usize {
1283        assert_eq!(len, std::mem::size_of::<usize>());
1284        if len != size_of::<Self::Output>() {
1285            return Self::default();
1286        }
1287        unsafe {
1288            usize::from_be(*(ptr as *mut [u8; std::mem::size_of::<Self::Output>()] as *mut usize))
1289        }
1290    }
1291}
1292
1293impl RaxKey for i16 {
1294    type Output = i16;
1295
1296    #[inline]
1297    fn encode(self) -> Self::Output {
1298        self.to_be()
1299    }
1300
1301    #[inline]
1302    fn to_buf(&self) -> (*const u8, usize) {
1303        (self as *const _ as *const u8, size_of::<Self::Output>())
1304    }
1305
1306    #[inline]
1307    unsafe fn from_buf(ptr: *const u8, len: usize) -> Self {
1308        if len != size_of::<Self::Output>() {
1309            return Self::default();
1310        }
1311        unsafe { i16::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut i16)) }
1312    }
1313}
1314
1315impl RaxKey for u16 {
1316    type Output = u16;
1317
1318    #[inline]
1319    fn encode(self) -> Self::Output {
1320        self.to_be()
1321    }
1322
1323    #[inline]
1324    fn to_buf(&self) -> (*const u8, usize) {
1325        (self as *const _ as *const u8, size_of::<Self::Output>())
1326    }
1327
1328    #[inline]
1329    unsafe fn from_buf(ptr: *const u8, len: usize) -> u16 {
1330        assert_eq!(len, std::mem::size_of::<u16>());
1331        if len != size_of::<Self::Output>() {
1332            return Self::default();
1333        }
1334        unsafe { u16::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut u16)) }
1335    }
1336}
1337
1338impl RaxKey for i32 {
1339    type Output = i32;
1340
1341    #[inline]
1342    fn encode(self) -> Self::Output {
1343        self.to_be()
1344    }
1345
1346    #[inline]
1347    fn to_buf(&self) -> (*const u8, usize) {
1348        (self as *const _ as *const u8, size_of::<Self::Output>())
1349    }
1350
1351    #[inline]
1352    unsafe fn from_buf(ptr: *const u8, len: usize) -> i32 {
1353        assert_eq!(len, std::mem::size_of::<i32>());
1354        if len != size_of::<Self::Output>() {
1355            return Self::default();
1356        }
1357        unsafe { i32::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut i32)) }
1358    }
1359}
1360
1361impl RaxKey for u32 {
1362    type Output = u32;
1363
1364    #[inline]
1365    fn encode(self) -> Self::Output {
1366        self.to_be()
1367    }
1368
1369    #[inline]
1370    fn to_buf(&self) -> (*const u8, usize) {
1371        (self as *const _ as *const u8, size_of::<Self::Output>())
1372    }
1373
1374    #[inline]
1375    unsafe fn from_buf(ptr: *const u8, len: usize) -> u32 {
1376        assert_eq!(len, std::mem::size_of::<u32>());
1377        if len != size_of::<Self::Output>() {
1378            return Self::default();
1379        }
1380        unsafe { u32::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut u32)) }
1381    }
1382}
1383
1384impl RaxKey for i64 {
1385    type Output = i64;
1386
1387    #[inline]
1388    fn encode(self) -> Self::Output {
1389        self.to_be()
1390    }
1391
1392    #[inline]
1393    fn to_buf(&self) -> (*const u8, usize) {
1394        (self as *const _ as *const u8, size_of::<Self::Output>())
1395    }
1396
1397    #[inline]
1398    unsafe fn from_buf(ptr: *const u8, len: usize) -> i64 {
1399        assert_eq!(len, std::mem::size_of::<i64>());
1400        if len != size_of::<Self::Output>() {
1401            return Self::default();
1402        }
1403        unsafe { i64::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut i64)) }
1404    }
1405}
1406
1407impl RaxKey for u64 {
1408    type Output = u64;
1409
1410    #[inline]
1411    fn encode(self) -> Self::Output {
1412        self.to_be()
1413    }
1414
1415    #[inline]
1416    fn to_buf(&self) -> (*const u8, usize) {
1417        (self as *const _ as *const u8, size_of::<Self::Output>())
1418    }
1419
1420    #[inline]
1421    unsafe fn from_buf(ptr: *const u8, len: usize) -> u64 {
1422        assert_eq!(len, std::mem::size_of::<u64>());
1423        if len != size_of::<Self::Output>() {
1424            return Self::default();
1425        }
1426        unsafe { u64::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut u64)) }
1427    }
1428}
1429
1430impl RaxKey for i128 {
1431    type Output = i128;
1432
1433    #[inline]
1434    fn encode(self) -> Self::Output {
1435        self.to_be()
1436    }
1437
1438    #[inline]
1439    fn to_buf(&self) -> (*const u8, usize) {
1440        (self as *const _ as *const u8, size_of::<Self::Output>())
1441    }
1442
1443    #[inline]
1444    unsafe fn from_buf(ptr: *const u8, len: usize) -> i128 {
1445        assert_eq!(len, std::mem::size_of::<i128>());
1446        if len != size_of::<Self::Output>() {
1447            return Self::default();
1448        }
1449        unsafe { i128::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut i128)) }
1450    }
1451}
1452
1453impl RaxKey for u128 {
1454    type Output = u128;
1455
1456    #[inline]
1457    fn encode(self) -> Self::Output {
1458        self.to_be()
1459    }
1460
1461    #[inline]
1462    fn to_buf(&self) -> (*const u8, usize) {
1463        (self as *const _ as *const u8, size_of::<Self::Output>())
1464    }
1465
1466    #[inline]
1467    unsafe fn from_buf(ptr: *const u8, len: usize) -> u128 {
1468        assert_eq!(len, std::mem::size_of::<u128>());
1469        if len != size_of::<Self::Output>() {
1470            return Self::default();
1471        }
1472        unsafe { u128::from_be(*(ptr as *mut [u8; size_of::<Self::Output>()] as *mut u128)) }
1473    }
1474}
1475
1476impl RaxKey for Vec<u8> {
1477    type Output = Vec<u8>;
1478
1479    #[inline]
1480    fn encode(self) -> Vec<u8> {
1481        self
1482    }
1483
1484    #[inline]
1485    fn to_buf(&self) -> (*const u8, usize) {
1486        (self.as_ptr(), self.len())
1487    }
1488
1489    #[inline]
1490    unsafe fn from_buf(ptr: *const u8, len: usize) -> Vec<u8> {
1491        unsafe { std::slice::from_raw_parts(ptr, len).to_vec() }
1492    }
1493}
1494
1495impl<'a> RaxKey for &'a [u8] {
1496    type Output = &'a [u8];
1497
1498    #[inline]
1499    fn encode(self) -> &'a [u8] {
1500        self
1501    }
1502
1503    #[inline]
1504    fn to_buf(&self) -> (*const u8, usize) {
1505        (self.as_ptr(), self.len())
1506    }
1507
1508    #[inline]
1509    unsafe fn from_buf(ptr: *const u8, len: usize) -> &'a [u8] {
1510        unsafe { std::slice::from_raw_parts(ptr, len) }
1511    }
1512}
1513
1514//impl RaxKey for SDS {
1515//    type Output = SDS;
1516//
1517//    #[inline]
1518//    fn encode(self) -> Self::Output {
1519//        self
1520//    }
1521//
1522//    #[inline]
1523//    fn to_buf(&self) -> (*const u8, usize) {
1524//        (self.as_ptr(), self.len())
1525//    }
1526//
1527//    #[inline]
1528//    unsafe fn from_buf(ptr: *const u8, len: usize) -> SDS {
1529//        SDS::from_ptr(ptr, len)
1530//    }
1531//}
1532
1533impl<'a> RaxKey for &'a str {
1534    type Output = &'a str;
1535
1536    #[inline]
1537    fn encode(self) -> Self::Output {
1538        self
1539    }
1540
1541    #[inline]
1542    fn to_buf(&self) -> (*const u8, usize) {
1543        ((*self).as_ptr(), self.len())
1544    }
1545
1546    #[inline]
1547    unsafe fn from_buf(ptr: *const u8, len: usize) -> &'a str {
1548        unsafe { std::str::from_utf8(std::slice::from_raw_parts(ptr, len)).unwrap_or_default() }
1549    }
1550}
1551
1552#[repr(C)]
1553pub struct RaxIterator<K: RaxKey, V> {
1554    pub flags: libc::c_int,
1555    pub rt: *mut rax,
1556    pub key: *mut u8,
1557    pub data: *mut libc::c_void,
1558    pub key_len: libc::size_t,
1559    pub key_max: libc::size_t,
1560    pub key_static_string: [u8; 128],
1561    pub node: *mut raxNode,
1562    pub stack: raxStack,
1563    pub node_cb: Option<raxNodeCallback>,
1564    _marker: std::marker::PhantomData<(K, V)>,
1565}
1566
1567/// Free up memory.
1568impl<K: RaxKey, V> Drop for RaxIterator<K, V> {
1569    fn drop(&mut self) {
1570        unsafe {
1571            // Fix key pointer if it still points at the (moved) inline buffer.
1572            if self.key_max == RAX_ITER_STATIC_LEN as usize {
1573                self.key = self.key_static_string.as_mut_ptr();
1574            }
1575
1576            // Fix stack pointer if it still points at the (moved) inline array.
1577            if self.stack.maxitems == RAX_STACK_STATIC_ITEMS as usize {
1578                self.stack.stack = self.stack.static_items.as_mut_ptr();
1579            }
1580
1581            raxStop(&raw mut *self as *mut raxIterator);
1582        }
1583    }
1584}
1585
1586/// Implement std::Iterator
1587impl<K: RaxKey, V: 'static> Iterator for RaxIterator<K, V> {
1588    type Item = (K, Option<&'static V>);
1589
1590    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
1591        unsafe {
1592            if raxNext(&raw mut *self as *mut raxIterator) == 1 {
1593                let data: *mut libc::c_void = self.data;
1594                if data.is_null() {
1595                    None
1596                } else {
1597                    let val = data as *const V;
1598                    if val.is_null() {
1599                        Some((self.key(), None))
1600                    } else {
1601                        Some((self.key(), Some(&*val)))
1602                    }
1603                }
1604            } else {
1605                None
1606            }
1607        }
1608    }
1609}
1610
1611/// Implement std::DoubleEndedIterator
1612impl<K: RaxKey, V: 'static> DoubleEndedIterator for RaxIterator<K, V> {
1613    fn next_back(&mut self) -> Option<<Self as Iterator>::Item> {
1614        unsafe {
1615            if raxPrev(&raw mut *self as *mut raxIterator) == 1 {
1616                let data: *mut libc::c_void = self.data;
1617                if data.is_null() {
1618                    None
1619                } else {
1620                    let val = data as *const V;
1621                    if val.is_null() {
1622                        Some((self.key(), None))
1623                    } else {
1624                        Some((self.key(), Some(&*val)))
1625                    }
1626                }
1627            } else {
1628                None
1629            }
1630        }
1631    }
1632}
1633
1634/// Core iterator implementation
1635impl<K: RaxKey, V> RaxIterator<K, V> {
1636    /// Create a new iterator for the given RaxMap.
1637    pub fn new(r: RaxMap<K, V>) -> RaxIterator<K, V> {
1638        // SAFETY: Iterator is initialized with raxStart before use.
1639        unsafe {
1640            let mut iter = MaybeUninit::<RaxIterator<K, V>>::uninit();
1641            raxStart(iter.as_mut_ptr() as *mut raxIterator, r.rax);
1642            let mut iter = iter.assume_init();
1643
1644            // Fix key pointer if it still points at the (moved) inline buffer.
1645            iter.key = iter.key_static_string.as_mut_ptr();
1646
1647            // Fix stack pointer if it still points at the (moved) inline array.
1648            iter.stack.stack = iter.stack.static_items.as_mut_ptr();
1649
1650            iter
1651        }
1652    }
1653
1654    pub fn print_ptr(&self) {
1655        println!("ptr = {:p}", self);
1656        println!("ptr = {:p}", self as *const _ as *const raxIterator);
1657    }
1658
1659    #[inline]
1660    pub fn seek_min(&mut self) -> bool {
1661        unsafe {
1662            if raxSeek(
1663                &raw mut *self as *mut raxIterator,
1664                BEGIN.as_ptr(),
1665                std::ptr::null(),
1666                0,
1667            ) == 1
1668            {
1669                self.forward()
1670            } else {
1671                false
1672            }
1673        }
1674    }
1675
1676    #[inline]
1677    pub fn seek_max(&mut self) -> bool {
1678        unsafe {
1679            if raxSeek(
1680                &raw mut *self as *mut raxIterator,
1681                END.as_ptr(),
1682                std::ptr::null(),
1683                0,
1684            ) == 1
1685            {
1686                self.back()
1687            } else {
1688                false
1689            }
1690        }
1691    }
1692
1693    #[inline]
1694    pub fn back(&mut self) -> bool {
1695        unsafe { raxPrev(&raw mut *self as *mut raxIterator) == 1 }
1696    }
1697
1698    #[inline]
1699    pub fn forward(&mut self) -> bool {
1700        unsafe { raxNext(&raw mut *self as *mut raxIterator) == 1 }
1701    }
1702
1703    /// Key at current position
1704    #[inline]
1705    pub fn key(&self) -> K {
1706        unsafe { K::from_buf(self.key, self.key_len) }
1707    }
1708
1709    /// Data at current position.
1710    #[inline]
1711    pub fn value(&self) -> Option<&V> {
1712        unsafe {
1713            let data: *mut libc::c_void = self.data;
1714            if data.is_null() {
1715                None
1716            } else {
1717                Some(&*(data as *const V))
1718            }
1719        }
1720    }
1721
1722    #[inline]
1723    pub fn lesser(&mut self, key: K) -> bool {
1724        self.seek(LESSER, key)
1725    }
1726
1727    #[inline]
1728    pub fn lesser_equal(&mut self, key: K) -> bool {
1729        self.seek(LESSER_EQUAL, key)
1730    }
1731
1732    #[inline]
1733    pub fn greater(&mut self, key: K) -> bool {
1734        self.seek(GREATER, key)
1735    }
1736
1737    #[inline]
1738    pub fn greater_equal(&mut self, key: K) -> bool {
1739        self.seek(GREATER_EQUAL, key)
1740    }
1741
1742    #[inline]
1743    pub fn seek(&mut self, op: &str, key: K) -> bool {
1744        unsafe {
1745            let k = key.encode();
1746            let (p, len) = k.to_buf();
1747            raxSeek(&raw mut *self as *mut raxIterator, op.as_ptr(), p, len) == 1
1748                && self.flags & RAX_ITER_EOF != 0
1749        }
1750    }
1751
1752    #[inline]
1753    pub fn seek_raw(&mut self, op: &str, key: K) -> i32 {
1754        unsafe {
1755            let k = key.encode();
1756            let (p, len) = k.to_buf();
1757            raxSeek(&raw mut *self as *mut raxIterator, op.as_ptr(), p, len)
1758        }
1759    }
1760
1761    #[inline]
1762    pub fn seek_bytes(&mut self, op: &str, ele: &[u8]) -> bool {
1763        unsafe {
1764            raxSeek(
1765                &raw mut *self as *mut raxIterator,
1766                op.as_ptr(),
1767                ele.as_ptr(),
1768                ele.len() as libc::size_t,
1769            ) == 1
1770        }
1771    }
1772
1773    /// Return if the iterator is in an EOF state. This happens when raxSeek()
1774    /// failed to seek an appropriate element, so that raxNext() or raxPrev()
1775    /// will return zero, or when an EOF condition was reached while iterating
1776    /// with next() and prev().
1777    #[inline]
1778    pub fn eof(&self) -> bool {
1779        self.flags & RAX_ITER_EOF != 0
1780    }
1781}
1782
1783impl fmt::Display for RaxError {
1784    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1785        match *self {
1786            // Both underlying errors already impl `Display`, so we defer to
1787            // their implementations.
1788            RaxError::Generic(ref err) => write!(f, "{}", err),
1789            RaxError::OutOfMemory() => write!(f, "out of memory"),
1790        }
1791    }
1792}
1793
1794impl error::Error for RaxError {
1795    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
1796        match *self {
1797            // N.B. Both of these implicitly cast `err` from their concrete
1798            // types (either `&io::Error` or `&num::ParseIntError`)
1799            // to a trait object `&Error`. This works because both error types
1800            // implement `Error`.
1801            RaxError::Generic(ref err) => Some(err),
1802            RaxError::OutOfMemory() => Some(self),
1803        }
1804    }
1805}
1806
1807#[derive(Debug)]
1808pub struct GenericError {
1809    message: String,
1810}
1811
1812impl GenericError {
1813    pub fn new(message: &str) -> GenericError {
1814        GenericError {
1815            message: String::from(message),
1816        }
1817    }
1818}
1819
1820impl fmt::Display for GenericError {
1821    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1822        write!(f, "Store error: {}", self.message)
1823    }
1824}
1825
1826impl error::Error for GenericError {}
1827
1828#[repr(C)]
1829pub struct rax;
1830
1831#[repr(C)]
1832pub struct raxNode;
1833
1834#[repr(C)]
1835pub struct raxStack {
1836    stack: *mut *mut libc::c_void,
1837    items: libc::size_t,
1838    maxitems: libc::size_t,
1839    static_items: [*mut libc::c_void; 32],
1840    oom: libc::c_int,
1841}
1842
1843#[repr(C)]
1844pub struct raxIterator;
1845
1846#[allow(non_snake_case)]
1847#[allow(non_camel_case_types)]
1848extern "C" fn RaxFreeWithCallbackWrapper<V>(v: *mut libc::c_void) {
1849    unsafe {
1850        // Re-box it so it can drop it immediately after it leaves this scope.
1851        drop(Box::from_raw(v as *mut V));
1852    }
1853}
1854
1855#[allow(non_camel_case_types)]
1856type raxNodeCallback = extern "C" fn(v: *mut libc::c_void);
1857
1858type RaxFreeCallback = extern "C" fn(v: *mut libc::c_void);
1859
1860#[allow(improper_ctypes)]
1861#[allow(non_snake_case)]
1862#[allow(non_camel_case_types)]
1863#[link(name = "rax", kind = "static")]
1864extern "C" {
1865    pub static raxNotFound: *mut u8;
1866
1867    pub static mut rax_malloc: extern "C" fn(size: libc::size_t) -> *mut u8;
1868    pub static mut rax_realloc:
1869        extern "C" fn(ptr: *mut libc::c_void, size: libc::size_t) -> *mut u8;
1870    pub static mut rax_free: extern "C" fn(ptr: *mut libc::c_void);
1871
1872    pub fn raxIteratorSize() -> libc::c_int;
1873
1874    pub fn raxNew() -> *mut rax;
1875
1876    pub fn raxFree(rax: *mut rax);
1877
1878    pub fn raxFreeWithCallback(rax: *mut rax, callback: RaxFreeCallback);
1879
1880    pub fn raxInsert(
1881        rax: *mut rax,
1882        s: *const u8,
1883        len: libc::size_t,
1884        data: *const u8,
1885        old: &mut *mut u8,
1886    ) -> libc::c_int;
1887
1888    pub fn raxTryInsert(
1889        rax: *mut rax,
1890        s: *const u8,
1891        len: libc::size_t,
1892        data: *const u8,
1893        old: *mut *mut u8,
1894    ) -> libc::c_int;
1895
1896    pub fn raxRemove(
1897        rax: *mut rax,
1898        s: *const u8,
1899        len: libc::size_t,
1900        old: &mut *mut u8,
1901    ) -> libc::c_int;
1902
1903    pub fn raxFind(rax: *mut rax, s: *const u8, len: libc::size_t) -> *mut u8;
1904
1905    pub fn raxIteratorNew(rt: *mut rax) -> *mut raxIterator;
1906
1907    pub fn raxStart(it: *mut raxIterator, rt: *mut rax);
1908
1909    pub fn raxSeek(
1910        it: *mut raxIterator,
1911        op: *const u8,
1912        ele: *const u8,
1913        len: libc::size_t,
1914    ) -> libc::c_int;
1915
1916    pub fn raxNext(it: *mut raxIterator) -> libc::c_int;
1917
1918    pub fn raxPrev(it: *mut raxIterator) -> libc::c_int;
1919
1920    pub fn raxRandomWalk(it: *mut raxIterator, steps: libc::size_t) -> libc::c_int;
1921
1922    pub fn raxCompare(
1923        it: *mut raxIterator,
1924        op: *const u8,
1925        key: *mut u8,
1926        key_len: libc::size_t,
1927    ) -> libc::c_int;
1928
1929    pub fn raxStop(it: *mut raxIterator);
1930
1931    pub fn raxEOF(it: *mut raxIterator) -> libc::c_int;
1932
1933    pub fn raxShow(rax: *mut rax);
1934
1935    pub fn raxSize(rax: *mut rax) -> u64;
1936}
1937
1938#[cfg(test)]
1939mod tests {
1940    extern crate test;
1941    use std::{
1942        self,
1943        default::Default,
1944        fmt,
1945        time::{Duration, Instant},
1946    };
1947
1948    use super::*;
1949
1950    extern "C" fn rax_malloc_hook(size: libc::size_t) -> *mut u8 {
1951        unsafe {
1952            println!("malloc");
1953            libc::malloc(size) as *mut u8
1954        }
1955    }
1956
1957    extern "C" fn rax_realloc_hook(ptr: *mut libc::c_void, size: libc::size_t) -> *mut u8 {
1958        unsafe {
1959            println!("realloc");
1960            libc::realloc(ptr, size) as *mut u8
1961        }
1962    }
1963
1964    extern "C" fn rax_free_hook(ptr: *mut libc::c_void) {
1965        unsafe {
1966            println!("free");
1967            libc::free(ptr)
1968        }
1969    }
1970
1971    pub struct MyMsg<'a>(&'a str);
1972
1973    impl<'a> Drop for MyMsg<'a> {
1974        fn drop(&mut self) {
1975            println!("dropped -> {}", self.0);
1976        }
1977    }
1978
1979    #[derive(Clone, Copy)]
1980    pub struct Stopwatch {
1981        start_time: Option<Instant>,
1982        elapsed: Duration,
1983    }
1984
1985    impl Default for Stopwatch {
1986        fn default() -> Stopwatch {
1987            Stopwatch {
1988                start_time: None,
1989                elapsed: Duration::from_secs(0),
1990            }
1991        }
1992    }
1993
1994    impl fmt::Display for Stopwatch {
1995        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1996            return write!(f, "{}ms", self.elapsed_ms());
1997        }
1998    }
1999
2000    #[expect(unused)]
2001    impl Stopwatch {
2002        pub fn new() -> Stopwatch {
2003            let sw: Stopwatch = Default::default();
2004            return sw;
2005        }
2006        pub fn start_new() -> Stopwatch {
2007            let mut sw = Stopwatch::new();
2008            sw.start();
2009            return sw;
2010        }
2011
2012        pub fn start(&mut self) {
2013            self.start_time = Some(Instant::now());
2014        }
2015        pub fn stop(&mut self) {
2016            self.elapsed = self.elapsed();
2017            self.start_time = None;
2018        }
2019        pub fn reset(&mut self) {
2020            self.elapsed = Duration::from_secs(0);
2021            self.start_time = None;
2022        }
2023        pub fn restart(&mut self) {
2024            self.reset();
2025            self.start();
2026        }
2027
2028        pub fn is_running(&self) -> bool {
2029            return self.start_time.is_some();
2030        }
2031
2032        pub fn elapsed(&self) -> Duration {
2033            match self.start_time {
2034                Some(t1) => {
2035                    return t1.elapsed() + self.elapsed;
2036                }
2037                None => {
2038                    return self.elapsed;
2039                }
2040            }
2041        }
2042        pub fn elapsed_ms(&self) -> i64 {
2043            let dur = self.elapsed();
2044            return (dur.as_secs() * 1000 + (dur.subsec_nanos() / 1000000) as u64) as i64;
2045        }
2046    }
2047
2048    #[test]
2049    fn bench() {
2050        let ops = 1000000;
2051        println!("{} operations per function", ops);
2052
2053        for _ in 0..2 {
2054            println!();
2055            println!("Gets...");
2056            {
2057                let r = &mut RaxSet::<u64>::new();
2058                for x in 0..2000 {
2059                    r.insert(x).expect("whoops!");
2060                }
2061
2062                let sw = Stopwatch::start_new();
2063                for _po in 0..ops {
2064                    r.exists(1601);
2065                }
2066
2067                println!("RaxSet::get       {}ms", sw.elapsed_ms());
2068            }
2069            {
2070                let r = &mut RaxMap::<u64, &str>::new();
2071                for x in 0..2000 {
2072                    r.insert_null(x).expect("whoops!");
2073                }
2074
2075                match r.find(1601) {
2076                    Some(v) => println!("{}", v),
2077                    None => {}
2078                }
2079
2080                let sw = Stopwatch::start_new();
2081
2082                for _po in 0..ops {
2083                    r.find(1601);
2084                }
2085
2086                println!("RaxMap::get       {}ms", sw.elapsed_ms());
2087            }
2088
2089            {
2090                let r = &mut RaxMap::<u64, &str>::new();
2091                for x in 0..2000 {
2092                    r.insert_null(x).expect("whoops!");
2093                }
2094                let sw = Stopwatch::start_new();
2095
2096                for _po in 0..ops {
2097                    r.iter(|_, iter| {
2098                        iter.seek(EQUAL, 1601);
2099                    });
2100                }
2101
2102                println!("RaxCursor:seek    {}ms", sw.elapsed_ms());
2103            }
2104            {
2105                let r = &mut std::collections::HashSet::<u64>::new();
2106                for x in 0..2000 {
2107                    r.insert(x);
2108                }
2109
2110                let sw = Stopwatch::start_new();
2111
2112                let xx = 300;
2113                for _po in 0..ops {
2114                    r.get(&xx);
2115                }
2116
2117                println!("HashSet::get      {}ms", sw.elapsed_ms());
2118            }
2119            {
2120                let r = &mut std::collections::HashMap::<u64, &str>::new();
2121                for x in 0..2000 {
2122                    r.insert(x, "");
2123                }
2124
2125                let sw = Stopwatch::start_new();
2126
2127                let xx = 300;
2128                for _po in 0..ops {
2129                    r.get(&xx);
2130                }
2131
2132                println!("HashMap::get      {}ms", sw.elapsed_ms());
2133            }
2134            {
2135                let r = &mut std::collections::BTreeSet::<u64>::new();
2136                for x in 0..2000 {
2137                    r.insert(x);
2138                }
2139
2140                let sw = Stopwatch::start_new();
2141
2142                let xx = 300;
2143                for _po in 0..ops {
2144                    r.get(&xx);
2145                }
2146
2147                println!("BTreeSet::get      {}ms", sw.elapsed_ms());
2148            }
2149            {
2150                let r = &mut std::collections::BTreeMap::<u64, &str>::new();
2151                for x in 0..2000 {
2152                    r.insert(x, "");
2153                }
2154
2155                let sw = Stopwatch::start_new();
2156
2157                let xx = 300;
2158                for _po in 0..ops {
2159                    r.get(&xx);
2160                }
2161
2162                println!("BTreeMap::get     {}ms", sw.elapsed_ms());
2163            }
2164
2165            println!();
2166            println!("Inserts...");
2167            {
2168                let r = &mut RaxMap::<u64, &str>::new();
2169                let sw = Stopwatch::start_new();
2170
2171                for x in 0..ops {
2172                    r.insert(x, Box::new("")).expect("whoops!");
2173                }
2174
2175                println!("RaxMap::insert        {}ms", sw.elapsed_ms());
2176            }
2177
2178            {
2179                let r = &mut RaxSet::<u64>::new();
2180                let sw = Stopwatch::start_new();
2181
2182                for x in 0..ops {
2183                    r.insert(x).expect("whoops!");
2184                }
2185
2186                println!("RaxSet::insert        {}ms", sw.elapsed_ms());
2187            }
2188
2189            {
2190                let r = &mut std::collections::BTreeSet::<u64>::new();
2191                let sw = Stopwatch::start_new();
2192
2193                for x in 0..ops {
2194                    r.insert(x);
2195                }
2196
2197                println!("BTreeSet::insert      {}ms", sw.elapsed_ms());
2198            }
2199            {
2200                let r = &mut std::collections::BTreeMap::<u64, &str>::new();
2201                let sw = Stopwatch::start_new();
2202
2203                for x in 0..ops {
2204                    r.insert(x, "");
2205                }
2206
2207                println!("BTreeMap::insert      {}ms", sw.elapsed_ms());
2208            }
2209
2210            {
2211                let r = &mut std::collections::HashMap::<u64, &str>::new();
2212                let sw = Stopwatch::start_new();
2213
2214                for x in 0..ops {
2215                    r.insert(x, "");
2216                }
2217
2218                println!("HashMap::insert       {}ms", sw.elapsed_ms());
2219            }
2220        }
2221    }
2222
2223    #[test]
2224    fn bench_rax_find() {
2225        for _ in 0..10 {
2226            let r = &mut RaxMap::<u64, &str>::new();
2227            for x in 0..2000 {
2228                r.insert_null(x).expect("whoops!");
2229            }
2230
2231            match r.find(1601) {
2232                Some(v) => println!("{}", v),
2233                None => {}
2234            }
2235
2236            let sw = Stopwatch::start_new();
2237
2238            for _po in 0..1000000 {
2239                r.find(1601);
2240            }
2241
2242            println!("Thing took {}ms", sw.elapsed_ms());
2243        }
2244    }
2245
2246    #[test]
2247    fn bench_rax_cur_find() {
2248        for _ in 0..10 {
2249            let r = &mut RaxMap::<u64, &str>::new();
2250            for x in 0..2000 {
2251                r.insert_null(x).expect("whoops!");
2252            }
2253
2254            match r.find(1601) {
2255                Some(v) => println!("{}", v),
2256                None => {}
2257            }
2258
2259            let sw = Stopwatch::start_new();
2260
2261            for _po in 0..1000000 {
2262                r.iter(|_, iter| {
2263                    iter.seek(EQUAL, 1601);
2264                });
2265            }
2266
2267            println!("RaxMap::cursor_find {}ms", sw.elapsed_ms());
2268        }
2269    }
2270
2271    #[test]
2272    fn bench_rax_insert() {
2273        for _ in 0..10 {
2274            let r = &mut RaxMap::<u64, &str>::new();
2275            //
2276            let sw = Stopwatch::start_new();
2277
2278            for x in 0..1000000 {
2279                r.insert(x, Box::new("")).expect("whoops!");
2280            }
2281
2282            println!("RaxMap::insert {}ms", sw.elapsed_ms());
2283            println!("Size {}", r.size());
2284        }
2285    }
2286
2287    #[test]
2288    fn bench_rax_insert_show() {
2289        let r = &mut RaxMap::<u64, &str>::new();
2290        //
2291        let sw = Stopwatch::start_new();
2292
2293        for x in 0..100 {
2294            r.insert(x, Box::new("")).expect("whoops!");
2295        }
2296
2297        r.show();
2298        println!("RaxMap::insert {}ms", sw.elapsed_ms());
2299        assert_eq!(r.size(), 100);
2300    }
2301
2302    #[test]
2303    fn bench_rax_replace() {
2304        let ops = 1000000;
2305        for _ in 0..2 {
2306            let r = &mut RaxMap::<u64, &str>::new();
2307            // Insert values
2308            for x in 0..ops {
2309                r.insert(x, Box::new("")).expect("whoops!");
2310            }
2311
2312            let sw = Stopwatch::start_new();
2313
2314            for x in 0..ops {
2315                // Replace existing key
2316                r.insert(x, Box::new("")).expect("whoops!");
2317            }
2318
2319            println!("RaxMap::replace {}ms", sw.elapsed_ms());
2320            assert_eq!(r.size(), ops);
2321        }
2322    }
2323
2324    #[test]
2325    fn key_str() {
2326        unsafe {
2327            set_allocator(rax_malloc_hook, rax_realloc_hook, rax_free_hook);
2328        }
2329
2330        let mut r = RaxMap::<&str, MyMsg>::new();
2331
2332        let key = "hello-way";
2333
2334        r.insert(key, Box::new(MyMsg("world 80"))).expect("whoops!");
2335        r.insert("hello-war", Box::new(MyMsg("world 80")))
2336            .expect("whoops!");
2337
2338        r.insert("hello-wares", Box::new(MyMsg("world 80")))
2339            .expect("whoops!");
2340        r.insert("hello", Box::new(MyMsg("world 100")))
2341            .expect("whoops!");
2342
2343        {
2344            match r.find("hello") {
2345                Some(v) => println!("Found {}", v.0),
2346                None => println!("Not Found"),
2347            }
2348        }
2349
2350        r.show();
2351
2352        r.iter(|_, iter| {
2353            if !iter.seek_min() {
2354                return;
2355            }
2356            while iter.forward() {
2357                println!("{}", iter.key());
2358            }
2359            if !iter.seek_max() {
2360                return;
2361            }
2362            while iter.back() {
2363                println!("{}", iter.key());
2364            }
2365        });
2366    }
2367
2368    #[test]
2369    fn key_f64() {
2370        println!("sizeof(Rax) {}", std::mem::size_of::<RaxMap<f64, MyMsg>>());
2371
2372        let mut r = RaxMap::<f64, MyMsg>::new();
2373
2374        r.insert(100.01, Box::new(MyMsg("world 100")))
2375            .expect("whoops!");
2376        r.insert(80.20, Box::new(MyMsg("world 80")))
2377            .expect("whoops!");
2378        r.insert(100.00, Box::new(MyMsg("world 200")))
2379            .expect("whoops!");
2380        r.insert(99.10, Box::new(MyMsg("world 1")))
2381            .expect("whoops!");
2382
2383        r.show();
2384
2385        r.iter(|_, iter| {
2386            //            for (k, v) in iter {
2387            //
2388            //            }
2389            iter.seek_min();
2390            while iter.forward() {
2391                println!("{}", iter.key());
2392            }
2393            iter.seek_max();
2394            while iter.back() {
2395                println!("{}", iter.key());
2396            }
2397        });
2398    }
2399
2400    #[test]
2401    fn key_u64() {
2402        println!("sizeof(Rax) {}", std::mem::size_of::<RaxMap<u64, MyMsg>>());
2403
2404        let mut r = RaxMap::<u64, MyMsg>::new();
2405
2406        r.insert(100, Box::new(MyMsg("world 100")))
2407            .expect("whoops!");
2408        r.insert(80, Box::new(MyMsg("world 80"))).expect("whoops!");
2409        r.insert(200, Box::new(MyMsg("world 200")))
2410            .expect("whoops!");
2411        r.insert(1, Box::new(MyMsg("world 1"))).expect("whoops!");
2412
2413        r.show();
2414
2415        //        let result = r.iter_result(move |it| {
2416        //
2417        //            if !it.seek(GREATER_EQUAL, 800) {
2418        //                println!("Not Found");
2419        //                return Ok("");
2420        //            }
2421        //
2422        //            if it.eof() {
2423        //                println!("Not Found");
2424        //                return Ok("");
2425        //            }
2426        //
2427        //            while it.forward() {
2428        //                println!("Key Len = {}", it.key());
2429        //                println!("Data = {}", it.data().unwrap().0);
2430        //            }
2431        //
2432        //            Ok("")
2433        //        });
2434
2435        //        r.seek(GREATER_EQUAL, 80, |_, iter| {
2436        //            for (key, value) in iter {
2437        //                println!("Key Len = {}", key);
2438        //                println!("Data = {}", value.unwrap().0);
2439        //            }
2440        //        });
2441
2442        //        r.seek_result(GREATER_EQUAL, 80, |_, iter| {
2443        //            for (key, value) in iter {
2444        //                println!("Key Len = {}", key);
2445        //                println!("Data = {}", value.unwrap().0);
2446        //            }
2447        //            Ok(())
2448        //        });
2449
2450        r.seek_min(|_, it| {
2451            for (key, value) in it.rev() {
2452                println!("Key Len = {}", key);
2453                println!("Data = {}", value.unwrap().0);
2454            }
2455        });
2456
2457        //        r.iter(move |it| {
2458        //            if !it.seek(GREATER_EQUAL, 800) {
2459        //                println!("Not Found");
2460        //                return;
2461        //            }
2462        //
2463        //
2464        //
2465        //            while it.forward() {
2466        //                println!("Key Len = {}", it.key());
2467        //                println!("Data = {}", it.data().unwrap().0);
2468        //            }
2469        //        });
2470
2471        //        let result = r.iter_apply(move |r, it| {
2472        //            if !it.seek(GREATER_EQUAL, 800) {
2473        //                println!("Out of Memory");
2474        //                return Ok("");
2475        //            }
2476        //
2477        //            r.insert(800, Box::new(MyMsg("moved")));
2478        //            it.seek(GREATER_EQUAL, 800);
2479        //
2480        //            if it.eof() {
2481        //                println!("Not Found");
2482        //                return Ok("");
2483        //            }
2484        //
2485        //            while it.back() {
2486        //                println!("Key Len = {}", it.key());
2487        //                println!("Data = {}", it.data().unwrap().0);
2488        //            }
2489        //
2490        //            Ok("")
2491        //        });
2492    }
2493}