unobtanium_segmenter/
subdivision_map.rs

1use std::marker::PhantomData;
2use std::mem;
3
4/// An owned iterator wrapper that allows shortcuts for common empty and one element cases. It is inteded for use with the [SubdivisionMap] iterator as the return type of the callback.
5///
6/// This implements an Iterator trait so it can be used like any other iterator.
7pub enum UseOrSubdivide<T, I: Iterator<Item = T>> {
8	/// Empty iterator.
9	Empty,
10
11	/// Exactly one element.
12	Use(T),
13
14	/// Multiple elements by wrapping an existing owned iterator.
15	///
16	/// For `I` something like [std::vec::IntoIter] is recommended.
17	Subdivide(I),
18}
19
20impl<T, I: Iterator<Item = T>> Iterator for UseOrSubdivide<T, I> {
21	type Item = T;
22
23	fn next(&mut self) -> Option<Self::Item> {
24		match self {
25			Self::Empty => {
26				return None;
27			}
28			Self::Use(_) => {
29				let taken_use_or_subdivide = mem::replace(self, UseOrSubdivide::Empty);
30				match taken_use_or_subdivide {
31					Self::Use(data) => {
32						return Some(data);
33					}
34					_ => unreachable!("this code only runs when UseOrSubdivide::Use"),
35				}
36			}
37			Self::Subdivide(iter) => {
38				return iter.next();
39			}
40		}
41	}
42}
43
44/// Iterator that allows subdividing each item into zero, one or multiple of itself.
45///
46/// It can do filtering, mapping and flattening of callback outputs.
47pub struct SubdivisionMap<
48	'i,
49	T: 'i,
50	I: Iterator<Item = T> + 'i,
51	R: Iterator<Item = T> + 'i,
52	F: Fn(T) -> UseOrSubdivide<T, R>,
53> {
54	_marker: PhantomData<&'i ()>,
55	iter: I,
56	buffer: Option<R>,
57	callback: F,
58}
59
60impl<
61		'i,
62		T,
63		I: Iterator<Item = T> + 'i,
64		R: Iterator<Item = T> + 'i,
65		F: Fn(T) -> UseOrSubdivide<T, R>,
66	> SubdivisionMap<'i, T, I, R, F>
67{
68	/// Create a new SubdivisionMap from an iterator and a callback
69	/// that does the splitting on the items returning a [UseOrSubdivide] instance.
70	pub fn new(iter: I, callback: F) -> Self {
71		Self {
72			_marker: PhantomData,
73			iter,
74			callback,
75			buffer: None,
76		}
77	}
78}
79
80impl<
81		'i,
82		T,
83		I: Iterator<Item = T> + 'i,
84		R: Iterator<Item = T> + 'i,
85		F: Fn(T) -> UseOrSubdivide<T, R>,
86	> Iterator for SubdivisionMap<'i, T, I, R, F>
87{
88	type Item = I::Item;
89
90	fn next(&mut self) -> Option<Self::Item> {
91		loop {
92			if let Some(buffer_iter) = &mut self.buffer {
93				if let Some(item) = buffer_iter.next() {
94					return Some(item);
95				}
96				self.buffer = None;
97			}
98			if let Some(raw_item) = self.iter.next() {
99				match (self.callback)(raw_item) {
100					UseOrSubdivide::Empty => return None,
101					UseOrSubdivide::Use(item) => return Some(item),
102					UseOrSubdivide::Subdivide(list) => {
103						self.buffer = Some(list.into_iter());
104						// Don't return and loop around
105					}
106				}
107			} else {
108				return None;
109			}
110		}
111	}
112}
unobtanium_segmenter/subdivision_map.rs

unobtanium_segmenter/
subdivision_map.rs
