_ Git - rust-lightning/blob - lightning/src/util/indexed_map.rs

   1 //! This module has a map which can be iterated in a deterministic order. See the [`IndexedMap`].
   2
   3 use crate::prelude::{HashMap, hash_map};
   4 use alloc::collections::{BTreeSet, btree_set};
   5 use core::hash::Hash;
   6 use core::cmp::Ord;
   7 use core::ops::RangeBounds;
   8
   9 /// A map which can be iterated in a deterministic order.
  10 ///
  11 /// This would traditionally be accomplished by simply using a [`BTreeMap`], however B-Trees
  12 /// generally have very slow lookups. Because we use a nodes+channels map while finding routes
  13 /// across the network graph, our network graph backing map must be as performant as possible.
  14 /// However, because peers expect to sync the network graph from us (and we need to support that
  15 /// without holding a lock on the graph for the duration of the sync or dumping the entire graph
  16 /// into our outbound message queue), we need an iterable map with a consistent iteration order we
  17 /// can jump to a starting point on.
  18 ///
  19 /// Thus, we have a custom data structure here - its API mimics that of Rust's [`BTreeMap`], but is
  20 /// actually backed by a [`HashMap`], with some additional tracking to ensure we can iterate over
  21 /// keys in the order defined by [`Ord`].
  22 ///
  23 /// [`BTreeMap`]: alloc::collections::BTreeMap
  24 #[derive(Clone, Debug, PartialEq, Eq)]
  25 pub struct IndexedMap<K: Hash + Ord, V> {
  26         map: HashMap<K, V>,
  27         // TODO: Explore swapping this for a sorted vec (that is only sorted on first range() call)
  28         keys: BTreeSet<K>,
  29 }
  30
  31 impl<K: Clone + Hash + Ord, V> IndexedMap<K, V> {
  32         /// Constructs a new, empty map
  33         pub fn new() -> Self {
  34                 Self {
  35                         map: HashMap::new(),
  36                         keys: BTreeSet::new(),
  37                 }
  38         }
  39
  40         #[inline(always)]
  41         /// Fetches the element with the given `key`, if one exists.
  42         pub fn get(&self, key: &K) -> Option<&V> {
  43                 self.map.get(key)
  44         }
  45
  46         /// Fetches a mutable reference to the element with the given `key`, if one exists.
  47         pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
  48                 self.map.get_mut(key)
  49         }
  50
  51         #[inline]
  52         /// Returns true if an element with the given `key` exists in the map.
  53         pub fn contains_key(&self, key: &K) -> bool {
  54                 self.map.contains_key(key)
  55         }
  56
  57         /// Removes the element with the given `key`, returning it, if one exists.
  58         pub fn remove(&mut self, key: &K) -> Option<V> {
  59                 let ret = self.map.remove(key);
  60                 if let Some(_) = ret {
  61                         assert!(self.keys.remove(key), "map and keys must be consistent");
  62                 }
  63                 ret
  64         }
  65
  66         /// Inserts the given `key`/`value` pair into the map, returning the element that was
  67         /// previously stored at the given `key`, if one exists.
  68         pub fn insert(&mut self, key: K, value: V) -> Option<V> {
  69                 let ret = self.map.insert(key.clone(), value);
  70                 if ret.is_none() {
  71                         assert!(self.keys.insert(key), "map and keys must be consistent");
  72                 }
  73                 ret
  74         }
  75
  76         /// Returns an [`Entry`] for the given `key` in the map, allowing access to the value.
  77         pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
  78                 match self.map.entry(key.clone()) {
  79                         hash_map::Entry::Vacant(entry) => {
  80                                 Entry::Vacant(VacantEntry {
  81                                         underlying_entry: entry,
  82                                         key,
  83                                         keys: &mut self.keys,
  84                                 })
  85                         },
  86                         hash_map::Entry::Occupied(entry) => {
  87                                 Entry::Occupied(OccupiedEntry {
  88                                         underlying_entry: entry,
  89                                         keys: &mut self.keys,
  90                                 })
  91                         }
  92                 }
  93         }
  94
  95         /// Returns an iterator which iterates over the keys in the map, in a random order.
  96         pub fn unordered_keys(&self) -> impl Iterator<Item = &K> {
  97                 self.map.keys()
  98         }
  99
 100         /// Returns an iterator which iterates over the `key`/`value` pairs in a random order.
 101         pub fn unordered_iter(&self) -> impl Iterator<Item = (&K, &V)> {
 102                 self.map.iter()
 103         }
 104
 105         /// Returns an iterator which iterates over the `key`s and mutable references to `value`s in a
 106         /// random order.
 107         pub fn unordered_iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
 108                 self.map.iter_mut()
 109         }
 110
 111         /// Returns an iterator which iterates over the `key`/`value` pairs in a given range.
 112         pub fn range<R: RangeBounds<K>>(&self, range: R) -> Range<K, V> {
 113                 Range {
 114                         inner_range: self.keys.range(range),
 115                         map: &self.map,
 116                 }
 117         }
 118
 119         /// Returns the number of `key`/`value` pairs in the map
 120         pub fn len(&self) -> usize {
 121                 self.map.len()
 122         }
 123
 124         /// Returns true if there are no elements in the map
 125         pub fn is_empty(&self) -> bool {
 126                 self.map.is_empty()
 127         }
 128 }
 129
 130 /// An iterator over a range of values in an [`IndexedMap`]
 131 pub struct Range<'a, K: Hash + Ord, V> {
 132         inner_range: btree_set::Range<'a, K>,
 133         map: &'a HashMap<K, V>,
 134 }
 135 impl<'a, K: Hash + Ord, V: 'a> Iterator for Range<'a, K, V> {
 136         type Item = (&'a K, &'a V);
 137         fn next(&mut self) -> Option<(&'a K, &'a V)> {
 138                 self.inner_range.next().map(|k| {
 139                         (k, self.map.get(k).expect("map and keys must be consistent"))
 140                 })
 141         }
 142 }
 143
 144 /// An [`Entry`] for a key which currently has no value
 145 pub struct VacantEntry<'a, K: Hash + Ord, V> {
 146         #[cfg(feature = "hashbrown")]
 147         underlying_entry: hash_map::VacantEntry<'a, K, V, hash_map::DefaultHashBuilder>,
 148         #[cfg(not(feature = "hashbrown"))]
 149         underlying_entry: hash_map::VacantEntry<'a, K, V>,
 150         key: K,
 151         keys: &'a mut BTreeSet<K>,
 152 }
 153
 154 /// An [`Entry`] for an existing key-value pair
 155 pub struct OccupiedEntry<'a, K: Hash + Ord, V> {
 156         #[cfg(feature = "hashbrown")]
 157         underlying_entry: hash_map::OccupiedEntry<'a, K, V, hash_map::DefaultHashBuilder>,
 158         #[cfg(not(feature = "hashbrown"))]
 159         underlying_entry: hash_map::OccupiedEntry<'a, K, V>,
 160         keys: &'a mut BTreeSet<K>,
 161 }
 162
 163 /// A mutable reference to a position in the map. This can be used to reference, add, or update the
 164 /// value at a fixed key.
 165 pub enum Entry<'a, K: Hash + Ord, V> {
 166         /// A mutable reference to a position within the map where there is no value.
 167         Vacant(VacantEntry<'a, K, V>),
 168         /// A mutable reference to a position within the map where there is currently a value.
 169         Occupied(OccupiedEntry<'a, K, V>),
 170 }
 171
 172 impl<'a, K: Hash + Ord, V> VacantEntry<'a, K, V> {
 173         /// Insert a value into the position described by this entry.
 174         pub fn insert(self, value: V) -> &'a mut V {
 175                 assert!(self.keys.insert(self.key), "map and keys must be consistent");
 176                 self.underlying_entry.insert(value)
 177         }
 178 }
 179
 180 impl<'a, K: Hash + Ord, V> OccupiedEntry<'a, K, V> {
 181         /// Remove the value at the position described by this entry.
 182         pub fn remove_entry(self) -> (K, V) {
 183                 let res = self.underlying_entry.remove_entry();
 184                 assert!(self.keys.remove(&res.0), "map and keys must be consistent");
 185                 res
 186         }
 187
 188         /// Get a reference to the value at the position described by this entry.
 189         pub fn get(&self) -> &V {
 190                 self.underlying_entry.get()
 191         }
 192
 193         /// Get a mutable reference to the value at the position described by this entry.
 194         pub fn get_mut(&mut self) -> &mut V {
 195                 self.underlying_entry.get_mut()
 196         }
 197
 198         /// Consume this entry, returning a mutable reference to the value at the position described by
 199         /// this entry.
 200         pub fn into_mut(self) -> &'a mut V {
 201                 self.underlying_entry.into_mut()
 202         }
 203 }