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 std::{
collections::{HashSet, VecDeque},
io::IoSliceMut,
};
use crate::{
api::{
JxlColorProfile, JxlDecoderOptions, JxlOutputBuffer, JxlPixelFormat,
inner::{
CodestreamParser,
box_parser::CodestreamInput,
codestream_parser::{ProcessMode, check_size_limit, validate_output_buffers},
process::SmallBuffer,
},
},
bit_reader::BitReader,
error::{Error, Result},
frame::{DecoderState, Frame, Section},
headers::{
FileHeader,
encodings::UnconditionalCoder,
frame_header::{Encoding, FrameHeader, FrameType},
toc::{IncrementalTocReader, Toc},
},
util::NewWithCapacity,
};
struct SectionBuffer {
len: usize,
data: Vec<u8>,
section: Section,
}
#[derive(Debug)]
struct SectionState {
lf_global_done: bool,
remaining_lf: usize,
hf_global_done: bool,
completed_passes: Vec<u8>,
lf_global_flush_len: usize,
}
impl SectionState {
fn new(num_lf_groups: usize, num_groups: usize) -> Self {
Self {
lf_global_done: false,
remaining_lf: num_lf_groups,
hf_global_done: false,
completed_passes: vec![0; num_groups],
lf_global_flush_len: 0,
}
}
}
pub struct FrameInfo {
frame_header: Option<FrameHeader>,
toc_parser: Option<IncrementalTocReader>,
frame: Option<Frame>,
// Keeps track of whether pixels have been modified.
pixels_dirty: bool,
// Section information.
sections: VecDeque<SectionBuffer>,
section_size: usize,
ready_section_data: 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>,
#[cfg(test)]
pub use_simple_pipeline: bool,
}
impl FrameInfo {
pub fn new() -> Self {
Self {
frame_header: None,
toc_parser: None,
frame: None,
sections: VecDeque::new(),
section_size: 0,
ready_section_data: 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(),
pixels_dirty: false,
#[cfg(test)]
use_simple_pipeline: false,
}
}
pub fn clear(&mut self, clear_frame: bool) {
self.frame_header = None;
self.toc_parser = None;
self.ready_section_data = 0;
if clear_frame {
self.frame = None;
}
// Clear sections
self.sections.clear();
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();
}
pub fn current_frame_header(&self) -> Option<&FrameHeader> {
self.frame_header
.as_ref()
.or(self.frame.as_ref().map(|x| x.header()))
}
pub fn parse_frame_header(
&mut self,
is_preview: bool,
decode_options: &JxlDecoderOptions,
file_header: &FileHeader,
br: &mut BitReader,
bits: &mut usize,
) -> Result<()> {
// For preview frames, use the preview dimensions instead of main image dimensions
let nonserialized = if is_preview {
file_header
.preview_frame_header_nonserialized()
.unwrap_or_else(|| file_header.frame_header_nonserialized())
} else {
file_header.frame_header_nonserialized()
};
let mut frame_header = FrameHeader::read_unconditional(&(), br, &nonserialized)?;
frame_header.postprocess(&nonserialized);
check_size_limit(
decode_options.sample_limit,
frame_header.size(),
frame_header.num_extra_channels as usize,
)?;
self.frame_header = Some(frame_header);
*bits = br.total_bits_read();
Ok(())
}
pub fn parse_toc(&mut self, br: &mut BitReader, bits: &mut usize) -> Result<Toc> {
if self.toc_parser.is_none() {
let num_toc_entries = self.frame_header.as_ref().unwrap().num_toc_entries();
self.toc_parser = Some(IncrementalTocReader::new(num_toc_entries as u32, br)?);
}
let toc_parser = self.toc_parser.as_mut().unwrap();
*bits = br.total_bits_read();
while !toc_parser.is_complete() {
match toc_parser.read_step(br) {
Ok(()) => *bits = br.total_bits_read(),
Err(Error::OutOfBounds(c)) => {
// Estimate >= 16 bits per remaining entry to read.
return Err(Error::OutOfBounds(
c + toc_parser.remaining_entries() as usize * 2,
));
}
Err(e) => return Err(e),
}
}
br.jump_to_byte_boundary()?;
*bits = br.total_bits_read();
Ok(self.toc_parser.take().unwrap().finalize())
}
#[allow(clippy::too_many_arguments)]
pub(super) fn make_frame(
&mut self,
cbuf: &mut SmallBuffer,
toc: Toc,
file_header: &FileHeader,
decode_options: &JxlDecoderOptions,
pixel_format: &JxlPixelFormat,
output_profile: &JxlColorProfile,
process_mode: ProcessMode,
) -> Result<()> {
self.section_size = toc.entries.iter().map(|x| *x as usize).sum();
self.ready_section_data = 0;
self.lf_global_section = None;
self.lf_sections.clear();
self.hf_global_section = None;
self.candidate_hf_sections.clear();
self.hf_sections.clear();
if !matches!(process_mode, ProcessMode::Skip(_)) {
// If we have a previous frame, compute the decoder state from there.
// Otherwise, compute a new decoder state.
// We finalize the previous frame here to allow progressive rendering
// to work properly if a flush is requested while we parse a frame
// header.
let decoder_state = self
.frame
.take()
.map(|x| x.finalize())
.transpose()?
.flatten()
.unwrap_or_else(|| DecoderState::new(file_header.clone(), decode_options));
let mut frame =
Frame::from_header_and_toc(self.frame_header.take().unwrap(), toc, decoder_state)?;
let num_groups = frame.header().num_groups();
let num_passes = frame.header().passes.num_passes as usize;
let mut hf_sections = Vec::new_with_capacity(num_groups)?;
for _ in 0..num_groups {
let mut row = Vec::new_with_capacity(num_passes)?;
for _ in 0..num_passes {
row.push(None);
}
hf_sections.push(row);
}
self.hf_sections = hf_sections;
let mut sections: Vec<_> = frame
.toc()
.entries
.iter()
.map(|x| SectionBuffer {
len: *x as usize,
data: vec![],
section: Section::LfGlobal, // will be fixed later
})
.collect();
let order = if frame.toc().permuted {
frame.toc().permutation.0.clone()
} else {
(0..sections.len() as u32).collect()
};
if sections.len() > 1 {
let base_sections = [Section::LfGlobal, Section::HfGlobal];
let lf_sections =
(0..frame.header().num_lf_groups()).map(|x| Section::Lf { group: x });
let hf_sections = (0..frame.header().passes.num_passes).flat_map(|p| {
(0..frame.header().num_groups()).map(move |g| Section::Hf {
group: g,
pass: p as usize,
})
});
for section in base_sections
.into_iter()
.chain(lf_sections)
.chain(hf_sections)
{
sections[order[frame.get_section_idx(section)] as usize].section = section;
}
}
self.sections = sections.into_iter().collect();
// Move data from the pre-section buffer into the sections.
for buf in self.sections.iter_mut() {
if cbuf.is_empty() {
break;
}
let mut data = Vec::new();
data.try_reserve_exact(buf.len)?;
data.resize(buf.len, 0);
buf.data = data;
let n = cbuf.take(&mut [IoSliceMut::new(&mut buf.data)]);
self.ready_section_data += n;
}
self.section_state =
SectionState::new(frame.header().num_lf_groups(), frame.header().num_groups());
frame.prepare_render_pipeline(pixel_format, output_profile)?;
self.frame = Some(frame);
} else {
let num = cbuf.len().min(self.section_size);
cbuf.consume(num);
self.ready_section_data += num;
}
Ok(())
}
#[cfg(test)]
pub fn frame(&mut self) -> Option<&mut Frame> {
self.frame.as_mut()
}
pub fn dequeue_ready_sections(&mut self) {
// Dequeue ready sections.
while self
.sections
.front()
.is_some_and(|s| s.len <= self.ready_section_data)
{
let s = self.sections.pop_front().unwrap();
self.ready_section_data -= s.len;
match s.section {
Section::LfGlobal => {
self.lf_global_section = Some(s);
}
Section::HfGlobal => {
self.hf_global_section = Some(s);
}
Section::Lf { .. } => {
self.lf_sections.push(s);
}
Section::Hf { group, pass } => {
self.hf_sections[group][pass] = Some(s);
self.candidate_hf_sections.insert(group);
}
}
}
}
pub fn skip_sections(
&mut self,
input: &mut CodestreamInput,
buf: &mut SmallBuffer,
) -> Result<()> {
let total_size = self.section_size;
let need_skip = total_size - self.ready_section_data;
let skipped = input.skip(need_skip)?;
buf.mark_consumed(skipped as u64);
self.ready_section_data += skipped;
if self.ready_section_data < total_size {
Err(Error::OutOfBounds(total_size - self.ready_section_data))
} else {
self.sections.clear();
Ok(())
}
}
pub fn fill_sections(
&mut self,
input: &mut CodestreamInput,
buf: &mut SmallBuffer,
) -> Result<()> {
// TODO(veluca): consider re-using `section_buffers` between different iterations.
let mut total = 0;
loop {
let available_codestream = input.available_bytes();
let mut available_codestream = match (available_codestream, total) {
(Ok(a), _) => a.max(1),
(Err(Error::OutOfBounds(_)), t) if t > 0 => return Ok(()),
(Err(e), _) => return Err(e),
};
let mut readable_section_data = self.ready_section_data + available_codestream;
// 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;
}
}
// Prepare buffers to read into.
let mut section_buffers = vec![];
let mut ready = self.ready_section_data;
for buf in self.sections.iter_mut() {
if buf.data.is_empty() && buf.len != 0 {
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 num = input.read(&mut section_buffers[..])?;
if num == 0 {
break;
}
self.ready_section_data += num;
total += num;
buf.mark_consumed(num as u64);
self.dequeue_ready_sections();
}
Ok(())
}
// Returns whether we modified the pixels.
fn process_single_section(
frame: &mut Frame,
buf: &[u8],
is_complete: bool,
output_buffers: &mut Option<&mut [JxlOutputBuffer<'_>]>,
output_profile: &JxlColorProfile,
pixel_format: &JxlPixelFormat,
do_flush: bool,
) -> Result<bool> {
let mut br = BitReader::new(buf);
frame.decode_lf_global(&mut br, !is_complete)?;
frame.decode_lf_group(0, &mut br)?;
frame.decode_hf_global(&mut br)?;
frame.finalize_lf()?;
frame.decode_and_render_hf_groups(
output_buffers,
pixel_format,
vec![(0, vec![(0, br)])],
do_flush,
output_profile,
)
}
// Returns:
// Ok(None) if the frame is complete
// Ok(Some(len)) if we need `len` bytes to make progress
// Err(_) if there was an error.
pub fn process_sections(
&mut self,
output_buffers: &mut Option<&mut [JxlOutputBuffer<'_>]>,
output_profile: &JxlColorProfile,
pixel_format: &JxlPixelFormat,
) -> Result<Option<usize>> {
let data_for_next_section = self
.sections
.front()
.map(|x| x.len - self.ready_section_data);
let frame = self.frame.as_mut().unwrap();
let frame_header = frame.header();
if frame_header.num_groups() == 1 && frame_header.passes.num_passes == 1 {
// Single-group special case.
let Some(buf) = self.lf_global_section.take() else {
return Ok(data_for_next_section);
};
assert!(self.sections.is_empty());
self.pixels_dirty |= Self::process_single_section(
frame,
&buf.data,
true,
output_buffers,
output_profile,
pixel_format,
false,
)?;
return Ok(None);
}
if let Some(buf) = self.lf_global_section.take() {
frame.decode_lf_global(&mut BitReader::new(&buf.data), false)?;
self.section_state.lf_global_done = true;
}
if !self.section_state.lf_global_done {
return Ok(data_for_next_section);
}
for lf_section in self.lf_sections.drain(..) {
let Section::Lf { group } = lf_section.section else {
unreachable!()
};
frame.decode_lf_group(group, &mut BitReader::new(&lf_section.data))?;
self.section_state.remaining_lf -= 1;
}
if self.section_state.remaining_lf != 0 {
return Ok(data_for_next_section);
}
if let Some(hf_global) = self.hf_global_section.take() {
frame.decode_hf_global(&mut BitReader::new(&hf_global.data))?;
frame.finalize_lf()?;
self.section_state.hf_global_done = true;
}
if !self.section_state.hf_global_done {
return Ok(data_for_next_section);
}
let mut group_readers = vec![];
let mut processed_groups = vec![];
let mut check_group = |g: usize| {
let mut sections = vec![];
for (pass, grp) in self.hf_sections[g]
.iter()
.enumerate()
.skip(self.section_state.completed_passes[g] as usize)
{
let Some(s) = &grp else {
break;
};
self.section_state.completed_passes[g] += 1;
sections.push((pass, BitReader::new(&s.data)));
}
if !sections.is_empty() {
group_readers.push((g, sections));
processed_groups.push(g);
}
};
if self.candidate_hf_sections.len() * 4 < self.hf_sections.len() {
for g in self.candidate_hf_sections.drain() {
check_group(g)
}
// Processing sections in order is more efficient because it lets us flush
// the pipeline faster.
group_readers.sort_by_key(|x| x.0);
} else {
for g in 0..self.hf_sections.len() {
if self.candidate_hf_sections.contains(&g) {
check_group(g);
}
}
self.candidate_hf_sections.clear();
}
self.pixels_dirty |= frame.decode_and_render_hf_groups(
output_buffers,
pixel_format,
group_readers,
false,
output_profile,
)?;
for g in processed_groups.into_iter() {
for i in 0..self.section_state.completed_passes[g] {
self.hf_sections[g][i as usize] = None;
}
}
Ok(data_for_next_section)
}
pub fn do_flush(
&mut self,
output_buffers: &mut [JxlOutputBuffer<'_>],
output_profile: &JxlColorProfile,
pixel_format: &JxlPixelFormat,
) -> Result<()> {
let Some(frame) = self.frame.as_mut() else {
return Ok(());
};
validate_output_buffers(output_buffers, Some(pixel_format))?;
let frame_header = frame.header();
let has_partial_lf = self
.sections
.front()
.is_some_and(|s| s.section == Section::LfGlobal)
&& 2 * self.ready_section_data > 3 * self.section_state.lf_global_flush_len
&& frame_header.encoding == Encoding::Modular
&& matches!(
frame_header.frame_type,
FrameType::RegularFrame | FrameType::LFFrame
);
if has_partial_lf {
self.section_state.lf_global_flush_len = self.ready_section_data;
let section = &self.sections[0].data[..self.ready_section_data];
// These are partial renders, so we ignore any errors.
if frame_header.num_groups() == 1 && frame_header.passes.num_passes == 1 {
if let Ok(dirty) = Self::process_single_section(
frame,
section,
false,
&mut Some(output_buffers),
output_profile,
pixel_format,
true,
) {
self.pixels_dirty |= dirty;
return Ok(());
}
} else {
let _ = frame.decode_lf_global(&mut BitReader::new(section), true);
}
}
if frame.can_do_early_rendering() || self.section_state.hf_global_done {
self.pixels_dirty |= frame.decode_and_render_hf_groups(
&mut Some(output_buffers),
pixel_format,
vec![],
true,
output_profile,
)?;
}
Ok(())
}
}
impl CodestreamParser {
pub fn get_and_clear_pixels_dirty(&mut self) -> bool {
let r = self.frame_info.pixels_dirty;
self.frame_info.pixels_dirty = false;
r
}
}