zencodec_coroutine/

decoder.rs

//! `IncrementalDecode` over a coroutine.

use alloc::boxed::Box;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::future::Future;
use core::marker::PhantomData;
use core::pin::Pin;
use core::task::{Context, Poll, Waker};

use zencodec::ImageInfo;
use zencodec::decode::{
    BoundaryAction, BoundaryId, ControlError, ControlledIncrementalDecode, ControlledProgress,
    DecodeBoundary, DecodeRowSink, Event, IncrementalDecode, Input, InputRequest, MetadataUpdate,
    OutputInfo, Progress, SinkError, StepError,
};
use zencodec::source::AccessPattern;
use zenpixels::PixelDescriptor;

use crate::context::{CoroutineContext, CoroutineError, Shared, SharedCell};

/// Default cap on fed-but-unread input, in bytes.
pub const DEFAULT_INPUT_CAP: usize = 16 * 1024 * 1024;

/// A boxed decode future, for decoders whose concrete future type is not
/// nameable (every `async` block).
pub type BoxFuture<E> = Pin<Box<dyn Future<Output = Result<OutputInfo, E>> + Send + 'static>>;

/// A [`CoroutineDecoder`] over a [`BoxFuture`]: the type to name in a
/// codec's `DecodeJob::dyn_incremental_decoder` or a struct field.
pub type BoxedCoroutineDecoder<E> = CoroutineDecoder<BoxFuture<E>, E>;

enum State {
    Running,
    /// The future returned while rows or metadata it deposited were still
    /// undelivered (it polled `write_rows` by hand and did not await it):
    /// they go out first, then `Done`.
    Finishing(OutputInfo),
    Done(OutputInfo),
    Failed,
}

/// An [`IncrementalDecode`] whose logic is an `async` block.
///
/// Build one with [`new`](Self::new) from a closure that receives the
/// [`CoroutineContext`] and returns the decode future. The future's error
/// type is the codec's; it must be constructible from [`CoroutineError`]
/// so the context's `?` operator works.
pub struct CoroutineDecoder<F, E> {
    shared: SharedCell,
    future: Pin<Box<F>>,
    state: State,
    begun: bool,
    /// The sink's `accepts_provisional`, mirrored here when `begin` ran so
    /// the `Rows` arm does not take the shared lock just to read it.
    accepts_provisional: bool,
    access: AccessPattern,
    need: Vec<InputRequest>,
    cached_info: Option<ImageInfo>,
    details: Option<zencodec::decode::DecodeDetails>,
    cached_info_generation: u64,
    /// A native batch may exceed the sink's capacity for one progress event.
    /// Keep its allocation until all strips have been delivered; do not poll
    /// the producer while any of its rows remain unpublished.
    pending_rows: Option<(crate::context::PendingRows, u32)>,
    _err: PhantomData<fn() -> E>,
}

// Stack-only cleanup: no extra Arc clone or lock on a successful step.
struct DetailsCleanup<'a, F, E>(&'a mut CoroutineDecoder<F, E>);
impl<F, E> Drop for DetailsCleanup<'_, F, E> {
    fn drop(&mut self) {
        if matches!(self.0.state, State::Failed) {
            let pending = self.0.shared.lock().details.take();
            drop(pending);
            self.0.details = None;
        }
    }
}

