From b49a94f4115efb34f91aa5dcc2fe576a94b12b3f Mon Sep 17 00:00:00 2001
From: Mark Poliakov <mark@alnyan.me>
Date: Sun, 31 Dec 2023 11:11:05 +0200
Subject: [PATCH] Add cut_with_origin() function

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

diff --git a/src/discrete_range_map.rs b/src/discrete_range_map.rs
index 03b6ebb..c673a27 100644
--- a/src/discrete_range_map.rs
+++ b/src/discrete_range_map.rs
@@ -487,6 +487,157 @@ where
 			))
 		}
 	}
+
+	/// Same as [DiscreteRangeMap::cut], but additionally returns the original bounds of the range
+	/// from which the requested range was cut.
+	pub fn cut_with_origin<'a>(
+		&'a mut self,
+		range: K,
+	) -> impl Iterator<Item = (K, K, V)>
+	where
+		V: Clone,
+	{
+		invalid_range_panic(range);
+
+		let start_comp = overlapping_comp(range.start());
+		let end_comp = overlapping_comp(range.end());
+
+		let left_overlapping = self
+			.inner
+			.get_key_value(start_comp)
+			.map(|(key, _)| key)
+			.copied();
+		let right_overlapping = self
+			.inner
+			.get_key_value(end_comp)
+			.map(|(key, _)| key)
+			.copied();
+
+		if let Some(left) = left_overlapping
+			&& let Some(right) = right_overlapping
+			&& left.start() == right.start()
+		{
+			Either::Left(self.cut_single_overlapping_with_origin(range, left))
+		} else {
+			Either::Right(self.cut_non_single_overlapping_with_origin(
+				range,
+				left_overlapping,
+				right_overlapping,
+			))
+		}
+	}
+
+	fn cut_single_overlapping_with_origin(
+		&mut self,
+		range: K,
+		single_overlapping_range: K,
+	) -> impl Iterator<Item = (K, K, V)>
+	where
+		V: Clone,
+	{
+		invalid_range_panic(range);
+
+		let cut_result = cut_range(single_overlapping_range, range);
+
+		let returning_before_cut = cut_result.before_cut.map(K::from);
+		let returning_after_cut = cut_result.after_cut.map(K::from);
+
+		let value = self.inner.remove(overlapping_comp(range.start())).unwrap();
+
+		if let Some(before) = returning_before_cut {
+			self.insert_unchecked(before, value.clone());
+		}
+		if let Some(after) = returning_after_cut {
+			self.insert_unchecked(after, value.clone());
+		}
+
+		once((
+			single_overlapping_range,
+			cut_result.inside_cut.map(K::from).unwrap(),
+			value,
+		))
+	}
+
+	fn cut_non_single_overlapping_with_origin<'a>(
+		&'a mut self,
+		range: K,
+		left_overlapping: Option<K>,
+		right_overlapping: Option<K>,
+	) -> impl Iterator<Item = (K, K, V)>
+	where
+		V: Clone,
+	{
+		invalid_range_panic(range);
+
+		let (origin_before, returning_before_cut, keeping_before) =
+			match left_overlapping {
+				Some(before) => {
+					let cut_result = cut_range(before, range);
+
+					(
+						Some(before),
+						cut_result.before_cut.map(K::from),
+						cut_result.inside_cut.map(K::from),
+					)
+				}
+				None => (None, None, None),
+			};
+		let (origin_after, returning_after_cut, keeping_after) =
+			match right_overlapping {
+				Some(after) => {
+					let cut_result = cut_range(after, range);
+
+					(
+						Some(after),
+						cut_result.after_cut.map(K::from),
+						cut_result.inside_cut.map(K::from),
+					)
+				}
+				None => (None, None, None),
+			};
+
+		let before_value = self.inner.remove(overlapping_comp(range.start()));
+		let after_value = self.inner.remove(overlapping_comp(range.end()));
+
+		if let Some(returning_before_cut) = returning_before_cut {
+			self.insert_unchecked(
+				returning_before_cut,
+				before_value.as_ref().cloned().unwrap(),
+			);
+		}
+		if let Some(returning_after_cut) = returning_after_cut {
+			self.insert_unchecked(
+				returning_after_cut,
+				after_value.as_ref().cloned().unwrap(),
+			);
+		}
+
+		let keeping_before_entry = keeping_before.map(|keeping_before| {
+			(
+				origin_before.unwrap(),
+				keeping_before,
+				before_value.unwrap(),
+			)
+		});
+		let keeping_after_entry = keeping_after.map(|keeping_after| {
+			(origin_after.unwrap(), keeping_after, after_value.unwrap())
+		});
+
+		return keeping_before_entry
+			.into_iter()
+			.chain(self.remove_overlapping(range).map(move |(key, value)| {
+				(
+					key,
+					K::from(InclusiveInterval {
+						start: key.start(),
+						end: key.end(),
+					}),
+					value,
+				)
+			}))
+			.chain(keeping_after_entry);
+	}
+
 	fn cut_single_overlapping<Q>(
 		&mut self,
 		range: Q,
@@ -514,6 +665,7 @@ where
 
 		once((cut_result.inside_cut.map(K::from).unwrap(), value))
 	}
+
 	fn cut_non_single_overlapping<'a, Q>(
 		&'a mut self,
 		range: Q,
@@ -2384,6 +2536,30 @@ mod tests {
 		}
 	}
 
+	#[test]
+	fn cut_range_with_origin_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.cut_with_origin(ie(5, 45)).collect();
+
+		assert_eq!(
+			&r,
+			&[
+				(ie(0, 10), ie(5, 10), "range0"),
+				(ie(20, 30), ie(20, 30), "range1"),
+				(ie(40, 50), ie(40, 45), "range2")
+			]
+		);
+	}
+
 	#[test]
 	fn cut_range_tests() {
 		for base in all_valid_test_bounds() {