From 10fd79828d2918bd079a11c2b2c623eede170c3f Mon Sep 17 00:00:00 2001
From: Mark Poliakov <mark@alnyan.me>
Date: Mon, 8 Jan 2024 12:46:41 +0200
Subject: [PATCH] Add `drain()` function to `DiscreteRangeMap`

---
 src/discrete_range_map.rs | 69 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/src/discrete_range_map.rs b/src/discrete_range_map.rs
index c673a27..b85e59d 100644
--- a/src/discrete_range_map.rs
+++ b/src/discrete_range_map.rs
@@ -1672,6 +1672,33 @@ impl<I, K, V> DiscreteRangeMap<I, K, V> {
 	pub fn last_entry(&self) -> Option<(&K, &V)> {
 		self.inner.last_key_value()
 	}
+
+	/// Removes all elements from the map, returning an owned iterator to these.
+	///
+	/// # Examples
+	/// ```
+	/// use discrete_range_map::DiscreteRangeMap;
+	/// use discrete_range_map::inclusive_interval::ie;
+	///
+	/// let mut map = DiscreteRangeMap::from_slice_strict([
+	/// 	(ie(1, 4), false),
+	/// 	(ie(4, 8), true),
+	/// 	(ie(8, 100), false),
+	/// ])
+	/// .unwrap();
+	///
+	/// assert_eq!(
+	/// 	&map.drain().collect::<Vec<_>>(),
+	/// 	&[
+	///         (ie(1, 4), false),
+	///         (ie(4, 8), true),
+	///         (ie(8, 100), false),
+	/// 	]
+	/// );
+	/// assert!(map.is_empty());
+	pub fn drain(&mut self) -> DrainIter<I, K, V> {
+		DrainIter { inner: self }
+	}
 }
 
 // Helper Functions ==========================
@@ -1827,6 +1854,22 @@ impl<I, K, V> IntoIterator for DiscreteRangeMap<I, K, V> {
 		};
 	}
 }
+
+/// An owning iterator over removed entries of a [`DiscreteRangeMap`].
+///
+/// This `struct` is created by the [`drain`] method on [`DiscreteRangeMap`].
+pub struct DrainIter<'a, I, K, V> {
+	inner: &'a mut DiscreteRangeMap<I, K, V>,
+}
+
+impl<'a, I, K, V> Iterator for DrainIter<'a, I, K, V> {
+	type Item = (K, V);
+
+	fn next(&mut self) -> Option<Self::Item> {
+		self.inner.inner.pop_first()
+	}
+}
+
 /// An owning iterator over the entries of a [`DiscreteRangeMap`].
 ///
 /// This `struct` is created by the [`into_iter`] method on
@@ -2536,6 +2579,32 @@ mod tests {
 		}
 	}
 
+	#[test]
+	fn drain_tests() {
+		let r0 = ie(0, 10);
+		let r1 = ie(20, 30);
+		let r2 = ie(40, 50);
+		let r3 = ie(60, 70);
+		let mut map = DiscreteRangeMap::new();
+		map.insert_strict(r0, "range0").unwrap();
+		map.insert_strict(r1, "range1").unwrap();
+		map.insert_strict(r2, "range2").unwrap();
+		map.insert_strict(r3, "range3").unwrap();
+
+		let r: Vec<_> = map.drain().collect();
+		assert_eq!(
+			&r,
+			&[
+				(ie(0, 10), "range0"),
+				(ie(20, 30), "range1"),
+				(ie(40, 50), "range2"),
+				(ie(60, 70), "range3"),
+			]
+		);
+
+		assert!(map.is_empty());
+	}
+
 	#[test]
 	fn cut_range_with_origin_tests() {
 		let r0 = ie(0, 10);