X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=lightning%2Fsrc%2Futil%2Findexed_map.rs;h=97788ffe68acbd5f90d67bd501710c3b5942ce51;hb=65efd92c4a931b3b3622c84c6055bc015152fde0;hp=841659714c68d73e9022e5f5e7d87a66247a7f71;hpb=1bd35367d8495d9a8b90f5de0f02b68014523e3b;p=rust-lightning diff --git a/lightning/src/util/indexed_map.rs b/lightning/src/util/indexed_map.rs index 84165971..97788ffe 100644 --- a/lightning/src/util/indexed_map.rs +++ b/lightning/src/util/indexed_map.rs @@ -1,9 +1,9 @@ //! 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 core::cmp::Ord; -use core::ops::RangeBounds; +use crate::prelude::*; +use alloc::slice::Iter; +use core::hash::Hash; +use core::ops::{Bound, RangeBounds}; /// A map which can be iterated in a deterministic order. /// @@ -19,17 +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, PartialEq, Eq)] -pub struct IndexedMap { - map: BTreeMap, +#[derive(Clone, Debug, Eq)] +pub struct IndexedMap { + map: HashMap, + keys: Vec, } -impl IndexedMap { +impl IndexedMap { /// Constructs a new, empty map pub fn new() -> Self { Self { - map: BTreeMap::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), } } @@ -44,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 { @@ -52,26 +69,38 @@ 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 { + let idx = self.keys.iter().position(|k| k == key).expect("map and keys must be consistent"); + self.keys.remove(idx); + } + 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() { + self.keys.push(key); + } + 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 +123,23 @@ 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>(&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[start..end].iter(), + map: &self.map, + } } /// Returns the number of `key`/`value` pairs in the map @@ -109,36 +153,71 @@ 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: Iter<'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>, +/// +/// 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> { + underlying_entry: VacantHashMapEntry<'a, K, V>, + key: K, + keys: &'a mut Vec, } /// An [`Entry`] for an existing key-value pair -pub struct OccupiedEntry<'a, K: Ord, V> { - underlying_entry: btree_map::OccupiedEntry<'a, K, V>, +/// +/// 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> { + 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. -pub enum Entry<'a, K: Ord, V> { +/// +/// 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>), /// 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 { + self.keys.push(self.key); 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(); + let idx = self.keys.iter().position(|k| k == &res.0).expect("map and keys must be consistent"); + self.keys.remove(idx); + res } /// Get a reference to the value at the position described by this entry.