quota 0.2.1

use crate::{Quota, QuotaId, QuotaPool, QuotaRef};
use std::hash::Hash;

pub struct OccupiedQuotaPoolEntry<'a, K: Eq + Hash> {
    pub(crate) key: K,
    pub(crate) id: QuotaId,
    pub(crate) pool: &'a mut QuotaPool<K>,
}

pub struct VacantQuotaPoolEntry<'a, K: Eq + Hash> {
    pub(crate) key: K,
    pub(crate) pool: &'a mut QuotaPool<K>,
}

pub enum QuotaEntry<'a, K: Eq + Hash> {
    /// An occupied Quota entry.
    Occupied(OccupiedQuotaPoolEntry<'a, K>),
    /// A vacant Quota entry.
    Vacant(VacantQuotaPoolEntry<'a, K>),
}

impl<'a, K: Eq + Hash> QuotaEntry<'a, K> {
    #[inline]
    pub fn or_insert(self, default: Quota) -> QuotaRef<'a> {
        match self {
            Self::Occupied(entry) => entry.into_mut(),
            Self::Vacant(entry) => entry.insert(default),
        }
    }

    #[inline]
    pub fn or_insert_with<F>(self, default: F) -> QuotaRef<'a>
    where
        F: FnOnce() -> Quota,
    {
        match self {
            Self::Occupied(entry) => entry.into_mut(),
            Self::Vacant(entry) => entry.insert(default()),
        }
    }

    #[inline]
    pub fn or_insert_with_key<F>(self, default: F) -> QuotaRef<'a>
    where
        F: FnOnce(&K) -> Quota,
    {
        match self {
            Self::Occupied(entry) => entry.into_mut(),
            Self::Vacant(entry) => {
                let quota = default(entry.key());
                entry.insert(quota)
            }
        }
    }

    #[inline]
    pub fn or_default(self) -> QuotaRef<'a> {
        self.or_insert_with(Quota::new)
    }

    #[inline]
    pub fn key(&self) -> &K {
        match self {
            Self::Occupied(entry) => entry.key(),
            Self::Vacant(entry) => entry.key(),
        }
    }

    #[inline]
    pub fn and_modify<F>(self, f: F) -> Self
    where
        F: FnOnce(QuotaRef<'_>),
    {
        match self {
            Self::Occupied(entry) => {
                f(entry.get());
                Self::Occupied(entry)
            }
            Self::Vacant(entry) => Self::Vacant(entry),
        }
    }

    #[inline]
    pub fn insert_entry(self, value: Quota) -> OccupiedQuotaPoolEntry<'a, K>
    where
        K: Clone,
    {
        match self {
            Self::Occupied(mut entry) => {
                entry.insert(value);
                entry
            }
            Self::Vacant(entry) => entry.insert_entry(value),
        }
    }
}

impl<'a, K: Eq + Hash> OccupiedQuotaPoolEntry<'a, K> {
    #[inline]
    pub fn key(&self) -> &K {
        &self.key
    }

    #[inline]
    pub fn remove_entry(self) -> (K, Quota) {
        let Self { pool, key, id } = self;
        let (key, removed_id) = pool
            .key_map
            .remove_entry(&key)
            .expect("QuotaPool occupied entry key was missing from the key map");

        debug_assert_eq!(removed_id, id);

        let quota = pool
            .buffer
            .take(id)
            .expect("QuotaPool key map pointed at a stale QuotaId");

        (key, quota)
    }

    #[inline]
    pub fn get(&self) -> QuotaRef<'_> {
        QuotaRef::new(&self.pool.buffer, self.id)
    }

    #[inline]
    pub fn get_mut(&mut self) -> QuotaRef<'_> {
        QuotaRef::new(&self.pool.buffer, self.id)
    }

    #[inline]
    pub fn into_mut(self) -> QuotaRef<'a> {
        let Self { pool, id, .. } = self;
        QuotaRef::new(&pool.buffer, id)
    }

    #[inline]
    pub fn insert(&mut self, value: Quota) -> Quota {
        self.pool
            .buffer
            .replace(self.id, value)
            .expect("QuotaPool key map pointed at a stale QuotaId")
    }

    #[inline]
    pub fn remove(self) -> Quota {
        self.remove_entry().1
    }
}

impl<'a, K: Eq + Hash> VacantQuotaPoolEntry<'a, K> {
    #[inline]
    pub fn key(&self) -> &K {
        &self.key
    }

    #[inline]
    pub fn into_key(self) -> K {
        self.key
    }

    #[inline]
    pub fn insert(self, value: Quota) -> QuotaRef<'a> {
        let id = self.pool.buffer.push(value);
        let old_id = self.pool.key_map.insert(self.key, id);

        debug_assert!(old_id.is_none());

        QuotaRef::new(&self.pool.buffer, id)
    }

    #[inline]
    pub fn insert_entry(self, value: Quota) -> OccupiedQuotaPoolEntry<'a, K>
    where
        K: Clone,
    {
        let key = self.key.clone();
        let id = self.pool.buffer.push(value);
        let old_id = self.pool.key_map.insert(self.key, id);

        debug_assert!(old_id.is_none());

        OccupiedQuotaPoolEntry {
            pool: self.pool,
            key,
            id,
        }
    }
}