use std::{num::NonZero, ops::ControlFlow};
use crate::{
buffer::{
self,
segments::{Many, Nonempty},
storage,
},
mem::Rejected,
};
#[must_use = "step or release the retained copy"]
pub struct Flatten<T: buffer::raw::Segment> {
source: Nonempty<T>,
offset: usize,
output: storage::Owned,
}
impl<T: buffer::raw::Segment> Flatten<T> {
pub fn try_new(source: Many<T>) -> Result<Self, Rejected<Many<T>, storage::AllocError>> {
let output = match storage::Owned::try_with(source.len().get()) {
Ok(output) => output,
Err(error) => return Err(Rejected::new(source, error)),
};
Ok(Self {
source: Nonempty::Many(source),
offset: 0,
output,
})
}
pub fn step(mut self, mut budget: NonZero<usize>) -> ControlFlow<storage::Shared, Self> {
loop {
let source = self.source.front().as_ref();
let remaining = source.len() - self.offset;
let amount = budget.get().min(remaining);
let prefix = self.output.len() as u32;
unsafe {
let source = source.get_unchecked(self.offset..self.offset + amount);
buffer::raw::Slices::new(NonZero::new_unchecked(amount as u32), [source])
.replace(&mut self.output, prefix);
}
let left = budget.get() - amount;
if amount < remaining {
self.offset += amount;
return ControlFlow::Continue(self);
}
self.source = match self.source {
Nonempty::One(_) => return ControlFlow::Break(self.output.freeze()),
Nonempty::Many(source) => source.take_front().1,
};
self.offset = 0;
let Some(left) = NonZero::new(left) else {
return ControlFlow::Continue(self);
};
budget = left;
}
}
}