mopenzl 0.1.0

//! Chunk headers and the stream graph they describe
//! (`decodeChunkHeader_internal`, `fillStoredStreams`).
//!
//! A chunk lists its codecs in decoding order, as parallel bit-packed and
//! varint lanes: types, IDs, private header sizes, variable-output counts,
//! regenerated-stream counts, then one distance per regenerated stream and
//! the size of every stored stream. Streams are numbered so that each codec's
//! inputs are the streams right after the previous codec's inputs, and each
//! regenerated stream lies `distance` streams past the end of its codec's
//! inputs. Streams no codec regenerates are stored, in index order, after the
//! codecs' private headers. The last stream is the chunk's output.

use alloc::vec::Vec;
use core::ops::Range;

use super::frame::FrameInfo;
use super::reader::{Parse, Reader};
use crate::decompressor::error::{DecodeError, InvalidDataKind, Unsupported};
use crate::shared::bitpack;
#[cfg(feature = "zstd")]
use crate::shared::wire::CODEC_ZSTD;
use crate::shared::wire::{INPUT_LIMIT, NODE_LIMIT, STREAM_LIMIT, codec_id_bits};

/// A codec this decoder runs; each is compiled with its own feature.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub(crate) enum Codec {
    /// `ZL_StandardTransformID_zstd`: one stream in, one stream out.
    #[cfg(feature = "zstd")]
    Zstd,
}

impl Codec {
    /// The codec of standard ID `id`.
    const fn standard(id: u32) -> Option<Self> {
        match id {
            #[cfg(feature = "zstd")]
            CODEC_ZSTD => Some(Self::Zstd),
            _ => None,
        }
    }

    /// Fixed inputs of the decoder (singleton outputs of the encoder).
    const fn fixed_inputs(self) -> usize {
        match self {
            #[cfg(feature = "zstd")]
            Self::Zstd => 1,
        }
    }

    /// Whether the codec takes variable inputs.
    const fn variable_inputs(self) -> bool {
        match self {
            #[cfg(feature = "zstd")]
            Self::Zstd => false,
        }
    }

    /// Streams the codec regenerates.
    const fn regenerated(self) -> usize {
        match self {
            #[cfg(feature = "zstd")]
            Self::Zstd => 1,
        }
    }
}

/// One codec of the chunk, ready to run.
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct Node {
    pub(crate) codec: Codec,
    /// Its private header, as a range of the chunk.
    pub(crate) header: Range<usize>,
    /// Its input streams.
    pub(crate) inputs: Range<usize>,
    /// The stream it regenerates.
    pub(crate) output: usize,
}

/// Where a stream comes from.
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum Source {
    /// Stored in the chunk at this range.
    Stored(Range<usize>),
    /// Regenerated by this node.
    Node(usize),
}

/// The parsed shape of one chunk, with reusable lanes.
#[derive(Debug, Default)]
pub(crate) struct Layout {
    pub(crate) nodes: Vec<Node>,
    pub(crate) streams: Vec<Source>,
    custom: Vec<bool>,
    ids: Vec<u32>,
    custom_ids: Vec<u32>,
    kinds: Vec<Codec>,
    producer: Vec<Option<usize>>,
    header_sizes: Vec<u64>,
    variable: Vec<u64>,
    regens: Vec<u64>,
    distances: Vec<u32>,
    stored: Vec<u64>,
    /// Bytes of the whole chunk, checksums included.
    pub(crate) len: usize,
    /// Offset of the content checksum, when present.
    pub(crate) content_checksum: Option<usize>,
    /// Offset of the compressed checksum, when present.
    pub(crate) compressed_checksum: Option<usize>,
}

/// Invalid chunk header data.
const fn bad_header() -> Parse {
    Parse::Failed(DecodeError::invalid(InvalidDataKind::ChunkHeader))
}

/// An inconsistent stream graph.
const fn bad_graph() -> Parse {
    Parse::Failed(DecodeError::invalid(InvalidDataKind::Graph))
}

/// Reads one flag lane, then a varint for every set flag, giving each entry
/// `base` when clear and `varint + offset` when set.
fn flagged_lane(
    reader: &mut Reader<'_>,
    count: usize,
    lane: &mut Vec<u64>,
    offset: u64,
    base: u64,
) -> Result<(), Parse> {
    lane.clear();
    reader.packed(count, 1, |flag| lane.push(u64::from(flag)))?;
    for entry in lane.iter_mut() {
        *entry = if *entry != 0 {
            reader
                .varint()?
                .checked_add(offset)
                .ok_or_else(bad_header)?
        } else {
            base
        };
    }
    Ok(())
}

impl Layout {
    /// Bytes the lanes keep allocated.
    pub(crate) fn retained_bytes(&self) -> usize {
        use core::mem::size_of;
        self.nodes.capacity() * size_of::<Node>()
            + self.streams.capacity() * size_of::<Source>()
            + self.custom.capacity()
            + self.kinds.capacity() * size_of::<Codec>()
            + self.producer.capacity() * size_of::<Option<usize>>()
            + (self.ids.capacity() + self.custom_ids.capacity() + self.distances.capacity()) * 4
            + (self.header_sizes.capacity()
                + self.variable.capacity()
                + self.regens.capacity()
                + self.stored.capacity())
                * 8
    }

