jxl 0.4.3

// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

use std::{
    collections::{HashSet, VecDeque},
    io::IoSliceMut,
};

use sections::SectionState;

#[cfg(test)]
use crate::api::FrameCallback;
use crate::{
    api::{
        JxlBasicInfo, JxlBitstreamInput, JxlColorEncoding, JxlColorProfile, JxlDataFormat,
        JxlDecoderOptions, JxlOutputBuffer, JxlPixelFormat, VisibleFrameInfo,
        VisibleFrameSeekTarget,
        inner::{box_parser::BoxParser, process::SmallBuffer},
    },
    error::{Error, Result},
    frame::{DecoderState, Frame, Section},
    headers::{Animation, FileHeader, frame_header::FrameHeader, toc::IncrementalTocReader},
    icc::IncrementalIccReader,
};

mod non_section;
mod sections;

struct SectionBuffer {
    len: usize,
    data: Vec<u8>,
    section: Section,
}

#[derive(Clone, Copy)]
struct FrameStartInfo {
    file_offset: usize,
    remaining_in_box: u64,
    visible_count_before: usize,
}

pub(super) struct CodestreamParser {
    // TODO(veluca): this would probably be cleaner with some kind of state enum.
    pub(super) file_header: Option<FileHeader>,
    notified_image_info: bool,
    icc_parser: Option<IncrementalIccReader>,
    // These fields are populated once image information is available.
    decoder_state: Option<DecoderState>,
    pub(super) basic_info: Option<JxlBasicInfo>,
    pub(super) animation: Option<Animation>,
    pub(super) embedded_color_profile: Option<JxlColorProfile>,
    pub(super) output_color_profile: Option<JxlColorProfile>,
    pub(super) pixel_format: Option<JxlPixelFormat>,
    xyb_encoded: bool,
    is_gray: bool,
    pub(super) output_color_profile_set_by_user: bool,

    // These fields are populated when starting to decode a frame, and cleared once
    // the frame is done.
    frame_header: Option<FrameHeader>,
    toc_parser: Option<IncrementalTocReader>,
    pub(super) frame: Option<Frame>,

    // Buffers.
    non_section_buf: SmallBuffer,
    non_section_bit_offset: u8,
    sections: VecDeque<SectionBuffer>,
    ready_section_data: usize,
    skip_sections: bool,
    // True when we need to process frames without copying them to output buffers, e.g. reference frames
    process_without_output: bool,
    // True once the preview frame has been processed (if there is one)
    preview_done: bool,
    /// Number of visible frames still to skip before returning to the caller.
    /// Set via `start_new_frame` when seeking to a non-keyframe.
    visible_frames_to_skip: usize,

    section_state: SectionState,

    // Or only section if in single section special case.
    lf_global_section: Option<SectionBuffer>,
    lf_sections: Vec<SectionBuffer>,
    hf_global_section: Option<SectionBuffer>,
    // indexed by group, then by pass.
    hf_sections: Vec<Vec<Option<SectionBuffer>>>,
    // group indices that *might* have new renderable data.
    candidate_hf_sections: HashSet<usize>,

    pub(super) has_more_frames: bool,

    header_needed_bytes: Option<u64>,

    // --- Frame info tracking (for frame scanning) ---
    /// Collected visible frame info entries.
    pub(super) scanned_frames: Vec<VisibleFrameInfo>,
    /// Zero-based visible frame index counter.
    visible_frame_index: usize,
    /// File offsets and visibility info for every non-preview frame (visible
    /// and non-visible), in parse order.
    frame_starts: Vec<FrameStartInfo>,
    /// For each reference slot, earliest frame index required to reconstruct
    /// the current contents of that slot.
    reference_slot_decode_start: [Option<usize>; DecoderState::MAX_STORED_FRAMES],
    /// For each LF slot, earliest frame index required to reconstruct the
    /// current contents of that slot.
    lf_slot_decode_start: [Option<usize>; DecoderState::NUM_LF_FRAMES],
    /// File byte offset where the current frame header parse started.
    /// Set when we begin parsing a frame header.
    current_frame_file_offset: usize,
    /// Remaining codestream bytes in the current box at frame start.
    /// Captured alongside `current_frame_file_offset`.
    current_frame_remaining_in_box: u64,

