From: Matt Corallo Date: Thu, 19 Jan 2023 17:59:10 +0000 (+0000) Subject: Swap `IndexedMap` implementation for a `HashMap`+B-Tree X-Git-Tag: v0.0.114-beta~45^2~2 X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=commitdiff_plain;h=039fa5255da65cd84c42af93e720e3db57a09b38;p=rust-lightning Swap `IndexedMap` implementation for a `HashMap`+B-Tree Our network graph has to be iterable in a deterministic order and with the ability to iterate over a specific range. Thus, historically, we've used a `BTreeMap` to do the iteration. This is fine, except our map needs to also provide high performance lookups in order to make route-finding fast. Sadly, `BTreeMap`s are quite slow due to the branching penalty. Here we replace the implementation of our `IndexedMap` with a `HashMap` to store the elements itself and a `BTreeSet` to store the keys set in sorted order for iteration. As of this commit on the same hardware as the above few commits, the benchmark results are: ``` test routing::router::benches::generate_mpp_routes_with_probabilistic_scorer ... bench: 109,544,993 ns/iter (+/- 27,553,574) test routing::router::benches::generate_mpp_routes_with_zero_penalty_scorer ... bench: 81,164,590 ns/iter (+/- 55,422,930) test routing::router::benches::generate_routes_with_probabilistic_scorer ... bench: 34,726,569 ns/iter (+/- 9,646,345) test routing::router::benches::generate_routes_with_zero_penalty_scorer ... bench: 22,772,355 ns/iter (+/- 9,574,418) ``` --- diff --git a/lightning/src/util/indexed_map.rs b/lightning/src/util/indexed_map.rs index 841659714..cccbfe7bc 100644 --- a/lightning/src/util/indexed_map.rs +++ b/lightning/src/util/indexed_map.rs @@ -1,7 +1,8 @@ //! This module has a map which can be iterated in a deterministic order. See the [`IndexedMap`]. -use crate::prelude::HashMap; -use alloc::collections::{BTreeMap, btree_map}; +use crate::prelude::{HashMap, hash_map}; +use alloc::collections::{BTreeSet, btree_set}; +use core::hash::Hash; use core::cmp::Ord; use core::ops::RangeBounds; @@ -20,16 +21,19 @@ use core::ops::RangeBounds; /// keys in the order defined by [`Ord`]. /// /// [`BTreeMap`]: alloc::collections::BTreeMap -#[derive(Clone, PartialEq, Eq)] -pub struct IndexedMap { - map: BTreeMap, +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct IndexedMap { + map: HashMap, + // TODO: Explore swapping this for a sorted vec (that is only sorted on first range() call) + keys: BTreeSet, } -impl IndexedMap { +impl IndexedMap { /// Constructs a new, empty map pub fn new() -> Self { Self { - map: BTreeMap::new(), + map: HashMap::new(), + keys: BTreeSet::new(), } } @@ -52,26 +56,37 @@ impl IndexedMap { /// Removes the element with the given `key`, returning it, if one exists. pub fn remove(&mut self, key: &K) -> Option { - self.map.remove(key) + let ret = self.map.remove(key); + if let Some(_) = ret { + assert!(self.keys.remove(key), "map and keys must be consistent"); + } + ret } /// Inserts the given `key`/`value` pair into the map, returning the element that was /// previously stored at the given `key`, if one exists. pub fn insert(&mut self, key: K, value: V) -> Option { - self.map.insert(key, value) + let ret = self.map.insert(key.clone(), value); + if ret.is_none() { + assert!(self.keys.insert(key), "map and keys must be consistent"); + } + ret } /// Returns an [`Entry`] for the given `key` in the map, allowing access to the value. pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { - match self.map.entry(key) { - btree_map::Entry::Vacant(entry) => { + match self.map.entry(key.clone()) { + hash_map::Entry::Vacant(entry) => { Entry::Vacant(VacantEntry { - underlying_entry: entry + underlying_entry: entry, + key, + keys: &mut self.keys, }) }, - btree_map::Entry::Occupied(entry) => { + hash_map::Entry::Occupied(entry) => { Entry::Occupied(OccupiedEntry { - underlying_entry: entry + underlying_entry: entry, + keys: &mut self.keys, }) } } @@ -94,8 +109,11 @@ impl IndexedMap { } /// Returns an iterator which iterates over the `key`/`value` pairs in a given range. - pub fn range>(&self, range: R) -> btree_map::Range { - self.map.range(range) + pub fn range>(&self, range: R) -> Range { + Range { + inner_range: self.keys.range(range), + map: &self.map, + } } /// Returns the number of `key`/`value` pairs in the map @@ -109,36 +127,62 @@ impl IndexedMap { } } +/// An iterator over a range of values in an [`IndexedMap`] +pub struct Range<'a, K: Hash + Ord, V> { + inner_range: btree_set::Range<'a, K>, + map: &'a HashMap, +} +impl<'a, K: Hash + Ord, V: 'a> Iterator for Range<'a, K, V> { + type Item = (&'a K, &'a V); + fn next(&mut self) -> Option<(&'a K, &'a V)> { + self.inner_range.next().map(|k| { + (k, self.map.get(k).expect("map and keys must be consistent")) + }) + } +} + /// An [`Entry`] for a key which currently has no value -pub struct VacantEntry<'a, K: Ord, V> { - underlying_entry: btree_map::VacantEntry<'a, K, V>, +pub struct VacantEntry<'a, K: Hash + Ord, V> { + #[cfg(feature = "hashbrown")] + underlying_entry: hash_map::VacantEntry<'a, K, V, hash_map::DefaultHashBuilder>, + #[cfg(not(feature = "hashbrown"))] + underlying_entry: hash_map::VacantEntry<'a, K, V>, + key: K, + keys: &'a mut BTreeSet, } /// An [`Entry`] for an existing key-value pair -pub struct OccupiedEntry<'a, K: Ord, V> { - underlying_entry: btree_map::OccupiedEntry<'a, K, V>, +pub struct OccupiedEntry<'a, K: Hash + Ord, V> { + #[cfg(feature = "hashbrown")] + underlying_entry: hash_map::OccupiedEntry<'a, K, V, hash_map::DefaultHashBuilder>, + #[cfg(not(feature = "hashbrown"))] + underlying_entry: hash_map::OccupiedEntry<'a, K, V>, + keys: &'a mut BTreeSet, } /// A mutable reference to a position in the map. This can be used to reference, add, or update the /// value at a fixed key. -pub enum Entry<'a, K: Ord, V> { +pub enum Entry<'a, K: Hash + Ord, V> { /// A mutable reference to a position within the map where there is no value. Vacant(VacantEntry<'a, K, V>), /// A mutable reference to a position within the map where there is currently a value. Occupied(OccupiedEntry<'a, K, V>), } -impl<'a, K: Ord, V> VacantEntry<'a, K, V> { +impl<'a, K: Hash + Ord, V> VacantEntry<'a, K, V> { /// Insert a value into the position described by this entry. pub fn insert(self, value: V) -> &'a mut V { + assert!(self.keys.insert(self.key), "map and keys must be consistent"); self.underlying_entry.insert(value) } } -impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> { +impl<'a, K: Hash + Ord, V> OccupiedEntry<'a, K, V> { /// Remove the value at the position described by this entry. pub fn remove_entry(self) -> (K, V) { - self.underlying_entry.remove_entry() + let res = self.underlying_entry.remove_entry(); + assert!(self.keys.remove(&res.0), "map and keys must be consistent"); + res } /// Get a reference to the value at the position described by this entry.