    /// Releases the lanes.
    pub(crate) fn release(&mut self) {
        *self = Self::default();
    }

    /// Parses the chunk at the front of `bytes` and lays out its streams.
    ///
    /// Fails with [`Parse::Truncated`] only while the chunk header itself is
    /// incomplete; the rest of the chunk need not be present yet.
    #[cfg_attr(all(feature = "hotpath", not(feature = "no_std")), hotpath::measure)]
    pub(crate) fn parse(&mut self, bytes: &[u8], frame: &FrameInfo) -> Result<(), Parse> {
        let mut reader = Reader::new(bytes, InvalidDataKind::ChunkHeader);
        let codecs = reader.varint()?.checked_sub(1).ok_or_else(bad_header)?;
        let stored = reader.varint()?;
        if codecs >= NODE_LIMIT || stored >= STREAM_LIMIT {
            return Err(bad_header());
        }
        let (codecs, stored) = (codecs as usize, stored as usize);

        self.custom.clear();
        reader.packed(codecs, 1, |flag| self.custom.push(flag != 0))?;
        let customs = self.custom.iter().filter(|&&custom| custom).count();
        self.ids.clear();
        reader.packed(codecs - customs, codec_id_bits(frame.version), |id| {
            self.ids.push(id)
        })?;
        self.custom_ids.clear();
        for _ in 0..customs {
            let id = u32::try_from(reader.varint()?).map_err(|_| bad_header())?;
            self.custom_ids.push(id);
        }
        flagged_lane(&mut reader, codecs, &mut self.header_sizes, 1, 0)?;
        if self
            .header_sizes
            .iter()
            .any(|&size| size >= u64::from(u32::MAX))
        {
            return Err(bad_header());
        }
        flagged_lane(&mut reader, codecs, &mut self.variable, 1, 0)?;
        flagged_lane(&mut reader, codecs, &mut self.regens, 2, 1)?;
        if self.regens.iter().any(|&count| count > INPUT_LIMIT as u64) {
            return Err(bad_header());
        }
        let regenerated = self.regens.iter().sum::<u64>() as usize;
        self.distances.clear();
        let bits = bitpack::bits_for((regenerated + stored) as u64);
        reader.packed(regenerated, bits, |distance| self.distances.push(distance))?;
        self.stored.clear();
        if bytes.len() - reader.pos() < stored {
            // The reference requires a byte per stored size up front.
            reader.take(stored)?;
        }
        for _ in 0..stored {
            let size = reader.varint()?;
            if size >= u64::from(u32::MAX) {
                return Err(bad_header());
            }
            self.stored.push(size);
        }
        let header_len = reader.pos();

        // The codecs must all be known before any stream is laid out.
        self.kinds.clear();
        let mut custom = self.custom_ids.iter();
        let mut standard = self.ids.iter();
        for &is_custom in &self.custom {
            if is_custom {
                let &id = custom.next().ok_or_else(bad_header)?;
                return Err(DecodeError::unsupported(Unsupported::CustomCodec { id }).into());
            }
            let &id = standard.next().ok_or_else(bad_header)?;
            let codec =
                Codec::standard(id).ok_or(DecodeError::unsupported(Unsupported::Codec { id }))?;
            self.kinds.push(codec);
        }
        self.lay_out(header_len, frame)
    }

