Skip to content

Instantly share code, notes, and snippets.

@timvermeulen
Created September 15, 2020 22:05
Show Gist options
  • Save timvermeulen/85eb5f3b8733f9a0b8d8ad3e7fdbb20f to your computer and use it in GitHub Desktop.
Save timvermeulen/85eb5f3b8733f9a0b8d8ad3e7fdbb20f to your computer and use it in GitHub Desktop.
Vec entry
use std::cmp::Ordering;
pub trait VecExt {
type T;
fn find_entry<F>(&mut self, f: F) -> Entry<'_, Self::T>
where
F: FnMut(&Self::T) -> bool;
fn first_entry_where<F>(&mut self, f: F) -> Entry<'_, Self::T>
where
F: FnOnce(&mut Self::T) -> bool;
fn last_entry_where<F>(&mut self, f: F) -> Entry<'_, Self::T>
where
F: FnOnce(&mut Self::T) -> bool;
fn binary_search_entry(&mut self, value: &Self::T) -> Entry<'_, Self::T>
where
Self::T: Ord;
fn binary_search_entry_by<F>(&mut self, f: F) -> Entry<'_, Self::T>
where
F: FnMut(&Self::T) -> Ordering;
fn binary_search_entry_by_key<B, F>(&mut self, b: &B, f: F) -> Entry<'_, Self::T>
where
B: Ord,
F: FnMut(&Self::T) -> B;
}
impl<T> VecExt for Vec<T> {
type T = T;
fn find_entry<F>(&mut self, f: F) -> Entry<'_, Self::T>
where
F: FnMut(&Self::T) -> bool,
{
let (index, occupied) = match self.iter().position(f) {
Some(index) => (index, true),
None => (self.len(), false),
};
Entry::new(self, index, occupied)
}
fn binary_search_entry(&mut self, value: &T) -> Entry<'_, T>
where
Self::T: Ord,
{
let (index, occupied) = match self.binary_search(value) {
Ok(index) => (index, true),
Err(index) => (index, false),
};
Entry::new(self, index, occupied)
}
fn binary_search_entry_by<F>(&mut self, f: F) -> Entry<'_, T>
where
F: FnMut(&T) -> Ordering,
{
let (index, occupied) = match self.binary_search_by(f) {
Ok(index) => (index, true),
Err(index) => (index, false),
};
Entry::new(self, index, occupied)
}
fn binary_search_entry_by_key<B, F>(&mut self, b: &B, f: F) -> Entry<'_, Self::T>
where
B: Ord,
F: FnMut(&Self::T) -> B,
{
let (index, occupied) = match self.binary_search_by_key(b, f) {
Ok(index) => (index, true),
Err(index) => (index, false),
};
Entry::new(self, index, occupied)
}
fn first_entry_where<F>(&mut self, f: F) -> Entry<'_, Self::T>
where
F: FnOnce(&mut Self::T) -> bool,
{
let occupied = self.first_mut().map_or(false, f);
Entry::new(self, 0, occupied)
}
fn last_entry_where<F>(&mut self, f: F) -> Entry<'_, Self::T>
where
F: FnOnce(&mut Self::T) -> bool,
{
let occupied = self.last_mut().map_or(false, f);
let index = if occupied { self.len() - 1 } else { self.len() };
Entry::new(self, index, occupied)
}
}
pub enum Entry<'a, T> {
Occupied(OccupiedEntry<'a, T>),
Vacant(VacantEntry<'a, T>),
}
impl<'a, T> Entry<'a, T> {
fn new(vec: &'a mut Vec<T>, index: usize, occupied: bool) -> Self {
if occupied {
Self::Occupied(OccupiedEntry { vec, index })
} else {
Self::Vacant(VacantEntry { vec, index })
}
}
pub fn and_modify(mut self, f: impl FnOnce(&mut T)) -> Self {
if let Entry::Occupied(entry) = &mut self {
f(entry.get_mut());
}
self
}
pub fn or_insert(self, value: T) -> &'a mut T {
self.or_insert_with(|| value)
}
pub fn or_insert_with(self, value: impl FnOnce() -> T) -> &'a mut T {
match self {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(value()),
}
}
pub fn remove(self) -> Option<T> {
match self {
Entry::Occupied(entry) => Some(entry.remove()),
Entry::Vacant(_entry) => None,
}
}
pub fn swap_remove(self) -> Option<T> {
match self {
Entry::Occupied(entry) => Some(entry.swap_remove()),
Entry::Vacant(_entry) => None,
}
}
pub fn index(&self) -> usize {
match self {
Entry::Occupied(entry) => entry.index(),
Entry::Vacant(entry) => entry.index(),
}
}
}
pub struct OccupiedEntry<'a, T> {
vec: &'a mut Vec<T>,
index: usize,
}
impl<'a, T> OccupiedEntry<'a, T> {
pub fn get(&self) -> &T {
&self.vec[self.index]
}
pub fn get_mut(&mut self) -> &mut T {
&mut self.vec[self.index]
}
pub fn into_mut(self) -> &'a mut T {
&mut self.vec[self.index]
}
pub fn remove(self) -> T {
self.vec.remove(self.index)
}
pub fn swap_remove(self) -> T {
self.vec.swap_remove(self.index)
}
pub fn index(&self) -> usize {
self.index
}
}
pub struct VacantEntry<'a, T> {
vec: &'a mut Vec<T>,
index: usize,
}
impl<'a, T> VacantEntry<'a, T> {
pub fn insert(self, value: T) -> &'a mut T {
self.vec.insert(self.index, value);
&mut self.vec[self.index]
}
pub fn index(&self) -> usize {
self.index
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment