해시맵 직접 구현하기

// FxHash 핵심: 곱셈 + XOR 두 번
// 이 상수는 φ(황금비)의 비트 표현 — 비트 패턴을 골고루 섞어줌
const K: u64 = 0x517cc1b727220a95;

fn fxhash(key: u64) -> u64 {
    key.rotate_left(5).bitxor(key).wrapping_mul(K)
}

fn h1(hash: u64, num_groups: usize) -> usize {
    (hash >> 7) as usize % num_groups
}

fn h2(hash: u64) -> u8 {
    // 하위 7비트만 (MSB는 항상 0 = FULL 상태를 의미)
    (hash & 0x7F) as u8
}

use std::alloc::{alloc, dealloc, Layout};
use std::ptr;

const EMPTY:   u8 = 0xFF;
const DELETED: u8 = 0x80;

pub struct SwissMap<K, V> {
    ctrl:    *mut u8,
    buckets: *mut (K, V),
    cap:     usize,
    len:     usize,
}

impl<K, V> SwissMap<K, V> {
    pub fn new() -> Self {
        Self::with_capacity(16)
    }

    pub fn with_capacity(cap: usize) -> Self {
        assert!(cap >= 16 && cap.is_power_of_two(),
                "cap은 16 이상의 2의 거듭제곱이어야 합니다");
        unsafe {
            // ctrl: 16바이트 정렬, 전부 EMPTY(0xFF)로 초기화
            let ctrl_layout = Layout::from_size_align(cap, 16).unwrap();
            let ctrl = alloc(ctrl_layout);
            ptr::write_bytes(ctrl, EMPTY, cap);

            // buckets: 초기화 하지 않음 — ctrl이 EMPTY인 슬롯은 절대 읽지 않음
            let bucket_layout = Layout::array::<(K, V)>(cap).unwrap();
            let buckets = alloc(bucket_layout) as *mut (K, V);

            Self { ctrl, buckets, cap, len: 0 }
        }
    }
}

// Drop: ctrl이 FULL인 슬롯의 값만 직접 drop
impl<K, V> Drop for SwissMap<K, V> {
    fn drop(&mut self) {
        unsafe {
            for i in 0..self.cap {
                // MSB가 0이면 FULL (0xFF나 0x80은 MSB가 1)
                if *self.ctrl.add(i) < 0x80 {
                    ptr::drop_in_place(self.buckets.add(i));
                }
            }
            let ctrl_layout = Layout::from_size_align(self.cap, 16).unwrap();
            dealloc(self.ctrl, ctrl_layout);

            let bucket_layout = Layout::array::<(K, V)>(self.cap).unwrap();
            dealloc(self.buckets as *mut u8, bucket_layout);
        }
    }
}

#![feature(portable_simd)]
use std::simd::prelude::*;

struct Group(*const u8);

impl Group {
    // 메모리에서 16개 ctrl 바이트를 SIMD 레지스터 하나로 로드
    unsafe fn load(&self) -> u8x16 {
        u8x16::from_slice(std::slice::from_raw_parts(self.0, 16))
    }

    // target과 같은 바이트의 위치를 비트마스크로 반환
    // 비트 i가 1 → i번 슬롯의 ctrl == target
    unsafe fn match_byte(&self, target: u8) -> u16 {
        self.load().simd_eq(u8x16::splat(target)).to_bitmask()
    }

    // EMPTY(0xFF) 슬롯 위치
    unsafe fn match_empty(&self) -> u16 {
        self.match_byte(EMPTY)
    }

    // EMPTY 또는 DELETED 슬롯 (삽입 가능한 위치)
    // 두 값 모두 MSB = 1이므로, ctrl & 0x80 != 0 이면 삽입 가능
    unsafe fn match_available(&self) -> u16 {
        (self.load() & u8x16::splat(0x80))
            .simd_ne(u8x16::splat(0))
            .to_bitmask()
    }
}

// 비트마스크의 set 비트 인덱스를 하나씩 꺼내는 이터레이터
fn each_set_bit(mut mask: u16) -> impl Iterator<Item = usize> {
    std::iter::from_fn(move || {
        if mask == 0 { return None; }
        let i = mask.trailing_zeros() as usize;
        mask &= mask - 1;  // 최하위 set 비트 제거 (비트 트릭)
        Some(i)
    })
}

