mopenzl 0.2.1

//! 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::codecs::{self, Desc};
use super::frame::FrameInfo;
use super::reader::{Parse, Reader};
use crate::decompressor::error::{DecodeError, InvalidDataKind, Unsupported};
use crate::shared::bitpack;
use crate::shared::wire::{INPUT_LIMIT, NODE_LIMIT, STREAM_LIMIT, codec_id_bits};

/// One codec of the chunk, ready to run.
#[derive(Clone, Debug)]
pub(crate) struct Node {
    pub(crate) desc: Desc,
    pub(crate) custom: bool,
    /// Its private header, as a range of the chunk.
    pub(crate) header: Range<usize>,
    /// Its input streams.
    pub(crate) inputs: Range<usize>,
    /// Its regenerated streams, as a range of [`Layout::regenerated`].
    pub(crate) regens: Range<usize>,
    /// The index of its dictionary in the frame's bundle.
    pub(crate) dict: Option<u32>,
}

impl PartialEq for Node {
    fn eq(&self, other: &Self) -> bool {
        self.desc.id == other.desc.id
            && self.custom == other.custom
            && self.header == other.header
            && self.inputs == other.inputs
            && self.regens == other.regens
            && self.dict == other.dict
    }
}

/// 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>,
    /// The streams each node regenerates, in the node's order.
    pub(crate) regenerated: Vec<usize>,
    custom: Vec<bool>,
    ids: Vec<u32>,
    custom_ids: Vec<u32>,
    kinds: Vec<Desc>,
    #[cfg(feature = "custom-codec")]
    pub(crate) custom_codecs: Vec<crate::custom_codec::Registered>,
    producer: Vec<Option<usize>>,
    header_sizes: Vec<u64>,
    variable: Vec<u64>,
    regens: Vec<u64>,
    dicts: Vec<Option<u32>>,
    dict_values: Vec<u32>,
    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>,
}

/// Streams one codec may read (`ZL_transformOutStreamsLimit` for format
/// versions 16 and later, so for every supported version).
const CODEC_INPUT_LIMIT: usize = 100_000;

/// 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(())
}

/// [`flagged_lane`] for the stream-count lanes, which the reference keeps
/// in 32 bits: a set entry is `(uint32_t)varint + offset`, wrapping
/// (`decompressNbVOs`, `decompressNbRegens`).
fn flagged_count_lane(
    reader: &mut Reader<'_>,
    count: usize,
    lane: &mut Vec<u64>,
    offset: u32,
    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 {
            // Truncation to 32 bits is the reference's arithmetic.
            u64::from((reader.varint()? as u32).wrapping_add(offset))
        } else {
            base
        };
    }
    Ok(())
}