impl<F, E> CoroutineDecoder<F, E>
where
    F: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    E: core::error::Error + Send + Sync + From<CoroutineError> + 'static,
{
    /// Create a decoder with the default input cap and
    /// [`AccessPattern::RandomAccess`].
    pub fn new(init: impl FnOnce(CoroutineContext) -> F) -> Self {
        Self::with_options(DEFAULT_INPUT_CAP, AccessPattern::RandomAccess, init)
    }

    /// Create a decoder requiring explicit interpretation decisions. No sink
    /// effects are allowed until `ctx.interpretation(...).await` has received
    /// adoption of final facts. Serve this through `step_controlled`; legacy
    /// `step` refuses before polling the future.
    pub fn controlled(
        input_cap: usize,
        access: AccessPattern,
        init: impl FnOnce(CoroutineContext) -> F,
    ) -> Self {
        let decoder = Self::with_options(input_cap, access, init);
        decoder.shared.lock().controlled = true;
        decoder
    }

    /// Controlled decoding with admitted input payload, range index and row
    /// staging. Admission runs before `init` is called. The input index holds at
    /// most 64 disjoint spans; excess fragmentation is an error, never eviction.
    /// Native workspace, image metadata, request bookkeeping and sink allocations
    /// are outside this scoped allowance. Reuse the same budget in native code
    /// to admit its additional owned buffers.
    pub fn controlled_with_storage(
        input_cap: usize,
        access: AccessPattern,
        storage: zencodec::storage::StorageBudget,
        init: impl FnOnce(CoroutineContext) -> F,
    ) -> Result<Self, CoroutineError> {
        let mut shared = Shared::with_storage(input_cap, storage)?;
        shared.controlled = true;
        Ok(Self::from_shared(shared, access, init))
    }

    /// Create a decoder that keeps at most `input_cap` unread bytes and
    /// declares `access` to drivers.
    pub fn with_options(
        input_cap: usize,
        access: AccessPattern,
        init: impl FnOnce(CoroutineContext) -> F,
    ) -> Self {
        Self::from_shared(Shared::new(input_cap), access, init)
    }

    fn from_shared(
        state: Shared,
        access: AccessPattern,
        init: impl FnOnce(CoroutineContext) -> F,
    ) -> Self {
        let shared = Arc::new(spin::Mutex::new(state));
        let future = Box::pin(init(CoroutineContext::new(Arc::clone(&shared))));
        Self {
            shared,
            future,
            state: State::Running,
            begun: false,
            accepts_provisional: false,
            access,
            need: Vec::new(),
            cached_info: None,
            details: None,
            cached_info_generation: 0,
            pending_rows: None,
            _err: PhantomData,
        }
    }

    /// Unread input currently retained, in bytes.
    pub fn retained_input(&self) -> usize {
        self.shared.lock().input.retained()
    }
}

impl<E> CoroutineDecoder<BoxFuture<E>, E>
where
    E: core::error::Error + Send + Sync + From<CoroutineError> + 'static,
{
    /// Like [`new`](Self::new), boxing the future so the decoder type is
    /// [`BoxedCoroutineDecoder`] regardless of the `async` block.
    pub fn boxed<Fut>(init: impl FnOnce(CoroutineContext) -> Fut) -> Self
    where
        Fut: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    {
        Self::new(|ctx| Box::pin(init(ctx)) as BoxFuture<E>)
    }

    /// Like [`with_options`](Self::with_options), boxing the future.
    pub fn boxed_with_options<Fut>(
        input_cap: usize,
        access: AccessPattern,
        init: impl FnOnce(CoroutineContext) -> Fut,
    ) -> Self
    where
        Fut: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    {
        Self::with_options(input_cap, access, |ctx| Box::pin(init(ctx)) as BoxFuture<E>)
    }
}