    /// Wires the streams of codecs `kinds` and computes the chunk's length.
    fn lay_out(&mut self, header_len: usize, frame: &FrameInfo) -> Result<(), Parse> {
        let total = self.stored.len() + self.distances.len();
        if total == 0 || total as u64 > STREAM_LIMIT {
            return Err(bad_graph());
        }
        self.producer.clear();
        self.producer.resize(total, None);
        let producer = &mut self.producer;
        self.nodes.clear();
        let mut next_stream = 0usize;
        let mut header_pos = header_len as u64;
        for (index, &codec) in self.kinds.iter().enumerate() {
            let variable = self.variable[index];
            if (variable > 0 && !codec.variable_inputs())
                || self.regens[index] != codec.regenerated() as u64
            {
                return Err(bad_graph());
            }
            let inputs = codec.fixed_inputs() + variable as usize;
            let end = next_stream + inputs;
            if end > total - 1 {
                // An input past the streams, or the output feeding a codec.
                return Err(bad_graph());
            }
            // Every codec here regenerates one stream, so codec `index`
            // owns distance `index`.
            let distance = self.distances[index] as usize;
            let output = end + distance;
            if output >= total || producer[output].is_some() {
                return Err(bad_graph());
            }
            producer[output] = Some(index);
            let header_end = header_pos + self.header_sizes[index];
            self.nodes.push(Node {
                codec,
                header: header_pos as usize..header_end as usize,
                inputs: next_stream..end,
                output,
            });
            header_pos = header_end;
            next_stream = end;
        }
        // A codec runs after every codec that regenerates one of its inputs.
        for (index, node) in self.nodes.iter().enumerate() {
            if node
                .inputs
                .clone()
                .any(|input| producer[input].is_some_and(|by| by >= index))
            {
                return Err(bad_graph());
            }
        }
        self.streams.clear();
        let mut pos = header_pos;
        let mut stored = self.stored.iter();
        for slot in producer.iter() {
            self.streams.push(match slot {
                Some(node) => Source::Node(*node),
                None => {
                    let &size = stored.next().ok_or_else(bad_graph)?;
                    let start = pos;
                    pos += size;
                    Source::Stored(start as usize..pos as usize)
                }
            });
        }
        if stored.next().is_some() {
            return Err(bad_graph());
        }
        self.content_checksum = frame.content_checksum.then_some(pos as usize);
        if frame.content_checksum {
            pos += 4;
        }
        self.compressed_checksum = frame.compressed_checksum.then_some(pos as usize);
        if frame.compressed_checksum {
            pos += 4;
        }
        // A length past the address space can never arrive whole.
        self.len = usize::try_from(pos).unwrap_or(usize::MAX);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn frame(checksums: bool) -> FrameInfo {
        FrameInfo {
            version: 27,
            content_checksum: checksums,
            compressed_checksum: checksums,
            size: 0,
            header_len: 0,
        }
    }

    fn parse(bytes: &[u8]) -> Result<Layout, Parse> {
        let mut layout = Layout::default();
        layout.parse(bytes, &frame(true)).map(|()| layout)
    }

    #[cfg(feature = "zstd")]
    #[test]
    fn a_zstd_chunk_wires_the_stored_stream_into_the_output() {
        let header = [0x02, 0x01, 0x00, 0x16, 0x00, 0x00, 0x00, 0x00, 0x15];
        let layout = parse(&header).unwrap();
        assert_eq!(
            layout.nodes,
            [Node {
                codec: Codec::Zstd,
                header: 9..9,
                inputs: 0..1,
                output: 1,
            }]
        );
        assert_eq!(layout.streams, [Source::Stored(9..30), Source::Node(0)]);
        assert_eq!(layout.content_checksum, Some(30));
        assert_eq!(layout.compressed_checksum, Some(34));
        assert_eq!(layout.len, 38);
        for len in 0..header.len() {
            assert_eq!(
                parse(&header[..len]).unwrap_err(),
                Parse::Truncated,
                "{len}"
            );
        }
        assert!(layout.retained_bytes() > 0);
    }

    #[test]
    fn a_store_chunk_is_its_stored_stream() {
        let mut layout = Layout::default();
        layout.parse(&[0x01, 0x01, 0x05], &frame(false)).unwrap();
        assert!(layout.nodes.is_empty());
        assert_eq!(layout.streams, [Source::Stored(3..8)]);
        assert_eq!((layout.len, layout.content_checksum), (8, None));
        layout.release();
        assert_eq!(layout.retained_bytes(), 0);
    }

    #[test]
    fn unknown_codecs_are_unsupported() {
        assert_eq!(
            parse(&[0x02, 0x01, 0x00, 0x31, 0x00, 0x00, 0x00, 0x00, 0x05]).unwrap_err(),
            Parse::Failed(DecodeError::unsupported(Unsupported::Codec { id: 49 }))
        );
        assert_eq!(
            parse(&[0x02, 0x01, 0x01, 0x80, 0x02, 0x00, 0x00, 0x00, 0x00, 0x05]).unwrap_err(),
            Parse::Failed(DecodeError::unsupported(Unsupported::CustomCodec {
                id: 256
            }))
        );
        // A custom ID beyond 32 bits.
        assert_eq!(
            parse(&[0x02, 0x01, 0x01, 0x80, 0x80, 0x80, 0x80, 0x10, 0x00]).unwrap_err(),
            bad_header()
        );
    }

    #[cfg(feature = "zstd")]
    #[test]
    fn malformed_headers_and_graphs_are_rejected() {
        // Zero codec count field; IDs only reach 127 in seven bits.
        assert_eq!(parse(&[0x00, 0x01, 0x00]).unwrap_err(), bad_header());
        assert_eq!(
            parse(&[0x02, 0x01, 0x00, 0x7f, 0x00, 0x00, 0x00, 0x00, 0x05]).unwrap_err(),
            Parse::Failed(DecodeError::unsupported(Unsupported::Codec { id: 127 }))
        );
        // No streams at all; a distance
        // past the streams; zstd with a variable input or two outputs.
        assert_eq!(parse(&[0x01, 0x00]).unwrap_err(), bad_graph());
        assert_eq!(
            parse(&[0x02, 0x01, 0x00, 0x16, 0x00, 0x00, 0x00, 0x01, 0x05]).unwrap_err(),
            bad_graph()
        );
        assert_eq!(
            parse(&[0x02, 0x01, 0x00, 0x16, 0x00, 0x01, 0x00, 0x00, 0x00, 0x05]).unwrap_err(),
            bad_graph()
        );
        assert_eq!(
            parse(&[
                0x02, 0x01, 0x00, 0x16, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x05
            ])
            .unwrap_err(),
            bad_graph()
        );
    }
}