impl<K: Eq, V> SwissMap<K, V> {
    pub fn get(&self, key: &K, hash: u64) -> Option<&V> {
        let h2val     = h2(hash);
        let num_groups = self.cap / 16;
        let start      = h1(hash, num_groups);

        for probe in 0..num_groups {
            let g        = (start + probe) % num_groups;
            let ctrl_ptr = unsafe { self.ctrl.add(g * 16) };
            let group    = Group(ctrl_ptr);

            // ① SIMD로 h2 일치 위치 탐색
            for slot in each_set_bit(unsafe { group.match_byte(h2val) }) {
                let idx = g * 16 + slot;
                // ② h2가 같아도 실제 키가 다를 수 있음 (hash 충돌)
                //    키를 직접 비교해서 최종 확인
                let (k, v) = unsafe { &*self.buckets.add(idx) };
                if k == key {
                    return Some(v);
                }
            }

            // ③ 이 그룹에 EMPTY가 하나라도 있으면
            //    삽입 시 여기서 멈췄을 것이므로 더 이상 없음
            if unsafe { group.match_empty() } != 0 {
                return None;
            }
            // ④ 이 그룹이 꽉 차있으면 다음 그룹으로
        }
        None
    }
}

impl<K: Eq, V> SwissMap<K, V> {
    pub fn insert(&mut self, key: K, value: V, hash: u64) -> Option<V> {
        // 87.5% = 7/8 이상이면 resize
        if self.len * 8 >= self.cap * 7 {
            self.resize();
        }

        let h2val      = h2(hash);
        let num_groups = self.cap / 16;
        let start      = h1(hash, num_groups);

        for probe in 0..num_groups {
            let g        = (start + probe) % num_groups;
            let ctrl_ptr = unsafe { self.ctrl.add(g * 16) };
            let group    = Group(ctrl_ptr);

            // ① 이미 같은 키가 있으면 값을 교체
            for slot in each_set_bit(unsafe { group.match_byte(h2val) }) {
                let idx = g * 16 + slot;
                let (k, v) = unsafe { &mut *self.buckets.add(idx) };
                if k == &key {
                    return Some(std::mem::replace(v, value));
                }
            }

            // ② 삽입 가능한 슬롯(EMPTY 또는 DELETED) 찾기
            let available = unsafe { group.match_available() };
            if available != 0 {
                let slot = available.trailing_zeros() as usize;
                let idx  = g * 16 + slot;
                unsafe {
                    // ctrl에 h2 기록 (이 슬롯이 FULL임을 표시)
                    *self.ctrl.add(idx) = h2val;
                    // 버킷에 값 기록 (이전엔 미초기화 메모리였음)
                    self.buckets.add(idx).write((key, value));
                }
                self.len += 1;
                return None;
            }
            // ③ 이 그룹이 꽉 차있으면 다음 그룹으로
        }
        unreachable!("resize가 됐는데 빈 슬롯이 없음 — 버그");
    }

    fn resize(&mut self) {
        let new_cap = self.cap * 2;
        let mut new_map = SwissMap::with_capacity(new_cap);

        // 모든 FULL 슬롯을 새 맵으로 이동
        for i in 0..self.cap {
            unsafe {
                if *self.ctrl.add(i) < 0x80 {
                    // read()는 소유권을 이동시킴 (원본은 초기화 안 된 상태로 남김)
                    let (k, v) = self.buckets.add(i).read();
                    // 새 해시를 계산해서 재삽입
                    let hash = fxhash_of(&k);
                    new_map.insert(k, v, hash);
                }
            }
        }
        // self와 new_map을 교체
        // swap 후 new_map이 drop → 이전 ctrl/buckets 메모리 해제
        std::mem::swap(self, &mut new_map);
    }
}

impl<K: Eq, V> SwissMap<K, V> {
    pub fn remove(&mut self, key: &K, hash: u64) -> Option<V> {
        let h2val      = h2(hash);
        let num_groups = self.cap / 16;
        let start      = h1(hash, num_groups);

        for probe in 0..num_groups {
            let g        = (start + probe) % num_groups;
            let ctrl_ptr = unsafe { self.ctrl.add(g * 16) };
            let group    = Group(ctrl_ptr);

            for slot in each_set_bit(unsafe { group.match_byte(h2val) }) {
                let idx = g * 16 + slot;
                let (k, _) = unsafe { &*self.buckets.add(idx) };
                if k == key {
                    unsafe {
                        // DELETED 마킹 (EMPTY로 두면 뒤에 있는 항목을 못 찾음)
                        *self.ctrl.add(idx) = DELETED;
                        // 값의 소유권을 꺼내서 반환
                        let (_, v) = self.buckets.add(idx).read();
                        self.len -= 1;
                        return Some(v);
                    }
                }
            }
            if unsafe { group.match_empty() } != 0 { return None; }
        }
        None
    }
}

// 우리 구현: 탐색 2번 (없으면 삽입, 있으면 수정)
let hash = fxhash_of(&key);
if let Some(v) = map.get(&key, hash) {
    map.insert(key, v + 1, hash);  // 탐색 또 한 번
} else {
    map.insert(key, 1, hash);      // 탐색 또 한 번
}

// hashbrown / std HashMap: 탐색 1번
// Entry가 "이미 찾은 슬롯 포인터"를 들고 있음
*map.entry(key).or_insert(0) += 1;