impl<F, E> CoroutineDecoder<F, E>
where
    F: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    E: core::error::Error + Send + Sync + From<CoroutineError> + 'static,
{
    /// Copy into the store the bytes of `input[from..]` that the
    /// outstanding reads need, and report how far into `input` that took
    /// (a prefix: needed bytes plus any unneeded bytes before them).
    fn absorb(
        shared: &mut Shared,
        input: &Input<'_>,
        from: usize,
    ) -> Result<usize, CoroutineError> {
        if let Some(t) = input.total_len {
            shared.input.set_total(t);
        }
        let front = input
            .offset
            .checked_add(from as u64)
            .ok_or(CoroutineError::InputRangeOverflow)?;
        input
            .offset
            .checked_add(input.bytes.len() as u64)
            .ok_or(CoroutineError::InputRangeOverflow)?;
        let end = input.end();
        if front >= end {
            return Ok(from);
        }
        let mut taken_to = front;
        let needs = core::mem::take(&mut shared.need);
        for r in &needs {
            let (lo, hi, sequential) = match *r {
                InputRequest::AtLeast { offset, .. } => (offset, end, true),
                InputRequest::Range(range) => (range.offset, range.end(), false),
                _ => continue,
            };
            let lo = lo.max(front);
            let hi = hi.min(end);
            if lo >= hi {
                continue;
            }
            let slice = &input.bytes[(lo - input.offset) as usize..(hi - input.offset) as usize];
            let stored = shared.input.feed(lo, slice, false)?;
            let reached = lo + stored as u64;
            if stored == slice.len() || sequential {
                taken_to = taken_to.max(reached);
            } else {
                // The store is full: take only up to where storing stopped.
                taken_to = taken_to.max(reached);
                break;
            }
        }
        shared.need = needs;
        Ok((taken_to - input.offset) as usize)
    }

    fn begin_if_ready(&mut self, sink: &mut dyn DecodeRowSink) -> Result<(), StepError<E>> {
        self.check_interpretation()?;
        if self.begun {
            return Ok(());
        }
        let (dims, desc) = {
            let shared = self.shared.lock();
            (
                shared.info.as_ref().map(|i| (i.width, i.height)),
                shared.descriptor,
            )
        };
        if let (Some((w, h)), Some(desc)) = (dims, desc) {
            sink.begin(w, h, desc).map_err(StepError::Sink)?;
            let accepts = sink.accepts_provisional();
            self.shared.lock().accepts_provisional = accepts;
            self.accepts_provisional = accepts;
            self.begun = true;
        }
        Ok(())
    }

    fn check_interpretation(&self) -> Result<(), StepError<E>> {
        let shared = self.shared.lock();
        if let Some(error) = shared.control_failure {
            return Err(StepError::Control(error));
        }
        if shared.controlled && (shared.boundary.is_some() || shared.adopted.is_none()) {
            return Err(StepError::Control(ControlError::OutputBeforeInterpretation));
        }
        Ok(())
    }
}

