raii-map 0.1.0

//! [`Map`] and related items.

use alloc::collections::{
    btree_map::{Iter, IterMut, Keys, Values, ValuesMut},
    BTreeMap,
};
use core::{marker::PhantomData, usize};
use owner_monad::{Owner, OwnerMut};

/// A map in which the presence of an entry is tied to the lifetime of a RAII
/// handle returned by the insertion.
#[derive(Debug, PartialEq, Eq)]
pub struct Map<K, V>(BTreeMap<K, V>);

impl<K: Ord, V> Map<K, V> {
    /// Creates a new, empty map.
    pub fn new() -> Self {
        Self::default()
    }

    /// Gets the number of items in the map.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Checks if the map is empty.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Gets an iterator over the contents of the map.
    pub fn iter(&self) -> Iter<K, V> {
        self.0.iter()
    }

    /// Gets a mutable iterator over the contents of the map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        self.0.iter_mut()
    }

    /// Gets an iterator over the keys of the map.
    pub fn keys(&self) -> Keys<K, V> {
        self.0.keys()
    }

    /// Gets an iterator over the values of the map.
    pub fn values(&self) -> Values<K, V> {
        self.0.values()
    }

    /// Gets a mutable iterator over the values of the map.
    pub fn values_mut(&mut self) -> ValuesMut<K, V> {
        self.0.values_mut()
    }
}

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

impl<K: Ord, V> IntoIterator for Map<K, V> {
    type Item = (K, V);

    type IntoIter = <BTreeMap<K, V> as IntoIterator>::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<'a, K: Ord, V> IntoIterator for &'a Map<K, V> {
    type Item = (&'a K, &'a V);

    type IntoIter = <&'a BTreeMap<K, V> as IntoIterator>::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        #[allow(clippy::into_iter_on_ref)]
        (&self.0).into_iter()
    }
}

impl<'a, K: Ord, V> IntoIterator for &'a mut Map<K, V> {
    type Item = (&'a K, &'a mut V);

    type IntoIter = <&'a mut BTreeMap<K, V> as IntoIterator>::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        #[allow(clippy::into_iter_on_ref)]
        (&mut self.0).into_iter()
    }
}

/// A RAII handle to data in a [`Map`].
pub struct MapHandle<K: Ord, V, O: OwnerMut<Map<K, V>>> {
    key: K,
    owner: O,
    _phantom: PhantomData<V>,
}

impl<K: Ord, V, O: OwnerMut<Map<K, V>>> Owner<O> for MapHandle<K, V, O> {
    fn with<'a, U>(&'a self, f: impl FnOnce(&O) -> U) -> Option<U>
    where
        O: 'a,
    {
        Some(f(&self.owner))
    }
}

impl<K: Ord, V, O: OwnerMut<Map<K, V>>> Drop for MapHandle<K, V, O> {
    fn drop(&mut self) {
        let key = &self.key;
        self.owner.with(|map| map.0.remove(key));
    }
}

/// Inserts a key-value pair into a map, returning a handle to the data if
/// successful. A result of `None` indicates that either the map could not be
/// accessed or the key already exists.
pub fn insert<K: Clone + Ord, V, O: OwnerMut<Map<K, V>>>(
    mut owner: O,
    key: K,
    value: V,
) -> Option<MapHandle<K, V, O>> {
    let r = {
        let key = key.clone();
        owner.with(move |map| map.0.insert(key, value))?
    };
    if r.is_none() {
        Some(MapHandle {
            key,
            owner,
            _phantom: PhantomData,
        })
    } else {
        None
    }
}