skip_list/
lib.rs

1//! Implementing a skip list with Rust. The SkipList supports `insert`,
2//! `get`, `delete` and iterator such as `iter`, `iter_mut`, `into_iter`.
3//! The default max level of skip list is 12 when use SkipList::default().
4//! The max level of skip list can be customized by SkipList::new(max_level: usize).
5//!
6//! # Example
7//! ```rust
8//! use skip_list::SkipList;
9//! 
10//! let mut skip_list = SkipList::default();
11// insert
12//! assert_eq!(skip_list.insert(1, 10), None); // there is no value with key with 1
13//! assert_eq!(skip_list.insert(2, 20), None);
14//! assert_eq!(skip_list.insert(3, 30), None);
15//! 
16//! // get
17//! assert_eq!(skip_list.get(&1), Some(&10));
18//! assert_eq!(skip_list.get(&2), Some(&20));
19//! assert_eq!(skip_list.get(&3), Some(&30));
20//! 
21//! // update
22//! assert_eq!(skip_list.insert(1, 100), Some(10));
23//! assert_eq!(skip_list.insert(2, 200), Some(20));
24//! assert_eq!(skip_list.insert(3, 300), Some(30));
25//! 
26//! // iterator
27//! for (k, v) in skip_list.iter() {
28//!     let value = k * 100;
29//!     assert_eq!(*v, value);
30//! }
31//! 
32//! // delete
33//! assert_eq!(skip_list.delete(&1), Some(100));
34//! assert_eq!(skip_list.delete(&10), None);
35//! assert_eq!(skip_list.get(&1), None);
36//! ```
37
38use std::{marker::PhantomData, ptr::NonNull};
39
40use rand::Rng;
41
42struct Node<K, V> {
43    key: std::mem::MaybeUninit<K>,
44    value: std::mem::MaybeUninit<V>,
45    level: usize,
46    next: Vec<Option<NonNull<Node<K, V>>>>,
47}
48
49impl<K, V> Node<K, V> {
50    fn new(key: K, value: V, level: usize, max_level: usize) -> Self {
51        Self {
52            key: std::mem::MaybeUninit::new(key),
53            value: std::mem::MaybeUninit::new(value),
54            level,
55            next: vec![None; max_level],
56        }
57    }
58
59    fn sigil(max_level: usize) -> Self {
60        Self {
61            key: std::mem::MaybeUninit::uninit(),
62            value: std::mem::MaybeUninit::uninit(),
63            level: 0,
64            next: vec![None; max_level],
65        }
66    }
67}
68
69pub struct SkipList<K, V> {
70    head: NonNull<Node<K, V>>,
71    len: usize,
72    level: usize,
73    max_level: usize,
74    marker: PhantomData<Node<K, V>>,
75}
76
77pub struct Iter<'a, K: 'a, V: 'a> {
78    len: usize,
79    head: Option<NonNull<Node<K, V>>>,
80    marker: PhantomData<&'a Node<K, V>>,
81}
82
83pub struct IterMut<'a, K: 'a, V: 'a> {
84    len: usize,
85    head: Option<NonNull<Node<K, V>>>,
86    marker: PhantomData<&'a Node<K, V>>,
87}
88
89pub struct IntoIter<K, V> {
90    len: usize,
91    head: Option<NonNull<Node<K, V>>>,
92    marker: PhantomData<Node<K, V>>,
93}
94
95impl<'a, K, V> Iterator for Iter<'a, K, V> {
96    type Item = (&'a K, &'a V);
97    fn next(&mut self) -> Option<Self::Item> {
98        self.head.map(|node| unsafe {
99            self.head = node.as_ref().next[0];
100            self.len -= 1;
101            let k = node.as_ref().key.assume_init_ref();
102            let v = node.as_ref().value.assume_init_ref();
103            (k, v)
104        })
105    }
106
107    fn size_hint(&self) -> (usize, Option<usize>) {
108        (self.len, Some(self.len))
109    }
110
111    fn count(self) -> usize
112    where
113        Self: Sized,
114    {
115        self.len
116    }
117}
118
119impl<'a, K, V> Iterator for IterMut<'a, K, V> {
120    type Item = (&'a K, &'a mut V);
121
122    fn next(&mut self) -> Option<Self::Item> {
123        self.head.map(|mut node| unsafe {
124            self.head = node.as_ref().next[0];
125            self.len -= 1;
126            let k = node.as_ref().key.assume_init_ref();
127            let v = &mut *node.as_mut().value.as_mut_ptr();
128            (k, v)
129        })
130    }
131
132    fn size_hint(&self) -> (usize, Option<usize>) {
133        (self.len, Some(self.len))
134    }
135
136    fn count(self) -> usize
137    where
138        Self: Sized,
139    {
140        self.len
141    }
142}
143
144impl<K, V> Iterator for IntoIter<K, V> {
145    type Item = (K, V);
146    fn next(&mut self) -> Option<Self::Item> {
147        self.head.map(|node| unsafe {
148            let node = Box::from_raw(node.as_ptr());
149
150            self.head = node.next[0];
151            self.len -= 1;
152            let k = node.key.assume_init();
153            let v = node.value.assume_init();
154
155            (k, v)
156        })
157    }
158
159    fn size_hint(&self) -> (usize, Option<usize>) {
160        (self.len, Some(self.len))
161    }
162
163    fn count(self) -> usize
164    where
165        Self: Sized,
166    {
167        self.len
168    }
169}
170
171impl<K, V> Default for SkipList<K, V> {
172    /// Create a skip list with max level(12)
173    /// 
174    /// # Example
175    /// 
176    /// ```rust
177    /// use skip_list::SkipList;
178    /// let mut skiplist: SkipList<i32, i32> = SkipList::default();
179    /// ```
180    fn default() -> Self {
181        let max_level = 12;
182        let node = Box::leak(Box::new(Node::sigil(max_level))).into();
183        Self {
184            head: node,
185            len: 0,
186            level: 0,
187            max_level,
188            marker: PhantomData,
189        }
190    }
191}
192
193impl<K: Ord, V> SkipList<K, V> {
194    /// Create a skip list with max level
195    /// 
196    /// # Example
197    /// 
198    /// ```rust
199    /// use skip_list::SkipList;
200    /// let mut skiplist: SkipList<i32, i32> = SkipList::new(12);
201    /// ```
202    pub fn new(max_level: usize) -> Self {
203        let node = Box::leak(Box::new(Node::sigil(max_level))).into();
204        Self {
205            head: node,
206            len: 0,
207            level: 0,
208            max_level,
209            marker: PhantomData,
210        }
211    }
212
213    /// Returns a reference to the value of the key in skip list or None if
214    /// not exist.
215    /// 
216    /// # Example
217    /// 
218    /// ```rust
219    /// use skip_list::SkipList;
220    /// 
221    /// let mut skip_list = SkipList::default();
222    /// 
223    /// skip_list.insert(1, "a");
224    /// 
225    /// assert_eq!(skip_list.get(&1), Some(&"a"));
226    /// ```
227    pub fn get(&self, k: &K) -> Option<&V> {
228        let mut node = self.head;
229        for l in (0..self.level).rev() {
230            unsafe {
231                while let Some(next) = node.as_ref().next[l] {
232                    let key = &*next.as_ref().key.as_ptr();
233                    if key == k {
234                        return Some(&*next.as_ref().value.as_ptr());
235                    }
236                    if key < k {
237                        node = next;
238                    } else {
239                        break;
240                    }
241                }
242            }
243        }
244        None
245    }
246
247    /// Insert a key-value pair into skip list. If the key already exists,
248    /// updates key's value and return old value. Otherwise, `None` is returned.
249    /// 
250    /// # Example
251    /// 
252    /// ```rust
253    /// use skip_list::SkipList;
254    /// 
255    /// let mut skip_list = SkipList::default();
256    /// 
257    /// assert_eq!(skip_list.insert(1, "a"), None);
258    /// assert_eq!(skip_list.get(&1), Some(&"a"));
259    /// assert_eq!(skip_list.insert(1, "aa"), Some("a"));
260    /// assert_eq!(skip_list.get(&1), Some(&"aa"));
261    /// 
262    /// ```
263    pub fn insert(&mut self, k: K, mut v: V) -> Option<V> {
264        let mut node = self.head;
265        let mut updates = vec![None; self.max_level];
266
267        for l in (0..self.level).rev() {
268            unsafe {
269                while let Some(mut next) = node.as_ref().next[l] {
270                    let key = &*next.as_ref().key.as_ptr();
271                    if key == &k {
272                        let value = &mut *next.as_mut().value.as_mut_ptr();
273                        std::mem::swap(value, &mut v);
274                        return Some(v);
275                    }
276                    if key < &k {
277                        node = next;
278                    } else {
279                        break;
280                    }
281                }
282            }
283            updates[l] = Some(node);
284        }
285
286        let level = self.random_level();
287        if level > self.level {
288            for node in updates.iter_mut().take(level).skip(self.level) {
289                node.replace(self.head);
290            }
291            self.level = level;
292        }
293
294        let mut node: NonNull<Node<K, V>> =
295            Box::leak(Box::new(Node::new(k, v, level, self.max_level))).into();
296        for (l, ln) in updates.iter_mut().enumerate().take(level) {
297            if let Some(ln) = ln {
298                unsafe {
299                    node.as_mut().next[l] = ln.as_ref().next[l];
300                    ln.as_mut().next[l] = Some(node);
301                }
302            }
303        }
304        self.len += 1;
305        None
306    }
307
308    /// Deletes and returns the key's value from skip list or `None` if not exist.
309    /// 
310    /// # Example
311    /// 
312    /// ```rust
313    /// use skip_list::SkipList;
314    /// 
315    /// let mut skip_list = SkipList::default();
316    /// 
317    /// assert_eq!(skip_list.insert(1, "a"), None);
318    /// assert_eq!(skip_list.get(&1), Some(&"a"));
319    /// assert_eq!(skip_list.delete(&1), Some("a"));
320    /// assert_eq!(skip_list.get(&1), None);
321    /// ```
322    /// 
323    pub fn delete(&mut self, k: &K) -> Option<V> {
324        let mut node = self.head;
325        let mut updates = vec![None; self.max_level];
326
327        let mut target = None;
328        for l in (0..self.level).rev() {
329            unsafe {
330                while let Some(next) = node.as_ref().next[l] {
331                    let key = &*next.as_ref().key.as_ptr();
332                    if key == k {
333                        target = Some(next);
334                        break;
335                    }
336                    if key < k {
337                        node = next;
338                    } else {
339                        break;
340                    }
341                }
342            }
343            updates[l] = Some(node);
344        }
345
346        if let Some(node) = target {
347            unsafe {
348                for (l, ln) in updates.iter().enumerate().take(node.as_ref().level) {
349                    if let Some(mut ln) = ln {
350                        ln.as_mut().next[l] = node.as_ref().next[l];
351                    }
352                }
353                self.len -= 1;
354                let mut node = Box::from_raw(node.as_ptr());
355                node.key.assume_init_drop();
356                return Some(node.value.assume_init());
357            }
358        }
359        None
360    }
361
362    /// Visit all key-value pairs in the order of keys
363    /// The Iterator element type is (&K, &V).
364    /// 
365    /// # Example
366    /// 
367    /// ```rust
368    /// use skip_list::SkipList;
369    /// 
370    /// let mut skip_list = SkipList::default();
371    /// assert_eq!(skip_list.insert(3, "c"), None);
372    /// assert_eq!(skip_list.insert(4, "d"), None);
373    /// assert_eq!(skip_list.insert(2, "b"), None);
374    /// assert_eq!(skip_list.insert(5, "e"), None);
375    /// assert_eq!(skip_list.insert(1, "a"), None);
376    /// 
377    /// // visit all key-value paris in the order of keys
378    /// let mut keys = vec![];
379    /// let mut values = vec![];
380    /// for (k, v) in skip_list.iter() {
381    ///     keys.push(k);
382    ///     values.push(v);
383    /// }
384    /// 
385    /// // check key sorted
386    /// assert_eq!(keys, vec![&1, &2, &3, &4, &5]);
387    /// assert_eq!(values, vec![&"a", &"b", &"c", &"d", &"e"]);
388    /// ```
389    pub fn iter(&self) -> Iter<'_, K, V> {
390        Iter {
391            len: self.len,
392            head: unsafe { self.head.as_ref().next[0] },
393            marker: PhantomData,
394        }
395    }
396
397    /// Visit all key-value pairs in the order of keys
398    /// The Iterator element type is (&K, &mut V).
399    /// The value is mut, can be update;
400    /// 
401    /// # Example
402    /// 
403    /// ```rust
404    /// use skip_list::SkipList;
405    /// 
406    /// let mut skip_list = SkipList::default();
407    /// assert_eq!(skip_list.insert(3, 3), None);
408    /// assert_eq!(skip_list.insert(4, 4), None);
409    /// assert_eq!(skip_list.insert(2, 2), None);
410    /// assert_eq!(skip_list.insert(5, 5), None);
411    /// assert_eq!(skip_list.insert(1, 1), None);
412    /// 
413    /// for (k, v) in skip_list.iter_mut() {
414    ///     *v = k * 10;
415    /// }
416    /// // visit all key-value paris in the order of keys
417    /// let mut keys = vec![];
418    /// let mut values = vec![];
419    /// for (k, v) in skip_list.iter() {
420    ///     keys.push(k);
421    ///     values.push(v);
422    /// }
423    /// 
424    /// // check key sorted
425    /// assert_eq!(keys, vec![&1, &2, &3, &4, &5]);
426    /// assert_eq!(values, vec![&10, &20, &30, &40, &50]);
427    /// ```
428    pub fn iter_mut(&self) -> IterMut<'_, K, V> {
429        IterMut {
430            len: self.len,
431            head: unsafe { self.head.as_ref().next[0] },
432            marker: PhantomData,
433        }
434    }
435
436    fn random_level(&self) -> usize {
437        let mut rng = rand::thread_rng();
438        rng.gen_range(1..self.max_level)
439    }
440}
441
442impl<K, V> IntoIterator for SkipList<K, V> {
443    type Item = (K, V);
444    type IntoIter = IntoIter<K, V>;
445
446    /// Visit all key-value pairs in the order of keys
447    /// The Iterator element type is (K, V).
448    /// 
449    /// # Example
450    /// 
451    /// ```rust
452    /// use skip_list::SkipList;
453    /// 
454    /// let mut skip_list = SkipList::default();
455    /// assert_eq!(skip_list.insert(3, "c"), None);
456    /// assert_eq!(skip_list.insert(4, "d"), None);
457    /// assert_eq!(skip_list.insert(2, "b"), None);
458    /// assert_eq!(skip_list.insert(5, "e"), None);
459    /// assert_eq!(skip_list.insert(1, "a"), None);
460    /// 
461    /// // visit all key-value paris in the order of keys
462    /// let mut keys = vec![];
463    /// let mut values = vec![];
464    /// for (k, v) in skip_list.into_iter() {
465    ///     keys.push(k);
466    ///     values.push(v);
467    /// }
468    /// 
469    /// // check key sorted
470    /// assert_eq!(keys, vec![1, 2, 3, 4, 5]);
471    /// assert_eq!(values, vec!["a", "b", "c", "d", "e"]);
472    /// ```
473    fn into_iter(mut self) -> Self::IntoIter {
474        let node = unsafe { self.head.as_ref().next[0] };
475        unsafe {
476            self.head.as_mut().next[0] = None;
477        }
478        IntoIter {
479            len: self.len,
480            head: node,
481            marker: PhantomData,
482        }
483    }
484}
485
486impl<K, V> Drop for SkipList<K, V> {
487    fn drop(&mut self) {
488        unsafe {
489            let mut node = self.head.as_mut().next[0];
490
491            while let Some(n) = node {
492                let mut n = Box::from_raw(n.as_ptr());
493                node = n.next[0];
494                n.key.assume_init_drop();
495                n.value.assume_init_drop();
496            }
497
498            Box::from_raw(self.head.as_ptr());
499        }
500    }
501}
502
503#[cfg(test)]
504mod tests {
505    use super::SkipList;
506    #[test]
507    fn test_of_skip_list() {
508        let mut skip_list = SkipList::default();
509        for i in 0..100 {
510            assert_eq!(skip_list.insert(i, i), None);
511        }
512
513        for i in 0..100 {
514            assert_eq!(skip_list.insert(i, 10 * i), Some(i));
515        }
516
517        for i in 0..100 {
518            let v = i * 10;
519            assert_eq!(skip_list.get(&i), Some(&v))
520        }
521
522        for i in 0..50 {
523            let v = 10 * i;
524            assert_eq!(skip_list.delete(&i), Some(v));
525        }
526
527        for i in 0..50 {
528            assert_eq!(skip_list.get(&i), None);
529        }
530
531        for i in 50..100 {
532            let v = i * 10;
533            assert_eq!(skip_list.get(&i), Some(&v));
534        }
535
536        for (k, v) in skip_list.iter_mut() {
537            *v = *k * 20;
538        }
539
540        let mut key = 50;
541        for (k, v) in skip_list.iter() {
542            assert_eq!(*k, key);
543            assert_eq!(*v, key * 20);
544            key += 1;
545        }
546
547        key = 50;
548        for (k, v) in skip_list.into_iter() {
549            assert_eq!(k, key);
550            assert_eq!(v, key * 20);
551            key += 1;
552        }
553    }
554
555    #[test]
556    fn test_sl() {
557        let mut skip_list = SkipList::default();
558        // insert
559        assert_eq!(skip_list.insert(1, 10), None); // there is no value with key with 1
560        assert_eq!(skip_list.insert(2, 20), None);
561        assert_eq!(skip_list.insert(3, 30), None);
562
563        // get
564        assert_eq!(skip_list.get(&1), Some(&10));
565        assert_eq!(skip_list.get(&2), Some(&20));
566        assert_eq!(skip_list.get(&3), Some(&30));
567
568        // update
569        assert_eq!(skip_list.insert(1, 100), Some(10));
570        assert_eq!(skip_list.insert(2, 200), Some(20));
571        assert_eq!(skip_list.insert(3, 300), Some(30));
572
573        // iterator
574        for (k, v) in skip_list.iter() {
575            let value = k * 100;
576            assert_eq!(*v, value);
577        }
578
579        // delete
580        assert_eq!(skip_list.delete(&1), Some(100));
581        assert_eq!(skip_list.delete(&10), None);
582        assert_eq!(skip_list.get(&1), None);
583    }
584}