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