raii_map/
map.rs

1//! [`Map`] and related items.
2
3use alloc::collections::{
4    btree_map::{Iter, IterMut, Keys, Values, ValuesMut},
5    BTreeMap,
6};
7use core::{marker::PhantomData, usize};
8use owner_monad::{Owner, OwnerMut};
9
10/// A map in which the presence of an entry is tied to the lifetime of a RAII
11/// handle returned by the insertion.
12#[derive(Debug, PartialEq, Eq)]
13pub struct Map<K, V>(BTreeMap<K, V>);
14
15impl<K: Ord, V> Map<K, V> {
16    /// Creates a new, empty map.
17    pub fn new() -> Self {
18        Self::default()
19    }
20
21    /// Gets the number of items in the map.
22    pub fn len(&self) -> usize {
23        self.0.len()
24    }
25
26    /// Checks if the map is empty.
27    pub fn is_empty(&self) -> bool {
28        self.0.is_empty()
29    }
30
31    /// Gets an iterator over the contents of the map.
32    pub fn iter(&self) -> Iter<K, V> {
33        self.0.iter()
34    }
35
36    /// Gets a mutable iterator over the contents of the map.
37    pub fn iter_mut(&mut self) -> IterMut<K, V> {
38        self.0.iter_mut()
39    }
40
41    /// Gets an iterator over the keys of the map.
42    pub fn keys(&self) -> Keys<K, V> {
43        self.0.keys()
44    }
45
46    /// Gets an iterator over the values of the map.
47    pub fn values(&self) -> Values<K, V> {
48        self.0.values()
49    }
50
51    /// Gets a mutable iterator over the values of the map.
52    pub fn values_mut(&mut self) -> ValuesMut<K, V> {
53        self.0.values_mut()
54    }
55}
56
57impl<K: Ord, V> Default for Map<K, V> {
58    fn default() -> Self {
59        Self(Default::default())
60    }
61}
62
63impl<K: Ord, V> IntoIterator for Map<K, V> {
64    type Item = (K, V);
65
66    type IntoIter = <BTreeMap<K, V> as IntoIterator>::IntoIter;
67
68    fn into_iter(self) -> Self::IntoIter {
69        self.0.into_iter()
70    }
71}
72
73impl<'a, K: Ord, V> IntoIterator for &'a Map<K, V> {
74    type Item = (&'a K, &'a V);
75
76    type IntoIter = <&'a BTreeMap<K, V> as IntoIterator>::IntoIter;
77
78    fn into_iter(self) -> Self::IntoIter {
79        #[allow(clippy::into_iter_on_ref)]
80        (&self.0).into_iter()
81    }
82}
83
84impl<'a, K: Ord, V> IntoIterator for &'a mut Map<K, V> {
85    type Item = (&'a K, &'a mut V);
86
87    type IntoIter = <&'a mut BTreeMap<K, V> as IntoIterator>::IntoIter;
88
89    fn into_iter(self) -> Self::IntoIter {
90        #[allow(clippy::into_iter_on_ref)]
91        (&mut self.0).into_iter()
92    }
93}
94
95/// A RAII handle to data in a [`Map`].
96pub struct MapHandle<K: Ord, V, O: OwnerMut<Map<K, V>>> {
97    key: K,
98    owner: O,
99    _phantom: PhantomData<V>,
100}
101
102impl<K: Ord, V, O: OwnerMut<Map<K, V>>> Owner<O> for MapHandle<K, V, O> {
103    fn with<'a, U>(&'a self, f: impl FnOnce(&O) -> U) -> Option<U>
104    where
105        O: 'a,
106    {
107        Some(f(&self.owner))
108    }
109}
110
111impl<K: Ord, V, O: OwnerMut<Map<K, V>>> Drop for MapHandle<K, V, O> {
112    fn drop(&mut self) {
113        let key = &self.key;
114        self.owner.with(|map| map.0.remove(key));
115    }
116}
117
118/// Inserts a key-value pair into a map, returning a handle to the data if
119/// successful. A result of `None` indicates that either the map could not be
120/// accessed or the key already exists.
121pub fn insert<K: Clone + Ord, V, O: OwnerMut<Map<K, V>>>(
122    mut owner: O,
123    key: K,
124    value: V,
125) -> Option<MapHandle<K, V, O>> {
126    let r = {
127        let key = key.clone();
128        owner.with(move |map| map.0.insert(key, value))?
129    };
130    if r.is_none() {
131        Some(MapHandle {
132            key,
133            owner,
134            _phantom: PhantomData,
135        })
136    } else {
137        None
138    }
139}