1 //! This module has a map which can be iterated in a deterministic order. See the [`IndexedMap`].
4 use alloc::slice::Iter;
6 use core::ops::{Bound, RangeBounds};
8 /// A map which can be iterated in a deterministic order.
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.
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`].
22 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
24 /// [`BTreeMap`]: alloc::collections::BTreeMap
25 #[derive(Clone, Debug, Eq)]
26 pub struct IndexedMap<K: Hash + Ord, V> {
31 impl<K: Clone + Hash + Ord, V> IndexedMap<K, V> {
32 /// Constructs a new, empty map
33 pub fn new() -> Self {
40 /// Constructs a new, empty map with the given capacity pre-allocated
41 pub fn with_capacity(capacity: usize) -> Self {
43 map: hash_map_with_capacity(capacity),
44 keys: Vec::with_capacity(capacity),
49 /// Fetches the element with the given `key`, if one exists.
50 pub fn get(&self, key: &K) -> Option<&V> {
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> {
60 /// Returns true if an element with the given `key` exists in the map.
61 pub fn contains_key(&self, key: &K) -> bool {
62 self.map.contains_key(key)
65 /// Removes the element with the given `key`, returning it, if one exists.
66 pub fn remove(&mut self, key: &K) -> Option<V> {
67 let ret = self.map.remove(key);
68 if let Some(_) = ret {
69 let idx = self.keys.iter().position(|k| k == key).expect("map and keys must be consistent");
70 self.keys.remove(idx);
75 /// Inserts the given `key`/`value` pair into the map, returning the element that was
76 /// previously stored at the given `key`, if one exists.
77 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
78 let ret = self.map.insert(key.clone(), value);
85 /// Returns an [`Entry`] for the given `key` in the map, allowing access to the value.
86 pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
87 match self.map.entry(key.clone()) {
88 hash_map::Entry::Vacant(entry) => {
89 Entry::Vacant(VacantEntry {
90 underlying_entry: entry,
95 hash_map::Entry::Occupied(entry) => {
96 Entry::Occupied(OccupiedEntry {
97 underlying_entry: entry,
104 /// Returns an iterator which iterates over the keys in the map, in a random order.
105 pub fn unordered_keys(&self) -> impl Iterator<Item = &K> {
109 /// Returns an iterator which iterates over the `key`/`value` pairs in a random order.
110 pub fn unordered_iter(&self) -> impl Iterator<Item = (&K, &V)> {
114 /// Returns an iterator which iterates over the `key`s and mutable references to `value`s in a
116 pub fn unordered_iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
120 /// Returns an iterator which iterates over the `key`/`value` pairs in a given range.
121 pub fn range<R: RangeBounds<K>>(&mut self, range: R) -> Range<K, V> {
122 self.keys.sort_unstable();
123 let start = match range.start_bound() {
124 Bound::Unbounded => 0,
125 Bound::Included(key) => self.keys.binary_search(key).unwrap_or_else(|index| index),
126 Bound::Excluded(key) => self.keys.binary_search(key).and_then(|index| Ok(index + 1)).unwrap_or_else(|index| index),
128 let end = match range.end_bound() {
129 Bound::Unbounded => self.keys.len(),
130 Bound::Included(key) => self.keys.binary_search(key).and_then(|index| Ok(index + 1)).unwrap_or_else(|index| index),
131 Bound::Excluded(key) => self.keys.binary_search(key).unwrap_or_else(|index| index),
135 inner_range: self.keys[start..end].iter(),
140 /// Returns the number of `key`/`value` pairs in the map
141 pub fn len(&self) -> usize {
145 /// Returns true if there are no elements in the map
146 pub fn is_empty(&self) -> bool {
151 impl<K: Hash + Ord + PartialEq, V: PartialEq> PartialEq for IndexedMap<K, V> {
152 fn eq(&self, other: &Self) -> bool {
153 self.map == other.map
157 /// An iterator over a range of values in an [`IndexedMap`]
159 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
160 pub struct Range<'a, K: Hash + Ord, V> {
161 inner_range: Iter<'a, K>,
162 map: &'a HashMap<K, V>,
164 impl<'a, K: Hash + Ord, V: 'a> Iterator for Range<'a, K, V> {
165 type Item = (&'a K, &'a V);
166 fn next(&mut self) -> Option<(&'a K, &'a V)> {
167 self.inner_range.next().map(|k| {
168 (k, self.map.get(k).expect("map and keys must be consistent"))
173 /// An [`Entry`] for a key which currently has no value
175 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
176 pub struct VacantEntry<'a, K: Hash + Ord, V> {
177 underlying_entry: VacantHashMapEntry<'a, K, V>,
179 keys: &'a mut Vec<K>,
182 /// An [`Entry`] for an existing key-value pair
184 /// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly.
185 pub struct OccupiedEntry<'a, K: Hash + Ord, V> {
186 underlying_entry: OccupiedHashMapEntry<'a, K, V>,
187 keys: &'a mut Vec<K>,
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.
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>),
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)
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);
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()
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()
228 /// Consume this entry, returning a mutable reference to the value at the position described by
230 pub fn into_mut(self) -> &'a mut V {
231 self.underlying_entry.into_mut()