Vec entry

vec_entry.rs

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
    }
}