lace_sans_io/iltp/

output.rs

//! The outgoing half of one carrier codec: the accepted output batch, produced
//! progressively into one bounded chunk.
//!
//! Prefaces and control lines are rendered one line at a time and records are
//! read from their canonical bytes after their label line, followed by their
//! framing LF. Everything after a
//! raw preface passes through
//! the direction's 065 encoder into a chunk of at most the codec's chunk size,
//! grown only as far as the batch can use. An `identity` direction hands a
//! record that does not fit the current chunk to the writer in place, at most
//! one chunk of it at a time, instead of copying it. The writer takes the current bytes, reports how many it
//! wrote, and asks again; partial writes keep the same chunk and offset.
//!
//! A record without contiguous bytes (a file range) is read in pieces of at
//! most one chunk: the cursor hands its owner a [`RecordRead`], the owner
//! performs it, possibly on a blocking worker, and returns it before asking
//! for more bytes.

use super::control::{LineCursor, next_line};
use super::encoded_stream::{Incoming, Outgoing, install};
use super::framing::{IltpError, IltpItem, write_preface};
use super::stream_encoding::{ENCODE_MIN_OUTPUT, EncodeMode, EncodeProgress};
use crate::records::Record;

/// Chunk room beyond a batch's plaintext that encoder framing, flushes and
/// end marks may use.
const ENCODING_SLACK: usize = 64;

pub(super) struct Output {
    chunk_bytes: usize,
    /// Zero-filled chunk storage, grown lazily up to `chunk_bytes`.
    chunk: Vec<u8>,
    filled: usize,
    /// Bytes the writer took from the chunk, or from the in-place record.
    taken: usize,
    /// An `identity` record handed to the writer without copying.
    in_place: Option<Record>,
    /// The current preface or control line.
    line: Vec<u8>,
    /// Bytes read from a file-range record and the count encoded so far.
    read: Vec<u8>,
    read_used: usize,
    read_state: ReadState,
    batch: Option<Batch>,
}

enum ReadState {
    Idle,
    /// Prepared and not yet handed to the owner.
    Requested(RecordRead),
    /// With the owner.
    Outstanding,
}

/// The next bytes of accepted output; see [`IltpCodec::output_bytes`](super::IltpCodec::output_bytes).
#[derive(Debug)]
pub enum OutputBytes<'a> {
    /// Bytes the writer must take next.
    Write(&'a [u8]),
    /// Record bytes the owner must read and return with
    /// [`IltpCodec::complete_read`](super::IltpCodec::complete_read) before
    /// more output exists.
    Read(RecordRead),
    /// The accepted output is complete, or none is accepted.
    Complete,
}

/// A read of `buffer.len()` canonical bytes of `record` at `offset`.
#[derive(Debug)]
pub struct RecordRead {
    record: Record,
    offset: usize,
    buffer: Vec<u8>,
}

impl RecordRead {
    /// Fill the buffer from the record's source. This is blocking file IO.
    pub fn run(mut self) -> std::io::Result<Self> {
        let mut filled = 0;
        while filled < self.buffer.len() {
            match self
                .record
                .read_at(self.offset + filled, &mut self.buffer[filled..])?
            {
                0 => return Err(std::io::ErrorKind::UnexpectedEof.into()),
                count => filled += count,
            }
        }
        Ok(self)
    }
}

struct Batch {
    items: Vec<IltpItem>,
    /// Canonical plaintext bytes of all items, which bound useful chunk size.
    plaintext: usize,
    step: Step,
    /// Offset into the current line or record.
    offset: usize,
    /// Encoded items follow the last preface, so the batch ends with a flush.
    flush: bool,
}

#[derive(Debug, Clone, Copy)]
enum Step {
    /// Start item `index`; past the last item, finish the batch.
    Item(usize),
    /// End the previous logical session's encoded stream before preface `index`.
    EndBeforePreface(usize),
    /// The raw bytes of preface `index`, held in `line`.
    Preface(usize),
    /// The current line of control item `index`, held in `line`.
    Line(usize, LineCursor),
    Record(usize),
    /// The LF after record `index`.
    RecordEnd(usize),
    Flush,
    End,
    Done,
}

impl Output {
    pub(super) fn new(chunk_bytes: usize) -> Self {
        Self {
            chunk_bytes,
            chunk: Vec::new(),
            filled: 0,
            taken: 0,
            in_place: None,
            line: Vec::new(),
            read: Vec::new(),
            read_used: 0,
            read_state: ReadState::Idle,
            batch: None,
        }
    }

    /// Bytes of chunk storage allocated so far.
    #[cfg(test)]
    pub(super) fn chunk_len(&self) -> usize {
        self.chunk.len()
    }

    pub(super) fn is_pending(&self) -> bool {
        self.batch.is_some()
    }

    /// Accept validated items holding `plaintext` canonical bytes.
    pub(super) fn accept(&mut self, items: Vec<IltpItem>, plaintext: usize) {
        self.start(items, plaintext, Step::Item(0));
    }

