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