unobtanium_segmenter/
subdivision_map.rs

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

unobtanium_segmenter/
subdivision_map.rs