impl<F, E> CoroutineDecoder<F, E>
where
    F: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    E: core::error::Error + Send + Sync + From<CoroutineError> + 'static,
{
    fn with_details_cleanup<T>(&mut self, work: impl FnOnce(&mut Self) -> T) -> T {
        let guard = DetailsCleanup(self);
        work(&mut *guard.0)
    }

    fn advance(
        &mut self,
        input: Input<'_>,
        sink: &mut dyn DecodeRowSink,
        pulse: &dyn zencodec::Pulse,
    ) -> Result<ControlledProgress, StepError<E>> {
        match &self.state {
            State::Done(info) => return Ok(Progress::new(0, Event::Done(info.clone())).into()),
            State::Failed => return Err(StepError::Finished),
            State::Running | State::Finishing(_) => {}
        }
        if let Err(reason) = pulse.check() {
            self.pending_rows = None;
            self.shared.lock().cancelled = Some(reason);
            self.state = State::Failed;
            return Err(StepError::Cancelled(reason));
        }
        if let Some(boundary) = self.shared.lock().boundary {
            return Ok(ControlledProgress::Suspended {
                consumed: 0,
                boundary,
            });
        }
        if let Some((rows, offset)) = self.pending_rows.take() {
            return self.deliver_rows(rows, offset, 0, sink);
        }
        if let State::Finishing(info) = &self.state {
            return self.finish(sink, info.clone(), 0).map(Into::into);
        }
        self.state = State::Failed;
        let mut consumed = Self::absorb(&mut self.shared.lock(), &input, 0)
            .map_err(|e| StepError::Codec(E::from(e)))?;
        self.state = State::Running;
        let waker = Waker::noop();
        let mut cx = Context::from_waker(waker);
        loop {
            // A caught panic never permits polling the half-mutated future again.
            self.state = State::Failed;
            let polled = self.future.as_mut().poll(&mut cx);
            self.state = State::Running;
            if let Err(reason) = pulse.check() {
                self.shared.lock().cancelled = Some(reason);
                self.state = State::Failed;
                return Err(StepError::Cancelled(reason));
            }
            let event = {
                let mut shared = self.shared.lock();
                if shared.info_generation != self.cached_info_generation {
                    self.cached_info = shared.info.clone();
                    self.cached_info_generation = shared.info_generation;
                }
                if let Some(error) = shared.control_failure {
                    self.state = State::Failed;
                    return Err(StepError::Control(error));
                }
                if let Some(boundary) = shared.boundary {
                    if shared.rows.is_some() {
                        self.state = State::Failed;
                        return Err(StepError::Control(ControlError::OutputBeforeInterpretation));
                    }
                    Pending::Boundary(boundary)
                } else if let Some(rows) = shared.rows.take() {
                    Pending::Rows(rows)
                } else if let Some(update) = shared.metadata.take() {
                    Pending::Metadata(update)
                } else if !shared.need.is_empty() {
                    Pending::Need
                } else {
                    Pending::Stalled
                }
            };
            match polled {
                Poll::Ready(Ok(info)) => {
                    if matches!(event, Pending::Boundary(_)) {
                        self.state = State::Failed;
                        return Err(StepError::Control(ControlError::InconsistentFacts));
                    }
                    if matches!(event, Pending::Rows(_) | Pending::Metadata(_)) {
                        // Deliver what was deposited before finishing; the
                        // future is complete and must not be polled again.
                        self.state = State::Finishing(info);
                        self.need.clear();
                    } else {
                        return self.finish(sink, info, consumed).map(Into::into);
                    }
                }
                Poll::Ready(Err(e)) => {
                    self.state = State::Failed;
                    self.need.clear();
                    return Err(StepError::Codec(e));
                }
                Poll::Pending => {}
            }
            match event {
                Pending::Boundary(boundary) => {
                    self.need.clear();
                    return Ok(ControlledProgress::Suspended { consumed, boundary });
                }
                Pending::Rows(rows) => {
                    let inconsistent = {
                        let shared = self.shared.lock();
                        shared.controlled
                            && (shared.descriptor != Some(rows.descriptor)
                                || shared
                                    .info
                                    .as_ref()
                                    .is_none_or(|info| info.width != rows.width))
                    };
                    if inconsistent {
                        self.state = State::Failed;
                        return Err(StepError::Control(ControlError::InconsistentFacts));
                    }
                    if let Err(e) = self.begin_if_ready(sink) {
                        self.state = State::Failed;
                        return Err(e);
                    }
                    if !self.begun {
                        self.state = State::Failed;
                        return Err(StepError::Codec(E::from(CoroutineError::RowsBeforeHeader)));
                    }
                    return self.deliver_rows(rows, 0, consumed, sink);
                }
                Pending::Metadata(update) => {
                    let controlled = self.shared.lock().controlled;
                    if !controlled
                        && update.is_header_ready()
                        && let Err(e) = self.begin_if_ready(sink)
                    {
                        self.state = State::Failed;
                        return Err(e);
                    }
                    return Ok(Progress::new(consumed, Event::Metadata(update)).into());
                }
                Pending::Need => {
                    // The same input may hold what the coroutine now asks
                    // for: take it and poll again rather than send the
                    // driver back for bytes it already offered.
                    let mut shared = self.shared.lock();
                    self.state = State::Failed;
                    let more = Self::absorb(&mut shared, &input, consumed)
                        .map_err(|e| StepError::Codec(E::from(e)))?;
                    self.state = State::Running;
                    if more > consumed {
                        consumed = more;
                        continue;
                    }
                    self.need.clear();
                    self.need.extend_from_slice(&shared.need);
                    return Ok(Progress::need(consumed).into());
                }
                Pending::Stalled => {
                    self.state = State::Failed;
                    return Err(StepError::Codec(E::from(CoroutineError::Stalled)));
                }
            }
        }
    }
}

