lock_free/

hashmap_fast.rs

//! Ultra-fast lock-free HashMap using open addressing
//! 
//! This implementation uses linear probing with atomic operations
//! for maximum performance on modern CPUs.

use std::sync::atomic::{AtomicU64, AtomicPtr, Ordering};
use std::ptr::null_mut;
use std::hash::{Hash, Hasher};
use std::collections::hash_map::DefaultHasher;
use std::marker::PhantomData;

/// Entry states
const EMPTY: u64 = 0;
const TOMBSTONE: u64 = 1;
const OCCUPIED_FLAG: u64 = 2;

/// Cache line size
const CACHE_LINE: usize = 128;

/// Default capacity (must be power of 2)
const DEFAULT_CAPACITY: usize = 1024;

/// Entry in the hashmap
#[repr(align(64))]
struct Entry<K, V> {
    /// Combined hash and state (2 bits state, 62 bits hash)
    hash_state: AtomicU64,
    /// Key storage
    key: AtomicPtr<K>,
    /// Value storage  
    value: AtomicPtr<V>,
}

impl<K, V> Entry<K, V> {
    fn new() -> Self {
        Self {
            hash_state: AtomicU64::new(EMPTY),
            key: AtomicPtr::new(null_mut()),
            value: AtomicPtr::new(null_mut()),
        }
    }
}

/// Ultra-fast lock-free HashMap
pub struct FastHashMap<K, V> {
    /// Entry array
    entries: Box<[Entry<K, V>]>,
    /// Capacity mask (capacity - 1)
    mask: usize,
    /// Size counter
    size: AtomicU64,
    /// Type markers
    _phantom: PhantomData<(K, V)>,
}

impl<K: Hash + Eq + Clone, V: Clone> FastHashMap<K, V> {
    /// Create new hashmap
    pub fn new() -> Self {
        Self::with_capacity(DEFAULT_CAPACITY)
    }

    /// Create with specified capacity
    pub fn with_capacity(capacity: usize) -> Self {
        let capacity = capacity.next_power_of_two();
        let mut entries = Vec::with_capacity(capacity);
        
        for _ in 0..capacity {
            entries.push(Entry::new());
        }
        
        Self {
            entries: entries.into_boxed_slice(),
            mask: capacity - 1,
            size: AtomicU64::new(0),
            _phantom: PhantomData,
        }
    }