/// Reads the dictionary lane (`decompressDictIdxs`): a has-dictionary flag
/// per codec, then, behind a width byte, the index of every flagged codec.
fn dictionary_lane(
    reader: &mut Reader<'_>,
    count: usize,
    lane: &mut Vec<Option<u32>>,
    values: &mut Vec<u32>,
) -> Result<(), Parse> {
    lane.clear();
    reader.packed(count, 1, |flag| lane.push((flag != 0).then_some(0)))?;
    let flagged = lane.iter().flatten().count();
    if flagged == 0 {
        return Ok(());
    }
    let bits = u32::from(reader.byte()?);
    if bits > 16 {
        return Err(bad_header());
    }
    values.clear();
    reader.packed(flagged, bits, |value| values.push(value))?;
    // The reference also rejects values above 65 535, which 16 bits cannot
    // hold.
    let mut values = values.iter();
    for slot in lane.iter_mut().filter(|slot| slot.is_some()) {
        *slot = values.next().copied();
    }
    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.regenerated.capacity() * size_of::<usize>()
            + self.kinds.capacity() * size_of::<Desc>()
            + self.dicts.capacity() * size_of::<Option<u32>>()
            + self.dict_values.capacity() * 4
            + 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) {
        #[cfg(feature = "custom-codec")]
        let custom_codecs = core::mem::take(&mut self.custom_codecs);
        *self = Self::default();
        #[cfg(feature = "custom-codec")]
        {
            self.custom_codecs = custom_codecs;
        }
    }

    /// 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_count_lane(&mut reader, codecs, &mut self.variable, 1, 0)?;
        flagged_count_lane(&mut reader, codecs, &mut self.regens, 2, 1)?;
        if self.regens.iter().any(|&count| count > INPUT_LIMIT as u64) {
            return Err(bad_header());
        }
        if frame.bundle.is_some() {
            dictionary_lane(&mut reader, codecs, &mut self.dicts, &mut self.dict_values)?;
        } else {
            self.dicts.clear();
            self.dicts.resize(codecs, None);
        }
        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 {
            let desc = if is_custom {
                let &id = custom.next().ok_or_else(bad_header)?;
                #[cfg(feature = "custom-codec")]
                let desc = self
                    .custom_codecs
                    .iter()
                    .find(|codec| codec.id == id)
                    .map(crate::custom_codec::descriptor);
                #[cfg(not(feature = "custom-codec"))]
                let desc: Option<Desc> = None;
                desc.ok_or(DecodeError::unsupported(Unsupported::CustomCodec { id }))?
            } else {
                let &id = standard.next().ok_or_else(bad_header)?;
                *codecs::standard(id).ok_or(DecodeError::unsupported(Unsupported::Codec { id }))?
            };
            if frame.version < desc.min_version {
                return Err(bad_header());
            }
            self.kinds.push(desc);
        }
        self.lay_out(header_len, frame)
    }

    /// Wires the streams of codecs `kinds` and computes the chunk's length
    /// (`fillStoredStreams`, `DCTX_validateNodeStatic`).
    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 || frame.outputs > total {
            return Err(bad_graph());
        }
        // The last `outputs` streams are the chunk's outputs.
        let first_output = total - frame.outputs;
        self.producer.clear();
        self.producer.resize(total, None);
        let producer = &mut self.producer;
        self.nodes.clear();
        self.regenerated.clear();
        let mut next_stream = 0usize;
        let mut next_distance = 0usize;
        let mut header_pos = header_len as u64;
        for (index, &desc) in self.kinds.iter().enumerate() {
            let variable = self.variable[index];
            let regens = self.regens[index] as usize;
            if (variable > 0 && desc.variable_inputs == 0)
                || regens == 0
                || !desc.regens_fit(regens)
                || variable > STREAM_LIMIT
            {
                return Err(bad_graph());
            }
            let inputs = desc.inputs.len() + variable as usize;
            if inputs > CODEC_INPUT_LIMIT {
                return Err(bad_graph());
            }
            let end = next_stream + inputs;
            if end > first_output {
                // An input past the streams, or an output feeding a codec.
                return Err(bad_graph());
            }
            let first = self.regenerated.len();
            for &distance in &self.distances[next_distance..next_distance + regens] {
                let output = end + distance as usize;
                if output >= total || producer[output].is_some() {
                    return Err(bad_graph());
                }
                producer[output] = Some(index);
                self.regenerated.push(output);
            }
            next_distance += regens;
            let header_end = header_pos + self.header_sizes[index];
            self.nodes.push(Node {
                desc,
                custom: self.custom[index],
                header: header_pos as usize..header_end as usize,
                inputs: next_stream..end,
                regens: first..self.regenerated.len(),
                dict: self.dicts[index],
            });
            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,
            outputs: 1,
            strings: 0,
            bundle: None,
            comment: None,
        }
    }

    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 {
                desc: super::codecs::zstd::ZSTD,
                custom: false,
                header: 9..9,
                inputs: 0..1,
                regens: 0..1,
                dict: None,
            }]
        );
        assert_eq!(layout.regenerated, [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);
    }

    #[cfg(feature = "zstd")]
    #[test]
    fn the_dictionary_lane_assigns_indices_to_flagged_codecs() {
        let bundle = FrameInfo {
            bundle: Some(crate::dict::BundleId::from([1; 32])),
            ..frame(true)
        };
        let parse = |bytes: &[u8]| {
            let mut layout = Layout::default();
            layout.parse(bytes, &bundle).map(|()| layout)
        };
        // One zstd codec: flag 1, width 3 bits, index 5.
        let with = [
            0x02, 0x01, 0x00, 0x16, 0x00, 0x00, 0x00, 0x01, 0x03, 0x05, 0x00, 0x15,
        ];
        let layout = parse(&with).unwrap();
        assert_eq!(layout.nodes[0].dict, Some(5));
        assert!(layout.retained_bytes() > 0);
        // The flag clear: no index follows.
        let without = [0x02, 0x01, 0x00, 0x16, 0x00, 0x00, 0x00, 0x00, 0x00, 0x15];
        assert_eq!(parse(&without).unwrap().nodes[0].dict, None);
        // Widths above 16 bits are corrupt; a cut lane is truncation.
        let wide = [
            0x02, 0x01, 0x00, 0x16, 0x00, 0x00, 0x00, 0x01, 17, 0x05, 0x00, 0x15,
        ];
        assert_eq!(parse(&wide).unwrap_err(), bad_header());
        for len in 7..with.len() - 2 {
            assert_eq!(parse(&with[..len]).unwrap_err(), Parse::Truncated, "{len}");
        }
        // Without a bundle in the frame the same bytes have no lane.
        assert_eq!(
            parse_plain(
                &without[..7]
                    .iter()
                    .chain(&[0x00, 0x15])
                    .copied()
                    .collect::<alloc::vec::Vec<_>>()
            )
            .unwrap()
            .nodes[0]
                .dict,
            None
        );
    }

    #[cfg(feature = "zstd")]
    fn parse_plain(bytes: &[u8]) -> Result<Layout, Parse> {
        parse(bytes)
    }

    #[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, 0x27, 0x00, 0x00, 0x00, 0x00, 0x05]).unwrap_err(),
            Parse::Failed(DecodeError::unsupported(Unsupported::Codec { id: 39 }))
        );
        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()
        );
    }
}