    /// Accept the end of the outgoing encoded stream.
    pub(super) fn accept_end(&mut self) {
        self.start(Vec::new(), 0, Step::End);
    }

    fn start(&mut self, items: Vec<IltpItem>, plaintext: usize, step: Step) {
        debug_assert!(self.batch.is_none());
        self.batch = Some(Batch {
            items,
            plaintext,
            step,
            offset: 0,
            flush: false,
        });
    }

    pub(super) fn discard(&mut self) {
        self.batch = None;
        self.in_place = None;
        self.filled = 0;
        self.taken = 0;
        self.read.clear();
        self.read_used = 0;
        self.read_state = ReadState::Idle;
    }

    /// The next bytes the writer must take, producing them when the previous
    /// ones are taken, or the record read that must precede them.
    pub(super) fn bytes(
        &mut self,
        outgoing: &mut Outgoing,
        incoming: &mut Incoming,
    ) -> Result<OutputBytes<'_>, IltpError> {
        if matches!(self.read_state, ReadState::Outstanding) {
            return Err(IltpError::Local("ILTP output read is outstanding"));
        }
        if self.in_place.is_none()
            && self.taken == self.filled
            && matches!(self.read_state, ReadState::Idle)
        {
            self.fill(outgoing, incoming)?;
        }
        if let Some(record) = &self.in_place {
            let bytes = record
                .contiguous_bytes()
                .ok_or(IltpError::Local("in-place record is not contiguous"))?;
            let rest = &bytes[self.taken..];
            return Ok(OutputBytes::Write(
                &rest[..rest.len().min(self.chunk_bytes)],
            ));
        }
        if self.taken < self.filled {
            return Ok(OutputBytes::Write(&self.chunk[self.taken..self.filled]));
        }
        Ok(
            match std::mem::replace(&mut self.read_state, ReadState::Outstanding) {
                ReadState::Requested(read) => OutputBytes::Read(read),
                _ => {
                    self.read_state = ReadState::Idle;
                    OutputBytes::Complete
                }
            },
        )
    }

    /// Accept the owner's completed [`RecordRead`].
    pub(super) fn complete_read(&mut self, read: RecordRead) -> Result<(), IltpError> {
        let expected = match &self.batch {
            Some(Batch {
                step: Step::Record(_),
                offset,
                ..
            }) => *offset,
            // Output discarded while the read ran.
            None => return Ok(()),
            _ => return Err(IltpError::Local("ILTP output read without its record")),
        };
        if !matches!(self.read_state, ReadState::Outstanding) || read.offset != expected {
            return Err(IltpError::Local(
                "ILTP output read does not match its request",
            ));
        }
        self.read = read.buffer;
        self.read_used = 0;
        self.read_state = ReadState::Idle;
        Ok(())
    }

    /// Accept `count` written bytes of the current [`bytes`](Self::bytes).
    pub(super) fn advance(&mut self, count: usize) {
        self.taken += count;
        match &self.in_place {
            Some(record) if self.taken >= record.canonical_len() => {
                self.in_place = None;
                self.taken = 0;
                self.filled = 0;
            }
            Some(_) => {}
            None => debug_assert!(self.taken <= self.filled),
        }
        self.finish_when_taken();
    }

    fn finish_when_taken(&mut self) {
        if self.in_place.is_none()
            && self.taken == self.filled
            && matches!(&self.batch, Some(batch) if matches!(batch.step, Step::Done))
        {
            self.batch = None;
        }
    }

    /// Produce the next chunk, or hand out the next record in place.
    fn fill(&mut self, outgoing: &mut Outgoing, incoming: &mut Incoming) -> Result<(), IltpError> {
        self.filled = 0;
        self.taken = 0;
        let Self {
            chunk_bytes,
            chunk,
            filled,
            taken,
            in_place,
            line,
            read,
            read_used,
            read_state,
            batch,
        } = self;
        let Some(batch) = batch else {
            return Ok(());
        };
        let useful = batch
            .plaintext
            .saturating_add(ENCODING_SLACK)
            .min(*chunk_bytes);
        if chunk.len() < useful {
            chunk.resize(useful, 0);
        }
        loop {
            match batch.step {
                Step::Item(index) => {
                    batch.offset = 0;
                    batch.step = match batch.items.get(index) {
                        None if batch.flush => Step::Flush,
                        None => Step::Done,
                        Some(IltpItem::SessionPreface(preface)) => {
                            line.clear();
                            write_preface(preface, line);
                            Step::EndBeforePreface(index)
                        }
                        Some(item) => {
                            batch.flush = true;
                            let mut cursor = LineCursor::default();
                            next_line(item, &mut cursor, line);
                            Step::Line(index, cursor)
                        }
                    };
                }
                Step::EndBeforePreface(index) => {
                    match encode_step(outgoing, chunk, filled, &[], EncodeMode::End)? {
                        Some(progress) if progress.done => batch.step = Step::Preface(index),
                        Some(_) => {}
                        None => break,
                    }
                }
                Step::Preface(index) => {
                    let count = copy(&line[batch.offset..], &mut chunk[*filled..]);
                    if count == 0 {
                        break;
                    }
                    *filled += count;
                    batch.offset += count;
                    if batch.offset == line.len() {
                        let IltpItem::SessionPreface(preface) = &batch.items[index] else {
                            unreachable!("preface step on another item");
                        };
                        *outgoing = Outgoing::Pending(preface.offers.clone());
                        if let Incoming::Pending(peer) = incoming {
                            let peer = peer.clone();
                            install(outgoing, incoming, &preface.offers, &peer, *chunk_bytes)?;
                        }
                        batch.flush = false;
                        batch.step = Step::Item(index + 1);
                    }
                }
                Step::Line(index, mut cursor) => {
                    let input = &line[batch.offset..];
                    match encode_step(outgoing, chunk, filled, input, EncodeMode::Write)? {
                        Some(progress) => batch.offset += progress.consumed,
                        None => break,
                    }
                    if batch.offset == line.len() {
                        batch.offset = 0;
                        let item = &batch.items[index];
                        batch.step = if next_line(item, &mut cursor, line) {
                            Step::Line(index, cursor)
                        } else if matches!(item, IltpItem::Record(..)) {
                            Step::Record(index)
                        } else {
                            Step::Item(index + 1)
                        };
                    }
                }
                Step::Record(index) => {
                    let IltpItem::Record(_, record) = &batch.items[index] else {
                        unreachable!("record step on another item");
                    };
                    let Some(bytes) = record.contiguous_bytes() else {
                        let rest = record.canonical_len() - batch.offset;
                        if rest == 0 {
                            batch.offset = 0;
                            batch.step = Step::RecordEnd(index);
                        } else if *read_used < read.len() {
                            let input = &read[*read_used..];
                            match encode_step(outgoing, chunk, filled, input, EncodeMode::Write)? {
                                Some(progress) => {
                                    batch.offset += progress.consumed;
                                    *read_used += progress.consumed;
                                }
                                None => break,
                            }
                        } else {
                            let mut buffer = std::mem::take(read);
                            buffer.clear();
                            buffer.resize(rest.min(*chunk_bytes), 0);
                            *read_used = 0;
                            *read_state = ReadState::Requested(RecordRead {
                                record: record.clone(),
                                offset: batch.offset,
                                buffer,
                            });
                            break;
                        }
                        continue;
                    };
                    let rest = &bytes[batch.offset..];
                    if rest.is_empty() {
                        batch.offset = 0;
                        batch.step = Step::RecordEnd(index);
                    } else if outgoing.is_identity() && rest.len() > chunk.len() - *filled {
                        if *filled > 0 {
                            break;
                        }
                        *taken = batch.offset;
                        *in_place = Some(record.clone());
                        batch.offset = 0;
                        batch.step = Step::RecordEnd(index);
                        return Ok(());
                    } else {
                        match encode_step(outgoing, chunk, filled, rest, EncodeMode::Write)? {
                            Some(progress) => batch.offset += progress.consumed,
                            None => break,
                        }
                    }
                }
                Step::RecordEnd(index) => {
                    match encode_step(outgoing, chunk, filled, b"\n", EncodeMode::Write)? {
                        Some(progress) if progress.consumed == 1 => {
                            batch.step = Step::Item(index + 1);
                        }
                        Some(_) => {}
                        None => break,
                    }
                }
                Step::Flush | Step::End => {
                    let mode = match batch.step {
                        Step::Flush => EncodeMode::Flush,
                        _ => EncodeMode::End,
                    };
                    match encode_step(outgoing, chunk, filled, &[], mode)? {
                        Some(progress) if progress.done => batch.step = Step::Done,
                        Some(_) => {}
                        None => break,
                    }
                }
                Step::Done => break,
            }
        }
        self.finish_when_taken();
        Ok(())
    }
}

fn copy(input: &[u8], room: &mut [u8]) -> usize {
    let count = input.len().min(room.len());
    room[..count].copy_from_slice(&input[..count]);
    count
}

/// One encoder step into the chunk's free room. `None` when the room is too
/// small, so the filled bytes go to the writer first.
fn encode_step(
    outgoing: &mut Outgoing,
    chunk: &mut [u8],
    filled: &mut usize,
    input: &[u8],
    mode: EncodeMode,
) -> Result<Option<EncodeProgress>, IltpError> {
    let room = &mut chunk[*filled..];
    if room.len() < ENCODE_MIN_OUTPUT {
        return Ok(None);
    }
    let progress = outgoing.encode(input, mode, room)?;
    *filled += progress.produced;
    if !progress.done && progress.consumed == 0 && progress.produced == 0 {
        return if *filled > 0 {
            Ok(None)
        } else {
            Err(IltpError::Local("ILTP stream encoder made no progress"))
        };
    }
    Ok(Some(progress))
}