    /// Remember whether we are decoding the file sequentially or we seeked.
    did_seek: bool,

    #[cfg(test)]
    pub frame_callback: Option<Box<FrameCallback>>,
    #[cfg(test)]
    pub decoded_frames: usize,
}

impl CodestreamParser {
    pub(super) fn new() -> Self {
        Self {
            file_header: None,
            notified_image_info: false,
            icc_parser: None,
            decoder_state: None,
            basic_info: None,
            animation: None,
            embedded_color_profile: None,
            output_color_profile: None,
            pixel_format: None,
            xyb_encoded: false,
            is_gray: false,
            output_color_profile_set_by_user: false,
            frame_header: None,
            toc_parser: None,
            frame: None,
            non_section_buf: SmallBuffer::new(4096),
            non_section_bit_offset: 0,
            sections: VecDeque::new(),
            ready_section_data: 0,
            skip_sections: false,
            process_without_output: false,
            preview_done: false,
            visible_frames_to_skip: 0,
            section_state: SectionState::new(0, 0),
            lf_global_section: None,
            lf_sections: vec![],
            hf_global_section: None,
            hf_sections: vec![],
            candidate_hf_sections: HashSet::new(),
            has_more_frames: true,
            header_needed_bytes: None,
            scanned_frames: Vec::new(),
            visible_frame_index: 0,
            frame_starts: Vec::new(),
            reference_slot_decode_start: [None; DecoderState::MAX_STORED_FRAMES],
            lf_slot_decode_start: [None; DecoderState::NUM_LF_FRAMES],
            current_frame_file_offset: 0,
            current_frame_remaining_in_box: u64::MAX,
            did_seek: false,
            #[cfg(test)]
            frame_callback: None,
            #[cfg(test)]
            decoded_frames: 0,
        }
    }

    fn has_visible_frame(&self) -> bool {
        if let Some(frame) = &self.frame {
            frame.header().is_visible()
        } else {
            false
        }
    }

