diff --git a/src/haystack/val.rs b/src/haystack/val.rs index 9676a5e..7bb7e67 100644 --- a/src/haystack/val.rs +++ b/src/haystack/val.rs @@ -81,6 +81,7 @@ pub mod datetime; pub mod dict; pub mod dis_macro; pub mod grid; +pub mod hybrid_map; pub mod kind; pub mod list; pub mod marker; @@ -102,6 +103,7 @@ pub use crate::haystack::val::datetime::*; pub use crate::haystack::val::dict::*; pub use crate::haystack::val::dis_macro::*; pub use crate::haystack::val::grid::*; +pub use crate::haystack::val::hybrid_map::*; pub use crate::haystack::val::list::*; pub use crate::haystack::val::marker::*; pub use crate::haystack::val::na::*; diff --git a/src/haystack/val/dict.rs b/src/haystack/val/dict.rs index 1152c88..2b0505a 100644 --- a/src/haystack/val/dict.rs +++ b/src/haystack/val/dict.rs @@ -15,15 +15,6 @@ use std::ops::Index; // Alias for the underlying Dict type pub(crate) type DictType = BTreeMap; -#[derive(Clone, Debug)] -enum DictRepr { - Small(Vec<(String, Value)>), - // Boxed so the discriminant can be packed into `Vec`'s pointer niche - // (`BTreeMap` doesn't expose an equivalent niche on its own), keeping - // `DictRepr`/`Dict` as small as possible. - Tree(Box), -} - /// A Haystack Dictionary /// /// Uses a hybrid back-store: a sorted small-vector for tiny dicts and a @@ -51,8 +42,7 @@ enum DictRepr { ///``` #[derive(Clone, Debug)] pub struct Dict { - value: DictRepr, - small_max_entries: usize, + entries: HybridMap, } /// Dictionary trait with utilities that help working with @@ -132,253 +122,149 @@ pub trait HaystackDict { impl Dict { /// Hint for the maximum number of entries for the small-vector back-store. - pub const SMALL_DICT_MAX_ENTRIES_HINT: usize = 32; + pub const SMALL_DICT_MAX_ENTRIES_HINT: usize = + HybridMap::::DEFAULT_SMALL_MAX_ENTRIES; /// Construct a new `Dict` with a threshold of 32 entries for the small-vector back-store. pub fn new() -> Dict { - Self::with_small_max_entries(Self::SMALL_DICT_MAX_ENTRIES_HINT) + Dict { + entries: HybridMap::new(), + } } /// Construct a new `Dict` with a custom small-store threshold. /// If `small_max_entries` is 0, the small-vector back-store is disabled /// and the dict will use the `BTreeMap` representation. pub fn with_small_max_entries(small_max_entries: usize) -> Dict { - let value = if small_max_entries == 0 { - DictRepr::Tree(Box::default()) - } else { - DictRepr::Small(Vec::new()) - }; Dict { - value, - small_max_entries, + entries: HybridMap::with_small_max_entries(small_max_entries), } } /// Return the active small-store threshold for this dict. + #[inline] pub fn small_max_entries(&self) -> usize { - self.small_max_entries - } - - fn small_search(entries: &[(String, Value)], key: &str) -> Result { - entries.binary_search_by(|(k, _)| k.as_str().cmp(key)) - } - - fn spill_to_tree(&mut self) { - if let DictRepr::Small(entries) = &mut self.value { - let map = entries.drain(..).collect::(); - self.value = DictRepr::Tree(Box::new(map)); - } + self.entries.small_max_entries() } + /// Returns the number of entries in the dict. + #[inline] pub fn len(&self) -> usize { - match &self.value { - DictRepr::Small(entries) => entries.len(), - DictRepr::Tree(map) => map.len(), - } + self.entries.len() } + /// Returns `true` if the dict contains no entries. + #[inline] pub fn is_empty(&self) -> bool { - self.len() == 0 + self.entries.is_empty() } + /// Removes all entries from the dict. + #[inline] pub fn clear(&mut self) { - match &mut self.value { - DictRepr::Small(entries) => entries.clear(), - DictRepr::Tree(map) => map.clear(), - } + self.entries.clear(); } + /// Returns `true` if the dict contains `key`. + #[inline] pub fn contains_key(&self, key: &str) -> bool { - self.get(key).is_some() + self.entries.contains_key(key) } + /// Returns a reference to the value for `key`, if present. + #[inline] pub fn get(&self, key: &str) -> Option<&Value> { - match &self.value { - DictRepr::Small(entries) => Self::small_search(entries, key) - .ok() - .map(|pos| &entries[pos].1), - DictRepr::Tree(map) => map.get(key), - } + self.entries.get(key) } + /// Returns a mutable reference to the value for `key`, if present. + #[inline] pub fn get_mut(&mut self, key: &str) -> Option<&mut Value> { - match &mut self.value { - DictRepr::Small(entries) => Self::small_search(entries, key) - .ok() - .map(|pos| &mut entries[pos].1), - DictRepr::Tree(map) => map.get_mut(key), - } + self.entries.get_mut(key) } + /// Inserts a key-value pair, returning the previous value for `key`, if any. pub fn insert(&mut self, key: String, value: Value) -> Option { - match &mut self.value { - DictRepr::Small(entries) => { - if entries.len() < self.small_max_entries - && entries - .last() - .is_none_or(|(last_key, _)| key.as_str() > last_key.as_str()) - { - entries.push((key, value)); - return None; - } - - match Self::small_search(entries, &key) { - Ok(pos) => Some(std::mem::replace(&mut entries[pos].1, value)), - Err(pos) => { - if entries.len() < self.small_max_entries { - entries.insert(pos, (key, value)); - None - } else { - self.spill_to_tree(); - match &mut self.value { - DictRepr::Tree(map) => map.insert(key, value), - DictRepr::Small(_) => None, - } - } - } - } - } - DictRepr::Tree(map) => map.insert(key, value), - } + self.entries.insert(key, value) } + /// Removes `key` from the dict, returning its value if it was present. + #[inline] pub fn remove(&mut self, key: &str) -> Option { - match &mut self.value { - DictRepr::Small(entries) => Self::small_search(entries, key) - .ok() - .map(|pos| entries.remove(pos).1), - DictRepr::Tree(map) => map.remove(key), - } + self.entries.remove(key) } + /// Removes and returns the first (lowest-keyed) entry, if any. + #[inline] pub fn pop_first(&mut self) -> Option<(String, Value)> { - match &mut self.value { - DictRepr::Small(entries) => { - if entries.is_empty() { - None - } else { - Some(entries.remove(0)) - } - } - DictRepr::Tree(map) => map.pop_first(), - } + self.entries.pop_first() } - /// Demote a `Tree`-backed dict back to `Small` when its entry count has - /// dropped to at or below the small-store threshold. + /// Demote a tree-backed dict back to the small-vector store when its + /// entry count has dropped to at or below the small-store threshold. /// /// This is the inverse of the automatic spill that happens in [`insert`](Self::insert). /// Call it after a burst of [`remove`](Self::remove) calls to recover the /// performance and memory advantages of the sorted-vector representation. /// - /// If the dict is already `Small`-backed this is a no-op. + /// If the dict is already small-backed this is a no-op. pub fn shrink_to_fit(&mut self) { - if let DictRepr::Tree(map) = &self.value - && map.len() <= self.small_max_entries - { - let entries = map.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); - self.value = DictRepr::Small(entries); - } + self.entries.shrink_to_fit(); } /// Retains only the entries for which the predicate returns `true`. /// /// Mirrors [`BTreeMap::retain`](std::collections::BTreeMap::retain). - /// When called on a `Tree`-backed dict and the surviving entry count drops + /// When called on a tree-backed dict and the surviving entry count drops /// to or below the small-store threshold, the storage is automatically /// downgraded back to the sorted-vector representation. pub fn retain(&mut self, mut f: F) where F: FnMut(&str, &mut Value) -> bool, { - match &mut self.value { - DictRepr::Small(entries) => { - entries.retain_mut(|(k, v)| f(k.as_str(), v)); - } - DictRepr::Tree(map) => { - map.retain(|k, v| f(k.as_str(), v)); - } - } - // Downgrade Tree -> Small when the surviving count drops to the threshold. - // No-op when already Small. - self.shrink_to_fit(); + self.entries.retain(|k, v| f(k.as_str(), v)); } + /// Returns an iterator over `(&String, &Value)` pairs, in key order. + #[inline] pub fn iter(&self) -> DictIter<'_> { - match &self.value { - DictRepr::Small(entries) => DictIter::Small(entries.iter()), - DictRepr::Tree(map) => DictIter::Tree(map.iter()), - } + self.entries.iter() } + /// Returns an iterator over `(&String, &mut Value)` pairs, in key order. + #[inline] pub fn iter_mut(&mut self) -> DictIterMut<'_> { - match &mut self.value { - DictRepr::Small(entries) => DictIterMut::Small(entries.iter_mut()), - DictRepr::Tree(map) => DictIterMut::Tree(map.iter_mut()), - } + self.entries.iter_mut() } + /// Returns an iterator over `&String` keys, in key order. + #[inline] pub fn keys(&self) -> DictKeys<'_> { - DictKeys { inner: self.iter() } + self.entries.keys() } + /// Returns an iterator over `&Value` values, in key order. + #[inline] pub fn values(&self) -> DictValues<'_> { - DictValues { inner: self.iter() } + self.entries.values() } + /// Returns an iterator over `&mut Value` values, in key order. + #[inline] pub fn values_mut(&mut self) -> DictValuesMut<'_> { - DictValuesMut { - inner: self.iter_mut(), - } - } - - /// Returns `None` when the size hint signals the entry count will exceed - /// the small-vec threshold (callers should build a `Tree` directly), or - /// `Some(dict)` with a `Small`-backed dict pre-allocated to the hinted - /// capacity. - fn prepare_from_hint(lower: usize, upper: Option) -> Option { - if lower > Self::SMALL_DICT_MAX_ENTRIES_HINT - || upper.is_some_and(|upper| upper > Self::SMALL_DICT_MAX_ENTRIES_HINT) - { - return None; - } - let mut dict = Dict::new(); - if lower > 0 - && let DictRepr::Small(entries) = &mut dict.value - { - entries.reserve(lower.min(dict.small_max_entries)); - } - Some(dict) + self.entries.values_mut() } /// Constructs a `Dict` from a fallible iterator of `(String, Value)` pairs. /// - /// Applies the same size-hint optimisation as [`FromIterator`]: when the - /// iterator reports more than `small_max_entries` items the backing store - /// starts as a `Tree` directly, skipping the small-vec stage. - /// /// The first `Err` item short-circuits collection and is returned - /// immediately, leaving any remaining items unconsumed. + /// immediately, leaving any remaining items unconsumed. See + /// [`HybridMap::try_from_iter`] for details on the size-hint optimisation applied. pub fn try_from_iter(iter: I) -> Result where I: IntoIterator>, { - let iter = iter.into_iter(); - let (lower, upper) = iter.size_hint(); - - let Some(mut dict) = Dict::prepare_from_hint(lower, upper) else { - let map = iter.collect::>()?; - return Ok(Dict { - value: DictRepr::Tree(Box::new(map)), - small_max_entries: Dict::SMALL_DICT_MAX_ENTRIES_HINT, - }); - }; - - for result in iter { - let (k, v) = result?; - dict.insert(k, v); - } - Ok(dict) + HybridMap::try_from_iter(iter).map(|entries| Dict { entries }) } } @@ -405,90 +291,21 @@ impl Hash for Dict { } } -pub enum DictIter<'a> { - Small(std::slice::Iter<'a, (String, Value)>), - Tree(std::collections::btree_map::Iter<'a, String, Value>), -} +/// Iterator over `(&String, &Value)` pairs, in key order. +pub type DictIter<'a> = HybridIter<'a, String, Value, DictType>; -impl<'a> Iterator for DictIter<'a> { - type Item = (&'a String, &'a Value); +/// Iterator over `(&String, &mut Value)` pairs, in key order. +pub type DictIterMut<'a> = HybridIterMut<'a, String, Value, DictType>; - fn next(&mut self) -> Option { - match self { - DictIter::Small(iter) => iter.next().map(|(k, v)| (k, v)), - DictIter::Tree(iter) => iter.next(), - } - } - - fn size_hint(&self) -> (usize, Option) { - match self { - DictIter::Small(iter) => iter.size_hint(), - DictIter::Tree(iter) => iter.size_hint(), - } - } -} - -impl ExactSizeIterator for DictIter<'_> {} - -pub enum DictIterMut<'a> { - Small(std::slice::IterMut<'a, (String, Value)>), - Tree(std::collections::btree_map::IterMut<'a, String, Value>), -} - -impl<'a> Iterator for DictIterMut<'a> { - type Item = (&'a String, &'a mut Value); - - fn next(&mut self) -> Option { - match self { - DictIterMut::Small(iter) => iter.next().map(|(k, v)| (&*k, v)), - DictIterMut::Tree(iter) => iter.next(), - } - } - - fn size_hint(&self) -> (usize, Option) { - match self { - DictIterMut::Small(iter) => iter.size_hint(), - DictIterMut::Tree(iter) => iter.size_hint(), - } - } -} - -impl ExactSizeIterator for DictIterMut<'_> {} - -pub enum DictIntoIter { - Small(std::vec::IntoIter<(String, Value)>), - Tree(std::collections::btree_map::IntoIter), -} - -impl Iterator for DictIntoIter { - type Item = (String, Value); - - fn next(&mut self) -> Option { - match self { - DictIntoIter::Small(iter) => iter.next(), - DictIntoIter::Tree(iter) => iter.next(), - } - } - - fn size_hint(&self) -> (usize, Option) { - match self { - DictIntoIter::Small(iter) => iter.size_hint(), - DictIntoIter::Tree(iter) => iter.size_hint(), - } - } -} - -impl ExactSizeIterator for DictIntoIter {} +/// Owning iterator over `(String, Value)` pairs, in key order. +pub type DictIntoIter = HybridIntoIter; impl IntoIterator for Dict { type Item = (String, Value); type IntoIter = DictIntoIter; fn into_iter(self) -> Self::IntoIter { - match self.value { - DictRepr::Small(entries) => DictIntoIter::Small(entries.into_iter()), - DictRepr::Tree(map) => DictIntoIter::Tree(map.into_iter()), - } + self.entries.into_iter() } } @@ -519,59 +336,14 @@ impl Index<&str> for Dict { } } -pub struct DictKeys<'a> { - inner: DictIter<'a>, -} - -impl<'a> Iterator for DictKeys<'a> { - type Item = &'a String; - - fn next(&mut self) -> Option { - self.inner.next().map(|(k, _)| k) - } - - fn size_hint(&self) -> (usize, Option) { - self.inner.size_hint() - } -} - -impl ExactSizeIterator for DictKeys<'_> {} - -pub struct DictValues<'a> { - inner: DictIter<'a>, -} - -impl<'a> Iterator for DictValues<'a> { - type Item = &'a Value; - - fn next(&mut self) -> Option { - self.inner.next().map(|(_, v)| v) - } - - fn size_hint(&self) -> (usize, Option) { - self.inner.size_hint() - } -} - -impl ExactSizeIterator for DictValues<'_> {} +/// Iterator over `&String` keys, in key order. +pub type DictKeys<'a> = HybridKeys<'a, String, Value, DictType>; -pub struct DictValuesMut<'a> { - inner: DictIterMut<'a>, -} +/// Iterator over `&Value` values, in key order. +pub type DictValues<'a> = HybridValues<'a, String, Value, DictType>; -impl<'a> Iterator for DictValuesMut<'a> { - type Item = &'a mut Value; - - fn next(&mut self) -> Option { - self.inner.next().map(|(_, v)| v) - } - - fn size_hint(&self) -> (usize, Option) { - self.inner.size_hint() - } -} - -impl ExactSizeIterator for DictValuesMut<'_> {} +/// Iterator over `&mut Value` values, in key order. +pub type DictValuesMut<'a> = HybridValuesMut<'a, String, Value, DictType>; /// A newtype wrapper around any `IntoIterator` whose items are /// `Result<(String, Value), E>`, used as the source type for @@ -609,20 +381,9 @@ where /// Allows constructing a `Dict` from a `(String, Value)` tuple iterator impl FromIterator<(String, Value)> for Dict { fn from_iter>(iter: T) -> Self { - let mut iter = iter.into_iter(); - let (lower, upper) = iter.size_hint(); - - let Some(mut dict) = Dict::prepare_from_hint(lower, upper) else { - return Dict { - value: DictRepr::Tree(Box::new(iter.collect())), - small_max_entries: Dict::SMALL_DICT_MAX_ENTRIES_HINT, - }; - }; - - for (k, v) in iter.by_ref() { - dict.insert(k, v); + Dict { + entries: iter.into_iter().collect(), } - dict } } @@ -736,17 +497,8 @@ impl HaystackDict for Dict { /// Converts from `DictType` to a `Dict` impl From for Dict { fn from(from: DictType) -> Self { - let small_max_entries = Dict::SMALL_DICT_MAX_ENTRIES_HINT; - if from.len() <= small_max_entries { - Dict { - value: DictRepr::Small(from.into_iter().collect()), - small_max_entries, - } - } else { - Dict { - value: DictRepr::Tree(Box::new(from)), - small_max_entries, - } + Dict { + entries: HybridMap::from(from), } } } @@ -754,10 +506,7 @@ impl From for Dict { /// Converts from `Dict` to a `DictType` impl From for DictType { fn from(dict: Dict) -> Self { - match dict.value { - DictRepr::Small(entries) => entries.into_iter().collect(), - DictRepr::Tree(map) => *map, - } + dict.entries.into_inner() } } @@ -960,7 +709,6 @@ fn decode_str_from_value(val: &'_ Value) -> Cow<'_, str> { mod test { use std::borrow::Cow; - use crate::val::dict::DictRepr; use crate::val::{Dict, HaystackDict, Value, dict_to_dis}; fn get_localized<'a>(key: &str) -> Option> { @@ -1107,7 +855,7 @@ mod test { /// Returns true if the dict is backed by the Small (Vec) repr. fn is_small(d: &Dict) -> bool { - matches!(d.value, DictRepr::Small(_)) + d.entries.is_small() } // -- with_small_max_entries ------------------------------------------------ diff --git a/src/haystack/val/hybrid_map.rs b/src/haystack/val/hybrid_map.rs new file mode 100644 index 0000000..4747f1a --- /dev/null +++ b/src/haystack/val/hybrid_map.rs @@ -0,0 +1,877 @@ +// Copyright (C) 2020 - 2026, J2 Innovations + +//! A generic hybrid small-vector / tree-map backing store. +//! +//! [`HybridMap`] factors out the storage strategy used internally by +//! [`Dict`](crate::val::Dict): entries are kept in a sorted `Vec` while the +//! map stays small (avoiding a heap-heavy tree allocation for the common +//! case of a handful of tags), and are only promoted to a heavier ordered +//! map (by default a [`BTreeMap`]) once the entry count exceeds a +//! configurable threshold. +//! +//! Projects that wrap [`Value`](crate::val::Value) with their own extra +//! variants (and so can't reuse `Dict` itself) can reuse this type directly +//! with their own key/value/map types instead of reimplementing the hybrid +//! storage strategy, e.g. `HybridMap`. + +use std::borrow::Borrow; +use std::collections::BTreeMap; + +/// Trait implemented by the "heavy" backing map used as the fallback tier of +/// a [`HybridMap`] once its entry count exceeds the small-vector threshold. +/// +/// Implemented out of the box for [`BTreeMap`]; other ordered-map +/// implementations can implement this trait to plug into a `HybridMap`. +pub trait TreeMap: Default + FromIterator<(K, V)> { + /// Iterator returned by [`TreeMap::iter`]. + type Iter<'a>: Iterator + ExactSizeIterator + where + Self: 'a, + K: 'a, + V: 'a; + + /// Iterator returned by [`TreeMap::iter_mut`]. + type IterMut<'a>: Iterator + ExactSizeIterator + where + Self: 'a, + K: 'a, + V: 'a; + + /// Iterator returned by [`TreeMap::into_iter`]. + type IntoIter: Iterator + ExactSizeIterator; + + /// Returns the number of entries in the map. + fn len(&self) -> usize; + + /// Returns `true` if the map contains no entries. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Removes all entries from the map. + fn clear(&mut self); + + /// Inserts a key-value pair, returning the previous value for `key`, if any. + fn insert(&mut self, key: K, value: V) -> Option; + + /// Removes and returns the first (lowest-keyed) entry, if any. + fn pop_first(&mut self) -> Option<(K, V)>; + + /// Retains only the entries for which the predicate returns `true`. + fn retain bool>(&mut self, f: F); + + /// Returns a reference to the value for `key`, if present. + fn get(&self, key: &Q) -> Option<&V> + where + K: Borrow + Ord, + Q: Ord + ?Sized; + + /// Returns a mutable reference to the value for `key`, if present. + fn get_mut(&mut self, key: &Q) -> Option<&mut V> + where + K: Borrow + Ord, + Q: Ord + ?Sized; + + /// Removes `key` from the map, returning its value if it was present. + fn remove(&mut self, key: &Q) -> Option + where + K: Borrow + Ord, + Q: Ord + ?Sized; + + /// Returns an iterator over `(&K, &V)` pairs, in key order. + fn iter(&self) -> Self::Iter<'_>; + + /// Returns an iterator over `(&K, &mut V)` pairs, in key order. + fn iter_mut(&mut self) -> Self::IterMut<'_>; + + /// Converts the map into an owning iterator over `(K, V)` pairs, in key order. + fn into_iter(self) -> Self::IntoIter; +} + +impl TreeMap for BTreeMap { + type Iter<'a> + = std::collections::btree_map::Iter<'a, K, V> + where + K: 'a, + V: 'a; + type IterMut<'a> + = std::collections::btree_map::IterMut<'a, K, V> + where + K: 'a, + V: 'a; + type IntoIter = std::collections::btree_map::IntoIter; + + fn len(&self) -> usize { + self.len() + } + + fn clear(&mut self) { + self.clear(); + } + + fn insert(&mut self, key: K, value: V) -> Option { + self.insert(key, value) + } + + fn pop_first(&mut self) -> Option<(K, V)> { + self.pop_first() + } + + fn retain bool>(&mut self, f: F) { + self.retain(f); + } + + fn get(&self, key: &Q) -> Option<&V> + where + K: Borrow + Ord, + Q: Ord + ?Sized, + { + self.get(key) + } + + fn get_mut(&mut self, key: &Q) -> Option<&mut V> + where + K: Borrow + Ord, + Q: Ord + ?Sized, + { + self.get_mut(key) + } + + fn remove(&mut self, key: &Q) -> Option + where + K: Borrow + Ord, + Q: Ord + ?Sized, + { + self.remove(key) + } + + fn iter(&self) -> Self::Iter<'_> { + self.iter() + } + + fn iter_mut(&mut self) -> Self::IterMut<'_> { + self.iter_mut() + } + + fn into_iter(self) -> Self::IntoIter { + IntoIterator::into_iter(self) + } +} + +#[derive(Clone, Debug)] +enum HybridRepr { + Small(Vec<(K, V)>), + // Boxed so the discriminant can be packed into `Vec`'s pointer niche. + Tree(Box), +} + +/// A hybrid map: a sorted small-vector back-store for tiny maps, spilling +/// over to a heavier ordered map `M` (a [`BTreeMap`] by default) once +/// the entry count exceeds a threshold. +/// +/// Iteration is always in key order, regardless of the active back-store. +#[derive(Clone, Debug)] +pub struct HybridMap> { + repr: HybridRepr, + small_max_entries: usize, +} + +impl HybridMap +where + K: Ord, + M: TreeMap, +{ + /// Default threshold for the small-vector back-store. + pub const DEFAULT_SMALL_MAX_ENTRIES: usize = 32; + + /// Construct a new `HybridMap` with the default small-store threshold. + pub fn new() -> Self { + Self::with_small_max_entries(Self::DEFAULT_SMALL_MAX_ENTRIES) + } + + /// Construct a new `HybridMap` with a custom small-store threshold. + /// If `small_max_entries` is 0, the small-vector back-store is disabled + /// and the map will use the `M` representation right away. + pub fn with_small_max_entries(small_max_entries: usize) -> Self { + let repr = if small_max_entries == 0 { + HybridRepr::Tree(Box::default()) + } else { + HybridRepr::Small(Vec::new()) + }; + HybridMap { + repr, + small_max_entries, + } + } + + /// Return the active small-store threshold for this map. + #[inline] + pub fn small_max_entries(&self) -> usize { + self.small_max_entries + } + + /// True if this map is currently backed by the small-vector store. + #[inline] + pub fn is_small(&self) -> bool { + matches!(self.repr, HybridRepr::Small(_)) + } + + fn small_search(entries: &[(K, V)], key: &Q) -> Result + where + K: Borrow, + Q: Ord + ?Sized, + { + entries.binary_search_by(|(k, _)| k.borrow().cmp(key)) + } + + fn spill_to_tree(&mut self) { + if let HybridRepr::Small(entries) = &mut self.repr { + let map = entries.drain(..).collect::(); + self.repr = HybridRepr::Tree(Box::new(map)); + } + } + + #[inline] + pub fn len(&self) -> usize { + match &self.repr { + HybridRepr::Small(entries) => entries.len(), + HybridRepr::Tree(map) => map.len(), + } + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn clear(&mut self) { + match &mut self.repr { + HybridRepr::Small(entries) => entries.clear(), + HybridRepr::Tree(map) => map.clear(), + } + } + + #[inline] + pub fn contains_key(&self, key: &Q) -> bool + where + K: Borrow, + Q: Ord + ?Sized, + { + self.get(key).is_some() + } + + #[inline] + pub fn get(&self, key: &Q) -> Option<&V> + where + K: Borrow, + Q: Ord + ?Sized, + { + match &self.repr { + HybridRepr::Small(entries) => Self::small_search(entries, key) + .ok() + .map(|pos| &entries[pos].1), + HybridRepr::Tree(map) => map.get(key), + } + } + + #[inline] + pub fn get_mut(&mut self, key: &Q) -> Option<&mut V> + where + K: Borrow, + Q: Ord + ?Sized, + { + match &mut self.repr { + HybridRepr::Small(entries) => Self::small_search(entries, key) + .ok() + .map(|pos| &mut entries[pos].1), + HybridRepr::Tree(map) => map.get_mut(key), + } + } + + pub fn insert(&mut self, key: K, value: V) -> Option { + match &mut self.repr { + HybridRepr::Small(entries) => { + if entries.len() < self.small_max_entries + && entries.last().is_none_or(|(last_key, _)| &key > last_key) + { + entries.push((key, value)); + return None; + } + + match Self::small_search(entries, &key) { + Ok(pos) => Some(std::mem::replace(&mut entries[pos].1, value)), + Err(pos) => { + if entries.len() < self.small_max_entries { + entries.insert(pos, (key, value)); + None + } else { + self.spill_to_tree(); + match &mut self.repr { + HybridRepr::Tree(map) => map.insert(key, value), + HybridRepr::Small(_) => None, + } + } + } + } + } + HybridRepr::Tree(map) => map.insert(key, value), + } + } + + #[inline] + pub fn remove(&mut self, key: &Q) -> Option + where + K: Borrow, + Q: Ord + ?Sized, + { + match &mut self.repr { + HybridRepr::Small(entries) => Self::small_search(entries, key) + .ok() + .map(|pos| entries.remove(pos).1), + HybridRepr::Tree(map) => map.remove(key), + } + } + + #[inline] + pub fn pop_first(&mut self) -> Option<(K, V)> { + match &mut self.repr { + HybridRepr::Small(entries) => { + if entries.is_empty() { + None + } else { + Some(entries.remove(0)) + } + } + HybridRepr::Tree(map) => map.pop_first(), + } + } + + /// Demote a tree-backed map back to the small-vector store once its + /// entry count has dropped to at or below the small-store threshold. + /// + /// This is the inverse of the automatic spill that happens in + /// [`insert`](Self::insert). If the map is already small-backed this is + /// a no-op. + pub fn shrink_to_fit(&mut self) { + if let HybridRepr::Tree(map) = &mut self.repr + && map.len() <= self.small_max_entries + { + // Take ownership without cloning; `M: Default` leaves a valid empty map behind. + let owned = std::mem::take(map.as_mut()); + self.repr = HybridRepr::Small(owned.into_iter().collect()); + } + } + + /// Retains only the entries for which the predicate returns `true`. + /// + /// When called on a tree-backed map and the surviving entry count drops + /// to or below the small-store threshold, the storage is automatically + /// downgraded back to the small-vector representation. + pub fn retain(&mut self, mut f: F) + where + F: FnMut(&K, &mut V) -> bool, + { + match &mut self.repr { + HybridRepr::Small(entries) => entries.retain_mut(|(k, v)| f(k, v)), + HybridRepr::Tree(map) => map.retain(|k, v| f(k, v)), + } + self.shrink_to_fit(); + } + + #[inline] + pub fn iter(&self) -> HybridIter<'_, K, V, M> { + match &self.repr { + HybridRepr::Small(entries) => HybridIter::Small(entries.iter()), + HybridRepr::Tree(map) => HybridIter::Tree(map.iter()), + } + } + + #[inline] + pub fn iter_mut(&mut self) -> HybridIterMut<'_, K, V, M> { + match &mut self.repr { + HybridRepr::Small(entries) => HybridIterMut::Small(entries.iter_mut()), + HybridRepr::Tree(map) => HybridIterMut::Tree(map.iter_mut()), + } + } + + #[inline] + pub fn keys(&self) -> HybridKeys<'_, K, V, M> { + HybridKeys { inner: self.iter() } + } + + #[inline] + pub fn values(&self) -> HybridValues<'_, K, V, M> { + HybridValues { inner: self.iter() } + } + + #[inline] + pub fn values_mut(&mut self) -> HybridValuesMut<'_, K, V, M> { + HybridValuesMut { + inner: self.iter_mut(), + } + } + + /// Returns `None` when the size hint signals the entry count will exceed + /// the small-vec threshold (callers should build a tree-backed map + /// directly), or `Some(map)` with a small-backed map pre-allocated to + /// the hinted capacity. + fn prepare_from_hint(lower: usize, upper: Option) -> Option { + if lower > Self::DEFAULT_SMALL_MAX_ENTRIES + || upper.is_some_and(|upper| upper > Self::DEFAULT_SMALL_MAX_ENTRIES) + { + return None; + } + let mut map = Self::new(); + if lower > 0 + && let HybridRepr::Small(entries) = &mut map.repr + { + entries.reserve(lower.min(map.small_max_entries)); + } + Some(map) + } + + /// Constructs a `HybridMap` from a fallible iterator of `(K, V)` pairs. + /// + /// Applies the same size-hint optimisation as [`FromIterator`]: when the + /// iterator reports more than `small_max_entries` items the backing + /// store starts as a tree directly, skipping the small-vec stage. + /// + /// The first `Err` item short-circuits collection and is returned + /// immediately, leaving any remaining items unconsumed. + pub fn try_from_iter(iter: I) -> Result + where + I: IntoIterator>, + { + let iter = iter.into_iter(); + let (lower, upper) = iter.size_hint(); + + let Some(mut map) = Self::prepare_from_hint(lower, upper) else { + let tree = iter.collect::>()?; + return Ok(HybridMap { + repr: HybridRepr::Tree(Box::new(tree)), + small_max_entries: Self::DEFAULT_SMALL_MAX_ENTRIES, + }); + }; + + for result in iter { + let (k, v) = result?; + map.insert(k, v); + } + Ok(map) + } +} + +impl Default for HybridMap +where + K: Ord, + M: TreeMap, +{ + fn default() -> Self { + Self::new() + } +} + +/// Implement FromIterator for `HybridMap` +impl FromIterator<(K, V)> for HybridMap +where + K: Ord, + M: TreeMap, +{ + fn from_iter>(iter: T) -> Self { + let mut iter = iter.into_iter(); + let (lower, upper) = iter.size_hint(); + + let Some(mut map) = Self::prepare_from_hint(lower, upper) else { + return HybridMap { + repr: HybridRepr::Tree(Box::new(iter.collect())), + small_max_entries: Self::DEFAULT_SMALL_MAX_ENTRIES, + }; + }; + + for (k, v) in iter.by_ref() { + map.insert(k, v); + } + map + } +} + +/// Converts from the heavy map type `M` to a `HybridMap` +impl From for HybridMap +where + K: Ord, + M: TreeMap, +{ + fn from(from: M) -> Self { + let small_max_entries = Self::DEFAULT_SMALL_MAX_ENTRIES; + if from.len() <= small_max_entries { + HybridMap { + repr: HybridRepr::Small(from.into_iter().collect()), + small_max_entries, + } + } else { + HybridMap { + repr: HybridRepr::Tree(Box::new(from)), + small_max_entries, + } + } + } +} + +impl HybridMap +where + K: Ord, + M: TreeMap, +{ + /// Converts this `HybridMap` back into the heavy map type `M`. + /// + /// A free-standing `From> for M` impl isn't possible + /// here since `M` is a type parameter, not a local type (orphan rules). + pub fn into_inner(self) -> M { + match self.repr { + HybridRepr::Small(entries) => entries.into_iter().collect(), + HybridRepr::Tree(map) => *map, + } + } +} + +pub enum HybridIter<'a, K: 'a, V: 'a, M: TreeMap + 'a> { + Small(std::slice::Iter<'a, (K, V)>), + Tree(M::Iter<'a>), +} + +impl<'a, K, V, M: TreeMap + 'a> Iterator for HybridIter<'a, K, V, M> { + type Item = (&'a K, &'a V); + + fn next(&mut self) -> Option { + match self { + HybridIter::Small(iter) => iter.next().map(|(k, v)| (k, v)), + HybridIter::Tree(iter) => iter.next(), + } + } + + fn size_hint(&self) -> (usize, Option) { + match self { + HybridIter::Small(iter) => iter.size_hint(), + HybridIter::Tree(iter) => iter.size_hint(), + } + } +} + +impl<'a, K, V, M: TreeMap + 'a> ExactSizeIterator for HybridIter<'a, K, V, M> {} + +pub enum HybridIterMut<'a, K: 'a, V: 'a, M: TreeMap + 'a> { + Small(std::slice::IterMut<'a, (K, V)>), + Tree(M::IterMut<'a>), +} + +impl<'a, K, V, M: TreeMap + 'a> Iterator for HybridIterMut<'a, K, V, M> { + type Item = (&'a K, &'a mut V); + + fn next(&mut self) -> Option { + match self { + HybridIterMut::Small(iter) => iter.next().map(|(k, v)| (&*k, v)), + HybridIterMut::Tree(iter) => iter.next(), + } + } + + fn size_hint(&self) -> (usize, Option) { + match self { + HybridIterMut::Small(iter) => iter.size_hint(), + HybridIterMut::Tree(iter) => iter.size_hint(), + } + } +} + +impl<'a, K, V, M: TreeMap + 'a> ExactSizeIterator for HybridIterMut<'a, K, V, M> {} + +pub enum HybridIntoIter> { + Small(std::vec::IntoIter<(K, V)>), + Tree(M::IntoIter), +} + +impl> Iterator for HybridIntoIter { + type Item = (K, V); + + fn next(&mut self) -> Option { + match self { + HybridIntoIter::Small(iter) => iter.next(), + HybridIntoIter::Tree(iter) => iter.next(), + } + } + + fn size_hint(&self) -> (usize, Option) { + match self { + HybridIntoIter::Small(iter) => iter.size_hint(), + HybridIntoIter::Tree(iter) => iter.size_hint(), + } + } +} + +impl> ExactSizeIterator for HybridIntoIter {} + +impl> IntoIterator for HybridMap { + type Item = (K, V); + type IntoIter = HybridIntoIter; + + fn into_iter(self) -> Self::IntoIter { + match self.repr { + HybridRepr::Small(entries) => HybridIntoIter::Small(entries.into_iter()), + HybridRepr::Tree(map) => HybridIntoIter::Tree(map.into_iter()), + } + } +} + +impl<'a, K: Ord, V, M: TreeMap + 'a> IntoIterator for &'a HybridMap { + type Item = (&'a K, &'a V); + type IntoIter = HybridIter<'a, K, V, M>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl<'a, K: Ord, V, M: TreeMap + 'a> IntoIterator for &'a mut HybridMap { + type Item = (&'a K, &'a mut V); + type IntoIter = HybridIterMut<'a, K, V, M>; + + fn into_iter(self) -> Self::IntoIter { + self.iter_mut() + } +} + +pub struct HybridKeys<'a, K: 'a, V: 'a, M: TreeMap + 'a> { + inner: HybridIter<'a, K, V, M>, +} + +impl<'a, K, V, M: TreeMap + 'a> Iterator for HybridKeys<'a, K, V, M> { + type Item = &'a K; + + fn next(&mut self) -> Option { + self.inner.next().map(|(k, _)| k) + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} + +impl<'a, K, V, M: TreeMap + 'a> ExactSizeIterator for HybridKeys<'a, K, V, M> {} + +pub struct HybridValues<'a, K: 'a, V: 'a, M: TreeMap + 'a> { + inner: HybridIter<'a, K, V, M>, +} + +impl<'a, K, V, M: TreeMap + 'a> Iterator for HybridValues<'a, K, V, M> { + type Item = &'a V; + + fn next(&mut self) -> Option { + self.inner.next().map(|(_, v)| v) + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} + +impl<'a, K, V, M: TreeMap + 'a> ExactSizeIterator for HybridValues<'a, K, V, M> {} + +pub struct HybridValuesMut<'a, K: 'a, V: 'a, M: TreeMap + 'a> { + inner: HybridIterMut<'a, K, V, M>, +} + +impl<'a, K, V, M: TreeMap + 'a> Iterator for HybridValuesMut<'a, K, V, M> { + type Item = &'a mut V; + + fn next(&mut self) -> Option { + self.inner.next().map(|(_, v)| v) + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} + +impl<'a, K, V, M: TreeMap + 'a> ExactSizeIterator for HybridValuesMut<'a, K, V, M> {} + +#[cfg(test)] +mod test { + use super::*; + + /// A minimal custom `TreeMap` impl (a sorted `Vec`, distinct from + /// `BTreeMap`) used to verify that `HybridMap`/`TreeMap` are genuinely + /// generic over the heavy map type, not implicitly tied to `BTreeMap`. + #[derive(Clone, Debug)] + struct SortedVecMap(Vec<(K, V)>); + + impl Default for SortedVecMap { + fn default() -> Self { + SortedVecMap(Vec::new()) + } + } + + impl FromIterator<(K, V)> for SortedVecMap { + fn from_iter>(iter: T) -> Self { + let mut entries: Vec<(K, V)> = iter.into_iter().collect(); + entries.sort_by(|a, b| a.0.cmp(&b.0)); + SortedVecMap(entries) + } + } + + impl TreeMap for SortedVecMap { + type Iter<'a> + = std::iter::Map, fn(&'a (K, V)) -> (&'a K, &'a V)> + where + K: 'a, + V: 'a; + type IterMut<'a> + = std::iter::Map< + std::slice::IterMut<'a, (K, V)>, + fn(&'a mut (K, V)) -> (&'a K, &'a mut V), + > + where + K: 'a, + V: 'a; + type IntoIter = std::vec::IntoIter<(K, V)>; + + fn len(&self) -> usize { + self.0.len() + } + + fn clear(&mut self) { + self.0.clear(); + } + + fn insert(&mut self, key: K, value: V) -> Option { + match self.0.binary_search_by(|(k, _)| k.cmp(&key)) { + Ok(pos) => Some(std::mem::replace(&mut self.0[pos].1, value)), + Err(pos) => { + self.0.insert(pos, (key, value)); + None + } + } + } + + fn pop_first(&mut self) -> Option<(K, V)> { + if self.0.is_empty() { + None + } else { + Some(self.0.remove(0)) + } + } + + fn retain bool>(&mut self, mut f: F) { + self.0.retain_mut(|(k, v)| f(k, v)); + } + + fn get(&self, key: &Q) -> Option<&V> + where + K: Borrow + Ord, + Q: Ord + ?Sized, + { + self.0 + .binary_search_by(|(k, _)| k.borrow().cmp(key)) + .ok() + .map(|pos| &self.0[pos].1) + } + + fn get_mut(&mut self, key: &Q) -> Option<&mut V> + where + K: Borrow + Ord, + Q: Ord + ?Sized, + { + match self.0.binary_search_by(|(k, _)| k.borrow().cmp(key)) { + Ok(pos) => Some(&mut self.0[pos].1), + Err(_) => None, + } + } + + fn remove(&mut self, key: &Q) -> Option + where + K: Borrow + Ord, + Q: Ord + ?Sized, + { + self.0 + .binary_search_by(|(k, _)| k.borrow().cmp(key)) + .ok() + .map(|pos| self.0.remove(pos).1) + } + + fn iter(&self) -> Self::Iter<'_> { + self.0.iter().map(|(k, v)| (k, v)) + } + + fn iter_mut(&mut self) -> Self::IterMut<'_> { + self.0.iter_mut().map(|(k, v)| (&*k, v)) + } + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } + } + + /// Insert `n` ordered keys "k00".."kNN" into a `HybridMap` backed by + /// `SortedVecMap` with the given threshold. + fn make_hybrid( + n: usize, + threshold: usize, + ) -> HybridMap> { + let mut map = HybridMap::with_small_max_entries(threshold); + for i in 0..n { + map.insert(format!("k{i:02}"), i as i32); + } + map + } + + #[test] + fn custom_tree_map_stays_small_below_threshold() { + let map = make_hybrid(4, 8); + assert!(map.is_small()); + assert_eq!(map.len(), 4); + } + + #[test] + fn custom_tree_map_spills_to_tree_at_threshold() { + let mut map = make_hybrid(8, 8); + assert!(map.is_small()); + map.insert("z_extra".into(), 99); + assert!(!map.is_small()); + assert_eq!(map.len(), 9); + assert_eq!(map.get("z_extra"), Some(&99)); + } + + #[test] + fn custom_tree_map_get_remove_and_ordering() { + let mut map = make_hybrid(10, 4); // spills to SortedVecMap + assert!(!map.is_small()); + assert_eq!(map.get("k05"), Some(&5)); + assert_eq!(map.remove("k05"), Some(5)); + assert!(map.get("k05").is_none()); + assert_eq!(map.len(), 9); + + let keys: Vec<&String> = map.keys().collect(); + let mut expected = keys.clone(); + expected.sort(); + assert_eq!(keys, expected); + } + + #[test] + fn custom_tree_map_shrink_to_fit_round_trip() { + let mut map = make_hybrid(10, 8); // spills at 9 entries + for i in 8..10 { + map.remove(&format!("k{i:02}")); + } + assert!(!map.is_small()); + map.shrink_to_fit(); + assert!(map.is_small()); + assert_eq!(map.len(), 8); + } + + #[test] + fn custom_tree_map_into_inner_round_trip() { + let map = make_hybrid(20, 4); // tree-backed + let inner: SortedVecMap = map.into_inner(); + assert_eq!(inner.len(), 20); + assert_eq!(inner.get("k10"), Some(&10)); + } +}