btree_indexmap/

multi_set.rs

use core::fmt;
use equivalent::Comparable;
use std::{
	cmp::Ordering,
	hash::{Hash, Hasher},
	iter::FusedIterator,
	ops::Index,
};

use crate::{
	multi_map::{self, BTreeIndexMultiMap},
	util::entries_as_slice,
};

/// BTree-based index multi-set.
///
/// This is a thin wrapper around
/// [`BTreeIndexMultiMap<T, ()>`](BTreeIndexMultiMap) providing a multi-set
/// API: a value can occur any number of times, every occurrence keeps a
/// stable index until it (or an earlier entry) is removed, and occurrences
/// can be iterated either in insertion order or in sorted order.
///
/// # Example
///
/// ```
/// use btree_indexmap::BTreeIndexMultiSet;
///
/// let mut set = BTreeIndexMultiSet::new();
///
/// // Multiple occurrences of the same value, each with its own stable
/// // index.
/// set.push_back("apple");
/// set.push_back("carrot");
/// set.push_back("apple");
///
/// assert_eq!(set.count("apple"), 2);
///
/// let apples: Vec<_> = set.get("apple").copied().collect();
/// assert_eq!(apples, ["apple", "apple"]);
///
/// // Iterate in sorted order, occurrences in position order.
/// let sorted: Vec<_> = set.iter_sorted().copied().collect();
/// assert_eq!(sorted, ["apple", "apple", "carrot"]);
/// ```
#[derive(Clone)]
pub struct BTreeIndexMultiSet<T> {
	map: BTreeIndexMultiMap<T, ()>,
}

impl<T> Default for BTreeIndexMultiSet<T> {
	fn default() -> Self {
		Self::new()
	}
}

impl<T> BTreeIndexMultiSet<T> {
	/// Creates a new, empty multi-set.
	pub fn new() -> Self {
		Self {
			map: BTreeIndexMultiMap::new(),
		}
	}

	/// Creates a new, empty multi-set with space reserved for at least
	/// `capacity` occurrences.
	pub fn with_capacity(capacity: usize) -> Self {
		Self {
			map: BTreeIndexMultiMap::with_capacity(capacity),
		}
	}

	/// Creates a new multi-set containing the single `value`.
	pub fn singleton(value: T) -> Self
	where
		T: Ord,
	{
		Self {
			map: BTreeIndexMultiMap::singleton(value, ()),
		}
	}

	/// Removes all occurrences from the set.
	///
	/// This keeps the capacity allocated for the occurrences.
	pub fn clear(&mut self) {
		self.map.clear();
	}

	/// Returns the number of occurrences the set can hold without
	/// reallocating.
	pub fn capacity(&self) -> usize {
		self.map.capacity()
	}

	/// Returns `true` if the set contains no occurrences.
	pub fn is_empty(&self) -> bool {
		self.map.is_empty()
	}

	/// Returns the total number of occurrences in the set.
	pub fn len(&self) -> usize {
		self.map.len()
	}

	/// Returns the values of the set as a slice, in insertion (index)
	/// order.
	pub fn as_slice(&self) -> &[T] {
		entries_as_slice(self.map.as_entries())
	}

	/// Returns a reference to the first occurrence in the set (in insertion
	/// order), if any.
	pub fn first(&self) -> Option<&T> {
		self.map.first().map(|(t, ())| t)
	}

	/// Returns a reference to the last occurrence in the set (in insertion
	/// order), if any.
	pub fn last(&self) -> Option<&T> {
		self.map.last().map(|(t, ())| t)
	}