    /// Record frame info for the just-parsed frame.
    /// Called after process_non_section() creates a Frame, for frame scanning.
    fn record_frame_info(&mut self) {
        let frame = match self.frame.as_ref() {
            Some(f) => f,
            None => return,
        };
        let header = frame.header();

        let current_frame_index = self.frame_starts.len();
        let is_visible = header.is_visible();
        self.frame_starts.push(FrameStartInfo {
            file_offset: self.current_frame_file_offset,
            remaining_in_box: self.current_frame_remaining_in_box,
            visible_count_before: self.visible_frame_index,
        });

        let mut decode_start_frame_index = current_frame_index;

        // Track frame dependencies through reference slots. For blending we know
        // exactly which slots are used. For patches we conservatively assume any
        // reference slot may be used.
        let mut used_reference_slots = [false; DecoderState::MAX_STORED_FRAMES];
        if header.needs_blending() {
            for blending_info in header
                .ec_blending_info
                .iter()
                .chain(std::iter::once(&header.blending_info))
            {
                let source = blending_info.source as usize;
                assert!(
                    source < DecoderState::MAX_STORED_FRAMES,
                    "invalid blending source slot {source}, max {}",
                    DecoderState::MAX_STORED_FRAMES - 1
                );
                used_reference_slots[source] = true;
            }
        }
        if header.has_patches() {
            used_reference_slots.fill(true);
        }

        for (slot, used) in used_reference_slots.iter().enumerate() {
            if *used && let Some(dep_start) = self.reference_slot_decode_start[slot] {
                decode_start_frame_index = decode_start_frame_index.min(dep_start);
            }
        }

        if header.has_lf_frame() {
            let lf_slot = header.lf_level as usize;
            assert!(
                lf_slot < DecoderState::NUM_LF_FRAMES,
                "invalid lf slot {lf_slot}, max {}",
                DecoderState::NUM_LF_FRAMES - 1
            );
            if let Some(dep_start) = self.lf_slot_decode_start[lf_slot] {
                decode_start_frame_index = decode_start_frame_index.min(dep_start);
            }
        }

        if is_visible {
            let duration_ticks = header.duration;
            let duration_ms = if let Some(ref anim) = self.animation {
                if anim.tps_numerator > 0 {
                    (duration_ticks as f64) * 1000.0 * (anim.tps_denominator as f64)
                        / (anim.tps_numerator as f64)
                } else {
                    0.0
                }
            } else {
                0.0
            };

            let decode_start = self.frame_starts[decode_start_frame_index];
            let seek_target = VisibleFrameSeekTarget {
                decode_start_file_offset: decode_start.file_offset,
                remaining_in_box: decode_start.remaining_in_box,
                visible_frames_to_skip: self
                    .visible_frame_index
                    .saturating_sub(decode_start.visible_count_before),
            };
            let is_keyframe = seek_target.visible_frames_to_skip == 0;

            self.scanned_frames.push(VisibleFrameInfo {
                index: self.visible_frame_index,
                duration_ms,
                duration_ticks,
                file_offset: self.current_frame_file_offset,
                is_last: header.is_last,
                is_keyframe,
                seek_target,
                name: header.name.clone(),
            });

            self.visible_frame_index += 1;
        }

        // Update slot dependency origins after processing this frame.
        if header.can_be_referenced {
            let slot = header.save_as_reference as usize;
            assert!(
                slot < DecoderState::MAX_STORED_FRAMES,
                "invalid save_as_reference slot {slot}, max {}",
                DecoderState::MAX_STORED_FRAMES - 1
            );
            self.reference_slot_decode_start[slot] = Some(decode_start_frame_index);
        }

        if header.lf_level != 0 {
            let slot = (header.lf_level - 1) as usize;
            assert!(
                slot < DecoderState::NUM_LF_FRAMES,
                "invalid lf save slot {slot}, max {}",
                DecoderState::NUM_LF_FRAMES - 1
            );
            self.lf_slot_decode_start[slot] = Some(decode_start_frame_index);
        }
    }

    /// Returns the number of passes that are fully completed across all groups.
    pub(super) fn num_completed_passes(&self) -> usize {
        self.section_state.num_completed_passes()
    }

    #[cfg(test)]
    pub(crate) fn set_use_simple_pipeline(&mut self, u: bool) {
        self.decoder_state
            .as_mut()
            .unwrap()
            .set_use_simple_pipeline(u);
    }

    /// Rewinds for animation loop replay, keeping pixel_format setting.
    pub(super) fn rewind(&mut self) -> Option<JxlPixelFormat> {
        let pixel_format = self.pixel_format.take();
        *self = Self::new();
        self.pixel_format = pixel_format.clone();
        pixel_format
    }

    /// Resets frame-level state for seeking to a new frame.
    ///
    /// Preserves: file_header, decoder_state (including reference frames),
    /// basic_info, animation, color profiles, pixel_format, xyb_encoded,
    /// is_gray, output_color_profile_set_by_user, preview_done.
    ///
    /// Clears: frame_header, toc_parser, frame, all section buffers,
    /// non_section_buf, and processing flags.
    pub(super) fn start_new_frame(&mut self, visible_frames_to_skip: usize) {
        self.frame_header = None;
        self.toc_parser = None;
        self.frame = None;
        self.non_section_buf = SmallBuffer::new(4096);
        self.non_section_bit_offset = 0;
        self.sections.clear();
        self.ready_section_data = 0;
        self.skip_sections = false;
        self.process_without_output = false;
        self.visible_frames_to_skip = visible_frames_to_skip;
        self.section_state = SectionState::new(0, 0);
        self.lf_global_section = None;
        self.lf_sections.clear();
        self.hf_global_section = None;
        self.hf_sections.clear();
        self.candidate_hf_sections.clear();
        self.has_more_frames = true;
        self.header_needed_bytes = None;
        self.did_seek = true;
    }

