use alloc::{boxed::Box, vec::Vec};
use arbitrary::Arbitrary;
use crate::{prelude::*, queues::Queue};
#[derive(Debug, PartialEq, Eq, Arbitrary, Clone, Copy)]
pub enum ConsumerOperation {
Consume,
Flush,
}
fn do_operations_make_progress(ops: &[ConsumerOperation]) -> bool {
ops.iter().any(|op| !matches!(op, ConsumerOperation::Flush))
}
#[derive(Debug)]
pub struct Scrambled<C, Q> {
inner: C,
buffer: Q,
ops: Box<[ConsumerOperation]>,
op_index: usize,
}
impl<C, Q> Scrambled<C, Q> {
pub(crate) fn new(inner: C, buffer: Q, mut ops: Vec<ConsumerOperation>) -> Self {
if !do_operations_make_progress(&ops) {
ops.push(ConsumerOperation::Consume)
}
Self {
inner,
buffer,
ops: ops.into_boxed_slice(),
op_index: 0,
}
}
pub fn into_inner(self) -> C {
self.inner
}
}
impl<C, Q> AsRef<C> for Scrambled<C, Q> {
fn as_ref(&self) -> &C {
&self.inner
}
}
impl<C, Q> Scrambled<C, Q>
where
C: Consumer<Item: Clone>,
Q: Queue<Item = C::Item>,
{
async fn write_buffer_to_inner(&mut self) -> Result<(), C::Error> {
while !self.buffer.is_empty() {
match self.ops[self.op_index] {
ConsumerOperation::Consume => {
let item = self
.buffer
.dequeue()
.expect("Dequeueing from a non-empty queue must always succeed.");
self.inner.consume_item(item).await?;
}
ConsumerOperation::Flush => self.inner.flush().await?,
}
if self.op_index == self.ops.len() - 1 {
self.op_index = 0;
} else {
self.op_index += 1;
}
}
debug_assert!(!self.buffer.is_full());
Ok(())
}
}
impl<C, Q> Consumer for Scrambled<C, Q>
where
C: Consumer<Item: Clone>,
Q: Queue<Item = C::Item>,
{
type Item = C::Item;
type Final = C::Final;
type Error = C::Error;
async fn consume(&mut self, val: Either<Self::Item, Self::Final>) -> Result<(), Self::Error> {
match val {
Left(item) => match self.buffer.enqueue(item) {
None => Ok(()),
Some(item) => {
self.write_buffer_to_inner().await?;
let res = self.buffer.enqueue(item);
debug_assert!(
res.is_none(),
"Enqueueing into an empty queue must always succeed."
);
Ok(())
}
},
Right(fin) => {
self.write_buffer_to_inner().await?;
self.inner.consume_final(fin).await
}
}
}
async fn flush(&mut self) -> Result<(), Self::Error> {
self.write_buffer_to_inner().await?;
self.inner.flush().await
}
}
impl<C, Q> BulkConsumer for Scrambled<C, Q>
where
C: Consumer<Item: Clone>,
Q: Queue<Item = C::Item>,
{
async fn expose_slots_gracefully<F, R>(&mut self, f: F) -> Result<R, (F, Self::Error)>
where
F: AsyncFnOnce(&mut [Self::Item]) -> (usize, R),
{
if self.buffer.is_full() {
if let Err(err) = self.write_buffer_to_inner().await {
return Err((f, err));
}
}
Ok(self.buffer.expose_slots(async |buffer_slots| {
debug_assert!(!buffer_slots.is_empty(), "A non-full queue must expose at least one item slot when expose_slots is invoked.");
f(buffer_slots).await
}).await)
}
}