feat: generic collection

src/lib.rs

@@ -1,7 +1,7 @@
 //! ```rust
 //! use median_accumulator::*;
 //!
-//! let mut acc = MedianAcc::new();
+//! let mut acc = vec::MedianAcc::new();
 //!
 //! assert_eq!(acc.get_median(), None);
 //! acc.push(7);
@@ -14,14 +14,29 @@
 //!
 //! In doc comments, _N_ represents the number of samples, _D_ represents the number of different values taken by the samples.
 
-use std::cmp::Ordering;
+#![cfg_attr(feature = "nostd", no_std)]
+
+mod traits;
+
+pub use traits::*;
+
+use core::{cmp::Ordering, ops::DerefMut};
 
 /// Accumulator for computing median
 #[derive(Clone, Debug, Default)]
-pub struct MedianAcc<T: Clone + Ord> {
-	samples: Vec<(T, u32)>,
+pub struct MedianAcc<
+	T: Clone + Ord,
+	V: DerefMut<Target = [(T, u32)]> + cc_traits::VecMut<(T, u32)> + InsertIndex,
+> {
+	samples: V,
 	median_index: Option<usize>,
 	median_subindex: u32,
+	_t: core::marker::PhantomData<T>,
+}
+
+#[cfg(not(feature = "nostd"))]
+pub mod vec {
+	pub type MedianAcc<T> = crate::MedianAcc<T, Vec<(T, u32)>>;
 }
 
 /// Computed median
@@ -34,17 +49,41 @@ pub enum MedianResult<T: Clone + Ord> {
 	Two(T, T),
 }
 
-impl<T: Clone + Ord> MedianAcc<T> {
+impl<
+		T: Clone + Ord,
+		V: DerefMut<Target = [(T, u32)]> + cc_traits::VecMut<(T, u32)> + InsertIndex,
+	> MedianAcc<T, V>
+{
 	/// Create an empty accumulator
 	///
 	/// _O(1)_
 	///
-	/// Does not allocate until the first push.
-	pub fn new() -> Self {
+	/// If using `std::vec::Vec`, does not allocate until the first push.
+	pub fn new() -> Self
+	where
+		V: Default,
+	{
 		Self {
-			samples: Vec::new(),
+			samples: Default::default(),
 			median_index: None,
 			median_subindex: 0,
+			_t: Default::default(),
+		}
+	}
+
+	/// Create an empty accumulator from an existing (empty) collection
+	///
+	/// _O(1)_
+	///
+	/// Useful when using fixed-length collections or to avoid allocations.
+	pub fn new_from(collection: V) -> Self {
+		assert!(collection.is_empty(), "the collection must be empty");
+
+		Self {
+			samples: collection,
+			median_index: None,
+			median_subindex: 0,
+			_t: Default::default(),
 		}
 	}
 
@@ -85,7 +124,7 @@ impl<T: Clone + Ord> MedianAcc<T> {
 					}
 				}
 				Err(sample_index) => {
-					self.samples.insert(sample_index, (sample, 1));
+					self.samples.insert_index(sample_index, (sample, 1));
 					if *median_index >= sample_index {
 						if self.median_subindex == 0 {
 							self.median_subindex =
@@ -105,7 +144,7 @@ impl<T: Clone + Ord> MedianAcc<T> {
 				}
 			}
 		} else {
-			self.samples.push((sample, 1));
+			self.samples.push_back((sample, 1));
 			self.median_index = Some(0);
 		}
 	}
@@ -148,19 +187,22 @@ impl<T: Clone + Ord> MedianAcc<T> {
 	}
 
 	/// Clear the data
-	pub fn clear(&mut self) {
+	pub fn clear(&mut self)
+	where
+		V: cc_traits::Clear,
+	{
 		self.samples.clear();
 		self.median_index = None;
 		self.median_subindex = 0;
 	}
 
-	/// Access the underlying vec
+	/// Access the underlying collection
 	///
 	/// Just in case you need finer allocation management.
 	///
 	/// # Safety
 	/// Leaving the vector in an invalid state may cause invalid result or panic (but no UB).
-	pub unsafe fn get_samples_mut(&mut self) -> &mut Vec<(T, u32)> {
+	pub unsafe fn get_samples_mut(&mut self) -> &mut V {
 		&mut self.samples
 	}
 }
@@ -198,7 +240,24 @@ mod tests {
 			let len: usize = rng.gen_range(0..100);
 			let mut samples: Vec<i32> = (0..len).map(|_| rng.gen_range(-100..100)).collect();
 
-			let mut median = MedianAcc::new();
+			let mut median = vec::MedianAcc::new();
+			for sample in samples.iter() {
+				median.push(*sample);
+			}
+
+			assert_eq!(median.get_median(), naive_median(&mut samples));
+		}
+	}
+
+	#[test]
+	fn correctness_smallvec() {
+		let mut rng = rand::thread_rng();
+
+		for _ in 0..100_000 {
+			let len: usize = rng.gen_range(0..64);
+			let mut samples: Vec<i32> = (0..len).map(|_| rng.gen_range(-100..100)).collect();
+
+			let mut median = MedianAcc::<i32, smallvec::SmallVec<[(i32, u32); 64]>>::new();
 			for sample in samples.iter() {
 				median.push(*sample);
 			}

src/traits.rs

@@ -0,0 +1,28 @@
+/// Collection where an item can be inserted at a given index.
+pub trait InsertIndex: cc_traits::Collection {
+	type Output;
+
+	fn insert_index(
+		&mut self,
+		index: usize,
+		element: <Self as cc_traits::Collection>::Item,
+	) -> Self::Output;
+}
+
+#[cfg(not(feature = "nostd"))]
+impl<T> InsertIndex for Vec<T> {
+	type Output = ();
+
+	fn insert_index(&mut self, index: usize, element: T) {
+		self.insert(index, element)
+	}
+}
+
+#[cfg(feature = "smallvec")]
+impl<T, A: smallvec::Array<Item = T>> InsertIndex for smallvec::SmallVec<A> {
+	type Output = ();
+
+	fn insert_index(&mut self, index: usize, element: T) {
+		self.insert(index, element)
+	}
+}