lock_free/

skiplist.rs

//! Simplified high-performance lock-free skip list.

use std::sync::atomic::{AtomicPtr, AtomicBool, Ordering};
use std::ptr;
use std::cmp::Ordering as CmpOrdering;

const MAX_HEIGHT: usize = 32;

struct Node<K, V> {
    key: K,
    value: V,
    next: [AtomicPtr<Node<K, V>>; MAX_HEIGHT],
    height: usize,
    marked: AtomicBool,
}

impl<K, V> Node<K, V> {
    fn new(key: K, value: V, height: usize) -> Self {
        let mut next: [AtomicPtr<Node<K, V>>; MAX_HEIGHT] = Default::default();
        for i in 0..MAX_HEIGHT {
            next[i] = AtomicPtr::new(ptr::null_mut());
        }
        
        Node {
            key,
            value,
            next,
            height,
            marked: AtomicBool::new(false),
        }
    }
}

/// A lock-free skip list implementation
pub struct SkipList<K, V> {
    head: Box<Node<K, V>>,
}

impl<K: Ord + Default + Clone, V: Default> SkipList<K, V> {
    /// Creates a new skip list
    pub fn new() -> Self {
        let head = Box::new(Node::new(K::default(), V::default(), MAX_HEIGHT));
        SkipList { head }
    }

    fn random_height() -> usize {
        let mut height = 1;
        let mut rng = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .subsec_nanos();
        
        while height < MAX_HEIGHT && rng % 2 == 0 {
            height += 1;
            rng = rng.wrapping_mul(1664525).wrapping_add(1013904223);
        }
        height
    }

    fn find(&self, key: &K) -> ([*mut Node<K, V>; MAX_HEIGHT], [*mut Node<K, V>; MAX_HEIGHT]) {
        let mut preds = [ptr::null_mut(); MAX_HEIGHT];
        let mut succs = [ptr::null_mut(); MAX_HEIGHT];
        
        let mut pred = &*self.head as *const Node<K, V> as *mut Node<K, V>;
        
        for level in (0..MAX_HEIGHT).rev() {
            let mut curr = unsafe { (*pred).next[level].load(Ordering::Acquire) };
            
            while !curr.is_null() {
                let curr_ref = unsafe { &*curr };
                
                if curr_ref.marked.load(Ordering::Acquire) {
                    curr = curr_ref.next[level].load(Ordering::Acquire);
                    continue;
                }
                
                match curr_ref.key.cmp(key) {
                    CmpOrdering::Less => {
                        pred = curr;
                        curr = curr_ref.next[level].load(Ordering::Acquire);
                    }
                    _ => break,
                }
            }
            
            preds[level] = pred;
            succs[level] = curr;
        }
        
        (preds, succs)
    }

    /// Inserts a key-value pair
    pub fn insert(&self, key: K, value: V) -> bool {
        let height = Self::random_height();
        let new_node = Box::into_raw(Box::new(Node::new(key, value, height)));
        
        loop {
            let (preds, succs) = self.find(unsafe { &(*new_node).key });
            
            // Check if already exists
            if !succs[0].is_null() {
                let succ_ref = unsafe { &*succs[0] };
                if succ_ref.key == unsafe { (*new_node).key.clone() } && !succ_ref.marked.load(Ordering::Acquire) {
                    unsafe { drop(Box::from_raw(new_node)) };
                    return false;
                }
            }
            
            // Set next pointers
            for level in 0..height {
                unsafe {
                    (*new_node).next[level].store(succs[level], Ordering::Relaxed);
                }
            }
            
            // Try to link at bottom level
            let pred_next = unsafe { &(*preds[0]).next[0] };
            match pred_next.compare_exchange(
                succs[0],
                new_node,
                Ordering::Release,
                Ordering::Acquire,
            ) {
                Ok(_) => {
                    // Link at higher levels
                    for level in 1..height {
                        loop {
                            let pred = preds[level];
                            let succ = succs[level];
                            let pred_next = unsafe { &(*pred).next[level] };
                            
                            if pred_next.compare_exchange(
                                succ,
                                new_node,
                                Ordering::Release,
                                Ordering::Acquire,
                            ).is_ok() {
                                break;
                            }
                            
                            let (_new_preds, new_succs) = self.find(unsafe { &(*new_node).key });
                            if !new_succs[0].is_null() && unsafe { &*new_succs[0] }.key == unsafe { (*new_node).key.clone() } {
                                break;
                            }
                        }
                    }
                    return true;
                }
                Err(_) => continue,
            }
        }
    }

    /// Checks if a key exists
    pub fn contains(&self, key: &K) -> bool {
        let (_, succs) = self.find(key);
        
        if !succs[0].is_null() {
            let succ_ref = unsafe { &*succs[0] };
            succ_ref.key == *key && !succ_ref.marked.load(Ordering::Acquire)
        } else {
            false
        }
    }

    /// Gets a value by key
    pub fn get(&self, key: &K) -> Option<V> 
    where 
        V: Clone 
    {
        let (_, succs) = self.find(key);
        
        if !succs[0].is_null() {
            let succ_ref = unsafe { &*succs[0] };
            if succ_ref.key == *key && !succ_ref.marked.load(Ordering::Acquire) {
                Some(succ_ref.value.clone())
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Removes a key
    pub fn remove(&self, key: &K) -> bool {
        loop {
            let (mut preds, succs) = self.find(key);
            
            if succs[0].is_null() {
                return false;
            }
            
            let node_to_remove = succs[0];
            let node_ref = unsafe { &*node_to_remove };
            
            if node_ref.key != *key {
                return false;
            }
            
            // Mark the node
            if !node_ref.marked.compare_exchange(
                false,
                true,
                Ordering::Release,
                Ordering::Acquire,
            ).is_ok() {
                return false;
            }
            
            // Try to physically remove
            for level in (0..node_ref.height).rev() {
                let succ = node_ref.next[level].load(Ordering::Acquire);
                loop {
                    let pred_next = unsafe { &(*preds[level]).next[level] };
                    if pred_next.compare_exchange(
                        node_to_remove,
                        succ,
                        Ordering::Release,
                        Ordering::Acquire,
                    ).is_ok() {
                        break;
                    }
                    
                    let (new_preds, _) = self.find(key);
                    if level < node_ref.height {
                        preds[level] = new_preds[level];
                    }
                }
            }
            
            return true;
        }
    }
}

impl<K: Default + Ord + Clone, V: Default> Default for SkipList<K, V> {
    fn default() -> Self {
        Self::new()
    }
}

unsafe impl<K: Send, V: Send> Send for SkipList<K, V> {}
unsafe impl<K: Send + Sync, V: Send + Sync> Sync for SkipList<K, V> {}