    /// Hash a key
    #[inline]
    fn hash_key(key: &K) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }

    /// Insert key-value pair
    pub fn insert(&self, key: K, value: V) -> Option<V> {
        let hash = Self::hash_key(&key);
        let hash_with_flag = hash | OCCUPIED_FLAG;
        let mut index = (hash as usize) & self.mask;
        let start_index = index;
        
        loop {
            let entry = &self.entries[index];
            let current_hash = entry.hash_state.load(Ordering::Acquire);
            
            if current_hash == EMPTY || current_hash == TOMBSTONE {
                // Try to claim this slot
                match entry.hash_state.compare_exchange(
                    current_hash,
                    hash_with_flag,
                    Ordering::AcqRel,
                    Ordering::Acquire
                ) {
                    Ok(_) => {
                        // We got it! Store key and value
                        let key_ptr = Box::into_raw(Box::new(key));
                        let value_ptr = Box::into_raw(Box::new(value));
                        
                        entry.key.store(key_ptr, Ordering::Release);
                        entry.value.store(value_ptr, Ordering::Release);
                        
                        if current_hash == EMPTY {
                            self.size.fetch_add(1, Ordering::Relaxed);
                        }
                        return None;
                    }
                    Err(_) => {
                        // Someone else got it, continue
                    }
                }
            } else if current_hash == hash_with_flag {
                // Check if same key
                let key_ptr = entry.key.load(Ordering::Acquire);
                if !key_ptr.is_null() {
                    let stored_key = unsafe { &*key_ptr };
                    if stored_key == &key {
                        // Update value
                        let new_value_ptr = Box::into_raw(Box::new(value));
                        let old_value_ptr = entry.value.swap(new_value_ptr, Ordering::AcqRel);
                        
                        if !old_value_ptr.is_null() {
                            let old_value = unsafe { Box::from_raw(old_value_ptr) };
                            return Some(*old_value);
                        }
                        return None;
                    }
                }
            }
            
            // Linear probe to next slot
            index = (index + 1) & self.mask;
            
            // Check if we've wrapped around completely
            if index == start_index {
                // Table is full, can't insert
                return None;
            }
        }
    }

    /// Get value by key
    pub fn get(&self, key: &K) -> Option<V> {
        let hash = Self::hash_key(key);
        let hash_with_flag = hash | OCCUPIED_FLAG;
        let mut index = (hash as usize) & self.mask;
        let start_index = index;
        
        loop {
            let entry = &self.entries[index];
            let current_hash = entry.hash_state.load(Ordering::Acquire);
            
            if current_hash == EMPTY {
                return None;
            }
            
            if current_hash == hash_with_flag {
                let key_ptr = entry.key.load(Ordering::Acquire);
                if !key_ptr.is_null() {
                    let stored_key = unsafe { &*key_ptr };
                    if stored_key == key {
                        let value_ptr = entry.value.load(Ordering::Acquire);
                        if !value_ptr.is_null() {
                            let value = unsafe { &*value_ptr };
                            return Some(value.clone());
                        }
                    }
                }
            }
            
            index = (index + 1) & self.mask;
            if index == start_index {
                return None;
            }
        }
    }

    /// Remove key-value pair
    pub fn remove(&self, key: &K) -> Option<V> {
        let hash = Self::hash_key(key);
        let hash_with_flag = hash | OCCUPIED_FLAG;
        let mut index = (hash as usize) & self.mask;
        let start_index = index;
        
        loop {
            let entry = &self.entries[index];
            let current_hash = entry.hash_state.load(Ordering::Acquire);
            
            if current_hash == EMPTY {
                return None;
            }
            
            if current_hash == hash_with_flag {
                let key_ptr = entry.key.load(Ordering::Acquire);
                if !key_ptr.is_null() {
                    let stored_key = unsafe { &*key_ptr };
                    if stored_key == key {
                        // Mark as tombstone
                        match entry.hash_state.compare_exchange(
                            hash_with_flag,
                            TOMBSTONE,
                            Ordering::AcqRel,
                            Ordering::Acquire
                        ) {
                            Ok(_) => {
                                // Remove key and value
                                let key_ptr = entry.key.swap(null_mut(), Ordering::AcqRel);
                                let value_ptr = entry.value.swap(null_mut(), Ordering::AcqRel);
                                
                                self.size.fetch_sub(1, Ordering::Relaxed);
                                
                                if !key_ptr.is_null() {
                                    unsafe { let _ = Box::from_raw(key_ptr); }
                                }
                                
                                if !value_ptr.is_null() {
                                    let value = unsafe { Box::from_raw(value_ptr) };
                                    return Some(*value);
                                }
                            }
                            Err(_) => {
                                // Someone else removed it
                                return None;
                            }
                        }
                    }
                }
            }
            
            index = (index + 1) & self.mask;
            if index == start_index {
                return None;
            }
        }
    }

    /// Check if key exists
    pub fn contains_key(&self, key: &K) -> bool {
        let hash = Self::hash_key(key);
        let hash_with_flag = hash | OCCUPIED_FLAG;
        let mut index = (hash as usize) & self.mask;
        let start_index = index;
        
        loop {
            let entry = &self.entries[index];
            let current_hash = entry.hash_state.load(Ordering::Acquire);
            
            if current_hash == EMPTY {
                return false;
            }
            
            if current_hash == hash_with_flag {
                let key_ptr = entry.key.load(Ordering::Acquire);
                if !key_ptr.is_null() {
                    let stored_key = unsafe { &*key_ptr };
                    if stored_key == key {
                        return true;
                    }
                }
            }
            
            index = (index + 1) & self.mask;
            if index == start_index {
                return false;
            }
        }
    }

    /// Get number of elements
    pub fn len(&self) -> usize {
        self.size.load(Ordering::Relaxed) as usize
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

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

impl<K, V> Drop for FastHashMap<K, V> {
    fn drop(&mut self) {
        for entry in self.entries.iter() {
            let key_ptr = entry.key.load(Ordering::Acquire);
            let value_ptr = entry.value.load(Ordering::Acquire);
            
            if !key_ptr.is_null() {
                unsafe { let _ = Box::from_raw(key_ptr); }
            }
            
            if !value_ptr.is_null() {
                unsafe { let _ = Box::from_raw(value_ptr); }
            }
        }
    }
}