/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

#![forbid(unsafe_code)]

use fxhash::FxHashMap;
use malloc_size_of::{MallocShallowSizeOf, MallocSizeOfOps};
use std::hash::Hash;
use std::mem;

pub(super) struct Map<K, V> {
    inner: MapInner<K, V>,
}

enum MapInner<K, V> {
    Empty,
    One(V),
    Map(Box<FxHashMap<K, V>>),
}

pub(super) struct MapIter<'a, K, V> {
    inner: MapIterInner<'a, K, V>,
}

enum MapIterInner<'a, K, V> {
    One(std::option::IntoIter<&'a V>),
    Map(std::collections::hash_map::Values<'a, K, V>),
}

impl<K, V> Default for Map<K, V> {
    fn default() -> Self {
        Map {
            inner: MapInner::Empty,
        }
    }
}

impl<'a, K, V> IntoIterator for &'a Map<K, V> {
    type Item = &'a V;
    type IntoIter = MapIter<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        MapIter {
            inner: match &self.inner {
                MapInner::Empty => MapIterInner::One(None.into_iter()),
                MapInner::One(one) => MapIterInner::One(Some(one).into_iter()),
                MapInner::Map(map) => MapIterInner::Map(map.values()),
            },
        }
    }
}

impl<'a, K, V> Iterator for MapIter<'a, K, V> {
    type Item = &'a V;

    fn next(&mut self) -> Option<Self::Item> {
        match &mut self.inner {
            MapIterInner::One(one_iter) => one_iter.next(),
            MapIterInner::Map(map_iter) => map_iter.next(),
        }
    }
}

impl<K, V> Map<K, V>
where
    K: Eq + Hash,
{
    pub(super) fn is_empty(&self) -> bool {
        match &self.inner {
            MapInner::Empty => true,
            MapInner::One(_) => false,
            MapInner::Map(map) => map.is_empty(),
        }
    }

    #[cfg(debug_assertions)]
    pub(super) fn len(&self) -> usize {
        match &self.inner {
            MapInner::Empty => 0,
            MapInner::One(_) => 1,
            MapInner::Map(map) => map.len(),
        }
    }

    pub(super) fn get(&self, key: &K, key_from_value: impl FnOnce(&V) -> K) -> Option<&V> {
        match &self.inner {
            MapInner::One(one) if *key == key_from_value(one) => Some(one),
            MapInner::Map(map) => map.get(key),
            MapInner::Empty | MapInner::One(_) => None,
        }
    }

    pub(super) fn get_or_insert_with(
        &mut self,
        key: K,
        key_from_value: impl FnOnce(&V) -> K,
        new_value: impl FnOnce() -> V,
    ) -> &mut V {
        match self.inner {
            MapInner::Empty => {
                self.inner = MapInner::One(new_value());
                match &mut self.inner {
                    MapInner::One(one) => one,
                    _ => unreachable!(),
                }
            },
            MapInner::One(_) => {
                let one = match mem::replace(&mut self.inner, MapInner::Empty) {
                    MapInner::One(one) => one,
                    _ => unreachable!(),
                };
                // If this panics, the child `one` will be lost.
                let one_key = key_from_value(&one);
                // Same for the equality test.
                if key == one_key {
                    self.inner = MapInner::One(one);
                    match &mut self.inner {
                        MapInner::One(one) => return one,
                        _ => unreachable!(),
                    }
                }
                self.inner = MapInner::Map(Box::new(FxHashMap::with_capacity_and_hasher(
                    2,
                    Default::default(),
                )));
                let map = match &mut self.inner {
                    MapInner::Map(map) => map,
                    _ => unreachable!(),
                };
                map.insert(one_key, one);
                // But it doesn't matter if f panics, by this point
                // the map is as before but represented as a map instead
                // of a single value.
                map.entry(key).or_insert_with(new_value)
            },
            MapInner::Map(ref mut map) => map.entry(key).or_insert_with(new_value),
        }
    }

    pub(super) fn remove(&mut self, key: &K, key_from_value: impl FnOnce(&V) -> K) -> Option<V> {
        match &mut self.inner {
            MapInner::One(one) if *key == key_from_value(one) => {
                match mem::replace(&mut self.inner, MapInner::Empty) {
                    MapInner::One(one) => Some(one),
                    _ => unreachable!(),
                }
            },
            MapInner::Map(map) => map.remove(key),
            MapInner::Empty | MapInner::One(_) => None,
        }
    }
}

impl<K, V> MallocShallowSizeOf for Map<K, V>
where
    K: Eq + Hash,
{
    fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
        match &self.inner {
            MapInner::Map(m) => {
                // We want to account for both the box and the hashmap.
                m.shallow_size_of(ops) + (**m).shallow_size_of(ops)
            },
            MapInner::One(_) | MapInner::Empty => 0,
        }
    }
}