use std::{fmt::Debug, iter::FusedIterator, mem::replace};
#[derive(Clone)]
pub struct SlabMap<T> {
entries: Vec<Entry<T>>,
next_vacant_idx: usize,
len: usize,
non_optimized: usize,
}
const INVALID_INDEX: usize = usize::MAX;
#[derive(Clone)]
enum Entry<T> {
Occupied(T),
VacantHead { vacant_body_len: usize },
VacantTail { next_vacant_idx: usize },
}
impl<T> SlabMap<T> {
#[inline]
pub fn new() -> Self {
Self {
entries: Vec::new(),
next_vacant_idx: INVALID_INDEX,
len: 0,
non_optimized: 0,
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self {
entries: Vec::with_capacity(capacity),
next_vacant_idx: INVALID_INDEX,
len: 0,
non_optimized: 0,
}
}
#[inline]
pub fn capacity(&self) -> usize {
self.entries.capacity()
}
#[inline]
pub fn reserve(&mut self, additional: usize) {
self.entries.reserve(self.entries_additional(additional));
}
#[inline]
pub fn reserve_exact(&mut self, additional: usize) {
self.entries
.reserve_exact(self.entries_additional(additional));
}
#[inline]
fn entries_additional(&self, additional: usize) -> usize {
additional.saturating_sub(self.entries.len() - self.len)
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn get(&self, key: usize) -> Option<&T> {
if let Entry::Occupied(value) = self.entries.get(key)? {
Some(value)
} else {
None
}
}
#[inline]
pub fn get_mut(&mut self, key: usize) -> Option<&mut T> {
if let Entry::Occupied(value) = self.entries.get_mut(key)? {
Some(value)
} else {
None
}
}
#[inline]
pub fn contains_key(&self, key: usize) -> bool {
self.get(key).is_some()
}
pub fn insert(&mut self, value: T) -> usize {
self.insert_with_key(|_| value)
}
#[inline]
pub fn insert_with_key(&mut self, f: impl FnOnce(usize) -> T) -> usize {
let idx;
if self.next_vacant_idx < self.entries.len() {
idx = self.next_vacant_idx;
self.next_vacant_idx = match self.entries[idx] {
Entry::VacantHead { vacant_body_len } => {
if vacant_body_len > 0 {
self.entries[idx + 1] = Entry::VacantHead {
vacant_body_len: vacant_body_len - 1,
};
}
idx + 1
}
Entry::VacantTail { next_vacant_idx } => next_vacant_idx,
Entry::Occupied(_) => unreachable!(),
};
self.entries[idx] = Entry::Occupied(f(idx));
self.non_optimized = self.non_optimized.saturating_sub(1);
} else {
idx = self.entries.len();
self.entries.push(Entry::Occupied(f(idx)));
}
self.len += 1;
idx
}
pub fn remove(&mut self, key: usize) -> Option<T> {
let is_last = key + 1 == self.entries.len();
let e = self.entries.get_mut(key)?;
if !matches!(e, Entry::Occupied(..)) {
return None;
}
self.len -= 1;
let e = if is_last {
self.entries.pop().unwrap()
} else {
let e = replace(
e,
Entry::VacantTail {
next_vacant_idx: self.next_vacant_idx,
},
);
self.next_vacant_idx = key;
self.non_optimized += 1;
e
};
if self.is_empty() {
self.clear();
}
if let Entry::Occupied(value) = e {
Some(value)
} else {
unreachable!()
}
}
pub fn clear(&mut self) {
self.entries.clear();
self.len = 0;
self.next_vacant_idx = INVALID_INDEX;
self.non_optimized = 0;
}
pub fn drain(&mut self) -> Drain<T> {
let len = self.len;
self.len = 0;
self.next_vacant_idx = INVALID_INDEX;
self.non_optimized = 0;
Drain {
iter: self.entries.drain(..),
len,
}
}
pub fn retain(&mut self, f: impl FnMut(usize, &mut T) -> bool) {
let mut f = f;
let mut idx = 0;
let mut idx_vacant_start = 0;
self.next_vacant_idx = INVALID_INDEX;
while let Some(e) = self.entries.get_mut(idx) {
match e {
Entry::VacantTail { .. } => {
idx += 1;
}
Entry::VacantHead { vacant_body_len } => {
idx += *vacant_body_len + 2;
}
Entry::Occupied(value) => {
if f(idx, value) {
self.merge_vacant(idx_vacant_start, idx);
idx += 1;
idx_vacant_start = idx;
} else {
self.entries[idx] = Entry::VacantTail {
next_vacant_idx: INVALID_INDEX,
};
idx += 1;
}
}
}
}
self.entries.truncate(idx_vacant_start);
self.non_optimized = 0;
}
pub fn optimize(&mut self) {
if !self.is_optimized() {
self.retain(|_, _| true);
}
}
#[inline]
fn is_optimized(&self) -> bool {
self.non_optimized == 0
}
fn merge_vacant(&mut self, start: usize, end: usize) {
if start < end {
if start < end - 1 {
self.entries[start] = Entry::VacantHead {
vacant_body_len: end - start - 2,
}
}
self.entries[end - 1] = Entry::VacantTail {
next_vacant_idx: self.next_vacant_idx,
};
self.next_vacant_idx = start;
}
}
#[inline]
pub fn iter(&self) -> Iter<T> {
Iter {
iter: self.entries.iter().enumerate(),
len: self.len,
}
}
#[inline]
pub fn iter_mut(&mut self) -> IterMut<T> {
IterMut {
iter: self.entries.iter_mut().enumerate(),
len: self.len,
}
}
#[inline]
pub fn keys(&self) -> Keys<T> {
Keys(self.iter())
}
#[inline]
pub fn values(&self) -> Values<T> {
Values(self.iter())
}
#[inline]
pub fn values_mut(&mut self) -> ValuesMut<T> {
ValuesMut(self.iter_mut())
}
}
impl<T> Default for SlabMap<T> {
#[inline]
fn default() -> Self {
Self::new()
}
}
impl<T: Debug> Debug for SlabMap<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_map().entries(self.iter()).finish()
}
}
impl<T> std::ops::Index<usize> for SlabMap<T> {
type Output = T;
#[inline]
fn index(&self, index: usize) -> &Self::Output {
self.get(index).expect("out of index.")
}
}
impl<T> std::ops::IndexMut<usize> for SlabMap<T> {
#[inline]
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
self.get_mut(index).expect("out of index.")
}
}
impl<T> IntoIterator for SlabMap<T> {
type Item = T;
type IntoIter = IntoIter<T>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
IntoIter {
iter: self.entries.into_iter(),
len: self.len,
}
}
}
impl<'a, T> IntoIterator for &'a SlabMap<T> {
type Item = (usize, &'a T);
type IntoIter = Iter<'a, T>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, T> IntoIterator for &'a mut SlabMap<T> {
type Item = (usize, &'a mut T);
type IntoIter = IterMut<'a, T>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
pub struct IntoIter<T> {
iter: std::vec::IntoIter<Entry<T>>,
len: usize,
}
impl<T> Iterator for IntoIter<T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let mut e_opt = self.iter.next();
while let Some(e) = e_opt {
e_opt = match e {
Entry::Occupied(value) => {
self.len -= 1;
return Some(value);
}
Entry::VacantHead { vacant_body_len } => self.iter.nth(vacant_body_len + 1),
Entry::VacantTail { .. } => self.iter.next(),
}
}
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.len
}
}
pub struct Drain<'a, T> {
iter: std::vec::Drain<'a, Entry<T>>,
len: usize,
}
impl<'a, T> Iterator for Drain<'a, T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let mut e_opt = self.iter.next();
while let Some(e) = e_opt {
e_opt = match e {
Entry::Occupied(value) => {
self.len -= 1;
return Some(value);
}
Entry::VacantHead { vacant_body_len } => self.iter.nth(vacant_body_len + 1),
Entry::VacantTail { .. } => self.iter.next(),
}
}
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.len
}
}
pub struct Iter<'a, T> {
iter: std::iter::Enumerate<std::slice::Iter<'a, Entry<T>>>,
len: usize,
}
impl<'a, T> Iterator for Iter<'a, T> {
type Item = (usize, &'a T);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let mut e_opt = self.iter.next();
while let Some(e) = e_opt {
e_opt = match e {
(key, Entry::Occupied(value)) => {
self.len -= 1;
return Some((key, value));
}
(_, Entry::VacantHead { vacant_body_len }) => self.iter.nth(*vacant_body_len + 1),
(_, Entry::VacantTail { .. }) => self.iter.next(),
}
}
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.len
}
}
impl<'a, T> FusedIterator for Iter<'a, T> {}
impl<'a, T> ExactSizeIterator for Iter<'a, T> {}
pub struct IterMut<'a, T> {
iter: std::iter::Enumerate<std::slice::IterMut<'a, Entry<T>>>,
len: usize,
}
impl<'a, T> Iterator for IterMut<'a, T> {
type Item = (usize, &'a mut T);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let mut e_opt = self.iter.next();
while let Some(e) = e_opt {
e_opt = match e {
(key, Entry::Occupied(value)) => {
self.len -= 1;
return Some((key, value));
}
(_, Entry::VacantHead { vacant_body_len }) => self.iter.nth(*vacant_body_len + 1),
(_, Entry::VacantTail { .. }) => self.iter.next(),
}
}
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.len
}
}
impl<'a, T> FusedIterator for IterMut<'a, T> {}
impl<'a, T> ExactSizeIterator for IterMut<'a, T> {}
pub struct Keys<'a, T>(Iter<'a, T>);
impl<'a, T> Iterator for Keys<'a, T> {
type Item = usize;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|(k, _)| k)
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.0.count()
}
}
impl<'a, T> FusedIterator for Keys<'a, T> {}
impl<'a, T> ExactSizeIterator for Keys<'a, T> {}
pub struct Values<'a, T>(Iter<'a, T>);
impl<'a, T> Iterator for Values<'a, T> {
type Item = &'a T;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|(_, v)| v)
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.0.count()
}
}
impl<'a, T> FusedIterator for Values<'a, T> {}
impl<'a, T> ExactSizeIterator for Values<'a, T> {}
pub struct ValuesMut<'a, T>(IterMut<'a, T>);
impl<'a, T> Iterator for ValuesMut<'a, T> {
type Item = &'a mut T;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|(_, v)| v)
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
#[inline]
fn count(self) -> usize
where
Self: Sized,
{
self.0.count()
}
}
impl<'a, T> FusedIterator for ValuesMut<'a, T> {}
impl<'a, T> ExactSizeIterator for ValuesMut<'a, T> {}