[no description]
init
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | Cargo.toml | 13 | ||||
| -rw-r--r-- | rustfmt.toml | 7 | ||||
| -rw-r--r-- | src/lib.rs | 139 |
4 files changed, 160 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..03314f7 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +Cargo.lock diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..479068b --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "btree_multiset" +version = "0.1.0" +edition = "2024" +description = "btreeset with more capacity" +authors = ["bendn <[email protected]>"] +license = "MIT" +repository = "https://git.bendn.org/btree_multiset" +keywords = ["btree", "multiset"] +rust-version = "1.95.0" + +[dependencies] +smolvec = "1.0.0" diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..c39da03 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,7 @@ +max_width = 75 +format_strings = true +group_imports = "StdExternalCrate" +imports_granularity = "Module" +match_arm_blocks = false +style_edition = "2024" +use_small_heuristics = "Max" diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..472df7e --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,139 @@ +#![feature(allocator_api, extend_one, btreemap_alloc)] +//! [`BTreeSet`](std::collections::BTreeSet) but with duplicates. +use std::alloc::{Allocator, Global}; +use std::collections::BTreeMap; +use std::collections::btree_map::{self, Entry}; +use std::fmt::Debug; +use std::iter::FlatMap; + +use smolvec::SmolVec; + +#[derive(Clone)] +/// Like a btree set, but with duplicates. +/// Currently immutable cause i havent thought about it yet. +pub struct BTreeMultiset<K, T, F, A: Allocator + Clone> { + inner: BTreeMap<K, SmolVec<T>, A>, + f: F, +} +impl<K: Ord, T: Debug, F: Fn(&T) -> K, A: Allocator + Clone> Debug + for BTreeMultiset<K, T, F, A> +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_set().entries(self.iter()).finish() + } +} +impl<K, T, F: Fn(&T) -> K, A: Allocator + Clone> IntoIterator + for BTreeMultiset<K, T, F, A> +{ + type Item = T; + + type IntoIter = FlatMap< + btree_map::IntoIter<K, SmolVec<T>, A>, + SmolVec<T>, + fn((K, SmolVec<T>)) -> SmolVec<T>, + >; + + fn into_iter(self) -> Self::IntoIter { + self.inner.into_iter().flat_map(|x| x.1) + } +} +impl<K: Ord, T, F: Fn(&T) -> K, A: Allocator + Clone> Extend<T> + for BTreeMultiset<K, T, F, A> +{ + #[inline] + fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) { + iter.into_iter().for_each(move |x| { + self.insert(x); + }); + } + + #[inline] + fn extend_one(&mut self, x: T) { + self.insert(x); + } +} +impl<K: Ord, T, F: Fn(&T) -> K> BTreeMultiset<K, T, F, Global> { + pub fn new(f: F) -> Self { + Self { inner: BTreeMap::new_in(Global), f } + } +} +type Iter<'a, K, T> = FlatMap< + btree_map::Iter<'a, K, SmolVec<T>>, + &'a SmolVec<T>, + fn((&'a K, &'a SmolVec<T>)) -> &'a SmolVec<T>, +>; +type Range<'a, K, T> = FlatMap< + btree_map::Range<'a, K, SmolVec<T>>, + &'a SmolVec<T>, + fn((&'a K, &'a SmolVec<T>)) -> &'a SmolVec<T>, +>; +impl<K, T, F, A: Allocator + Clone> BTreeMultiset<K, T, F, A> { + pub fn into_inner(self) -> BTreeMap<K, SmolVec<T>, A> { + self.inner + } +} +impl<K: Ord, T, F: Fn(&T) -> K, A: Allocator + Clone> + BTreeMultiset<K, T, F, A> +{ + pub fn new_in(f: F, a: A) -> Self { + Self { inner: BTreeMap::new_in(a), f } + } + /// Gets an iterator that visits the elements in the [`BTreeMultiset`] in ascending order. + /// Similar to `range(..)`. + pub fn iter<'a>(&'a self) -> Iter<'a, K, T> { + self.inner.iter().flat_map(|x| x.1) + } + /// Adds a value to the multi set. + pub fn insert(&mut self, t: T) -> bool { + match self.inner.entry((self.f)(&t)) { + Entry::Vacant(vacant_entry) => { + vacant_entry.insert(SmolVec::from_iter([t])); + false + } + Entry::Occupied(occupied_entry) => { + occupied_entry.into_mut().push(t); + true + } + } + } + + /// Constructs a double-ended iterator over a sub-range of elements in the multiset. + /// The simplest way is to use the range syntax `min..max`, thus `range(min..max)` will + /// yield elements from min (inclusive) to max (exclusive). + /// The range may also be entered as `(Bound<T>, Bound<T>)`, so for example + /// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive + /// range from 4 to 10. + /// + /// # Panics + /// + /// Panics if range `start > end`. + /// Panics if range `start == end` and both bounds are `Excluded`. + pub fn range<'a>( + &'a self, + r: impl std::ops::RangeBounds<K>, + ) -> Range<'a, K, T> { + self.inner.range(r).flat_map(|x| x.1) + } + + // pub fn range_mut<'a, R>( + // &'a mut self, + // r: impl std::ops::RangeBounds<K>, + // ) -> FlatMap< + // RangeMut<'a, K, SmolVec<T>>, + // &'a mut SmolVec<T>, + // fn((&'a K, &'a mut SmolVec<T>)) -> &'a mut SmolVec<T>, + // > { + // self.inner.range_mut(r).flat_map(|x| x.1) + // } +} +#[test] +fn works() { + let mut x = BTreeMultiset::new(|x: &u32| *x); + + x.insert(4); + x.insert(3); + x.insert(7); + x.insert(3); + assert_eq!(x.iter().copied().collect::<Vec<_>>(), vec![3, 3, 4, 7]); + assert_eq!(x.range(0..=4).copied().collect::<Vec<_>>(), vec![3, 3, 4]); +} |