	/// Returns `true` if the set contains at least one occurrence of
	/// `value`.
	pub fn contains<Q>(&self, value: &Q) -> bool
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.contains_key(value)
	}

	/// Returns the number of occurrences of `value` in the set.
	pub fn count<Q>(&self, value: &Q) -> usize
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.key_occurences(value)
	}

	/// Returns an iterator over the indexes of every occurrence of `value`.
	pub fn indexes_of<Q>(&self, value: &Q) -> multi_map::IndexesIter<'_>
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.indexes_of(value)
	}

	/// Returns the index of the first occurrence of `value`.
	pub fn index_of<Q>(&self, value: &Q) -> Option<usize>
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.index_of(value)
	}

	/// Returns an iterator over `(index, &value)` for every occurrence of
	/// `value`.
	pub fn get_indexed<Q>(&self, value: &Q) -> GetIndexed<'_, T>
	where
		Q: ?Sized + Comparable<T>,
	{
		GetIndexed(self.map.get_indexed_entries(value))
	}

	/// Returns an iterator over a reference to `value` for every occurrence
	/// found in the set.
	pub fn get<Q>(&self, value: &Q) -> Get<'_, T>
	where
		Q: ?Sized + Comparable<T>,
	{
		Get(self.map.get_indexed_entries(value))
	}

	/// Inserts a new occurrence of `value` at the end of the set.
	///
	/// Returns the index of the new occurrence, and whether `value` was not
	/// already present in the set.
	pub fn push_back(&mut self, value: T) -> (usize, bool)
	where
		T: Ord,
	{
		self.map.push_back(value, ())
	}

	/// Inserts a new occurrence of `value` at the front of the set.
	///
	/// Returns whether `value` was not already present in the set.
	pub fn push_front(&mut self, value: T) -> bool
	where
		T: Ord,
	{
		self.map.push_front(value, ())
	}

	/// Inserts a new occurrence of `value` at the given `index`.
	///
	/// Returns whether `value` was not already present in the set.
	///
	/// # Panics
	///
	/// Panics if `index > len`.
	pub fn insert_at(&mut self, index: usize, value: T) -> bool
	where
		T: Ord,
	{
		self.map.insert_at(index, value, ())
	}

	/// Inserts a new occurrence of `value` in the set, without removing any
	/// existing occurrence.
	///
	/// This is an alias for [`Self::push_back`].
	pub fn insert(&mut self, value: T) -> (usize, bool)
	where
		T: Ord,
	{
		self.push_back(value)
	}

	/// Removes the occurrence at the given `index`, shifting every
	/// following occurrence one index down.
	pub fn shift_remove_at(&mut self, index: usize) -> Option<T>
	where
		T: Ord,
	{
		self.map.shift_remove_at(index).map(|(t, ())| t)
	}

	/// Removes the occurrence at the given `index`, moving the last
	/// occurrence into its place.
	pub fn swap_remove_at(&mut self, index: usize) -> Option<T>
	where
		T: Ord,
	{
		self.map.swap_remove_at(index).map(|(t, ())| t)
	}

	/// Removes every occurrence of `value`, returning an iterator over the
	/// removed occurrences.
	///
	/// Shifts the remaining occurrences down to fill the gaps.
	pub fn shift_remove<Q>(&mut self, value: &Q) -> ShiftRemove<'_, T>
	where
		Q: ?Sized + Comparable<T>,
	{
		ShiftRemove(self.map.shift_remove_entries(value))
	}

	/// Removes every occurrence of `value`, returning an iterator over the
	/// removed occurrences.
	///
	/// Moves occurrences from the end of the set into the gaps.
	pub fn swap_remove<Q>(&mut self, value: &Q) -> SwapRemove<'_, T>
	where
		Q: ?Sized + Comparable<T>,
		T: Ord,
	{
		SwapRemove(self.map.swap_remove_entries(value))
	}

	/// Returns an iterator over the values of the set, in insertion (index)
	/// order.
	pub fn iter(&self) -> Iter<'_, T> {
		self.as_slice().iter()
	}

	/// Returns an iterator over the values of the set, sorted according to
	/// `T`'s implementation of [`Ord`].
	pub fn iter_sorted(&self) -> IterSorted<'_, T> {
		IterSorted(self.map.iter_sorted())
	}

	/// Turns this set into an iterator over its values, sorted according to
	/// `T`'s implementation of [`Ord`].
	pub fn into_iter_sorted(self) -> IntoIterSorted<T> {
		IntoIterSorted(self.map.into_iter_sorted())
	}

	/// Sorts the entries of the set by value.
	///
	/// The relative order of equal values is unchanged.
	pub fn sort(&mut self)
	where
		T: Ord,
	{
		self.map.sort()
	}
}

impl<T> Index<usize> for BTreeIndexMultiSet<T> {
	type Output = T;

	/// # Panics
	///
	/// Panics if `index` is out of bounds.
	fn index(&self, index: usize) -> &Self::Output {
		&self.as_slice()[index]
	}
}

impl<T> IntoIterator for BTreeIndexMultiSet<T> {
	type IntoIter = IntoIter<T>;
	type Item = T;

	fn into_iter(self) -> Self::IntoIter {
		IntoIter(self.map.into_iter())
	}
}

impl<T: PartialEq> PartialEq for BTreeIndexMultiSet<T> {
	fn eq(&self, other: &Self) -> bool {
		self.iter_sorted().eq(other.iter_sorted())
	}
}

impl<T: Eq> Eq for BTreeIndexMultiSet<T> {}

impl<T: PartialOrd> PartialOrd for BTreeIndexMultiSet<T> {
	fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
		self.iter_sorted().partial_cmp(other.iter_sorted())
	}
}

impl<T: Ord> Ord for BTreeIndexMultiSet<T> {
	fn cmp(&self, other: &Self) -> Ordering {
		self.iter_sorted().cmp(other.iter_sorted())
	}
}

impl<T: Hash> Hash for BTreeIndexMultiSet<T> {
	fn hash<H: Hasher>(&self, state: &mut H) {
		state.write_usize(self.len());
		for item in self.iter_sorted() {
			item.hash(state);
		}
	}
}

impl<T: fmt::Debug> fmt::Debug for BTreeIndexMultiSet<T> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.debug_set().entries(self.iter()).finish()
	}
}

impl<T: Ord> Extend<T> for BTreeIndexMultiSet<T> {
	fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
		for value in iter {
			self.push_back(value);
		}
	}
}

impl<T: Ord> FromIterator<T> for BTreeIndexMultiSet<T> {
	fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
		let mut result = Self::new();
		result.extend(iter);
		result
	}
}

