Add jbrd box support for JPEG reconstruction

jxl/Cargo.toml

@@ -22,6 +22,7 @@ array-init = "2.0.0"
 tracing = { version = "0.1.40", optional = true }
 jxl_macros = { path = "../jxl_macros", version = "=0.3.0" }
 jxl_simd = { path = "../jxl_simd", version = "=0.3.0" }
+brotli = { version = "7.0", optional = true }
 
 [dev-dependencies]
 arbtest = "0.3.2"
@@ -37,6 +38,7 @@ sse42 = ["jxl_simd/sse42"]
 avx = ["jxl_simd/avx"]
 avx512 = ["jxl_simd/avx512"]
 neon = ["jxl_simd/neon"]
+jpeg-reconstruction = ["brotli"]
 
 [lints]
 workspace = true

jxl/src/api/decoder.rs

@@ -9,6 +9,8 @@ use super::{
 };
 #[cfg(test)]
 use crate::frame::Frame;
+#[cfg(feature = "jpeg-reconstruction")]
+use crate::jpeg::JpegReconstructionData;
 use crate::{api::JxlFrameHeader, error::Result};
 use states::*;
 use std::marker::PhantomData;
@@ -141,6 +143,21 @@ impl JxlDecoder<WithImageInfo> {
         self.inner.has_more_frames()
     }
 
