X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=lightning%2Fsrc%2Futil%2Findexed_map.rs;h=97788ffe68acbd5f90d67bd501710c3b5942ce51;hb=b783736356d17704852ea121a325b2978884ad9a;hp=cccbfe7bc7a43434a12b8c7608dff2d935e5c1dc;hpb=ccf92157620da45032d75f06b5972eaf142c1ce3;p=rust-lightning diff --git a/lightning/src/util/indexed_map.rs b/lightning/src/util/indexed_map.rs index cccbfe7b..97788ffe 100644 --- a/lightning/src/util/indexed_map.rs +++ b/lightning/src/util/indexed_map.rs @@ -1,10 +1,9 @@ //! This module has a map which can be iterated in a deterministic order. See the [`IndexedMap`]. -use crate::prelude::{HashMap, hash_map}; -use alloc::collections::{BTreeSet, btree_set}; +use crate::prelude::*; +use alloc::slice::Iter; use core::hash::Hash; -use core::cmp::Ord; -use core::ops::RangeBounds; +use core::ops::{Bound, RangeBounds}; /// A map which can be iterated in a deterministic order. /// @@ -20,20 +19,29 @@ use core::ops::RangeBounds; /// actually backed by a [`HashMap`], with some additional tracking to ensure we can iterate over /// keys in the order defined by [`Ord`]. /// +/// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly. +/// /// [`BTreeMap`]: alloc::collections::BTreeMap -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, Eq)] pub struct IndexedMap { map: HashMap, - // TODO: Explore swapping this for a sorted vec (that is only sorted on first range() call) - keys: BTreeSet, + keys: Vec, } impl IndexedMap { /// Constructs a new, empty map pub fn new() -> Self { Self { - map: HashMap::new(), - keys: BTreeSet::new(), + map: new_hash_map(), + keys: Vec::new(), + } + } + + /// Constructs a new, empty map with the given capacity pre-allocated + pub fn with_capacity(capacity: usize) -> Self { + Self { + map: hash_map_with_capacity(capacity), + keys: Vec::with_capacity(capacity), } } @@ -48,6 +56,11 @@ impl IndexedMap { self.map.get_mut(key) } + /// Fetches the key-value pair corresponding to the supplied key, if one exists. + pub fn get_key_value(&self, key: &K) -> Option<(&K, &V)> { + self.map.get_key_value(key) + } + #[inline] /// Returns true if an element with the given `key` exists in the map. pub fn contains_key(&self, key: &K) -> bool { @@ -58,7 +71,8 @@ impl IndexedMap { pub fn remove(&mut self, key: &K) -> Option { let ret = self.map.remove(key); if let Some(_) = ret { - assert!(self.keys.remove(key), "map and keys must be consistent"); + let idx = self.keys.iter().position(|k| k == key).expect("map and keys must be consistent"); + self.keys.remove(idx); } ret } @@ -68,7 +82,7 @@ impl IndexedMap { pub fn insert(&mut self, key: K, value: V) -> Option { let ret = self.map.insert(key.clone(), value); if ret.is_none() { - assert!(self.keys.insert(key), "map and keys must be consistent"); + self.keys.push(key); } ret } @@ -109,9 +123,21 @@ impl IndexedMap { } /// Returns an iterator which iterates over the `key`/`value` pairs in a given range. - pub fn range>(&self, range: R) -> Range { + pub fn range>(&mut self, range: R) -> Range { + self.keys.sort_unstable(); + let start = match range.start_bound() { + Bound::Unbounded => 0, + Bound::Included(key) => self.keys.binary_search(key).unwrap_or_else(|index| index), + Bound::Excluded(key) => self.keys.binary_search(key).and_then(|index| Ok(index + 1)).unwrap_or_else(|index| index), + }; + let end = match range.end_bound() { + Bound::Unbounded => self.keys.len(), + Bound::Included(key) => self.keys.binary_search(key).and_then(|index| Ok(index + 1)).unwrap_or_else(|index| index), + Bound::Excluded(key) => self.keys.binary_search(key).unwrap_or_else(|index| index), + }; + Range { - inner_range: self.keys.range(range), + inner_range: self.keys[start..end].iter(), map: &self.map, } } @@ -127,9 +153,17 @@ impl IndexedMap { } } +impl PartialEq for IndexedMap { + fn eq(&self, other: &Self) -> bool { + self.map == other.map + } +} + /// An iterator over a range of values in an [`IndexedMap`] +/// +/// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly. pub struct Range<'a, K: Hash + Ord, V> { - inner_range: btree_set::Range<'a, K>, + inner_range: Iter<'a, K>, map: &'a HashMap, } impl<'a, K: Hash + Ord, V: 'a> Iterator for Range<'a, K, V> { @@ -142,26 +176,26 @@ impl<'a, K: Hash + Ord, V: 'a> Iterator for Range<'a, K, V> { } /// An [`Entry`] for a key which currently has no value +/// +/// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly. 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>, + underlying_entry: VacantHashMapEntry<'a, K, V>, key: K, - keys: &'a mut BTreeSet, + keys: &'a mut Vec, } /// An [`Entry`] for an existing key-value pair +/// +/// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly. 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, + underlying_entry: OccupiedHashMapEntry<'a, K, V>, + keys: &'a mut Vec, } /// A mutable reference to a position in the map. This can be used to reference, add, or update the /// value at a fixed key. +/// +/// This is not exported to bindings users as bindings provide alternate accessors rather than exposing maps directly. 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>), @@ -172,7 +206,7 @@ pub enum Entry<'a, K: Hash + Ord, 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.keys.push(self.key); self.underlying_entry.insert(value) } } @@ -181,7 +215,8 @@ 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) { let res = self.underlying_entry.remove_entry(); - assert!(self.keys.remove(&res.0), "map and keys must be consistent"); + let idx = self.keys.iter().position(|k| k == &res.0).expect("map and keys must be consistent"); + self.keys.remove(idx); res }