Source code

Revision control

Copy as Markdown

Other Tools

// 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 crate::api::box_parser::CodestreamInput;
use crate::api::codestream_parser::frame_info::FrameInfo;
use crate::api::codestream_parser::frame_scan_info::FrameScanInfo;
pub(crate) use crate::api::codestream_parser::frame_scan_info::FrameStartInfo;
use crate::api::codestream_parser::image_info::ImageInfo;
use crate::api::{
BoxParserCheckpoint, Event, JxlColorProfile, JxlDecoderOptions, JxlOutputBuffer,
JxlParallelRunner, JxlPixelFormat, ProfileLevel, SmallBuffer, VisibleFrameSeekTarget,
};
use crate::error::{Error, Result};
use crate::headers::frame_header::FrameHeader;
mod frame_info;
mod frame_scan_info;
mod image_info;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProcessMode {
// Frame is processed regularly.
Process,
// Frame is an internal frame that is not shown to the user.
SkipOutput,
// Frame is skipped entirely; whether the user is notified
// depends on the argument.
Skip(bool),
}
impl ProcessMode {
fn notify_user(&self) -> bool {
*self == ProcessMode::Process || *self == ProcessMode::Skip(true)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParserState {
FileHeader,
ColorEncoding,
FrameHeader {
is_preview: bool,
},
Toc {
is_preview: bool,
},
Sections {
is_preview: bool,
process_mode: ProcessMode,
},
TrailingData,
Done,
}
fn level5_limits(input: &CodestreamInput, decode_options: &JxlDecoderOptions) -> bool {
decode_options.max_profile_level == ProfileLevel::Main5
|| input.box_parser().container_level() != Some(ProfileLevel::Main10)
}
fn check_size_limit(
pixel_limit: Option<usize>,
(xs, ys): (usize, usize),
num_ec: usize,
) -> Result<()> {
if let Some(limit) = pixel_limit {
let xs = xs.max(16); // xsize is always at least 64 bytes.
let total_pixels = xs.saturating_mul(ys).saturating_mul(3 + num_ec);
if total_pixels >= limit {
return Err(Error::ImageSizeTooLarge(xs, ys));
}
};
Ok(())
}
fn validate_output_buffers(
output_buffers: &[JxlOutputBuffer],
pixel_format: Option<&JxlPixelFormat>,
) -> Result<()> {
let Some(px) = pixel_format else {
return Err(Error::ApiUsageError(
"cannot provide output buffers before BasicInfo",
));
};
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));
}
Ok(())
}
pub(super) struct CodestreamParser {
state: ParserState,
// Image/frame information
pub image_info: ImageInfo,
pub frame_info: FrameInfo,
frame_scan_info: FrameScanInfo,
/// Total length of the file.
pub file_length: Option<u64>,
// Temporary buffer
local_buffer: SmallBuffer,
header_needed_bytes: Option<usize>,
// Output information
pub output_color_profile: Option<JxlColorProfile>,
pub pixel_format: Option<JxlPixelFormat>,
/// While seeking: the number of visible frames still to skip before the target.
visible_frames_to_skip: Option<usize>,
/// Index (in file order, previews not counted) of the next frame.
next_frame_index: usize,
/// For each stored frame of the decoder state (see `stored_slot`), the frame whose result it
/// is (as in a sequential decode, unless a seek has not reached its target yet).
slot_frames: [Option<usize>; 8],
/// Number of frames decoded (not skipped), for tests.
#[cfg(test)]
pub(super) frames_decoded: usize,
}
/// Where a frame is stored: reference frame 0-3, or LF frame 0-3 as 4-7.
pub(super) fn stored_slot(header: &FrameHeader) -> Option<usize> {
if header.can_be_referenced {
Some(header.save_as_reference as usize)
} else {
(header.lf_level != 0).then(|| 3 + header.lf_level as usize)
}
}
impl CodestreamParser {
pub(super) fn new() -> Self {
Self {
image_info: ImageInfo::new(),
state: ParserState::FileHeader,
output_color_profile: None,
pixel_format: None,
local_buffer: SmallBuffer::new(4096),
header_needed_bytes: None,
frame_info: FrameInfo::new(),
visible_frames_to_skip: None,
next_frame_index: 0,
slot_frames: [None; 8],
#[cfg(test)]
frames_decoded: 0,
frame_scan_info: FrameScanInfo::new(),
file_length: None,
}
}
pub(super) fn reset(&mut self) {
self.state = ParserState::FileHeader;
self.image_info = ImageInfo::new();
self.frame_info.reset((0, 0), [false; 8]);
self.frame_scan_info = FrameScanInfo::new();
self.file_length = None;
self.local_buffer.reset();
self.header_needed_bytes = None;
self.output_color_profile = None;
self.pixel_format = None;
self.visible_frames_to_skip = None;
self.next_frame_index = 0;
self.slot_frames = [None; 8];
#[cfg(test)]
{
self.frames_decoded = 0;
}
}
pub(super) fn start_new_frame(
&mut self,
seek_target: &VisibleFrameSeekTarget,
) -> BoxParserCheckpoint {
// On the way to a seek target, frames that do not matter for it can be decoded wrong,
// so stored frames are only usable when not in the middle of a frame or a seek.
let mid_frame = matches!(self.state, ParserState::Sections { .. });
let usable = !mid_frame && self.visible_frames_to_skip.is_none();
let have = if usable { self.slot_frames } else { [None; 8] };
// Keep the stored frames that the decoder already has, unless a frame decoded from `start`
// on would overwrite them, and start where the others can be reconstructed.
let kept = |s: usize, start: usize| {
seek_target.slots[s].is_some_and(|(f, _)| have[s] == Some(f) && f < start)
};
let mut start = seek_target.target;
while let Some(earlier) = (0..8)
.filter(|&s| !kept(s, start.frame_index))
.filter_map(|s| Some(seek_target.slots[s]?.1))
.filter(|s| s.frame_index < start.frame_index)
.min_by_key(|s| s.frame_index)
{
start = earlier;
}
let keep_slots = std::array::from_fn(|s| kept(s, start.frame_index));
self.slot_frames = std::array::from_fn(|s| have[s].filter(|_| keep_slots[s]));
self.frame_info.reset(start.frame_counters, keep_slots);
self.local_buffer.reset();
self.local_buffer
.mark_consumed(start.box_parser_checkpoint.consumed_codestream);
self.visible_frames_to_skip = (start.frame_index < seek_target.target.frame_index)
.then_some(seek_target.target.frame_counters.0 - start.frame_counters.0);
self.next_frame_index = start.frame_index;
self.state = ParserState::FrameHeader { is_preview: false };
self.header_needed_bytes = None;
start.box_parser_checkpoint
}
fn has_more_frames(&self) -> bool {
!matches!(self.state, ParserState::TrailingData | ParserState::Done)
}
fn refill_and_parse<T>(
&mut self,
input: &mut CodestreamInput,
mut do_parse: impl FnMut(&mut Self, &mut CodestreamInput) -> Result<T>,
) -> Result<T> {
loop {
let c = self.local_buffer.refill(|buf| input.read(buf))?;
if let Some(mut need) = self.header_needed_bytes.take() {
need = need.saturating_sub(c);
if need > 0 {
self.header_needed_bytes = Some(need);
let enlarge = !self.local_buffer.can_read_more();
if enlarge {
self.local_buffer.enlarge();
}
if c > 0 || enlarge {
continue;
}
return Err(Error::OutOfBounds(need));
}
}
let range = self.local_buffer.range();
let res = do_parse(self, input);
match res {
Ok(val) => {
self.header_needed_bytes = None;
return Ok(val);
}
Err(Error::OutOfBounds(n)) => {
let new_range = self.local_buffer.range();
let made_progress = new_range != range;
let enlarge = !made_progress && !self.local_buffer.can_read_more();
if enlarge {
self.local_buffer.enlarge();
}
self.header_needed_bytes = Some(n);
if c > 0 || made_progress || enlarge {
continue;
}
return Err(Error::OutOfBounds(n));
}
Err(e) => return Err(e),
}
}
}
pub(super) fn process(
&mut self,
input: &mut CodestreamInput,
decode_options: &JxlDecoderOptions,
output_buffers: Option<&mut [JxlOutputBuffer]>,
parallel_runner: &mut dyn JxlParallelRunner,
) -> Result<Event> {
match self.process_inner(input, decode_options, output_buffers, parallel_runner) {
Err(Error::OutOfBounds(_))
if self.state != ParserState::TrailingData
&& input.box_parser().is_codestream_complete() =>
{
Err(Error::UnexpectedCodestreamBoxEnd)
}
Err(Error::OutOfBounds(size_hint)) => Ok(Event::NeedMoreInput { size_hint }),
result => result,
}
}
fn process_inner(
&mut self,
input: &mut CodestreamInput,
decode_options: &JxlDecoderOptions,
mut output_buffers: Option<&mut [JxlOutputBuffer]>,
parallel_runner: &mut dyn JxlParallelRunner,
) -> Result<Event> {
if let Some(output_buffers) = &output_buffers {
validate_output_buffers(output_buffers, self.pixel_format.as_ref())?;
}
loop {
match self.state {
ParserState::FileHeader => {
self.refill_and_parse(input, |c, input| {
c.local_buffer.with_br(|br, bits| {
let level5_limits = level5_limits(input, decode_options);
c.image_info
.parse_file_header(decode_options, br, bits, level5_limits)
})
})?;
self.state = ParserState::ColorEncoding;
}
ParserState::ColorEncoding => {
self.refill_and_parse(input, |c, _| {
c.local_buffer
.with_br(|br, bits| c.image_info.parse_color_encoding(br, bits))
})?;
self.update_default_output_options();
self.state = ParserState::FrameHeader {
is_preview: self.image_info.has_preview(),
};
return Ok(Event::BasicInfo);
}
ParserState::FrameHeader { is_preview } => {
self.refill_and_parse(input, |c, input| {
c.frame_scan_info.set_current_frame_checkpoint(
input
.box_parser()
.state_checkpoint(c.local_buffer.consumed())?,
);
let level5_limits = level5_limits(input, decode_options);
c.local_buffer.with_br(|br, bits| {
c.frame_info.parse_frame_header(
is_preview,
decode_options,
c.image_info.file_header(),
br,
bits,
level5_limits,
)
})
})?;
self.state = ParserState::Toc { is_preview };
}
ParserState::Toc { is_preview } => {
let toc = self.refill_and_parse(input, |c, _| {
c.local_buffer
.with_br(|br, bits| c.frame_info.parse_toc(br, bits))
})?;
// Decide how much to decode this frame.
let mut process_mode = ProcessMode::Process;
if is_preview && decode_options.skip_preview {
process_mode = ProcessMode::Skip(false);
}
let header = self.frame_info.current_frame_header().unwrap();
let is_stored = stored_slot(header).is_some();
if !header.is_visible() {
process_mode = ProcessMode::SkipOutput;
} else if let Some(n) = self.visible_frames_to_skip.as_mut() {
if *n > 0 {
*n -= 1;
// A frame that is not stored does not affect later frames.
process_mode = if is_stored {
ProcessMode::SkipOutput
} else {
ProcessMode::Skip(false)
};
} else {
// The seek target: what is stored is now as in a sequential decode.
self.visible_frames_to_skip = None;
}
}
if !is_preview {
self.next_frame_index += 1;
}
if decode_options.scan_frames_only && process_mode == ProcessMode::Process {
process_mode = ProcessMode::Skip(true);
}
#[cfg(test)]
if !matches!(process_mode, ProcessMode::Skip(_)) {
self.frames_decoded += 1;
}
self.state = ParserState::Sections {
is_preview,
process_mode,
};
// Record frame for seeking purposes
if !is_preview {
self.frame_scan_info.record(
self.frame_info.current_frame_header().unwrap(),
&self.image_info.file_header().image_metadata.animation,
);
}
let level5_limits = level5_limits(input, decode_options);
self.frame_info.make_frame(
&mut self.local_buffer,
toc,
self.image_info.file_header(),
decode_options,
self.pixel_format.as_ref().unwrap(),
self.output_color_profile.as_ref().unwrap(),
process_mode,
level5_limits,
)?;
if !matches!(process_mode, ProcessMode::Skip(..)) {
self.frame_info.dequeue_ready_sections();
}
if process_mode.notify_user() {
return Ok(Event::FrameHeader);
}
}
ParserState::Sections {
is_preview,
process_mode,
} => {
if matches!(process_mode, ProcessMode::Skip(..)) {
self.frame_info
.skip_sections(input, &mut self.local_buffer)?;
} else {
self.frame_info
.fill_sections(input, &mut self.local_buffer)?;
match self.frame_info.process_sections(
(process_mode == ProcessMode::Process)
.then_some(output_buffers.as_deref_mut())
.flatten(),
self.output_color_profile.as_ref().unwrap(),
self.pixel_format.as_ref().unwrap(),
parallel_runner,
) {
Ok(None) => Ok(()),
Ok(Some(missing)) => Err(Error::OutOfBounds(missing)),
Err(Error::OutOfBounds(_)) => Err(Error::SectionTooShort),
Err(err) => Err(err),
}?;
}
let header = self.frame_info.current_frame_header().unwrap();
let is_last = header.is_last;
if !is_preview
&& !matches!(process_mode, ProcessMode::Skip(_))
&& let Some(slot) = stored_slot(header)
{
self.slot_frames[slot] = Some(self.next_frame_index - 1);
}
// After the last frame, the decoder state is dropped, except for a looping
// animation, which will likely be decoded again (seeking can reuse it).
let animation = &self.image_info.file_header().image_metadata.animation;
let looping = animation.as_ref().is_some_and(|a| a.num_loops != 1);
if !is_last || (looping && !is_preview) {
self.frame_info.clear();
} else {
self.frame_info.reset((0, 0), [false; 8]);
self.slot_frames = [None; 8];
}
if !is_last {
self.state = ParserState::FrameHeader { is_preview };
} else if is_preview {
self.state = ParserState::FrameHeader { is_preview: false };
} else {
self.state = ParserState::TrailingData;
self.file_length = Some(
input
.box_parser()
.total_bytes_consumed(self.local_buffer.consumed()),
);
}
if process_mode.notify_user() {
return Ok(Event::FrameComplete {
has_more_frames: self.has_more_frames(),
});
}
}
ParserState::TrailingData => {
input.consume_trailing_data()?;
self.state = ParserState::Done;
return Ok(Event::Complete);
}
ParserState::Done => {
return Err(Error::ApiUsageError("decoding is already complete"));
}
};
}
}
pub fn has_visible_frame(&self) -> bool {
matches!(
self.state,
ParserState::Sections { process_mode, .. } if process_mode.notify_user()
)
}
}