/* 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::collections::hash_map;
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>),
}
pub(super) enum Entry<'a, K, V> {
Occupied(&'a mut V),
Vacant(VacantEntry<'a, K, V>),
}
pub(super) struct VacantEntry<'a, K, V> {
inner: VacantEntryInner<'a, K, V>,
}
enum VacantEntryInner<'a, K, V> {
One(&'a mut MapInner<K, V>),
Map(hash_map::VacantEntry<'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 entry(
&mut self,
key: K,
key_from_value: impl FnOnce(&V) -> K,
) -> Entry<'_, K, V> {
match self.inner {
ref mut inner @ MapInner::Empty => Entry::Vacant(VacantEntry {
inner: VacantEntryInner::One(inner),
}),
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);
let one = match &mut self.inner {
MapInner::One(one) => one,
_ => unreachable!(),
};
return Entry::Occupied(one);
}
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);
match map.entry(key) {
hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
inner: VacantEntryInner::Map(entry),
}),
_ => unreachable!(),
}
},
MapInner::Map(ref mut map) => match map.entry(key) {
hash_map::Entry::Occupied(entry) => Entry::Occupied(entry.into_mut()),
hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
inner: VacantEntryInner::Map(entry),
}),
},
}
}
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<'a, K, V> VacantEntry<'a, K, V> {
pub(super) fn insert(self, value: V) -> &'a mut V {
match self.inner {
VacantEntryInner::One(map) => {
*map = MapInner::One(value);
match map {
MapInner::One(one) => one,
_ => unreachable!(),
}
},
VacantEntryInner::Map(entry) => entry.insert(value),
}
}
}
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,
}
}
}