    pub(super) fn process(
        &mut self,
        box_parser: &mut BoxParser,
        input: &mut dyn JxlBitstreamInput,
        decode_options: &JxlDecoderOptions,
        mut output_buffers: Option<&mut [JxlOutputBuffer]>,
        do_flush: bool,
    ) -> Result<()> {
        if let Some(output_buffers) = &output_buffers {
            let px = self.pixel_format.as_ref().unwrap();
            let expected_len = std::iter::once(&px.color_data_format)
                .chain(px.extra_channel_format.iter())
                .filter(|x| x.is_some())
                .count();
            if output_buffers.len() != expected_len {
                return Err(Error::WrongBufferCount(output_buffers.len(), expected_len));
            }
        }
        // If we have sections to read, read into sections; otherwise, read into the local buffer.
        loop {
            if !self.sections.is_empty() {
                let regular_frame = self.has_visible_frame();
                // Only skip sections if we don't need the frame data. Frames that can be
                // referenced must be decoded because they serve as sources for patches,
                // blending, or frame extension in subsequent frames.
                let can_be_referenced = self
                    .frame
                    .as_ref()
                    .is_some_and(|f| f.header().can_be_referenced);
                if decode_options.scan_frames_only
                    || (!self.process_without_output
                        && output_buffers.is_none()
                        && !can_be_referenced)
                {
                    self.skip_sections = true;
                }

                if !self.skip_sections {
                    // This is just an estimate as there could be box bytes in the middle.
                    let mut readable_section_data = (self.non_section_buf.len()
                        + input.available_bytes()?
                        + self.ready_section_data)
                        .max(1);
                    // Ensure enough section buffers are available for reading available data.
                    for buf in self.sections.iter_mut() {
                        if buf.data.is_empty() {
                            buf.data.resize(buf.len, 0);
                        }
                        readable_section_data =
                            readable_section_data.saturating_sub(buf.data.len());
                        if readable_section_data == 0 {
                            break;
                        }
                    }
                    // Read sections up to the end of the current box.
                    let mut available_codestream = match box_parser.get_more_codestream(input) {
                        Err(Error::OutOfBounds(_)) => 0,
                        Ok(c) => c as usize,
                        Err(e) => return Err(e),
                    };
                    let mut section_buffers = vec![];
                    let mut ready = self.ready_section_data;
                    for buf in self.sections.iter_mut() {
                        if buf.data.is_empty() {
                            break;
                        }
                        let len = buf.data.len();
                        if len > ready {
                            let readable = (available_codestream + ready).min(len);
                            section_buffers.push(IoSliceMut::new(&mut buf.data[ready..readable]));
                            available_codestream =
                                available_codestream.saturating_sub(readable - ready);
                            if available_codestream == 0 {
                                break;
                            }
                        }
                        ready = ready.saturating_sub(len);
                    }
                    let mut buffers = &mut section_buffers[..];
                    loop {
                        let num = if !box_parser.box_buffer.is_empty() {
                            box_parser.box_buffer.take(buffers)
                        } else {
                            let num = input.read(buffers)?;
                            box_parser.mark_file_consumed(num);
                            num
                        };
                        self.ready_section_data += num;
                        box_parser.consume_codestream(num as u64);
                        IoSliceMut::advance_slices(&mut buffers, num);
                        if num == 0 || buffers.is_empty() {
                            break;
                        }
                    }
                    match self.process_sections(decode_options, &mut output_buffers, do_flush) {
                        Ok(None) => Ok(()),
                        Ok(Some(missing)) => Err(Error::OutOfBounds(missing)),
                        Err(Error::OutOfBounds(_)) => Err(Error::SectionTooShort),
                        Err(err) => Err(err),
                    }?;
                } else {
                    let total_size = self.sections.iter().map(|x| x.len).sum::<usize>();
                    loop {
                        let to_skip = total_size - self.ready_section_data;
                        if to_skip == 0 {
                            break;
                        }
                        let available_codestream = box_parser.get_more_codestream(input)? as usize;
                        let to_skip = to_skip.min(available_codestream);
                        let skipped = if !box_parser.box_buffer.is_empty() {
                            box_parser.box_buffer.consume(to_skip)
                        } else {
                            let skipped = input.skip(to_skip)?;
                            box_parser.mark_file_consumed(skipped);
                            skipped
                        };
                        box_parser.consume_codestream(skipped as u64);
                        self.ready_section_data += skipped;
                        if skipped == 0 {
                            break;
                        }
                    }
                    if self.ready_section_data < total_size {
                        return Err(Error::OutOfBounds(total_size - self.ready_section_data));
                    } else {
                        self.sections.clear();
                        // Finalize the skipped frame, mirroring what process_sections does
                        let frame = self
                            .frame
                            .take()
                            .expect("frame must be set when skip_sections is true");
                        if let Some(decoder_state) = frame.finalize()? {
                            self.decoder_state = Some(decoder_state);
                        } else {
                            self.has_more_frames = false;
                            // Return immediately so we don't re-enter the outer loop and hit the
                            // API-misuse assertion below inside the same call.
                            return Ok(());
                        }
                        self.skip_sections = false;
                    }
                }
                if self.sections.is_empty() {
                    // Go back to parsing a new frame header, if any.
                    // Only return if this was a regular visible frame that was actually decoded
                    // (not a frame we were skipping like a preview frame)
                    let was_skipping = self.process_without_output;
                    self.process_without_output = false;
                    if regular_frame && !was_skipping {
                        return Ok(());
                    }
                    continue;
                }
            } else {
                // Trying to read a frame or a file header.
                assert!(self.frame.is_none());
                if !self.has_more_frames {
                    // If this is a flush request and the file is complete, we are done.
                    // Otherwise, this is an API usage error.
                    assert!(do_flush);
                    return Ok(());
                }

                // Capture frame-start metadata once before parsing the next
                // frame header. We do this after `get_more_codestream()` so we
                // are robust to the previous frame ending exactly at a box
                // boundary (BoxNeeded -> CodestreamBox transition).
                let mut capture_frame_start =
                    self.decoder_state.is_some() && self.frame_header.is_none();

                // Loop to handle incremental parsing (e.g. large ICC profiles) that may need
                // multiple buffer refills to complete.
                loop {
                    let available_codestream = match box_parser.get_more_codestream(input) {
                        Err(Error::OutOfBounds(_)) => 0,
                        Ok(c) => c,
                        Err(e) => return Err(e),
                    };

                    if capture_frame_start {
                        // total_file_consumed counts bytes read/skipped from
                        // raw input. non_section_buf and box_buffer contain
                        // unread bytes already accounted there.
                        self.current_frame_file_offset = (box_parser.total_file_consumed as usize)
                            .saturating_sub(self.non_section_buf.len())
                            .saturating_sub(box_parser.box_buffer.len());

                        // `available_codestream` includes bytes still in
                        // box_buffer and not yet in non_section_buf.
                        self.current_frame_remaining_in_box = if available_codestream > u64::MAX / 2
                        {
                            u64::MAX
                        } else {
                            available_codestream.saturating_add(self.non_section_buf.len() as u64)
                        };
                        capture_frame_start = false;
                    }

                    let c = self.non_section_buf.refill(
                        |buf| {
                            if !box_parser.box_buffer.is_empty() {
                                Ok(box_parser.box_buffer.take(buf))
                            } else {
                                let read = input.read(buf)?;
                                box_parser.mark_file_consumed(read);
                                Ok(read)
                            }
                        },
                        Some(available_codestream as usize),
                    )? as u64;
                    box_parser.consume_codestream(c);

                    // If we know that non-section parsing will require more bytes than what
                    // we added to the codestream, don't even try to parse non-section data.
                    if let Some(needed) = self.header_needed_bytes.as_mut() {
                        *needed = needed.saturating_sub(c);
                        if *needed > 0 {
                            if !self.non_section_buf.can_read_more() {
                                self.non_section_buf.enlarge();
                            }
                            // Check if input still has data - if so, refill and retry
                            if input.available_bytes().unwrap_or(0) > 0 {
                                continue;
                            } else {
                                return Err(Error::OutOfBounds(*needed as usize));
                            }
                        }
                    }

                    let range = self.non_section_buf.range();

                    match self.process_non_section(decode_options) {
                        Ok(()) => {
                            self.header_needed_bytes = None;
                            break;
                        }
                        Err(Error::OutOfBounds(n)) => {
                            let new_range = self.non_section_buf.range();
                            // If non-section parsing consumed no bytes, and the non-section buffer
                            // cannot accept more bytes, enlarge the buffer to allow to make progress.
                            if new_range == range && !self.non_section_buf.can_read_more() {
                                self.non_section_buf.enlarge();
                            }
                            self.header_needed_bytes = Some(n as u64);
                            // Check if input still has data - if so, refill and retry
                            if input.available_bytes().unwrap_or(0) > 0 {
                                continue;
                            } else {
                                return Err(Error::OutOfBounds(n));
                            }
                        }
                        Err(e) => return Err(e),
                    }
                }

                if (!self.notified_image_info && self.decoder_state.is_some())
                    && self.frame_header.is_none()
                {
                    self.notified_image_info = true;
                    // Return to caller if we found image info.
                    return Ok(());
                }
                if self.frame.is_some() {
                    // Check if this is a preview frame that should be skipped
                    let is_preview_frame = !self.preview_done
                        && self
                            .basic_info
                            .as_ref()
                            .is_some_and(|info| info.preview_size.is_some());
                    if is_preview_frame {
                        self.preview_done = true;
                        if decode_options.skip_preview {
                            self.process_without_output = true;
                            continue;
                        }
                    }

                    // Record frame info for scanning (after preview check).
                    if !is_preview_frame && !self.did_seek {
                        self.record_frame_info();
                    }

                    if self.has_visible_frame() {
                        if self.visible_frames_to_skip > 0 {
                            // Skip this visible frame without returning to the
                            // caller; decrement the counter and continue
                            // processing internally.
                            self.visible_frames_to_skip -= 1;
                            self.process_without_output = true;
                            continue;
                        }
                        // Return to caller if we found visible frame info.
                        return Ok(());
                    } else {
                        self.process_without_output = true;
                        continue;
                    }
                }
            }
        }
    }