impl<F, E> IncrementalDecode for CoroutineDecoder<F, E>
where
    F: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    E: core::error::Error + Send + Sync + From<CoroutineError> + 'static,
{
    type Error = E;

    fn step(
        &mut self,
        input: Input<'_>,
        sink: &mut dyn DecodeRowSink,
        pulse: Option<&dyn zencodec::Pulse>,
    ) -> Result<Progress, StepError<E>> {
        self.with_details_cleanup(|this| {
            let pulse = pulse.unwrap_or(&zencodec::NoPulse);
            if matches!(this.state, State::Failed) {
                return Err(StepError::Finished);
            }
            if this.shared.lock().controlled {
                this.state = State::Failed;
                return Err(StepError::Control(ControlError::UnhandledBoundary));
            }
            match this.advance(input, sink, pulse)? {
                ControlledProgress::Decode(progress) => Ok(progress),
                ControlledProgress::Suspended { .. } => {
                    this.state = State::Failed;
                    Err(StepError::Control(ControlError::UnhandledBoundary))
                }
            }
        })
    }

    fn requested(&self) -> &[InputRequest] {
        match self.state {
            State::Running => &self.need,
            _ => &[],
        }
    }

    fn info(&self) -> Option<&ImageInfo> {
        self.details
            .as_ref()
            .map(|d| &d.info)
            .or(self.cached_info.as_ref())
    }

    fn take_details(&mut self) -> Option<zencodec::decode::DecodeDetails> {
        if !matches!(self.state, State::Done(_)) {
            return None;
        }
        let details = self.details.take();
        if details.is_some() {
            self.cached_info = None;
        }
        details
    }

    fn output_descriptor(&self) -> Option<PixelDescriptor> {
        self.shared.lock().descriptor
    }

    fn access_pattern(&self) -> AccessPattern {
        self.access
    }
}

impl<F, E> ControlledIncrementalDecode for CoroutineDecoder<F, E>
where
    F: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    E: core::error::Error + Send + Sync + From<CoroutineError> + 'static,
{
    fn step_controlled(
        &mut self,
        input: Input<'_>,
        sink: &mut dyn DecodeRowSink,
        pulse: Option<&dyn zencodec::Pulse>,
    ) -> Result<ControlledProgress, StepError<E>> {
        self.with_details_cleanup(|this| {
            let pulse = pulse.unwrap_or(&zencodec::NoPulse);
            if matches!(this.state, State::Failed) {
                return Err(StepError::Finished);
            }
            if !this.shared.lock().controlled {
                this.state = State::Failed;
                return Err(StepError::Control(ControlError::NotControlled));
            }
            this.advance(input, sink, pulse)
        })
    }

    fn resume(
        &mut self,
        id: BoundaryId,
        action: BoundaryAction,
        pulse: Option<&dyn zencodec::Pulse>,
    ) -> Result<(), StepError<E>> {
        self.with_details_cleanup(|this| {
            let pulse = pulse.unwrap_or(&zencodec::NoPulse);
            let state = core::mem::replace(&mut this.state, State::Failed);
            if !matches!(state, State::Running) {
                return Err(StepError::Finished);
            }
            pulse.check()?;
            let mut shared = this.shared.lock();
            if !shared.controlled {
                return Err(StepError::Control(ControlError::NotControlled));
            }
            if let Some(error) = shared.control_failure {
                return Err(StepError::Control(error));
            }
            let Some(boundary) = shared.boundary.take() else {
                return Err(StepError::Control(ControlError::InvalidReply));
            };
            if boundary.id() != id || !action.permits(boundary.facts()) || shared.reply.is_some() {
                return Err(StepError::Control(ControlError::InvalidReply));
            }
            let facts = boundary.facts();
            if facts.descriptor != shared.descriptor
                || shared
                    .info
                    .as_ref()
                    .is_none_or(|info| info.width != facts.width || info.height != facts.height)
            {
                return Err(StepError::Control(ControlError::InconsistentFacts));
            }
            if action == BoundaryAction::AdoptInterpretation {
                shared.adopted = Some(facts);
            }
            shared.reply = Some((id, action));
            this.state = State::Running;
            Ok(())
        })
    }
}

enum Pending {
    Boundary(DecodeBoundary),
    Rows(crate::context::PendingRows),
    Metadata(MetadataUpdate),
    Need,
    Stalled,
}

