From 06dd7a5a9c1bc530c1e3350558ceadd58385c4b3 Mon Sep 17 00:00:00 2001 From: Tage Johansson Date: Fri, 15 Nov 2024 19:18:23 +0100 Subject: [PATCH] Add support for custom allocators by the allocator-api2 crate. This commit adds allocator-api2 (a drop-in replacement for the `Allocator` trait on stable) as a dependency. The underlying `Vec` is replaced by `allocator_api2::vec::Vec`, and the type signiture of `VecMap` is changed to `VecMap`. It adds the methods `new_in()` and `with_capacity_in()` which creates a `VecMap` with a custom allocator. The old `new()` and `with_capacity()` methods are not removed and defaults to the default global allocator. I think that this commit is backwards compatable. --- Cargo.toml | 2 + src/lib.rs | 136 +++++++++++++++++++++++++++++++++-------------------- 2 files changed, 88 insertions(+), 50 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f09dafd..6748036 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,8 +39,10 @@ readme = "README.md" exclude = ["/.travis.yml", "/deploy-docs.sh"] [features] +serde = ["dep:serde", "allocator-api2/serde"] # This feature is kept for backwards compatibility. Use feature "serde" instead. eders = [ "serde" ] [dependencies] +allocator-api2 = "0.2.20" serde = { version = "1.0", features = [ "derive" ], optional = true } diff --git a/src/lib.rs b/src/lib.rs index 209ef7c..e66fb1d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,10 @@ extern crate serde; use self::Entry::*; +use allocator_api2::{ + alloc::{Allocator, Global}, + vec::{self, Vec}, +}; use std::cmp::{max, Ordering}; use std::fmt; use std::hash::{Hash, Hasher}; @@ -29,7 +33,6 @@ use std::iter::{Enumerate, FilterMap, FromIterator}; use std::mem::{replace, swap}; use std::ops::{Index, IndexMut}; use std::slice; -use std::vec; /// A map optimized for small integer keys. /// @@ -63,30 +66,37 @@ use std::vec; /// months.clear(); /// assert!(months.is_empty()); /// ``` -#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] -pub struct VecMap { +#[cfg_attr( + feature = "serde", + derive(Serialize, Deserialize), + serde(bound( + deserialize = "V: serde::Deserialize<'de>, A: Default", + serialize = "V: serde::Serialize" + )) +)] +pub struct VecMap { n: usize, - v: Vec>, + v: Vec, A>, } /// A view into a single entry in a map, which may either be vacant or occupied. -pub enum Entry<'a, V> { +pub enum Entry<'a, V, A: Allocator = Global> { /// A vacant Entry - Vacant(VacantEntry<'a, V>), + Vacant(VacantEntry<'a, V, A>), /// An occupied Entry - Occupied(OccupiedEntry<'a, V>), + Occupied(OccupiedEntry<'a, V, A>), } /// A vacant Entry. -pub struct VacantEntry<'a, V> { - map: &'a mut VecMap, +pub struct VacantEntry<'a, V, A: Allocator = Global> { + map: &'a mut VecMap, index: usize, } /// An occupied Entry. -pub struct OccupiedEntry<'a, V> { - map: &'a mut VecMap, +pub struct OccupiedEntry<'a, V, A: Allocator = Global> { + map: &'a mut VecMap, index: usize, } @@ -97,7 +107,7 @@ impl Default for VecMap { } } -impl Hash for VecMap { +impl Hash for VecMap { fn hash(&self, state: &mut H) { // In order to not traverse the `VecMap` twice, count the elements // during iteration. @@ -120,7 +130,7 @@ impl VecMap { /// let mut map: VecMap<&str> = VecMap::new(); /// ``` pub fn new() -> Self { - VecMap { n: 0, v: vec![] } + VecMap::new_in(Default::default()) } /// Creates an empty `VecMap` with space for at least `capacity` @@ -133,9 +143,25 @@ impl VecMap { /// let mut map: VecMap<&str> = VecMap::with_capacity(10); /// ``` pub fn with_capacity(capacity: usize) -> Self { + VecMap::with_capacity_in(capacity, Default::default()) + } +} + +impl VecMap { + /// Creates an empty `VecMap` with a custom allocator. + pub fn new_in(alloc: A) -> Self { VecMap { n: 0, - v: Vec::with_capacity(capacity), + v: Vec::new_in(alloc), + } + } + + /// Creates an empty `VecMap` with space for at least `capacity` + /// elements before resizing, in a custom allocator. + pub fn with_capacity_in(capacity: usize, alloc: A) -> Self { + VecMap { + n: 0, + v: Vec::with_capacity_in(capacity, alloc), } } @@ -221,6 +247,11 @@ impl VecMap { self.v.shrink_to_fit() } + /// Returns a reference to the allocator for the `VecMap`. + pub fn allocator(&self) -> &A { + self.v.allocator() + } + /// Returns an iterator visiting all keys in ascending order of the keys. /// The iterator's element type is `usize`. pub fn keys(&self) -> Keys<'_, V> { @@ -355,8 +386,11 @@ impl VecMap { /// assert_eq!(b[3], "c"); /// assert_eq!(b[4], "d"); /// ``` - pub fn split_off(&mut self, at: usize) -> Self { - let mut other = VecMap::new(); + pub fn split_off(&mut self, at: usize) -> Self + where + A: Copy, + { + let mut other = VecMap::new_in(*self.allocator()); if at == 0 { // Move all elements to other @@ -414,7 +448,7 @@ impl VecMap { /// /// assert_eq!(vec, [(1, "a"), (2, "b"), (3, "c")]); /// ``` - pub fn drain(&mut self) -> Drain<'_, V> { + pub fn drain(&mut self) -> Drain<'_, V, A> { fn filter((i, v): (usize, Option)) -> Option<(usize, A)> { v.map(|v| (i, v)) } @@ -603,7 +637,7 @@ impl VecMap { /// /// assert_eq!(count[1], 3); /// ``` - pub fn entry(&mut self, key: usize) -> Entry<'_, V> { + pub fn entry(&mut self, key: usize) -> Entry<'_, V, A> { // FIXME(Gankro): this is basically the dumbest implementation of // entry possible, because weird non-lexical borrows issues make it // completely insane to do any other way. That said, Entry is a border-line @@ -651,7 +685,7 @@ impl VecMap { } } -impl<'a, V> Entry<'a, V> { +impl<'a, V, A: Allocator> Entry<'a, V, A> { /// Ensures a value is in the entry by inserting the default if empty, and /// returns a mutable reference to the value in the entry. pub fn or_insert(self, default: V) -> &'a mut V { @@ -672,7 +706,7 @@ impl<'a, V> Entry<'a, V> { } } -impl<'a, V> VacantEntry<'a, V> { +impl<'a, V, A: Allocator> VacantEntry<'a, V, A> { /// Sets the value of the entry with the VacantEntry's key, /// and returns a mutable reference to it. pub fn insert(self, value: V) -> &'a mut V { @@ -682,7 +716,7 @@ impl<'a, V> VacantEntry<'a, V> { } } -impl<'a, V> OccupiedEntry<'a, V> { +impl<'a, V, A: Allocator> OccupiedEntry<'a, V, A> { /// Gets a reference to the value in the entry. pub fn get(&self) -> &V { let index = self.index; @@ -715,7 +749,7 @@ impl<'a, V> OccupiedEntry<'a, V> { } } -impl Clone for VecMap { +impl Clone for VecMap { #[inline] fn clone(&self) -> Self { VecMap { @@ -731,29 +765,29 @@ impl Clone for VecMap { } } -impl PartialEq for VecMap { +impl PartialEq for VecMap { fn eq(&self, other: &Self) -> bool { self.n == other.n && self.iter().eq(other.iter()) } } -impl Eq for VecMap {} +impl Eq for VecMap {} -impl PartialOrd for VecMap { +impl PartialOrd for VecMap { #[inline] fn partial_cmp(&self, other: &Self) -> Option { self.iter().partial_cmp(other.iter()) } } -impl Ord for VecMap { +impl Ord for VecMap { #[inline] fn cmp(&self, other: &Self) -> Ordering { self.iter().cmp(other.iter()) } } -impl fmt::Debug for VecMap { +impl fmt::Debug for VecMap { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_map().entries(self).finish() } @@ -767,9 +801,9 @@ impl FromIterator<(usize, V)> for VecMap { } } -impl IntoIterator for VecMap { +impl IntoIterator for VecMap { type Item = (usize, T); - type IntoIter = IntoIter; + type IntoIter = IntoIter; /// Returns an iterator visiting all key-value pairs in ascending order of /// the keys, consuming the original `VecMap`. @@ -789,7 +823,7 @@ impl IntoIterator for VecMap { /// /// assert_eq!(vec, [(1, "a"), (2, "b"), (3, "c")]); /// ``` - fn into_iter(self) -> IntoIter { + fn into_iter(self) -> IntoIter { IntoIter { n: self.n, yielded: 0, @@ -798,7 +832,7 @@ impl IntoIterator for VecMap { } } -impl<'a, T> IntoIterator for &'a VecMap { +impl<'a, T, A: Allocator> IntoIterator for &'a VecMap { type Item = (usize, &'a T); type IntoIter = Iter<'a, T>; @@ -807,7 +841,7 @@ impl<'a, T> IntoIterator for &'a VecMap { } } -impl<'a, T> IntoIterator for &'a mut VecMap { +impl<'a, T, A: Allocator> IntoIterator for &'a mut VecMap { type Item = (usize, &'a mut T); type IntoIter = IterMut<'a, T>; @@ -816,7 +850,7 @@ impl<'a, T> IntoIterator for &'a mut VecMap { } } -impl Extend<(usize, V)> for VecMap { +impl Extend<(usize, V)> for VecMap { fn extend>(&mut self, iter: I) { for (k, v) in iter { self.insert(k, v); @@ -824,13 +858,13 @@ impl Extend<(usize, V)> for VecMap { } } -impl<'a, V: Copy> Extend<(usize, &'a V)> for VecMap { +impl<'a, V: Copy, A: Allocator> Extend<(usize, &'a V)> for VecMap { fn extend>(&mut self, iter: I) { self.extend(iter.into_iter().map(|(key, &value)| (key, value))); } } -impl Index for VecMap { +impl Index for VecMap { type Output = V; #[inline] @@ -839,7 +873,7 @@ impl Index for VecMap { } } -impl<'a, V> Index<&'a usize> for VecMap { +impl<'a, V, A: Allocator> Index<&'a usize> for VecMap { type Output = V; #[inline] @@ -848,14 +882,14 @@ impl<'a, V> Index<&'a usize> for VecMap { } } -impl IndexMut for VecMap { +impl IndexMut for VecMap { #[inline] fn index_mut(&mut self, i: usize) -> &mut V { self.get_mut(i).expect("key not present") } } -impl<'a, V> IndexMut<&'a usize> for VecMap { +impl<'a, V, A: Allocator> IndexMut<&'a usize> for VecMap { #[inline] fn index_mut(&mut self, i: &usize) -> &mut V { self.get_mut(*i).expect("key not present") @@ -985,21 +1019,21 @@ pub struct ValuesMut<'a, V> { } /// A consuming iterator over the key-value pairs of a map. -pub struct IntoIter { +pub struct IntoIter { n: usize, yielded: usize, - iter: Enumerate>>, + iter: Enumerate, A>>, } /// A draining iterator over the key-value pairs of a map. -pub struct Drain<'a, V> { +pub struct Drain<'a, V, A: Allocator = Global> { iter: FilterMap< - Enumerate>>, + Enumerate, A>>, fn((usize, Option)) -> Option<(usize, V)>, >, } -impl<'a, V> Iterator for Drain<'a, V> { +impl<'a, V, A: Allocator> Iterator for Drain<'a, V, A> { type Item = (usize, V); fn next(&mut self) -> Option<(usize, V)> { @@ -1010,9 +1044,9 @@ impl<'a, V> Iterator for Drain<'a, V> { } } -impl<'a, V> ExactSizeIterator for Drain<'a, V> {} +impl<'a, V, A: Allocator> ExactSizeIterator for Drain<'a, V, A> {} -impl<'a, V> DoubleEndedIterator for Drain<'a, V> { +impl<'a, V, A: Allocator> DoubleEndedIterator for Drain<'a, V, A> { fn next_back(&mut self) -> Option<(usize, V)> { self.iter.next_back() } @@ -1075,7 +1109,7 @@ impl<'a, V> DoubleEndedIterator for ValuesMut<'a, V> { } } -impl Iterator for IntoIter { +impl Iterator for IntoIter { type Item = (usize, V); fn next(&mut self) -> Option<(usize, V)> { @@ -1096,9 +1130,9 @@ impl Iterator for IntoIter { } } -impl ExactSizeIterator for IntoIter {} +impl ExactSizeIterator for IntoIter {} -impl DoubleEndedIterator for IntoIter { +impl DoubleEndedIterator for IntoIter { fn next_back(&mut self) -> Option<(usize, V)> { loop { match self.iter.next_back() { @@ -1112,11 +1146,13 @@ impl DoubleEndedIterator for IntoIter { #[allow(dead_code)] fn assert_properties() { - fn vec_map_covariant<'a, T>(map: VecMap<&'static T>) -> VecMap<&'a T> { + fn vec_map_covariant<'a, T, A: Allocator>(map: VecMap<&'static T, A>) -> VecMap<&'a T, A> { map } - fn into_iter_covariant<'a, T>(iter: IntoIter<&'static T>) -> IntoIter<&'a T> { + fn into_iter_covariant<'a, T, A: Allocator>( + iter: IntoIter<&'static T, A>, + ) -> IntoIter<&'a T, A> { iter }