bendns' repos / btree_multiset
[no description]
Diffstat (limited to 'src/lib.rs')
-rw-r--r--src/lib.rs139
1 files changed, 139 insertions, 0 deletions
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]);
+}