implement a skiplist in rust

main.rs
use rand::Rng;
use std::cmp::Ordering;

#[derive(Debug)]
struct SkipNode<T: Clone + PartialOrd> {
    value: T,
    next: Vec<Option<Box<SkipNode<T>>>>,
}

impl<T: Clone + PartialOrd> SkipNode<T> {
    fn new(value: T, height: usize) -> SkipNode<T> {
        SkipNode {
            value,
            next: vec![None; height],
        }
    }
}

#[derive(Debug)]
pub struct SkipList<T: Clone + PartialOrd> {
    head: SkipNode<T>,
    max_height: usize,
}

impl<T: Clone + PartialOrd> SkipList<T> {
    pub fn new() -> SkipList<T> {
        SkipList {
            head: SkipNode::new(Default::default(), 1),
            max_height: 1,
        }
    }

    pub fn insert(&mut self, value: T) -> bool {
        let mut update: Vec<Option<Box<SkipNode<T>>>> =
            vec![None; self.max_height];
        let mut current = &mut self.head;
        for height in (0..self.max_height).rev() {
            while let Some(ref next) = current.next[height] {
                match next.value.partial_cmp(&value) {
                    Some(Ordering::Equal) => return false,
                    Some(Ordering::Greater) => break,
                    _ => current = next,
                }
            }
            update[height] = Some(Box::new(current.clone()));
        }
        let height = Self::random_height(self.max_height);
        if height > self.max_height {
            update.resize_with(height, || None);
            self.max_height = height;
        }
        let node = SkipNode::new(value.clone(), height);
        for height in 0..height {
            if let Some(ref mut n) = update[height] {
                node.next[height] = n.next[height].take();
                n.next[height] = Some(Box::new(node.clone()));
            }
        }
        true
    }

    pub fn get(&self, value: T) -> Option<&T> {
        let mut current = &self.head;
        for height in (0..self.max_height).rev() {
            while let Some(ref next) = current.next[height] {
                match next.value.partial_cmp(&value) {
                    Some(Ordering::Equal) => return Some(&next.value),
                    Some(Ordering::Greater) => break,
                    _ => current = next,
                }
            }
        }
        None
    }

    pub fn remove(&mut self, value: T) -> bool {
        let mut removed = false;
        let mut update: Vec<Option<Box<SkipNode<T>>>> =
            vec![None; self.max_height];
        let mut current = &mut self.head;
        for height in (0..self.max_height).rev() {
            while let Some(ref mut next) = current.next[height] {
                match next.value.partial_cmp(&value) {
                    Some(Ordering::Equal) => {
                        removed = true;
                        let _ = next.clone().next
                            .iter()
                            .enumerate()
                            .filter(|(_, n)| n.is_some())
                            .for_each(|(i, n)| {
                                update[i] = next.next[i].take();
                            });
                        break;
                    }
                    Some(Ordering::Greater) => break,
                    _ => current = next,
                }
            }
            update[height] = Some(Box::new(current.clone()));
        }
        if removed {
            for height in (0..self.max_height).rev() {
                if let Some(ref mut n) = update[height] {
                    if let Some(ref mut next) = n.next[height] {
                        if next.value == value {
                            n.next[height] = next.next[height].take();
                        }
                    }
                }
            }
        }
        removed
    }

    fn random_height(max_height: usize) -> usize {
        let mut rng = rand::thread_rng();
        let mut height = 1;
        while rng.gen::<bool>() && height < max_height {
            height += 1;
        }
        height
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert() {
        let mut list: SkipList<i32> = SkipList::new();
        assert_eq!(list.insert(2), true);
        assert_eq!(list.insert(3), true);
        assert_eq!(list.insert(2), false);
    }

    #[test]
    fn test_remove() {
        let mut list: SkipList<i32> = SkipList::new();
        assert_eq!(list.insert(2), true);
        assert_eq!(list.remove(2), true);
        assert_eq!(list.get(2), None);
    }

    #[test]
    fn test_get() {
        let mut list: SkipList<i32> = SkipList::new();
        assert_eq!(list.insert(2), true);
        assert_eq!(list.get(2), Some(&2));
        assert_eq!(list.get(3), None);
    }
}
4692 chars
153 lines