]> git.bitcoin.ninja Git - rust-lightning/blob - lightning/src/util/indexed_map.rs
Merge pull request #2248 from TheBlueMatt/2023-04-gossip-check
[rust-lightning] / 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::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: HashMap::new(),
38                         keys: Vec::new(),
39                 }
40         }
41
42         #[inline(always)]
43         /// Fetches the element with the given `key`, if one exists.
44         pub fn get(&self, key: &K) -> Option<&V> {
45                 self.map.get(key)
46         }
47
48         /// Fetches a mutable reference to the element with the given `key`, if one exists.
49         pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
50                 self.map.get_mut(key)
51         }
52
53         #[inline]
54         /// Returns true if an element with the given `key` exists in the map.
55         pub fn contains_key(&self, key: &K) -> bool {
56                 self.map.contains_key(key)
57         }
58
59         /// Removes the element with the given `key`, returning it, if one exists.
60         pub fn remove(&mut self, key: &K) -> Option<V> {
61                 let ret = self.map.remove(key);
62                 if let Some(_) = ret {
63                         let idx = self.keys.iter().position(|k| k == key).expect("map and keys must be consistent");
64                         self.keys.remove(idx);
65                 }
66                 ret
67         }
68
69         /// Inserts the given `key`/`value` pair into the map, returning the element that was
70         /// previously stored at the given `key`, if one exists.
71         pub fn insert(&mut self, key: K, value: V) -> Option<V> {
72                 let ret = self.map.insert(key.clone(), value);
73                 if ret.is_none() {
74                         self.keys.push(key);
75                 }
76                 ret
77         }
78
79         /// Returns an [`Entry`] for the given `key` in the map, allowing access to the value.
80         pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
81                 match self.map.entry(key.clone()) {
82                         hash_map::Entry::Vacant(entry) => {
83                                 Entry::Vacant(VacantEntry {
84                                         underlying_entry: entry,
85                                         key,
86                                         keys: &mut self.keys,
87                                 })
88                         },
89                         hash_map::Entry::Occupied(entry) => {
90                                 Entry::Occupied(OccupiedEntry {
91                                         underlying_entry: entry,
92                                         keys: &mut self.keys,
93                                 })
94                         }
95                 }
96         }
97
98         /// Returns an iterator which iterates over the keys in the map, in a random order.
99         pub fn unordered_keys(&self) -> impl Iterator<Item = &K> {
100                 self.map.keys()
101         }
102
103         /// Returns an iterator which iterates over the `key`/`value` pairs in a random order.
104         pub fn unordered_iter(&self) -> impl Iterator<Item = (&K, &V)> {
105                 self.map.iter()
106         }
107
108         /// Returns an iterator which iterates over the `key`s and mutable references to `value`s in a
109         /// random order.
110         pub fn unordered_iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
111                 self.map.iter_mut()
112         }
113
114         /// Returns an iterator which iterates over the `key`/`value` pairs in a given range.
115         pub fn range<R: RangeBounds<K>>(&mut self, range: R) -> Range<K, V> {
116                 self.keys.sort_unstable();
117                 let start = match range.start_bound() {
118                         Bound::Unbounded => 0,
119                         Bound::Included(key) => self.keys.binary_search(key).unwrap_or_else(|index| index),
120                         Bound::Excluded(key) => self.keys.binary_search(key).and_then(|index| Ok(index + 1)).unwrap_or_else(|index| index),
121                 };
122                 let end = match range.end_bound() {
123                         Bound::Unbounded => self.keys.len(),
124                         Bound::Included(key) => self.keys.binary_search(key).and_then(|index| Ok(index + 1)).unwrap_or_else(|index| index),
125                         Bound::Excluded(key) => self.keys.binary_search(key).unwrap_or_else(|index| index),
126                 };
127
128                 Range {
129                         inner_range: self.keys[start..end].iter(),
130                         map: &self.map,
131                 }
132         }
133
134         /// Returns the number of `key`/`value` pairs in the map
135         pub fn len(&self) -> usize {
136                 self.map.len()
137         }
138
139         /// Returns true if there are no elements in the map
140         pub fn is_empty(&self) -> bool {
141                 self.map.is_empty()
142         }
143 }
144
145 impl<K: Hash + Ord + PartialEq, V: PartialEq> PartialEq for IndexedMap<K, V> {
146         fn eq(&self, other: &Self) -> bool {
147                 self.map == other.map
148         }
149 }
150
151 /// An iterator over a range of values in an [`IndexedMap`]
152 ///
153 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
154 pub struct Range<'a, K: Hash + Ord, V> {
155         inner_range: Iter<'a, K>,
156         map: &'a HashMap<K, V>,
157 }
158 impl<'a, K: Hash + Ord, V: 'a> Iterator for Range<'a, K, V> {
159         type Item = (&'a K, &'a V);
160         fn next(&mut self) -> Option<(&'a K, &'a V)> {
161                 self.inner_range.next().map(|k| {
162                         (k, self.map.get(k).expect("map and keys must be consistent"))
163                 })
164         }
165 }
166
167 /// An [`Entry`] for a key which currently has no value
168 ///
169 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
170 pub struct VacantEntry<'a, K: Hash + Ord, V> {
171         #[cfg(feature = "hashbrown")]
172         underlying_entry: hash_map::VacantEntry<'a, K, V, hash_map::DefaultHashBuilder>,
173         #[cfg(not(feature = "hashbrown"))]
174         underlying_entry: hash_map::VacantEntry<'a, K, V>,
175         key: K,
176         keys: &'a mut Vec<K>,
177 }
178
179 /// An [`Entry`] for an existing key-value pair
180 ///
181 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
182 pub struct OccupiedEntry<'a, K: Hash + Ord, V> {
183         #[cfg(feature = "hashbrown")]
184         underlying_entry: hash_map::OccupiedEntry<'a, K, V, hash_map::DefaultHashBuilder>,
185         #[cfg(not(feature = "hashbrown"))]
186         underlying_entry: hash_map::OccupiedEntry<'a, K, V>,
187         keys: &'a mut Vec<K>,
188 }
189
190 /// A mutable reference to a position in the map. This can be used to reference, add, or update the
191 /// value at a fixed key.
192 ///
193 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
194 pub enum Entry<'a, K: Hash + Ord, V> {
195         /// A mutable reference to a position within the map where there is no value.
196         Vacant(VacantEntry<'a, K, V>),
197         /// A mutable reference to a position within the map where there is currently a value.
198         Occupied(OccupiedEntry<'a, K, V>),
199 }
200
201 impl<'a, K: Hash + Ord, V> VacantEntry<'a, K, V> {
202         /// Insert a value into the position described by this entry.
203         pub fn insert(self, value: V) -> &'a mut V {
204                 self.keys.push(self.key);
205                 self.underlying_entry.insert(value)
206         }
207 }
208
209 impl<'a, K: Hash + Ord, V> OccupiedEntry<'a, K, V> {
210         /// Remove the value at the position described by this entry.
211         pub fn remove_entry(self) -> (K, V) {
212                 let res = self.underlying_entry.remove_entry();
213                 let idx = self.keys.iter().position(|k| k == &res.0).expect("map and keys must be consistent");
214                 self.keys.remove(idx);
215                 res
216         }
217
218         /// Get a reference to the value at the position described by this entry.
219         pub fn get(&self) -> &V {
220                 self.underlying_entry.get()
221         }
222
223         /// Get a mutable reference to the value at the position described by this entry.
224         pub fn get_mut(&mut self) -> &mut V {
225                 self.underlying_entry.get_mut()
226         }
227
228         /// Consume this entry, returning a mutable reference to the value at the position described by
229         /// this entry.
230         pub fn into_mut(self) -> &'a mut V {
231                 self.underlying_entry.into_mut()
232         }
233 }