btree_indexmap/

set.rs

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

use crate::{
	map::{self, BTreeIndexMap},
	util::entries_as_slice,
};

/// BTree-based index set.
///
/// This is a thin wrapper around [`BTreeIndexMap<T, ()>`](BTreeIndexMap)
/// providing a set-like API: every value is stored at most once, values keep
/// a stable index until they (or an earlier value) are removed, and values
/// can be iterated either in insertion order or in sorted order.
///
/// # Example
///
/// ```
/// use btree_indexmap::BTreeIndexSet;
///
/// let mut set = BTreeIndexSet::new();
///
/// // Values are stored in insertion order and assigned a stable index.
/// assert!(set.insert("banana"));
/// assert!(set.insert("apple"));
/// assert!(!set.insert("banana")); // already present
///
/// // O(1) access by index (insertion order).
/// assert_eq!(set.as_slice(), ["banana", "apple"]);
///
/// // O(log n) lookup, returning the stable index alongside the value.
/// assert_eq!(set.get_full("apple"), Some((1, &"apple")));
///
/// // Iterate in sorted order.
/// let sorted: Vec<_> = set.iter_sorted().copied().collect();
/// assert_eq!(sorted, ["apple", "banana"]);
/// ```
#[derive(Clone)]
pub struct BTreeIndexSet<T> {
	map: BTreeIndexMap<T, ()>,
}

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

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

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

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

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

	/// Returns the number of values 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 the index and a reference to the value equal to `value`, if
	/// any.
	pub fn get_full<Q>(&self, value: &Q) -> Option<(usize, &T)>
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.get_full(value).map(|(i, t, ())| (i, t))
	}

	/// Returns the index of `value` in the set, if any.
	pub fn get_index_of<Q>(&self, value: &Q) -> Option<usize>
	where
		Q: ?Sized + Comparable<T>,
	{
		self.get_full(value).map(|(i, _)| i)
	}

	/// Returns a reference to the value equal to `value`, if any.
	pub fn get<Q>(&self, value: &Q) -> Option<&T>
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.get_entry(value).map(|(t, ())| t)
	}

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

	/// Inserts `value` in the set, returning its index and whether it was
	/// newly inserted (`true`) or already present (`false`).
	pub fn insert_full(&mut self, value: T) -> (usize, bool)
	where
		T: Ord,
	{
		let (i, old) = self.map.insert_full(value, ());
		(i, old.is_none())
	}

	/// Inserts `value` in the set, returning whether it was newly inserted.
	pub fn insert(&mut self, value: T) -> bool
	where
		T: Ord,
	{
		self.insert_full(value).1
	}

	/// Removes `value` from the set, returning its former index and the
	/// value itself, if it was present.
	///
	/// Shifts every following value one index down.
	pub fn shift_remove_full<Q>(&mut self, value: &Q) -> Option<(usize, T)>
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.shift_remove_full(value).map(|(i, t, ())| (i, t))
	}

	/// Removes `value` from the set, returning the removed value, if any.
	///
	/// Shifts every following value one index down.
	pub fn shift_take<Q>(&mut self, value: &Q) -> Option<T>
	where
		Q: ?Sized + Comparable<T>,
	{
		self.shift_remove_full(value).map(|(_, t)| t)
	}

	/// Removes `value` from the set, returning whether the value was
	/// present.
	///
	/// Shifts every following value one index down.
	pub fn shift_remove<Q>(&mut self, value: &Q) -> bool
	where
		Q: ?Sized + Comparable<T>,
	{
		self.map.shift_remove(value).is_some()
	}

	/// Removes `value` from the set, returning its former index and the
	/// value itself, if it was present.
	///
	/// Moves the last value into the removed slot.
	pub fn swap_remove_full<Q>(&mut self, value: &Q) -> Option<(usize, T)>
	where
		Q: ?Sized + Comparable<T>,
		T: Ord,
	{
		self.map.swap_remove_full(value).map(|(i, t, ())| (i, t))
	}

	/// Removes `value` from the set, returning the removed value, if any.
	///
	/// Moves the last value into its place.
	pub fn swap_take<Q>(&mut self, value: &Q) -> Option<T>
	where
		Q: ?Sized + Comparable<T>,
		T: Ord,
	{
		self.swap_remove_full(value).map(|(_, t)| t)
	}

	/// Removes `value` from the set, returning whether the value was
	/// present.
	///
	/// Moves the last value into its place.
	pub fn swap_remove<Q>(&mut self, value: &Q) -> bool
	where
		Q: ?Sized + Comparable<T>,
		T: Ord,
	{
		self.map.swap_remove(value).is_some()
	}

	/// 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())
	}
}

impl<T> Index<usize> for BTreeIndexSet<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 BTreeIndexSet<T> {
	type IntoIter = IntoIter<T>;
	type Item = T;

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

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

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

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

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

impl<T: Hash> Hash for BTreeIndexSet<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 BTreeIndexSet<T> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.debug_set().entries(self.iter()).finish()
	}
}

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

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

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

impl<'a, T> IntoIterator for &'a BTreeIndexSet<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 [`BTreeIndexSet`], sorted according to
/// `T`'s implementation of [`Ord`].
pub struct IterSorted<'a, T>(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> DoubleEndedIterator for IterSorted<'_, T> {
	fn next_back(&mut self) -> Option<Self::Item> {
		self.0.next_back().map(|(t, ())| t)
	}
}

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

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

/// An owning iterator over the values of a [`BTreeIndexSet`], sorted
/// according to `T`'s implementation of [`Ord`].
pub struct IntoIterSorted<T>(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> DoubleEndedIterator for IntoIterSorted<T> {
	fn next_back(&mut self) -> Option<Self::Item> {
		self.0.next_back().map(|(t, ())| t)
	}
}

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

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

/// An owning iterator over the values of a [`BTreeIndexSet`], in insertion
/// order.
pub struct IntoIter<T>(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> {}