//! # Interning arena.
//!
//! Additionally, if same value was present on the arena, doesn't add new object, returns old.

use std::hash::{BuildHasherDefault, Hash};

use hashbrown::hash_map::{HashMap, RawEntryMut};

use super::{ptr::Ptr, wonly};

// No raw_entry API on hashset.
type Bi<P> = HashMap<Mapping<P>, (), BuildHasherDefault<hasher::NoHash>>;

#[derive(Debug, Default, Clone, Copy)]
struct Mapping<P> {
    hash: u64,
    key: P,
}

impl<P> Hash for Mapping<P> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        // [`hasher::NoHash`] implementation ensures that
        // hasher stores `write_u64` result only once.
        state.write_u64(self.hash);
    }
}

#[derive(Clone)]
pub struct Arena<T, P> {
    objects: wonly::Arena<T, P>,
    bi: Bi<P>,
}

impl<T, P> Default for Arena<T, P> {
    fn default() -> Self {
        Self {
            objects: wonly::Arena::new(),
            bi: Bi::default(),
        }
    }
}

impl<T, P> Arena<T, P> {
    pub fn modify(&mut self, ptr: P, f: impl FnOnce(&mut T)) -> bool
    where
        T: Hash + Eq,
        P: Ptr,
    {
        let Some(object) = self.objects.get(ptr.clone()) else {
            return false;
        };
        let hash = fxhash::hash64(object);

        let entry = self
            .bi
            .raw_entry_mut()
            .from_hash(hash, |m| self.objects.get(m.key.clone()) == Some(object));

        // SAFETY: if we got here, the object is in arena.
        let object = unsafe { self.objects.get_mut(ptr.clone()).unwrap_unchecked() };
        f(object);

        let hash = fxhash::hash64(object);
        entry.and_modify(|k, _| *k = Mapping { hash, key: ptr });

        true
    }

    pub fn get(&self, ptr: P) -> Option<&T>
    where
        P: Ptr,
    {
        self.objects.get(ptr)
    }

    pub fn insert(&mut self, object: T) -> P
    where
        P: Ptr,
        T: Hash + Eq,
    {
        let hash = fxhash::hash64(&object);
        match self.bi.raw_entry_mut().from_hash(hash, |m| {
            self.objects
                .get(m.key.clone())
                .map_or(false, |lhs| lhs == &object)
        }) {
            RawEntryMut::Occupied(occup) => occup.into_key().key.clone(),
            RawEntryMut::Vacant(vacant) => {
                let key = self.objects.insert(object);
                vacant.insert_hashed_nocheck(
                    hash,
                    Mapping {
                        hash,
                        key: key.clone(),
                    },
                    (),
                );

                key
            }
        }
    }

    pub fn new() -> Self {
        Self::default()
    }
}

mod hasher;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::arena::ptr::define;

    define! {
        struct Token;
    }

    #[test]
    fn modification_works() {
        // Ensure that when modifying the value,
        // pointer is still valid.

        let mut arena: Arena<u32, Token> = Arena::new();
        let key100to200 = arena.insert(100);
        let modified = arena.modify(key100to200, |x| *x = 200);

        assert!(modified);
        assert_eq!(arena.get(key100to200), Some(&200));

        let key100new = arena.insert(100);
        assert_ne!(key100to200, key100new);

        assert_eq!(arena.get(key100new), Some(&100));
        assert_eq!(arena.get(key100to200), Some(&200));
    }

    #[test]
    fn it_interns() {
        let mut arena: Arena<u32, Token> = Arena::new();
        let key100 = arena.insert(100);
        let key100again = arena.insert(100);

        assert_eq!(key100, key100again);
    }
}