/// An iterator over `(index, &value)` for every occurrence of a value in a
/// [`BTreeIndexMultiSet`].
///
/// See [`BTreeIndexMultiSet::get_indexed`].
pub struct GetIndexed<'a, T>(multi_map::GetIndexedEntries<'a, T, ()>);

impl<'a, T> Iterator for GetIndexed<'a, T> {
	type Item = (usize, &'a T);

	fn next(&mut self) -> Option<Self::Item> {
		self.0.next().map(|(i, t, ())| (i, t))
	}
}

/// An iterator over a reference to a value for every occurrence found in a
/// [`BTreeIndexMultiSet`].
///
/// See [`BTreeIndexMultiSet::get`].
pub struct Get<'a, T>(multi_map::GetIndexedEntries<'a, T, ()>);

impl<'a, T> Iterator for Get<'a, T> {
	type Item = &'a T;

	fn next(&mut self) -> Option<Self::Item> {
		self.0.next().map(|(_, t, ())| t)
	}
}

/// An iterator over the occurrences removed by
/// [`BTreeIndexMultiSet::shift_remove`].
pub struct ShiftRemove<'a, T>(multi_map::ShiftRemoveEntries<'a, T, ()>);

impl<'a, T> Iterator for ShiftRemove<'a, T> {
	type Item = T;

	fn size_hint(&self) -> (usize, Option<usize>) {
		self.0.size_hint()
	}

	fn next(&mut self) -> Option<Self::Item> {
		self.0.next().map(|(t, ())| t)
	}
}

impl<T> ExactSizeIterator for ShiftRemove<'_, T> {}

/// An iterator over the occurrences removed by
/// [`BTreeIndexMultiSet::swap_remove`].
pub struct SwapRemove<'a, T: Ord>(multi_map::SwapRemoveEntries<'a, T, ()>);

impl<'a, T: Ord> Iterator for SwapRemove<'a, T> {
	type Item = T;

	fn size_hint(&self) -> (usize, Option<usize>) {
		self.0.size_hint()
	}

	fn next(&mut self) -> Option<Self::Item> {
		self.0.next().map(|(t, ())| t)
	}
}

impl<T: Ord> ExactSizeIterator for SwapRemove<'_, T> {}

/// An iterator over the values of a [`BTreeIndexMultiSet`], in insertion
/// order.
///
/// Since the set's entries are exposed as a plain `&[T]` slice (see
/// [`BTreeIndexMultiSet::as_slice`]), this is just [`std::slice::Iter`].
pub type Iter<'a, T> = std::slice::Iter<'a, T>;

impl<'a, T> IntoIterator for &'a BTreeIndexMultiSet<T> {
	type IntoIter = Iter<'a, T>;
	type Item = &'a T;

	fn into_iter(self) -> Self::IntoIter {
		self.iter()
	}
}

/// An iterator over the values of a [`BTreeIndexMultiSet`], sorted
/// according to `T`'s implementation of [`Ord`].
pub struct IterSorted<'a, T>(multi_map::IterSorted<'a, T, ()>);

impl<T> Clone for IterSorted<'_, T> {
	fn clone(&self) -> Self {
		Self(self.0.clone())
	}
}

impl<'a, T> Iterator for IterSorted<'a, T> {
	type Item = &'a T;

	fn size_hint(&self) -> (usize, Option<usize>) {
		self.0.size_hint()
	}

	fn next(&mut self) -> Option<Self::Item> {
		self.0.next().map(|(t, ())| t)
	}
}

impl<T> ExactSizeIterator for IterSorted<'_, T> {}

impl<T> FusedIterator for IterSorted<'_, T> {}

/// An owning iterator over the values of a [`BTreeIndexMultiSet`], sorted
/// according to `T`'s implementation of [`Ord`].
pub struct IntoIterSorted<T>(multi_map::IntoIterSorted<T, ()>);

impl<T> Iterator for IntoIterSorted<T> {
	type Item = T;

	fn size_hint(&self) -> (usize, Option<usize>) {
		self.0.size_hint()
	}

	fn next(&mut self) -> Option<Self::Item> {
		self.0.next().map(|(t, ())| t)
	}
}

impl<T> ExactSizeIterator for IntoIterSorted<T> {}

impl<T> FusedIterator for IntoIterSorted<T> {}

/// An owning iterator over the values of a [`BTreeIndexMultiSet`], in
/// insertion order.
pub struct IntoIter<T>(multi_map::IntoIter<T, ()>);

impl<T> Iterator for IntoIter<T> {
	type Item = T;

	fn size_hint(&self) -> (usize, Option<usize>) {
		self.0.size_hint()
	}

	fn next(&mut self) -> Option<Self::Item> {
		self.0.next().map(|(t, ())| t)
	}
}

impl<T> DoubleEndedIterator for IntoIter<T> {
	fn next_back(&mut self) -> Option<Self::Item> {
		self.0.next_back().map(|(t, ())| t)
	}
}

impl<T> ExactSizeIterator for IntoIter<T> {}

impl<T> FusedIterator for IntoIter<T> {}