j2k 0.7.2

// SPDX-License-Identifier: MIT OR Apache-2.0

mod allocation;
mod boxes;
mod codestream;

use self::boxes::{extract_jp2_codestream_payload, parse_jp2};
use self::codestream::{parse_codestream, CodestreamInfo};
use crate::{J2kComponentInfo, J2kError, J2kFileMetadata, J2kSupportInfo};
use j2k_core::{
    Colorspace, CompressedPayloadKind, CompressedTransferSyntax, Info, TileLayout, Unsupported,
};

pub(crate) fn parse_info(input: &[u8]) -> Result<Info, J2kError> {
    Ok(parse_image_info(input)?.info)
}

pub(crate) fn parse_image_info(input: &[u8]) -> Result<ParsedImageInfo, J2kError> {
    if boxes::looks_like_jp2(input) {
        return parse_jp2(input);
    }
    if codestream::looks_like_codestream(input) {
        let parsed = parse_codestream(input, 0)?;
        let (info, transfer_syntax, components) = parsed.into_parts(None);
        return Ok(ParsedImageInfo {
            info,
            transfer_syntax,
            payload_kind: CompressedPayloadKind::Jpeg2000Codestream,
            components,
            file_metadata: None,
        });
    }
    Err(J2kError::Unsupported(Unsupported {
        what: "input is not a JP2 container or raw JPEG 2000 codestream",
    }))
}

/// Borrowed raw codestream slice extracted from raw JPEG 2000 bytes or a JP2/JPH wrapper.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct J2kCodestreamPayload<'a> {
    codestream: &'a [u8],
    payload_kind: CompressedPayloadKind,
    codestream_offset: usize,
}

impl<'a> J2kCodestreamPayload<'a> {
    /// Raw JPEG 2000 codestream bytes.
    #[must_use]
    pub const fn codestream(self) -> &'a [u8] {
        self.codestream
    }

    /// Encapsulation shape the codestream was extracted from.
    #[must_use]
    pub const fn payload_kind(self) -> CompressedPayloadKind {
        self.payload_kind
    }

    /// Byte offset where the codestream payload starts in the original input.
    #[must_use]
    pub const fn codestream_offset(self) -> usize {
        self.codestream_offset
    }
}

/// Return the raw JPEG 2000 codestream payload from raw codestream, JP2, or JPH input.
///
/// This helper validates only the wrapper framing needed to locate the borrowed
/// codestream slice. Full JP2 metadata and ordering validation remains part of
/// the decoder inspect path.
pub fn extract_j2k_codestream_payload(input: &[u8]) -> Result<J2kCodestreamPayload<'_>, J2kError> {
    if codestream::looks_like_codestream(input) {
        return Ok(J2kCodestreamPayload {
            codestream: input,
            payload_kind: CompressedPayloadKind::Jpeg2000Codestream,
            codestream_offset: 0,
        });
    }
    if boxes::looks_like_jp2(input) {
        let (payload_kind, codestream_offset, codestream) = extract_jp2_codestream_payload(input)?;
        return Ok(J2kCodestreamPayload {
            codestream,
            payload_kind,
            codestream_offset,
        });
    }
    Err(J2kError::Unsupported(Unsupported {
        what: "input is not a JP2 container or raw JPEG 2000 codestream",
    }))
}

#[derive(Debug, PartialEq, Eq)]
pub(crate) struct ParsedImageInfo {
    pub(crate) info: Info,
    pub(crate) transfer_syntax: CompressedTransferSyntax,
    pub(crate) payload_kind: CompressedPayloadKind,
    pub(crate) components: Vec<J2kComponentInfo>,
    pub(crate) file_metadata: Option<J2kFileMetadata>,
}

impl ParsedImageInfo {
    pub(crate) fn into_support_info(self) -> J2kSupportInfo {
        J2kSupportInfo {
            info: self.info,
            transfer_syntax: self.transfer_syntax,
            payload_kind: self.payload_kind,
            components: self.components,
            file_metadata: self.file_metadata,
        }
    }

    pub(crate) fn allocated_bytes(&self) -> Result<usize, J2kError> {
        let mut bytes = allocation::capacity_bytes::<J2kComponentInfo>(
            self.components.capacity(),
            "codestream component metadata",
        )?;
        if let Some(metadata) = &self.file_metadata {
            allocation::checked_add_bytes(
                &mut bytes,
                boxes::metadata_allocated_bytes(metadata)?,
                "JPEG 2000 inspection metadata",
            )?;
        }
        Ok(bytes)
    }
}