    pub(super) fn update_default_output_color_profile(&mut self) {
        // Only set default output_color_profile if not already configured by user
        if self.output_color_profile_set_by_user {
            return;
        }

        let embedded_color_profile = self.embedded_color_profile.as_ref().unwrap();
        let pixel_format = self.pixel_format.as_ref().unwrap();

        // Determine default output color profile following libjxl logic:
        // - For XYB: use embedded if can_output_to(), else:
        //   - if float samples are requested: linear sRGB,
        //   - else: sRGB
        // - For non-XYB: use embedded color profile
        let output_color_profile = if self.xyb_encoded {
            // Use embedded if we can output to it, otherwise fall back to sRGB
            let base_encoding = if embedded_color_profile.can_output_to() {
                match &embedded_color_profile {
                    JxlColorProfile::Simple(enc) => enc.clone(),
                    JxlColorProfile::Icc(_) => {
                        unreachable!("can_output_to returns false for ICC")
                    }
                }
            } else {
                let data_format = pixel_format
                    .color_data_format
                    .unwrap_or(JxlDataFormat::U8 { bit_depth: 8 });
                let is_float = matches!(
                    data_format,
                    JxlDataFormat::F32 { .. } | JxlDataFormat::F16 { .. }
                );
                if is_float {
                    JxlColorEncoding::linear_srgb(self.is_gray)
                } else {
                    JxlColorEncoding::srgb(self.is_gray)
                }
            };

            JxlColorProfile::Simple(base_encoding)
        } else {
            embedded_color_profile.clone()
        };
        self.output_color_profile = Some(output_color_profile);
    }
}