swiss_table/

lib.rs

use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hash, Hasher};
use std::mem;

const GROUP_SIZE: usize = 8;
const MAX_LOAD_FACTOR: f64 = 0.875;

// control bytes
const EMPTY: u8 = 0b1111_1111;
const DELETED: u8 = 0b1000_0000;

/// Extract h2 from a hash.
/// Lowest 7 bits.
fn h2(hash: u64) -> u8 {
    (hash & 0x7F) as u8
}

/// A group is 8 control bytes + 8 slots.
/// The control bytes summarize the state
/// of the group so we can do fast parallel matching.
struct Group<K, V> {
    // Each byte is either EMPTY, DELETED, or an h2 fingerprint
    ctrl: [u8; GROUP_SIZE],
    slots: [Option<(K, V)>; GROUP_SIZE],
}

impl<K, V> Group<K, V> {
    fn new() -> Self {
        Group {
            ctrl: [EMPTY; GROUP_SIZE],
            slots: std::array::from_fn(|_| None),
        }
    }

    /// SIMD-style parallel match.
    /// In a real implementation, this would be a single SIMD instruction
    /// Just simulating it here: return a bitmask of which slots have ctrl == needle.
    ///
    /// e.g., if slots 1 and 5 match, returns 0b00100010
    fn match_h2(&self, needle: u8) -> u8 {
        let mut mask = 0u8;
        for i in 0..GROUP_SIZE {
            if self.ctrl[i] == needle {
                mask |= 1 << i;
            }
        }
        mask
    }

    /// Return bitmask of empty slots.
    fn match_empty(&self) -> u8 {
        let mut mask = 0u8;
        for i in 0..GROUP_SIZE {
            if self.ctrl[i] == EMPTY {
                mask |= 1 << i;
            }
        }
        mask
    }

    /// Return bitmask of empty OR deleted slots (available for insertion).
    fn match_empty_or_deleted(&self) -> u8 {
        let mut mask = 0u8;
        for i in 0..GROUP_SIZE {
            // Both EMPTY and DELETED have the high bit set.
            // which will never happen for a valid h2 value.
            if self.ctrl[i] & 0x80 != 0 {
                mask |= 1 << i;
            }
        }
        mask
    }
}

pub struct SwissTable<K, V> {
    groups: Vec<Group<K, V>>,
    num_groups: usize,
    len: usize,
    hash_builder: RandomState,
}

