weak_map/
map.rs

1//! `BTreeMap` with weak references.
2
3use alloc::collections::btree_map;
4use core::{
5    borrow::Borrow,
6    fmt,
7    iter::FusedIterator,
8    sync::atomic::{AtomicUsize, Ordering},
9};
10
11use crate::{StrongRef, WeakRef};
12
13#[derive(Default)]
14struct OpsCounter(AtomicUsize);
15
16const OPS_THRESHOLD: usize = 1000;
17
18impl OpsCounter {
19    #[inline]
20    const fn new() -> Self {
21        Self(AtomicUsize::new(0))
22    }
23
24    #[inline]
25    fn add(&self, ops: usize) {
26        self.0.fetch_add(ops, Ordering::Relaxed);
27    }
28
29    #[inline]
30    fn bump(&self) {
31        self.add(1);
32    }
33
34    #[inline]
35    fn reset(&mut self) {
36        *self.0.get_mut() = 0;
37    }
38
39    #[inline]
40    fn get(&self) -> usize {
41        self.0.load(Ordering::Relaxed)
42    }
43
44    #[inline]
45    fn reach_threshold(&self) -> bool {
46        self.get() >= OPS_THRESHOLD
47    }
48}
49
50impl Clone for OpsCounter {
51    #[inline]
52    fn clone(&self) -> Self {
53        Self(AtomicUsize::new(self.get()))
54    }
55}
56
57/// Alias for `BTreeMap<K, V>`.
58pub type StrongMap<K, V> = btree_map::BTreeMap<K, V>;
59
60/// A B-Tree map that stores weak references to values.
61#[derive(Clone)]
62pub struct WeakMap<K, V> {
63    inner: btree_map::BTreeMap<K, V>,
64    ops: OpsCounter,
65}
66
67impl<K, V> WeakMap<K, V> {
68    /// Makes a new, empty `WeakMap`.
69    ///
70    /// Does not allocate anything on its own.
71    #[inline]
72    pub const fn new() -> Self {
73        Self {
74            inner: btree_map::BTreeMap::new(),
75            ops: OpsCounter::new(),
76        }
77    }
78}
79
80impl<K, V> Default for WeakMap<K, V> {
81    fn default() -> Self {
82        Self::new()
83    }
84}
85
86impl<K, V> From<btree_map::BTreeMap<K, V>> for WeakMap<K, V> {
87    #[inline]
88    fn from(inner: btree_map::BTreeMap<K, V>) -> Self {
89        Self {
90            inner,
91            ops: OpsCounter::new(),
92        }
93    }
94}
95
96impl<K, V> From<WeakMap<K, V>> for btree_map::BTreeMap<K, V> {
97    #[inline]
98    fn from(map: WeakMap<K, V>) -> Self {
99        map.inner
100    }
101}
102
103impl<K, V> WeakMap<K, V> {
104    /// Clears the map, removing all elements.
105    #[inline]
106    pub fn clear(&mut self) {
107        self.inner.clear();
108        self.ops.reset();
109    }
110
111    /// Returns the number of elements in the underlying map.
112    #[must_use]
113    pub fn raw_len(&self) -> usize {
114        self.inner.len()
115    }
116
117    /// Gets an iterator over the entries of the map, sorted by key.
118    #[inline]
119    pub fn iter(&self) -> Iter<'_, K, V> {
120        self.ops.add(self.inner.len());
121        Iter(self.inner.iter())
122    }
123
124    /// Gets an iterator over the keys of the map, in sorted order.
125    #[inline]
126    pub fn keys(&self) -> Keys<'_, K, V> {
127        Keys(self.iter())
128    }
129
130    /// Creates a consuming iterator visiting all the keys, in sorted order.
131    /// The map cannot be used after calling this.
132    #[inline]
133    pub fn into_keys(self) -> IntoKeys<K, V> {
134        IntoKeys(IntoIter(self.inner.into_iter()))
135    }
136
137    /// Gets an iterator over the values of the map, in order by key.
138    #[inline]
139    pub fn values(&self) -> Values<'_, K, V> {
140        Values(self.iter())
141    }
142
143    /// Creates a consuming iterator visiting all the values, in order by key.
144    /// The map cannot be used after calling this.
145    #[inline]
146    pub fn into_values(self) -> IntoValues<K, V> {
147        IntoValues(IntoIter(self.inner.into_iter()))
148    }
149}
150
151impl<K, V> WeakMap<K, V>
152where
153    K: Ord,
154    V: WeakRef,
155{
156    /// Cleans up the map by removing expired values.
157    ///
158    /// Usually you don't need to call this manually, as it is called
159    /// automatically when the number of operations reaches a threshold.
160    #[inline]
161    pub fn cleanup(&mut self) {
162        self.ops.reset();
163        self.inner.retain(|_, v| !v.is_expired());
164    }
165
166    #[inline]
167    fn try_bump(&mut self) {
168        self.ops.bump();
169        if self.ops.reach_threshold() {
170            self.cleanup();
171        }
172    }
173
174    /// Returns the number of elements in the map, excluding expired values.
175    ///
176    /// This is a linear operation, as it iterates over all elements in the map.
177    ///
178    /// The returned value may be less than the result of [`Self::raw_len`].
179    #[inline]
180    pub fn len(&self) -> usize {
181        self.iter().count()
182    }
183
184    /// Returns `true` if the map contains no valid elements.
185    #[inline]
186    pub fn is_empty(&self) -> bool {
187        self.len() == 0
188    }
189
190    /// Retains only the elements specified by the predicate.
191    #[inline]
192    pub fn retain<F>(&mut self, mut f: F)
193    where
194        F: FnMut(&K, V::Strong) -> bool,
195    {
196        self.ops.reset();
197        self.inner.retain(|k, v| {
198            if let Some(v) = v.upgrade() {
199                f(k, v)
200            } else {
201                false
202            }
203        });
204    }
205
206    /// Returns a reference to the value corresponding to the key.
207    ///
208    /// The key may be any borrowed form of the map's key type, but the ordering
209    /// on the borrowed form *must* match the ordering on the key type.
210    pub fn get<Q>(&self, key: &Q) -> Option<V::Strong>
211    where
212        K: Borrow<Q>,
213        Q: Ord + ?Sized,
214    {
215        self.ops.bump();
216        self.inner.get(key).and_then(V::upgrade)
217    }
218
219    /// Returns the key-value pair corresponding to the supplied key. This is
220    /// potentially useful:
221    /// - for key types where non-identical keys can be considered equal;
222    /// - for getting the `&K` stored key value from a borrowed `&Q` lookup key; or
223    /// - for getting a reference to a key with the same lifetime as the collection.
224    ///
225    /// The supplied key may be any borrowed form of the map's key type, but the ordering
226    /// on the borrowed form *must* match the ordering on the key type.
227    pub fn get_key_value<Q>(&self, key: &Q) -> Option<(&K, V::Strong)>
228    where
229        K: Borrow<Q>,
230        Q: Ord + ?Sized,
231    {
232        self.ops.bump();
233        self.inner
234            .get_key_value(key)
235            .and_then(|(k, v)| v.upgrade().map(|v| (k, v)))
236    }
237
238    /// Returns `true` if the map contains a value for the specified key.
239    ///
240    /// The key may be any borrowed form of the map's key type, but the ordering
241    /// on the borrowed form *must* match the ordering on the key type.
242    pub fn contains_key<Q>(&self, key: &Q) -> bool
243    where
244        K: Borrow<Q>,
245        Q: Ord + ?Sized,
246    {
247        self.ops.bump();
248        self.inner.get(key).is_some_and(|v| !v.is_expired())
249    }
250
251    /// Inserts a key-value pair into the map.
252    ///
253    /// If the map did not have this key present, `None` is returned.
254    ///
255    /// If the map did have this key present, the value is updated, and the old
256    /// value is returned. The key is not updated, though; this matters for
257    /// types that can be `==` without being identical. See the [module-level
258    /// documentation] for more.
259    ///
260    /// [module-level documentation]: https://doc.rust-lang.org/std/collections/index.html#insert-and-complex-keys
261    pub fn insert(&mut self, key: K, value: &V::Strong) -> Option<V::Strong> {
262        self.try_bump();
263        self.inner
264            .insert(key, V::Strong::downgrade(value))
265            .and_then(|v| v.upgrade())
266    }
267
268    /// Removes a key from the map, returning the value at the key if the key
269    /// was previously in the map.
270    ///
271    /// The key may be any borrowed form of the map's key type, but the ordering
272    /// on the borrowed form *must* match the ordering on the key type.
273    pub fn remove<Q>(&mut self, key: &Q) -> Option<V::Strong>
274    where
275        K: Borrow<Q>,
276        Q: Ord + ?Sized,
277    {
278        self.try_bump();
279        self.inner.remove(key).and_then(|v| v.upgrade())
280    }
281
282    /// Removes a key from the map, returning the stored key and value if the key
283    /// was previously in the map.
284    ///
285    /// The key may be any borrowed form of the map's key type, but the ordering
286    /// on the borrowed form *must* match the ordering on the key type.
287    pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V::Strong)>
288    where
289        K: Borrow<Q>,
290        Q: Ord + ?Sized,
291    {
292        self.try_bump();
293        self.inner
294            .remove_entry(key)
295            .and_then(|(k, v)| v.upgrade().map(|v| (k, v)))
296    }
297
298    /// Gets a mutable iterator over the entries of the map, sorted by key.
299    #[inline]
300    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
301        self.ops.add(self.inner.len());
302        if self.ops.reach_threshold() {
303            self.cleanup();
304        }
305        IterMut(self.inner.iter_mut())
306    }
307
308    /// Upgrade this `WeakMap` to a `StrongMap`.
309    pub fn upgrade(&self) -> StrongMap<K, V::Strong>
310    where
311        K: Clone,
312    {
313        self.ops.bump();
314        let mut map = StrongMap::new();
315        for (key, value) in self.iter() {
316            map.insert(key.clone(), value);
317        }
318        map
319    }
320}
321
322impl<K, V> PartialEq for WeakMap<K, V>
323where
324    K: Ord,
325    V: WeakRef,
326{
327    fn eq(&self, other: &Self) -> bool {
328        self.iter().all(|(key, value)| {
329            other
330                .get(key)
331                .is_some_and(|v| V::Strong::ptr_eq(&value, &v))
332        })
333    }
334}
335
336impl<K, V> Eq for WeakMap<K, V>
337where
338    K: Ord,
339    V: WeakRef,
340{
341}
342
343impl<K, V> fmt::Debug for WeakMap<K, V>
344where
345    K: fmt::Debug,
346    V: WeakRef,
347    V::Strong: fmt::Debug,
348{
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350        f.debug_map().entries(self.iter()).finish()
351    }
352}
353
354impl<'a, K, V> FromIterator<(K, &'a V::Strong)> for WeakMap<K, V>
355where
356    K: Ord,
357    V: WeakRef,
358{
359    #[inline]
360    fn from_iter<T: IntoIterator<Item = (K, &'a V::Strong)>>(iter: T) -> Self {
361        let iter = iter.into_iter();
362        let mut map = WeakMap::new();
363        for (key, value) in iter {
364            map.insert(key, value);
365        }
366        map
367    }
368}
369
370impl<K, V, const N: usize> From<[(K, &V::Strong); N]> for WeakMap<K, V>
371where
372    K: Ord,
373    V: WeakRef,
374{
375    #[inline]
376    fn from(array: [(K, &V::Strong); N]) -> Self {
377        array.into_iter().collect()
378    }
379}
380
381impl<K, V> From<&StrongMap<K, V::Strong>> for WeakMap<K, V>
382where
383    K: Ord + Clone,
384    V: WeakRef,
385{
386    fn from(value: &StrongMap<K, V::Strong>) -> Self {
387        let mut map = WeakMap::new();
388        for (key, value) in value.iter() {
389            map.insert(key.clone(), value);
390        }
391        map
392    }
393}
394
395/// An iterator over the entries of a `WeakMap`.
396#[must_use = "iterators are lazy and do nothing unless consumed"]
397pub struct Iter<'a, K, V>(btree_map::Iter<'a, K, V>);
398
399impl<'a, K, V> Iterator for Iter<'a, K, V>
400where
401    V: WeakRef,
402{
403    type Item = (&'a K, V::Strong);
404
405    fn next(&mut self) -> Option<Self::Item> {
406        for (key, value) in self.0.by_ref() {
407            if let Some(value) = value.upgrade() {
408                return Some((key, value));
409            }
410        }
411        None
412    }
413
414    #[inline]
415    fn size_hint(&self) -> (usize, Option<usize>) {
416        (0, Some(self.0.len()))
417    }
418}
419
420impl<K, V> FusedIterator for Iter<'_, K, V> where V: WeakRef {}
421
422impl<K, V> Default for Iter<'_, K, V> {
423    fn default() -> Self {
424        Iter(btree_map::Iter::default())
425    }
426}
427
428impl<K, V> Clone for Iter<'_, K, V> {
429    fn clone(&self) -> Self {
430        Iter(self.0.clone())
431    }
432}
433
434impl<K, V> fmt::Debug for Iter<'_, K, V>
435where
436    K: fmt::Debug,
437    V: WeakRef,
438    V::Strong: fmt::Debug,
439{
440    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
441        f.debug_list().entries(self.clone()).finish()
442    }
443}
444
445impl<'a, K, V> IntoIterator for &'a WeakMap<K, V>
446where
447    V: WeakRef,
448{
449    type IntoIter = Iter<'a, K, V>;
450    type Item = (&'a K, V::Strong);
451
452    fn into_iter(self) -> Self::IntoIter {
453        self.iter()
454    }
455}
456
457/// A mutable iterator over the entries of a `BTreeMap`.
458#[must_use = "iterators are lazy and do nothing unless consumed"]
459pub struct IterMut<'a, K, V>(btree_map::IterMut<'a, K, V>);
460
461impl<'a, K, V> Iterator for IterMut<'a, K, V> {
462    type Item = (&'a K, &'a mut V);
463
464    fn next(&mut self) -> Option<Self::Item> {
465        self.0.next()
466    }
467
468    #[inline]
469    fn size_hint(&self) -> (usize, Option<usize>) {
470        (0, Some(self.0.len()))
471    }
472}
473
474impl<K, V> ExactSizeIterator for IterMut<'_, K, V> {}
475
476impl<K, V> FusedIterator for IterMut<'_, K, V> {}
477
478impl<K, V> Default for IterMut<'_, K, V> {
479    fn default() -> Self {
480        IterMut(btree_map::IterMut::default())
481    }
482}
483
484impl<'a, K, V> IntoIterator for &'a mut WeakMap<K, V>
485where
486    K: Ord,
487    V: WeakRef,
488{
489    type IntoIter = IterMut<'a, K, V>;
490    type Item = (&'a K, &'a mut V);
491
492    fn into_iter(self) -> Self::IntoIter {
493        self.iter_mut()
494    }
495}
496
497/// An iterator over the keys of a `WeakMap`.
498#[must_use = "iterators are lazy and do nothing unless consumed"]
499pub struct Keys<'a, K, V>(Iter<'a, K, V>);
500
501impl<'a, K, V> Iterator for Keys<'a, K, V>
502where
503    V: WeakRef,
504{
505    type Item = &'a K;
506
507    fn next(&mut self) -> Option<Self::Item> {
508        self.0.next().map(|(key, _)| key)
509    }
510
511    fn size_hint(&self) -> (usize, Option<usize>) {
512        self.0.size_hint()
513    }
514}
515
516impl<K, V> FusedIterator for Keys<'_, K, V> where V: WeakRef {}
517
518impl<K, V> Default for Keys<'_, K, V> {
519    fn default() -> Self {
520        Keys(Iter::default())
521    }
522}
523
524impl<K, V> Clone for Keys<'_, K, V> {
525    fn clone(&self) -> Self {
526        Keys(self.0.clone())
527    }
528}
529
530impl<K, V> fmt::Debug for Keys<'_, K, V>
531where
532    K: fmt::Debug,
533    V: WeakRef,
534{
535    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
536        f.debug_list().entries(self.clone()).finish()
537    }
538}
539
540/// An iterator over the values of a `WeakMap`.
541#[must_use = "iterators are lazy and do nothing unless consumed"]
542pub struct Values<'a, K, V>(Iter<'a, K, V>);
543
544impl<K, V> Iterator for Values<'_, K, V>
545where
546    V: WeakRef,
547{
548    type Item = V::Strong;
549
550    fn next(&mut self) -> Option<Self::Item> {
551        self.0.next().map(|(_, value)| value)
552    }
553
554    fn size_hint(&self) -> (usize, Option<usize>) {
555        self.0.size_hint()
556    }
557}
558
559impl<K, V> FusedIterator for Values<'_, K, V> where V: WeakRef {}
560
561impl<K, V> Default for Values<'_, K, V> {
562    fn default() -> Self {
563        Values(Iter::default())
564    }
565}
566
567impl<K, V> Clone for Values<'_, K, V> {
568    fn clone(&self) -> Self {
569        Values(self.0.clone())
570    }
571}
572
573impl<K, V> fmt::Debug for Values<'_, K, V>
574where
575    V: WeakRef,
576    V::Strong: fmt::Debug,
577{
578    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
579        f.debug_list().entries(self.clone()).finish()
580    }
581}
582
583/// An owning iterator over the entries of a `WeakMap`.
584pub struct IntoIter<K, V>(btree_map::IntoIter<K, V>);
585
586impl<K, V> Iterator for IntoIter<K, V>
587where
588    V: WeakRef,
589{
590    type Item = (K, V::Strong);
591
592    fn next(&mut self) -> Option<Self::Item> {
593        for (key, value) in self.0.by_ref() {
594            if let Some(value) = value.upgrade() {
595                return Some((key, value));
596            }
597        }
598        None
599    }
600
601    fn size_hint(&self) -> (usize, Option<usize>) {
602        (0, Some(self.0.len()))
603    }
604}
605
606impl<K, V> FusedIterator for IntoIter<K, V> where V: WeakRef {}
607
608impl<K, V> Default for IntoIter<K, V> {
609    fn default() -> Self {
610        IntoIter(btree_map::IntoIter::default())
611    }
612}
613
614impl<K, V> IntoIterator for WeakMap<K, V>
615where
616    V: WeakRef,
617{
618    type IntoIter = IntoIter<K, V>;
619    type Item = (K, V::Strong);
620
621    fn into_iter(self) -> Self::IntoIter {
622        IntoIter(self.inner.into_iter())
623    }
624}
625
626/// An owning iterator over the keys of a `WeakMap`.
627pub struct IntoKeys<K, V>(IntoIter<K, V>);
628
629impl<K, V> Iterator for IntoKeys<K, V>
630where
631    V: WeakRef,
632{
633    type Item = K;
634
635    fn next(&mut self) -> Option<Self::Item> {
636        self.0.next().map(|(key, _)| key)
637    }
638
639    fn size_hint(&self) -> (usize, Option<usize>) {
640        self.0.size_hint()
641    }
642}
643
644impl<K, V> FusedIterator for IntoKeys<K, V> where V: WeakRef {}
645
646impl<K, V> Default for IntoKeys<K, V> {
647    fn default() -> Self {
648        IntoKeys(IntoIter::default())
649    }
650}
651
652/// An owning iterator over the values of a `WeakMap`.`
653pub struct IntoValues<K, V>(IntoIter<K, V>);
654
655impl<K, V> Iterator for IntoValues<K, V>
656where
657    V: WeakRef,
658{
659    type Item = V::Strong;
660
661    fn next(&mut self) -> Option<Self::Item> {
662        self.0.next().map(|(_, value)| value)
663    }
664
665    fn size_hint(&self) -> (usize, Option<usize>) {
666        self.0.size_hint()
667    }
668}
669
670impl<K, V> FusedIterator for IntoValues<K, V> where V: WeakRef {}
671
672impl<K, V> Default for IntoValues<K, V> {
673    fn default() -> Self {
674        IntoValues(IntoIter::default())
675    }
676}
677
678#[cfg(test)]
679mod tests {
680    use alloc::sync::{Arc, Weak};
681
682    use super::*;
683
684    #[test]
685    fn test_basic() {
686        let mut map = WeakMap::<u32, Weak<&str>>::new();
687
688        let elem1 = Arc::new("1");
689        map.insert(1, &elem1);
690
691        {
692            let elem2 = Arc::new("2");
693            map.insert(2, &elem2);
694        }
695
696        assert_eq!(map.len(), 1);
697        assert_eq!(map.get(&1), Some(elem1));
698        assert_eq!(map.get(&2), None);
699    }
700
701    #[test]
702    fn test_cleanup() {
703        let mut map = WeakMap::<usize, Weak<usize>>::new();
704
705        for i in 0..OPS_THRESHOLD * 10 {
706            let elem = Arc::new(i);
707            map.insert(i, &elem);
708        }
709
710        assert_eq!(map.len(), 0);
711        assert_eq!(map.raw_len(), 1);
712    }
713}
weak_map/map.rs

weak_map/
map.rs
