total_maps/
btree_map.rs

1//! Provides [TotalBTreeMap], an ordered map in which every possible key has an associated value.
2//! Only entries with *uncommon* values are actually stored in the map; all other keys are presumed
3//! to be associated with a *common* value.
4
5use std::{
6    borrow::Borrow,
7    cmp::Ordering,
8    collections::{BTreeMap, btree_map},
9    fmt::{self, Debug, Formatter},
10    hash::{Hash, Hasher},
11    iter::FusedIterator,
12    mem,
13    ops::{Deref, DerefMut, Index},
14    sync::OnceLock,
15};
16
17use crate::{Commonality, DefaultCommonality, PhantomPtr};
18
19// --------------------------------------------------------------------------
20
21/// An ordered map in which every possible key has an associated value. Only entries with *uncommon*
22/// values are actually stored in the map; all other keys are presumed to be associated with a
23/// *common* value.
24///
25/// See the [crate documentation](crate) for more information.
26///
27/// The API is more-or-less a subset of that of [BTreeMap]. However, methods that treat this type
28/// like a collection (for example, [`len()`](Self::len) and [`iter()`](Self::iter)) operate only on
29/// the *uncommon* entries.
30pub struct TotalBTreeMap<K, V, C = DefaultCommonality> {
31    inner: BTreeMap<K, V>,
32    common: OnceLock<V>, // need to store this so we can return references to it, e.g., in Self::get
33    _commonality: PhantomPtr<C>,
34}
35
36impl<K: Clone, V: Clone, C> Clone for TotalBTreeMap<K, V, C> {
37    fn clone(&self) -> Self {
38        Self {
39            inner: self.inner.clone(),
40            common: self.common.clone(),
41            _commonality: PhantomPtr::default(),
42        }
43    }
44}
45
46impl<K, V, C> Default for TotalBTreeMap<K, V, C> {
47    fn default() -> Self {
48        Self {
49            inner: BTreeMap::default(),
50            common: OnceLock::new(),
51            _commonality: PhantomPtr::default(),
52        }
53    }
54}
55impl<K, V, C> TotalBTreeMap<K, V, C> {
56    /// Constructs a `TotalBTreeMap` in which all keys are associated with the *common* value.
57    pub fn new() -> Self {
58        Self::default()
59    }
60}
61
62impl<K, V, C: Commonality<V>> TotalBTreeMap<K, V, C> {
63    fn common(&self) -> &V {
64        self.common.get_or_init(C::common)
65    }
66}
67
68impl<K, V, C> TotalBTreeMap<K, V, C> {
69    /// Returns the number of *uncommon* entries in the map.
70    pub fn len(&self) -> usize {
71        self.inner.len()
72    }
73    /// Returns true if the map contains no *uncommon* entries.
74    pub fn is_empty(&self) -> bool {
75        self.inner.is_empty()
76    }
77    /// Resets all entries in the map to the *common* value.
78    pub fn clear(&mut self) {
79        self.inner.clear()
80    }
81}
82
83// --------------------------------------------------------------------------
84// Element access
85
86impl<K: Ord, V, C> TotalBTreeMap<K, V, C> {
87    /// Returns true if the map contains an *uncommon* entry with the given key.
88    pub fn contains_key<Q>(&self, key: &Q) -> bool
89    where
90        K: Borrow<Q>,
91        Q: Ord + ?Sized,
92    {
93        self.inner.contains_key(key)
94    }
95}
96
97impl<K: Borrow<Q> + Ord, Q: Ord + ?Sized, V, C: Commonality<V>> Index<&Q>
98    for TotalBTreeMap<K, V, C>
99{
100    type Output = V;
101    fn index(&self, index: &Q) -> &Self::Output {
102        self.get(index)
103    }
104}
105
106impl<K: Ord, V, C: Commonality<V>> TotalBTreeMap<K, V, C> {
107    /// Returns a reference to the value associated with the given key.
108    pub fn get<Q>(&self, key: &Q) -> &V
109    where
110        K: Borrow<Q>,
111        Q: Ord + ?Sized,
112    {
113        self.inner.get(key).unwrap_or_else(|| self.common())
114    }
115
116    /// Associates a key with a value in the map, and returns the value previously associated with
117    /// that key.
118    pub fn insert(&mut self, key: K, value: V) -> V {
119        if C::is_common(&value) { self.inner.remove(&key) } else { self.inner.insert(key, value) }
120            .unwrap_or_else(C::common)
121    }
122
123    /// Associates a key with the *common* value in the map, and returns the value previously
124    /// associated with that key.
125    pub fn remove<Q>(&mut self, key: &Q) -> V
126    where
127        K: Borrow<Q>,
128        Q: Ord + ?Sized,
129    {
130        self.inner.remove(key).unwrap_or_else(C::common)
131    }
132
133    /// Gets the given key's associated entry in the map for in-place manipulation.
134    pub fn entry(&mut self, key: K) -> Entry<'_, K, K, V, C> {
135        Entry {
136            inner: match self.inner.entry(key) {
137                btree_map::Entry::Occupied(inner) => EntryInner::Occupied { inner },
138                btree_map::Entry::Vacant(inner) => EntryInner::Vacant { inner, value: C::common() },
139            },
140        }
141    }
142
143    /// Gets the given key's associated entry in the map if it has an *uncommon* value; otherwise
144    /// returns [None].
145    ///
146    /// In contrast with [Self::entry], this method accepts the key in borrowed form.
147    pub fn uncommon_entry<'a, Q>(&'a mut self, key: &'a Q) -> Option<Entry<'a, Q, K, V, C>>
148    where
149        K: Borrow<Q>,
150        Q: Ord + ?Sized,
151    {
152        let map = self as *mut _;
153        let value = self.inner.get_mut(key)?;
154        Some(Entry { inner: EntryInner::ByRef { map, key, value } })
155    }
156}
157
158/// A view into a single entry in a [TotalBTreeMap].
159///
160/// This view is constructed from [TotalBTreeMap::entry].
161pub struct Entry<'a, Q, K, V, C = DefaultCommonality>
162where
163    Q: Ord + ?Sized,
164    K: Ord + Borrow<Q>,
165    C: Commonality<V>,
166{
167    inner: EntryInner<'a, Q, K, V, C>,
168}
169
170impl<Q, K, V, C> Deref for Entry<'_, Q, K, V, C>
171where
172    Q: Ord + ?Sized,
173    K: Ord + Borrow<Q>,
174    C: Commonality<V>,
175{
176    type Target = V;
177    fn deref(&self) -> &Self::Target {
178        match &self.inner {
179            EntryInner::Occupied { inner } => inner.get(),
180            EntryInner::Vacant { value, .. } => value,
181            EntryInner::ByRef { value, .. } => value,
182            EntryInner::Dropping => unreachable!(),
183        }
184    }
185}
186impl<Q, K, V, C> DerefMut for Entry<'_, Q, K, V, C>
187where
188    Q: Ord + ?Sized,
189    K: Ord + Borrow<Q>,
190    C: Commonality<V>,
191{
192    fn deref_mut(&mut self) -> &mut Self::Target {
193        match &mut self.inner {
194            EntryInner::Occupied { inner } => inner.get_mut(),
195            EntryInner::Vacant { value, .. } => value,
196            EntryInner::ByRef { value, .. } => value,
197            EntryInner::Dropping => unreachable!(),
198        }
199    }
200}
201
202impl<Q, K, V, C> Drop for Entry<'_, Q, K, V, C>
203where
204    Q: Ord + ?Sized,
205    K: Ord + Borrow<Q>,
206    C: Commonality<V>,
207{
208    fn drop(&mut self) {
209        match mem::replace(&mut self.inner, EntryInner::Dropping) {
210            EntryInner::Occupied { inner } => {
211                if C::is_common(inner.get()) {
212                    inner.remove();
213                }
214            }
215            EntryInner::Vacant { inner, value } => {
216                if !C::is_common(&value) {
217                    inner.insert(value);
218                }
219            }
220            EntryInner::ByRef { map, key, value } => {
221                if C::is_common(value) {
222                    unsafe { &mut *map }.remove(key);
223                }
224            }
225            EntryInner::Dropping => unreachable!(),
226        }
227    }
228}
229
230impl<'a, Q, K, V, C> Debug for Entry<'a, Q, K, V, C>
231where
232    Q: Debug + Ord + ?Sized,
233    K: Debug + Ord + Borrow<Q>,
234    V: Debug,
235    C: Commonality<V>,
236{
237    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
238        let mut f = f.debug_tuple("Entry");
239        match &self.inner {
240            EntryInner::Occupied { inner } => f.field(inner.key()).field(inner.get()),
241            EntryInner::Vacant { inner, value } => f.field(inner.key()).field(value),
242            EntryInner::ByRef { key, value, .. } => f.field(key).field(value),
243            EntryInner::Dropping => &mut f,
244        };
245        f.finish()
246    }
247}
248
249enum EntryInner<'a, Q: ?Sized, K, V, C> {
250    Occupied { inner: btree_map::OccupiedEntry<'a, K, V> },
251    Vacant { inner: btree_map::VacantEntry<'a, K, V>, value: V },
252    ByRef { map: *mut TotalBTreeMap<K, V, C>, key: &'a Q, value: &'a mut V },
253    Dropping,
254}
255
256// --------------------------------------------------------------------------
257// Iteration
258
259impl<K, V, C> TotalBTreeMap<K, V, C> {
260    /// An iterator over all keys associated with *uncommon* values in the map, in sorted order.
261    pub fn keys(&self) -> Keys<'_, K, V> {
262        Keys(self.inner.keys())
263    }
264    /// Creates a consuming iterator over all keys associated with *uncommon* values in the map, in
265    /// sorted order.
266    pub fn into_keys(self) -> IntoKeys<K, V> {
267        IntoKeys(self.inner.into_keys())
268    }
269    /// An iterator over all *uncommon* values in the map, in sorted order.
270    pub fn values(&self) -> Values<'_, K, V> {
271        Values(self.inner.values())
272    }
273    /// Creates a consuming iterator over all *uncommon* values in the map, in sorted order.
274    pub fn into_values(self) -> IntoValues<K, V> {
275        IntoValues(self.inner.into_values())
276    }
277    /// An iterator over all *uncommon* entries in the map, in sorted order.
278    pub fn iter(&self) -> Iter<'_, K, V> {
279        Iter(self.inner.iter())
280    }
281}
282
283impl<K, V, C> IntoIterator for TotalBTreeMap<K, V, C> {
284    type Item = (K, V);
285    type IntoIter = IntoIter<K, V>;
286    fn into_iter(self) -> Self::IntoIter {
287        IntoIter(self.inner.into_iter())
288    }
289}
290impl<'a, K, V, C> IntoIterator for &'a TotalBTreeMap<K, V, C> {
291    type Item = (&'a K, &'a V);
292    type IntoIter = Iter<'a, K, V>;
293    fn into_iter(self) -> Self::IntoIter {
294        self.iter()
295    }
296}
297
298/// An iterator over the keys associated with *uncommon* values in a [TotalBTreeMap].
299///
300/// This iterator is created by [TotalBTreeMap::keys].
301pub struct Keys<'a, K, V>(btree_map::Keys<'a, K, V>);
302impl<K, V> Clone for Keys<'_, K, V> {
303    fn clone(&self) -> Self {
304        Self(self.0.clone())
305    }
306}
307impl<'a, K, V> Default for Keys<'a, K, V> {
308    fn default() -> Self {
309        Self(Default::default())
310    }
311}
312impl<'a, K, V> Iterator for Keys<'a, K, V> {
313    type Item = &'a K;
314    fn next(&mut self) -> Option<Self::Item> {
315        self.0.next()
316    }
317    fn size_hint(&self) -> (usize, Option<usize>) {
318        self.0.size_hint()
319    }
320}
321impl<K, V> DoubleEndedIterator for Keys<'_, K, V> {
322    fn next_back(&mut self) -> Option<Self::Item> {
323        self.0.next_back()
324    }
325}
326impl<K, V> ExactSizeIterator for Keys<'_, K, V> {
327    fn len(&self) -> usize {
328        self.0.len()
329    }
330}
331impl<K, V> FusedIterator for Keys<'_, K, V> {}
332impl<K: Debug, V: Debug> Debug for Keys<'_, K, V> {
333    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
334        f.debug_list().entries(self.clone()).finish()
335    }
336}
337
338/// An owning iterator over the keys associated with *uncommon* values in a [TotalBTreeMap].
339///
340/// This iterator is created by [TotalBTreeMap::into_keys].
341pub struct IntoKeys<K, V>(btree_map::IntoKeys<K, V>);
342impl<K, V> Default for IntoKeys<K, V> {
343    fn default() -> Self {
344        Self(Default::default())
345    }
346}
347impl<K, V> Iterator for IntoKeys<K, V> {
348    type Item = K;
349    fn next(&mut self) -> Option<Self::Item> {
350        self.0.next()
351    }
352    fn size_hint(&self) -> (usize, Option<usize>) {
353        self.0.size_hint()
354    }
355}
356impl<K, V> DoubleEndedIterator for IntoKeys<K, V> {
357    fn next_back(&mut self) -> Option<Self::Item> {
358        self.0.next_back()
359    }
360}
361impl<K, V> ExactSizeIterator for IntoKeys<K, V> {
362    fn len(&self) -> usize {
363        self.0.len()
364    }
365}
366impl<K, V> FusedIterator for IntoKeys<K, V> {}
367
368/// An iterator over the *uncommon* values in a [TotalBTreeMap].
369///
370/// This iterator is created by [TotalBTreeMap::values].
371pub struct Values<'a, K, V>(btree_map::Values<'a, K, V>);
372impl<K, V> Clone for Values<'_, K, V> {
373    fn clone(&self) -> Self {
374        Self(self.0.clone())
375    }
376}
377impl<'a, K, V> Default for Values<'a, K, V> {
378    fn default() -> Self {
379        Self(Default::default())
380    }
381}
382impl<'a, K, V> Iterator for Values<'a, K, V> {
383    type Item = &'a V;
384    fn next(&mut self) -> Option<Self::Item> {
385        self.0.next()
386    }
387    fn size_hint(&self) -> (usize, Option<usize>) {
388        self.0.size_hint()
389    }
390}
391impl<'a, K, V> DoubleEndedIterator for Values<'a, K, V> {
392    fn next_back(&mut self) -> Option<Self::Item> {
393        self.0.next_back()
394    }
395}
396impl<K, V> ExactSizeIterator for Values<'_, K, V> {
397    fn len(&self) -> usize {
398        self.0.len()
399    }
400}
401impl<K, V> FusedIterator for Values<'_, K, V> {}
402impl<K: Debug, V: Debug> Debug for Values<'_, K, V> {
403    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
404        f.debug_list().entries(self.clone()).finish()
405    }
406}
407
408/// An owning iterator over the *uncommon* values in a [TotalBTreeMap].
409///
410/// This iterator is created by [TotalBTreeMap::into_values].
411pub struct IntoValues<K, V>(btree_map::IntoValues<K, V>);
412impl<K, V> Default for IntoValues<K, V> {
413    fn default() -> Self {
414        Self(Default::default())
415    }
416}
417impl<K, V> Iterator for IntoValues<K, V> {
418    type Item = V;
419    fn next(&mut self) -> Option<Self::Item> {
420        self.0.next()
421    }
422    fn size_hint(&self) -> (usize, Option<usize>) {
423        self.0.size_hint()
424    }
425}
426impl<K, V> DoubleEndedIterator for IntoValues<K, V> {
427    fn next_back(&mut self) -> Option<Self::Item> {
428        self.0.next_back()
429    }
430}
431impl<K, V> ExactSizeIterator for IntoValues<K, V> {
432    fn len(&self) -> usize {
433        self.0.len()
434    }
435}
436impl<K, V> FusedIterator for IntoValues<K, V> {}
437
438/// An iterator over the *uncommon* entries in a [TotalBTreeMap].
439///
440/// This iterator is created by [TotalBTreeMap::iter].
441pub struct Iter<'a, K, V>(btree_map::Iter<'a, K, V>);
442impl<K, V> Clone for Iter<'_, K, V> {
443    fn clone(&self) -> Self {
444        Self(self.0.clone())
445    }
446}
447impl<'a, K, V> Default for Iter<'a, K, V> {
448    fn default() -> Self {
449        Self(Default::default())
450    }
451}
452impl<'a, K, V> Iterator for Iter<'a, K, V> {
453    type Item = (&'a K, &'a V);
454    fn next(&mut self) -> Option<Self::Item> {
455        self.0.next()
456    }
457    fn size_hint(&self) -> (usize, Option<usize>) {
458        self.0.size_hint()
459    }
460}
461impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V> {
462    fn next_back(&mut self) -> Option<Self::Item> {
463        self.0.next_back()
464    }
465}
466impl<K, V> ExactSizeIterator for Iter<'_, K, V> {
467    fn len(&self) -> usize {
468        self.0.len()
469    }
470}
471impl<K, V> FusedIterator for Iter<'_, K, V> {}
472impl<K: Debug, V: Debug> Debug for Iter<'_, K, V> {
473    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
474        f.debug_list().entries(self.clone()).finish()
475    }
476}
477
478/// An owning iterator over the *uncommon* entries in a [TotalBTreeMap].
479///
480/// This iterator is created by [TotalBTreeMap]'s implementation of [IntoIterator].
481pub struct IntoIter<K, V>(btree_map::IntoIter<K, V>);
482impl<K, V> Default for IntoIter<K, V> {
483    fn default() -> Self {
484        Self(Default::default())
485    }
486}
487impl<K, V> Iterator for IntoIter<K, V> {
488    type Item = (K, V);
489    fn next(&mut self) -> Option<Self::Item> {
490        self.0.next()
491    }
492    fn size_hint(&self) -> (usize, Option<usize>) {
493        self.0.size_hint()
494    }
495}
496impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
497    fn next_back(&mut self) -> Option<Self::Item> {
498        self.0.next_back()
499    }
500}
501impl<K, V> ExactSizeIterator for IntoIter<K, V> {
502    fn len(&self) -> usize {
503        self.0.len()
504    }
505}
506impl<K, V> FusedIterator for IntoIter<K, V> {}
507
508// --------------------------------------------------------------------------
509// Population from iterators
510
511impl<K: Ord, V, C: Commonality<V>> Extend<(K, V)> for TotalBTreeMap<K, V, C> {
512    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
513        for (key, value) in iter {
514            self.insert(key, value);
515        }
516    }
517}
518impl<K: Ord, V, C: Commonality<V>> FromIterator<(K, V)> for TotalBTreeMap<K, V, C> {
519    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
520        let mut this = Self::default();
521        this.extend(iter);
522        this
523    }
524}
525
526// --------------------------------------------------------------------------
527// Low-level access
528
529impl<K, V, C> TotalBTreeMap<K, V, C> {
530    /// Returns a view into the underlying [BTreeMap] of a [TotalBTreeMap], which contains the
531    /// *uncommon* entries.
532    pub fn as_btree_map(&self) -> &BTreeMap<K, V> {
533        &self.inner
534    }
535}
536
537impl<K: Ord, V, C: Commonality<V>> TotalBTreeMap<K, V, C> {
538    /// Returns a mutable view into the underlying [BTreeMap] of a [TotalBTreeMap], from which
539    /// mutating iterators can be obtained by calling [BTreeMap::values_mut] or
540    /// [BTreeMap::iter_mut].
541    ///
542    /// By directly mutating the underlying [BTreeMap], it is possible to store *uncommon* entries
543    /// in the map temporarily. When the returned view is dropped, all *uncommon* entries will be
544    /// removed, restoring the invariant of [TotalBTreeMap].
545    ///
546    /// You don't need this method if you are only mutating individual entries; use the
547    /// [entry][Self::entry] method instead.
548    pub fn as_btree_map_mut(&mut self) -> AsBTreeMapMut<'_, K, V, C> {
549        AsBTreeMapMut { map: &mut self.inner, _commonality: PhantomPtr::default() }
550    }
551}
552
553/// A mutable view into the underlying [BTreeMap] of a [TotalBTreeMap].
554///
555/// This view is created by [TotalBTreeMap::as_btree_map_mut].
556pub struct AsBTreeMapMut<'a, K: Ord, V, C: Commonality<V> = DefaultCommonality> {
557    map: &'a mut BTreeMap<K, V>,
558    _commonality: PhantomPtr<C>,
559}
560
561impl<K: Ord, V, C: Commonality<V>> Deref for AsBTreeMapMut<'_, K, V, C> {
562    type Target = BTreeMap<K, V>;
563    fn deref(&self) -> &Self::Target {
564        self.map
565    }
566}
567impl<K: Ord, V, C: Commonality<V>> DerefMut for AsBTreeMapMut<'_, K, V, C> {
568    fn deref_mut(&mut self) -> &mut Self::Target {
569        self.map
570    }
571}
572
573impl<K: Ord, V, C: Commonality<V>> Drop for AsBTreeMapMut<'_, K, V, C> {
574    fn drop(&mut self) {
575        self.map.retain(|_, value| !C::is_common(value));
576    }
577}
578
579impl<K: Ord, V: PartialEq, C: Commonality<V>> PartialEq for AsBTreeMapMut<'_, K, V, C> {
580    fn eq(&self, other: &Self) -> bool {
581        self.map == other.map
582        // deliberately ignoring commonality
583    }
584}
585impl<K: Ord, V: Eq, C: Commonality<V>> Eq for AsBTreeMapMut<'_, K, V, C> {}
586impl<K: Ord, V: PartialOrd, C: Commonality<V>> PartialOrd for AsBTreeMapMut<'_, K, V, C> {
587    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
588        self.map.partial_cmp(&other.map)
589        // deliberately ignoring commonality
590    }
591}
592impl<K: Ord, V: Ord, C: Commonality<V>> Ord for AsBTreeMapMut<'_, K, V, C> {
593    fn cmp(&self, other: &Self) -> Ordering {
594        self.map.cmp(&other.map)
595        // deliberately ignoring commonality
596    }
597}
598impl<K: Ord + Hash, V: Hash, C: Commonality<V>> Hash for AsBTreeMapMut<'_, K, V, C> {
599    fn hash<H: Hasher>(&self, state: &mut H) {
600        self.map.hash(state);
601        // deliberately ignoring commonality
602    }
603}
604impl<K: Ord + Debug, V: Debug, C: Commonality<V>> Debug for AsBTreeMapMut<'_, K, V, C> {
605    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
606        f.debug_tuple("AsBTreeMapMut").field(&self.map).finish()
607    }
608}
609
610// --------------------------------------------------------------------------
611// Miscellaneous traits
612
613impl<K: PartialEq, V: PartialEq, C> PartialEq for TotalBTreeMap<K, V, C> {
614    fn eq(&self, other: &Self) -> bool {
615        self.inner == other.inner
616    }
617}
618impl<K: Eq, V: Eq, C> Eq for TotalBTreeMap<K, V, C> {}
619
620impl<K: PartialOrd, V: PartialOrd, C> PartialOrd for TotalBTreeMap<K, V, C> {
621    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
622        self.inner.partial_cmp(&other.inner)
623    }
624}
625impl<K: Ord, V: Ord, C> Ord for TotalBTreeMap<K, V, C> {
626    fn cmp(&self, other: &Self) -> Ordering {
627        self.inner.cmp(&other.inner)
628    }
629}
630
631impl<K: Hash, V: Hash, C> Hash for TotalBTreeMap<K, V, C> {
632    fn hash<H: Hasher>(&self, state: &mut H) {
633        self.inner.hash(state);
634    }
635}
636
637impl<K: Debug, V: Debug, C> Debug for TotalBTreeMap<K, V, C> {
638    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
639        f.debug_map().entries(self.iter()).finish_non_exhaustive()
640    }
641}
total_maps/btree_map.rs

total_maps/
btree_map.rs