/// Route a batch to the sink: final rows through `provide_next_buffer`,
/// provisional rows through `provide_pass_buffer`, and only to a sink that
/// accepts them. The coroutine API (`accepts_provisional` on the context)
/// lets a decoder skip provisional passes nobody wants; sending one anyway
/// is an error rather than a silent drop.
fn copy_rows(
    sink: &mut dyn DecodeRowSink,
    rows: &crate::context::PendingRows,
    offset: u32,
    count: u32,
    accepts_provisional: bool,
) -> Result<(), SinkError> {
    if !rows.meta.is_final && !accepts_provisional {
        return Err("decoder wrote provisional rows to a sink that does not accept them".into());
    }
    let pixels = zenpixels::PixelSlice::new(
        rows.data.as_slice(),
        rows.width,
        rows.rows,
        rows.stride,
        rows.descriptor,
    )?;
    rows.meta
        .y
        .checked_add(rows.rows)
        .ok_or_else(|| SinkError::from("copied rows exceed the row coordinate range"))?;
    if count == 0 {
        return Ok(());
    }
    let mut meta = rows.meta;
    meta.y += offset;
    let mut destination = if meta.is_final {
        sink.provide_next_buffer(meta.y, count, rows.width, rows.descriptor)?
    } else {
        sink.provide_pass_buffer(meta, count, rows.width, rows.descriptor)?
    };
    if destination.width() != rows.width
        || destination.rows() != count
        || destination.descriptor() != rows.descriptor
    {
        return Err("sink returned a buffer with a different geometry or descriptor".into());
    }
    for row in 0..count {
        destination
            .row_mut(row)
            .copy_from_slice(pixels.row(offset + row));
    }
    Ok(())
}

impl<F, E> CoroutineDecoder<F, E>
where
    F: Future<Output = Result<OutputInfo, E>> + Send + 'static,
    E: core::error::Error + Send + Sync + From<CoroutineError> + 'static,
{
    fn deliver_rows(
        &mut self,
        rows: crate::context::PendingRows,
        offset: u32,
        consumed: usize,
        sink: &mut dyn DecodeRowSink,
    ) -> Result<ControlledProgress, StepError<E>> {
        // Keep terminal behavior even if a sink callback unwinds. The batch is
        // local until copying succeeds, so failure releases it immediately.
        let state = core::mem::replace(&mut self.state, State::Failed);
        self.check_interpretation()?;
        let count = sink.max_buffer_rows().get().min(rows.rows - offset);
        copy_rows(sink, &rows, offset, count, self.accepts_provisional).map_err(StepError::Sink)?;
        let mut meta = rows.meta;
        meta.y += offset;
        if offset + count < rows.rows {
            self.pending_rows = Some((rows, offset + count));
        } else {
            self.shared.lock().spare = rows.data;
        }
        self.state = state;
        self.need.clear();
        Ok(Progress::new(consumed, Event::Rows { meta, rows: count }).into())
    }

    /// The future completed with `info`: begin if the header never was,
    /// finish the sink, and report `Done`.
    fn finish(
        &mut self,
        sink: &mut dyn DecodeRowSink,
        info: OutputInfo,
        consumed: usize,
    ) -> Result<Progress, StepError<E>> {
        // Final sink callbacks may unwind. A caught panic must never resume a
        // completed future or expose its pending details as successful output.
        self.state = State::Failed;
        let inconsistent = {
            let shared = self.shared.lock();
            shared.details.as_ref().is_some_and(|details| {
                details.info.width != info.width
                    || details.info.height != info.height
                    || (shared.adopted.is_some()
                        && shared.info.as_ref().is_some_and(|old| {
                            old.orientation != details.info.orientation
                                || old.source_color != details.info.source_color
                        }))
            }) || (shared.controlled
                && shared.adopted.is_some_and(|facts| {
                    facts.width != info.width
                        || facts.height != info.height
                        || facts.descriptor != Some(info.native_format)
                }))
        };
        if inconsistent {
            self.state = State::Failed;
            return Err(StepError::Control(ControlError::InconsistentFacts));
        }
        if let Err(e) = self.begin_if_ready(sink) {
            self.state = State::Failed;
            return Err(e);
        }
        if self.begun
            && let Err(e) = sink.finish()
        {
            self.state = State::Failed;
            return Err(StepError::Sink(e));
        }
        self.state = State::Done(info.clone());
        self.need.clear();
        let mut shared = self.shared.lock();
        shared.input.clear();
        shared.spare = Default::default();
        self.details = shared.details.take();
        if self.details.is_some() {
            // The owned final result supersedes earlier metadata snapshots.
            self.cached_info = None;
            shared.info = None;
        }
        Ok(Progress::new(consumed, Event::Done(info)))
    }
}