fn infer_colorspace(components: u16, has_mct: bool, reversible: bool) -> Colorspace {
    match (components, has_mct, reversible) {
        (1, _, _) => Colorspace::SGray,
        (3, false, _) => Colorspace::Rgb,
        (3, true, false) => Colorspace::Ict,
        (3, true, true) => Colorspace::Rct,
        _ => Colorspace::IccTagged,
    }
}

#[derive(Debug)]
struct ParsedSiz {
    dimensions: (u32, u32),
    components: u16,
    bit_depth: u8,
    tile_layout: TileLayout,
    component_info: Vec<J2kComponentInfo>,
}

#[derive(Debug, Clone, Copy)]
struct ParsedCod {
    resolution_levels: u8,
    has_mct: bool,
    reversible: bool,
    high_throughput: bool,
}

impl CodestreamInfo {
    fn into_parts(
        self,
        colorspace: Option<Colorspace>,
    ) -> (Info, CompressedTransferSyntax, Vec<J2kComponentInfo>) {
        let Self { siz, cod } = self;
        let ParsedSiz {
            dimensions,
            components,
            bit_depth,
            tile_layout,
            component_info,
        } = siz;
        let info = Info {
            dimensions,
            components,
            colorspace: colorspace
                .unwrap_or_else(|| infer_colorspace(components, cod.has_mct, cod.reversible)),
            bit_depth,
            tile_layout: Some(tile_layout),
            coded_unit_layout: None,
            restart_interval: None,
            resolution_levels: cod.resolution_levels,
        };
        let transfer_syntax = match (cod.high_throughput, cod.reversible) {
            (false, true) => CompressedTransferSyntax::Jpeg2000Lossless,
            (false, false) => CompressedTransferSyntax::Jpeg2000Lossy,
            (true, true) => CompressedTransferSyntax::HtJpeg2000Lossless,
            (true, false) => CompressedTransferSyntax::HtJpeg2000Lossy,
        };
        (info, transfer_syntax, component_info)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{J2kColorSpec, J2kPaletteColumn, J2kPaletteMetadata};
    use core::mem::size_of;

    #[test]
    fn parsed_metadata_retained_bytes_count_nested_actual_capacities_exactly() {
        let mut components = Vec::new();
        components
            .try_reserve_exact(2)
            .expect("test component metadata");
        components.push(J2kComponentInfo {
            bit_depth: 8,
            signed: false,
            x_rsiz: 1,
            y_rsiz: 1,
        });

        let mut profile = Vec::new();
        profile.try_reserve_exact(5).expect("test ICC owner");
        profile.push(1);
        let mut color_specs = Vec::new();
        color_specs
            .try_reserve_exact(2)
            .expect("test COLR metadata");
        color_specs.push(J2kColorSpec::IccProfile { profile });

        let mut columns = Vec::new();
        columns.try_reserve_exact(2).expect("test palette columns");
        columns.push(J2kPaletteColumn {
            bit_depth: 8,
            signed: false,
        });
        let mut row = Vec::new();
        row.try_reserve_exact(3).expect("test palette row");
        row.push(7);
        let mut entries = Vec::new();
        entries.try_reserve_exact(2).expect("test palette rows");
        entries.push(row);

        let parsed = ParsedImageInfo {
            info: Info {
                dimensions: (1, 1),
                components: 1,
                colorspace: Colorspace::IccTagged,
                bit_depth: 8,
                tile_layout: None,
                coded_unit_layout: None,
                restart_interval: None,
                resolution_levels: 1,
            },
            transfer_syntax: CompressedTransferSyntax::Jpeg2000Lossless,
            payload_kind: CompressedPayloadKind::Jp2File,
            components,
            file_metadata: Some(J2kFileMetadata {
                bits_per_component: Vec::new(),
                color_specs,
                palette: Some(J2kPaletteMetadata { columns, entries }),
                component_mappings: Vec::new(),
                channel_definitions: Vec::new(),
                has_palette: true,
                has_component_mapping: false,
                has_channel_definition: false,
            }),
        };

        let metadata = parsed.file_metadata.as_ref().expect("metadata");
        let palette = metadata.palette.as_ref().expect("palette");
        let profile_capacity = match &metadata.color_specs[0] {
            J2kColorSpec::IccProfile { profile } => profile.capacity(),
            _ => 0,
        };
        let expected = parsed.components.capacity() * size_of::<J2kComponentInfo>()
            + metadata.color_specs.capacity() * size_of::<J2kColorSpec>()
            + profile_capacity
            + palette.columns.capacity() * size_of::<J2kPaletteColumn>()
            + palette.entries.capacity() * size_of::<Vec<u64>>()
            + palette.entries[0].capacity() * size_of::<u64>();
        assert_eq!(parsed.allocated_bytes().expect("retained bytes"), expected);
    }
}