+    /// Returns the JPEG reconstruction data if present in the file.
+    ///
+    /// This data is available after reading a jbrd box from a JXL file
+    /// that was created by losslessly recompressing a JPEG.
+    #[cfg(feature = "jpeg-reconstruction")]
+    pub fn jpeg_reconstruction_data(&self) -> Option<&JpegReconstructionData> {
+        self.inner.jpeg_reconstruction_data()
+    }
+
+    /// Returns true if the file contains JPEG reconstruction data.
+    #[cfg(feature = "jpeg-reconstruction")]
+    pub fn has_jpeg_reconstruction(&self) -> bool {
+        self.inner.has_jpeg_reconstruction()
+    }
+
     #[cfg(test)]
     pub(crate) fn set_use_simple_pipeline(&mut self, u: bool) {
         self.inner.set_use_simple_pipeline(u);

jxl/src/api/inner/box_parser.rs

@@ -4,17 +4,23 @@
 // license that can be found in the LICENSE file.
 
 use crate::error::{Error, Result};
+#[cfg(feature = "jpeg-reconstruction")]
+use crate::jpeg::JpegReconstructionData;
 
 use crate::api::{
     JxlBitstreamInput, JxlSignatureType, check_signature_internal, inner::process::SmallBuffer,
 };
+#[cfg(feature = "jpeg-reconstruction")]
+use std::io::IoSliceMut;
 
 #[derive(Clone)]
 enum ParseState {
     SignatureNeeded,
     BoxNeeded,
     CodestreamBox(u64),
     SkippableBox(u64),
+    #[cfg(feature = "jpeg-reconstruction")]
+    JbrdBox(u64),
 }
 
 enum CodestreamBoxType {
@@ -28,6 +34,12 @@ pub(super) struct BoxParser {
     pub(super) box_buffer: SmallBuffer,
     state: ParseState,
     box_type: CodestreamBoxType,
+    /// Buffer for accumulating jbrd box data
+    #[cfg(feature = "jpeg-reconstruction")]
+    jbrd_buffer: Vec<u8>,
+    /// Parsed JPEG reconstruction data (available after jbrd box is fully read)
+    #[cfg(feature = "jpeg-reconstruction")]
+    pub(super) jpeg_reconstruction: Option<JpegReconstructionData>,
 }
 
 impl BoxParser {
@@ -36,6 +48,10 @@ impl BoxParser {
             box_buffer: SmallBuffer::new(128),
             state: ParseState::SignatureNeeded,
             box_type: CodestreamBoxType::None,
+            #[cfg(feature = "jpeg-reconstruction")]
+            jbrd_buffer: Vec::new(),
+            #[cfg(feature = "jpeg-reconstruction")]
+            jpeg_reconstruction: None,
         }
     }
 
@@ -83,6 +99,48 @@ impl BoxParser {
                         self.state = ParseState::SkippableBox(s);
                     }
                 }
+                #[cfg(feature = "jpeg-reconstruction")]
+                ParseState::JbrdBox(mut remaining) => {
+                    // Accumulate jbrd box data for parsing
+                    let num = remaining.min(usize::MAX as u64) as usize;
+                    let read_count = if !self.box_buffer.is_empty() {
+                        let to_read = num.min(self.box_buffer.len());
+                        self.jbrd_buffer
+                            .extend_from_slice(&self.box_buffer[..to_read]);
+                        self.box_buffer.consume(to_read);
+                        to_read
+                    } else {
+                        // Read data into jbrd_buffer using IoSliceMut
+                        let start = self.jbrd_buffer.len();
+                        self.jbrd_buffer.resize(start + num, 0);
+                        let read =
+                            input.read(&mut [IoSliceMut::new(&mut self.jbrd_buffer[start..])])?;
+                        if read < num {
+                            self.jbrd_buffer.truncate(start + read);
+                        }
+                        read
+                    };
+                    if read_count == 0 {
+                        return Err(Error::OutOfBounds(num));
+                    }
+                    remaining -= read_count as u64;
+                    if remaining == 0 {
+                        // Parse the jbrd data
+                        match JpegReconstructionData::parse(&self.jbrd_buffer) {
+                            Ok(data) => {
+                                self.jpeg_reconstruction = Some(data);
+                            }
+                            Err(_e) => {
+                                // Parsing failed - leave jpeg_reconstruction as None
+                            }
+                        }
+                        self.jbrd_buffer.clear();
+                        self.jbrd_buffer.shrink_to_fit();
+                        self.state = ParseState::BoxNeeded;
+                    } else {
+                        self.state = ParseState::JbrdBox(remaining);
+                    }
+                }
                 ParseState::BoxNeeded => {
                     self.box_buffer.refill(|b| input.read(b), None)?;
                     let min_len = match &self.box_buffer[..] {
@@ -148,6 +206,12 @@ impl BoxParser {
                             };
                             self.state = ParseState::CodestreamBox(content_len);
                         }
+                        #[cfg(feature = "jpeg-reconstruction")]
+                        b"jbrd" => {
+                            // JPEG reconstruction data box - accumulate for later parsing
+                            self.jbrd_buffer.clear();
+                            self.state = ParseState::JbrdBox(content_len);
+                        }
                         _ => {
                             self.state = ParseState::SkippableBox(content_len);
                         }

jxl/src/api/inner/codestream_parser/mod.rs

@@ -245,7 +245,18 @@ impl CodestreamParser {
                             break;
                         }
                     }
-                    match self.process_sections(decode_options, &mut output_buffers, do_flush) {
+                    #[cfg(feature = "jpeg-reconstruction")]
+                    let result = self.process_sections(
+                        decode_options,
+                        &mut output_buffers,
+                        do_flush,
+                        box_parser,
+                    );
+                    #[cfg(not(feature = "jpeg-reconstruction"))]
+                    let result =
+                        self.process_sections(decode_options, &mut output_buffers, do_flush);
+
+                    match result {
                         Ok(None) => Ok(()),
                         Ok(Some(missing)) => Err(Error::OutOfBounds(missing)),
                         Err(Error::OutOfBounds(_)) => Err(Error::SectionTooShort),

jxl/src/api/inner/codestream_parser/non_section.rs

@@ -279,7 +279,8 @@ impl CodestreamParser {
         // Save file_header before creating frame (for preview frame recovery)
         self.saved_file_header = self.decoder_state.as_ref().map(|ds| ds.file_header.clone());
 
-        let frame = Frame::from_header_and_toc(
+        #[cfg_attr(not(feature = "jpeg-reconstruction"), allow(unused_mut))]
+        let mut frame = Frame::from_header_and_toc(
             self.frame_header.take().unwrap(),
             toc,
             self.decoder_state.take().unwrap(),
@@ -341,6 +342,12 @@ impl CodestreamParser {
         self.section_state =
             SectionState::new(frame.header().num_lf_groups(), frame.header().num_groups());
 
+        // Enable JPEG coefficient preservation if requested
+        #[cfg(feature = "jpeg-reconstruction")]
+        if decode_options.preserve_jpeg_coefficients {
+            frame.set_preserve_jpeg_coefficients(true);
+        }
+
         self.frame = Some(frame);
 
         Ok(())

jxl/src/api/inner/codestream_parser/sections.rs

@@ -3,10 +3,10 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+#[cfg(feature = "jpeg-reconstruction")]
+use crate::api::inner::box_parser::BoxParser;
 use crate::{
-    api::{JxlDecoderOptions, JxlOutputBuffer},
-    bit_reader::BitReader,
-    error::Result,
+    api::JxlDecoderOptions, api::JxlOutputBuffer, bit_reader::BitReader, error::Result,
     frame::Section,
 };
 
@@ -42,9 +42,9 @@ impl CodestreamParser {
         decode_options: &JxlDecoderOptions,
         output_buffers: &mut Option<&mut [JxlOutputBuffer<'_>]>,
         do_flush: bool,
+        #[cfg(feature = "jpeg-reconstruction")] box_parser: &mut BoxParser,
     ) -> Result<Option<usize>> {
         let frame = self.frame.as_mut().unwrap();
-        let frame_header = frame.header();
 
         // Dequeue ready sections.
         while self
@@ -75,6 +75,7 @@ impl CodestreamParser {
         let mut processed_section = false;
         let pixel_format = self.pixel_format.as_ref().unwrap();
         'process: {
+            let frame_header = frame.header();
             if frame_header.num_groups() == 1 && frame_header.passes.num_passes == 1 {
                 // Single-group special case.
                 let Some(sec) = self.lf_global_section.take() else {
@@ -228,6 +229,58 @@ impl CodestreamParser {
             .is_some_and(|info| info.preview_size.is_some());
         let might_be_preview = self.process_without_output && has_preview;
 
+        // Extract JPEG coefficients before finalizing the frame
+        #[cfg(feature = "jpeg-reconstruction")]
+        if let Some(frame) = self.frame.as_mut()
+            && let Some(coeffs) = frame.take_jpeg_coefficients()
+        {
+            let do_ycbcr = frame.header().do_ycbcr;
+            // Merge coefficients into the jpeg_reconstruction data
+            if let Some(ref mut jpeg_data) = box_parser.jpeg_reconstruction {
+                jpeg_data.dct_coefficients = Some(coeffs);
+                if let Some((qtable, qtable_den)) = frame.jpeg_raw_quant_table() {
+                    jpeg_data.update_quant_tables_from_raw(qtable, qtable_den, do_ycbcr)?;
+                }
+                {
+                    let header = frame.header();
+                    let is_gray = jpeg_data.is_gray || jpeg_data.components.len() == 1;
+                    let component_map = if is_gray {
+                        [1usize, 1, 1]
+                    } else {
+                        [1usize, 0, 2]
+                    };
+                    let mut max_hshift = 0usize;
+                    let mut max_vshift = 0usize;
+                    let chans = if is_gray {
+                        &[1usize][..]
+                    } else {
+                        &[0usize, 1, 2][..]
+                    };
+                    for &c in chans {
+                        max_hshift = max_hshift.max(header.hshift(c));
+                        max_vshift = max_vshift.max(header.vshift(c));
+                    }
+                    for (jpeg_idx, &vardct_chan) in component_map
+                        .iter()
+                        .enumerate()
+                        .take(jpeg_data.components.len())
+                    {
+                        let hshift = header.hshift(vardct_chan);
+                        let vshift = header.vshift(vardct_chan);
+                        jpeg_data.components[jpeg_idx].h_samp_factor =
+                            1u8 << (max_hshift.saturating_sub(hshift) as u8);
+                        jpeg_data.components[jpeg_idx].v_samp_factor =
+                            1u8 << (max_vshift.saturating_sub(vshift) as u8);
+                    }
+                }
+                if let Some(profile) = self.embedded_color_profile.as_ref()
+                    && let Some(icc) = profile.try_as_icc()
+                {
+                    jpeg_data.fill_icc_app_markers(icc.as_ref())?;
+                }
+            }
+        }
+
         let decoder_state = self.frame.take().unwrap().finalize()?;
         if let Some(state) = decoder_state {
             self.decoder_state = Some(state);

jxl/src/api/inner/mod.rs

@@ -5,6 +5,8 @@
 
 #[cfg(test)]
 use crate::api::FrameCallback;
+#[cfg(feature = "jpeg-reconstruction")]
+use crate::jpeg::JpegReconstructionData;
 use crate::{
     api::JxlFrameHeader,
     error::{Error, Result},
@@ -135,6 +137,21 @@ impl JxlDecoderInner {
         self.codestream_parser.has_more_frames
     }
 
+    /// Returns the JPEG reconstruction data if present in the file.
+    ///
+    /// This data is available after reading a jbrd box from a JXL file
+    /// that was created by losslessly recompressing a JPEG.
+    #[cfg(feature = "jpeg-reconstruction")]
+    pub fn jpeg_reconstruction_data(&self) -> Option<&JpegReconstructionData> {
+        self.box_parser.jpeg_reconstruction.as_ref()
+    }
+
+    /// Returns true if the file contains JPEG reconstruction data.
+    #[cfg(feature = "jpeg-reconstruction")]
+    pub fn has_jpeg_reconstruction(&self) -> bool {
+        self.box_parser.jpeg_reconstruction.is_some()
+    }
+
     #[cfg(test)]
     pub(crate) fn set_use_simple_pipeline(&mut self, u: bool) {
         self.codestream_parser.set_use_simple_pipeline(u);

jxl/src/api/mod.rs

@@ -15,6 +15,8 @@ mod signature;
 mod xyb_constants;
 
 pub use crate::image::JxlOutputBuffer;
+#[cfg(feature = "jpeg-reconstruction")]
+pub use crate::jpeg::{JpegDctCoefficients, JpegReconstructionData};
 pub use color::*;
 pub use data_types::*;
 pub use decoder::*;