impl<K, V> SwissTable<K, V>
where
    K: Eq + Hash,
    V: Clone,
{
    pub fn new() -> Self {
        let num_groups = 1;
        SwissTable {
            groups: vec![Group::new()],
            num_groups,
            len: 0,
            hash_builder: RandomState::new(),
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

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

    pub fn capacity(&self) -> usize {
        self.num_groups * GROUP_SIZE
    }

    fn hash_key(&self, key: &K) -> u64 {
        let mut hasher = self.hash_builder.build_hasher();
        key.hash(&mut hasher);
        hasher.finish()
    }

    fn find(&self, key: &K) -> Option<(usize, usize)> {
        let hash = self.hash_key(key);
        let h2_val = h2(hash);
        let start_group = (hash >> 7) as usize % self.num_groups;

        for i in 0..self.num_groups {
            let gi = (start_group + i) % self.num_groups;
            let group = &self.groups[gi];

            let mut candidates = group.match_h2(h2_val);
            while candidates != 0 {
                let slot = candidates.trailing_zeros() as usize;
                candidates &= candidates - 1;

                if let Some((ref k, _)) = group.slots[slot] {
                    if k == key {
                        return Some((gi, slot));
                    }
                }
            }

            if group.match_empty() != 0 {
                return None;
            }
        }
        None
    }

    pub fn get(&self, key: &K) -> Option<&V> {
        let (gi, si) = self.find(key)?;
        self.groups[gi].slots[si].as_ref().map(|(_, v)| v)
    }

    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        let (gi, si) = self.find(key)?;
        self.groups[gi].slots[si].as_mut().map(|(_, v)| v)
    }

    pub fn contains_key(&self, key: &K) -> bool {
        self.find(key).is_some()
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        if self.len >= (self.capacity() as f64 * MAX_LOAD_FACTOR) as usize {
            self.grow();
        }

        let hash = self.hash_key(&key);
        let h2_val = h2(hash);
        let start_group = (hash >> 7) as usize % self.num_groups;

        for i in 0..self.num_groups {
            let gi = (start_group + i) % self.num_groups;
            let group = &mut self.groups[gi];

            let mut candidates = group.match_h2(h2_val);
            while candidates != 0 {
                let slot = candidates.trailing_zeros() as usize;
                candidates &= candidates - 1;

                if let Some((ref k, ref mut v)) = group.slots[slot] {
                    if *k == key {
                        let old = mem::replace(v, value);
                        return Some(old);
                    }
                }
            }

            let available = group.match_empty_or_deleted();
            if available != 0 {
                let slot = available.trailing_zeros() as usize;
                group.ctrl[slot] = h2_val;
                group.slots[slot] = Some((key, value));
                self.len += 1;
                return None;
            }
        }

        unreachable!("table should have grown before reaching here");
    }

    pub fn remove(&mut self, key: &K) -> Option<V> {
        let (gi, si) = self.find(key)?;
        let group = &mut self.groups[gi];
        group.ctrl[si] = DELETED;
        let (_, v) = group.slots[si].take().unwrap();
        self.len -= 1;
        Some(v)
    }

    fn grow(&mut self) {
        let new_num_groups = self.num_groups * 2;
        let old_groups = mem::replace(
            &mut self.groups,
            (0..new_num_groups).map(|_| Group::new()).collect(),
        );
        let old_len = self.len;
        self.num_groups = new_num_groups;
        self.len = 0;

        for group in old_groups {
            for slot in group.slots {
                if let Some((k, v)) = slot {
                    self.insert(k, v);
                }
            }
        }

        debug_assert_eq!(self.len, old_len);
    }
}

impl<K, V> Default for SwissTable<K, V>
where
    K: Eq + Hash,
    V: Clone,
{
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn empty_table() {
        let table: SwissTable<String, i32> = SwissTable::new();
        assert_eq!(table.len(), 0);
        assert!(table.is_empty());
        assert_eq!(table.get(&"hello".to_string()), None);
    }

    #[test]
    fn insert_and_get() {
        let mut table = SwissTable::new();
        table.insert("foo", 1);
        table.insert("bar", 2);
        table.insert("baz", 3);

        assert_eq!(table.get(&"foo"), Some(&1));
        assert_eq!(table.get(&"bar"), Some(&2));
        assert_eq!(table.get(&"baz"), Some(&3));
        assert_eq!(table.get(&"missing"), None);
        assert_eq!(table.len(), 3);
    }

    #[test]
    fn insert_returns_old_value() {
        let mut table = SwissTable::new();
        assert_eq!(table.insert("key", 10), None);
        assert_eq!(table.insert("key", 20), Some(10));
        assert_eq!(table.get(&"key"), Some(&20));
        assert_eq!(table.len(), 1);
    }

    #[test]
    fn remove() {
        let mut table = SwissTable::new();
        table.insert("a", 1);
        table.insert("b", 2);
        table.insert("c", 3);

        assert_eq!(table.remove(&"b"), Some(2));
        assert_eq!(table.get(&"b"), None);
        assert!(!table.contains_key(&"b"));
        assert_eq!(table.len(), 2);

        assert_eq!(table.remove(&"b"), None);
    }

    #[test]
    fn remove_doesnt_break_probe_chain() {
        let mut table = SwissTable::new();
        for i in 0..20 {
            table.insert(i, i * 100);
        }

        table.remove(&5);
        table.remove(&10);
        table.remove(&15);

        for i in 0..20 {
            if i == 5 || i == 10 || i == 15 {
                assert_eq!(table.get(&i), None);
            } else {
                assert_eq!(table.get(&i), Some(&(i * 100)));
            }
        }
    }

    #[test]
    fn contains_key() {
        let mut table = SwissTable::new();
        table.insert("present", 42);

        assert!(table.contains_key(&"present"));
        assert!(!table.contains_key(&"absent"));
    }

    #[test]
    fn get_mut() {
        let mut table = SwissTable::new();
        table.insert("counter", 0);

        if let Some(v) = table.get_mut(&"counter") {
            *v += 10;
        }

        assert_eq!(table.get(&"counter"), Some(&10));
    }

    #[test]
    fn grow_under_load() {
        let mut table = SwissTable::new();

        for i in 0..100 {
            table.insert(i, i.to_string());
        }

        assert_eq!(table.len(), 100);

        for i in 0..100 {
            assert_eq!(table.get(&i), Some(&i.to_string()));
        }
    }

    #[test]
    fn string_keys() {
        let mut table = SwissTable::new();

        for i in 0..50 {
            table.insert(format!("key_{}", i), i);
        }

        for i in 0..50 {
            let key = format!("key_{}", i);
            assert_eq!(table.get(&key), Some(&i));
        }
    }

    #[test]
    fn insert_after_remove() {
        let mut table = SwissTable::new();
        table.insert("reuse", 1);
        table.remove(&"reuse");
        table.insert("reuse", 2);

        assert_eq!(table.get(&"reuse"), Some(&2));
        assert_eq!(table.len(), 1);
    }

    #[test]
    fn large_scale_insert_remove() {
        let mut table = SwissTable::new();

        for i in 0..1000 {
            table.insert(i, i);
        }
        assert_eq!(table.len(), 1000);

        for i in (0..1000).step_by(2) {
            table.remove(&i);
        }
        assert_eq!(table.len(), 500);

        for i in 0..1000 {
            if i % 2 == 0 {
                assert_eq!(table.get(&i), None);
            } else {
                assert_eq!(table.get(&i), Some(&i));
            }
        }
    }
}