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