diff --git a/jpegxl-rs/src/decode/event.rs b/jpegxl-rs/src/decode/event.rs index a010cd8..1413a34 100644 --- a/jpegxl-rs/src/decode/event.rs +++ b/jpegxl-rs/src/decode/event.rs @@ -1,21 +1,37 @@ use std::ffi::c_int; +use jpegxl_sys::common::types::JxlPixelFormat; /// Target of the color profile. pub use jpegxl_sys::decode::JxlColorProfileTarget as ColorProfileTarget; use super::{ColorEncodingConfig, Config}; +/// Events that can be subscribed to. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Event { + /// Basic information. BasicInfo, + /// Color encoding. ColorEncoding(ColorEncodingConfig), + /// Preview image. + PreviewImage { + /// Pixel format. + pixel_format: JxlPixelFormat, + }, + /// JPEG reconstruction. + JpegReconstruction { + /// Initial buffer size. Increase it to reduce the number of reallocations. + init_buffer_size: usize, + }, } impl From for c_int { fn from(value: Event) -> Self { match value { Event::BasicInfo => 0x40, - Event::ColorEncoding { .. } => 0x100, + Event::ColorEncoding(_) => 0x100, + Event::PreviewImage { .. } => 0x200, + Event::JpegReconstruction { .. } => 0x2000, } } } @@ -25,12 +41,13 @@ where I: IntoIterator, { iter.into_iter() - .fold((0, Config::default()), |(flag, config), x| { + .fold((0, Config::default()), |(flag, mut config), x| { let flag = flag | c_int::from(x); let config = match x { - Event::ColorEncoding(val) => Config { - color_profile: Some(val), - }, + Event::ColorEncoding(val) => { + config.color_profile = Some(val); + config + } _ => config, }; (flag, config) diff --git a/jpegxl-rs/src/decode/session.rs b/jpegxl-rs/src/decode/session.rs index 6a961c3..da1eef8 100644 --- a/jpegxl-rs/src/decode/session.rs +++ b/jpegxl-rs/src/decode/session.rs @@ -1,34 +1,40 @@ -use std::{mem::MaybeUninit, sync::Arc}; +use std::{ffi::CString, mem::MaybeUninit, sync::Arc}; use jpegxl_sys::{ - color_encoding::JxlColorEncoding, - decode::{ - JxlDecoderGetColorAsEncodedProfile, JxlDecoderGetColorAsICCProfile, - JxlDecoderGetICCProfileSize, JxlDecoderStatus, - }, + color::color_encoding::JxlColorEncoding, common::types::JxlPixelFormat, decode as d, }; -use super::{BasicInfo, ColorProfileTarget, Event, JxlDecoder}; +use super::{BasicInfo, ColorProfileTarget, Event, JxlDecoder, Pixels}; use crate::{decode::parse_events, errors::check_dec_status, DecodeError}; /// Represents the state of the session. pub enum State { - /// Initial state of the session. - Init, + /// Waiting for the next event. + Continue, /// Basic information such as image dimensions and extra channels. /// This event occurs max once per image. BasicInfo(Arc), /// ICC color profile . IccProfile(Vec), - /// - ColorProfile(JxlColorEncoding), + /// Color profile. + ColorProfile(Box), + /// Preview image. Dimensions can be accessed from [`BasicInfo::preview`] + PreviewImage(Pixels), + /// Begining of a frame + Frame, + /// JPEG reconstruction. + JpegReconstruction(Vec), } #[derive(Debug, Default)] pub(crate) struct Config { pub color_profile: Option, + pub preview: Option, + pub frame: Option, + pub jpeg_reconstruction: Option, } +/// Configuration for color encoding. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ColorEncodingConfig { target: ColorProfileTarget, @@ -39,6 +45,7 @@ pub struct ColorEncodingConfig { pub struct Session<'dec, 'pr, 'mm> { dec: &'dec mut JxlDecoder<'pr, 'mm>, basic_info: Option>, + jpeg_buffer: Vec, config: Config, state: State, } @@ -86,74 +93,161 @@ impl<'dec, 'pr, 'mm> Session<'dec, 'pr, 'mm> { Ok(Self { dec, basic_info: None, + jpeg_buffer: Vec::new(), config, - state: State::Init, + state: State::Continue, }) } - fn step(&mut self, status: JxlDecoderStatus) -> Result { - use jpegxl_sys::decode::{JxlDecoderGetBasicInfo, JxlDecoderStatus as s}; + fn step(&mut self, status: d::JxlDecoderStatus) -> Result { + use jpegxl_sys::decode::JxlDecoderStatus as s; match status { s::Success => panic!("Unexpected success status"), s::Error => Err(DecodeError::GenericError), - s::BasicInfo => { - let mut info = MaybeUninit::uninit(); + s::BasicInfo => self.get_basic_info(), + s::ColorEncoding => self.get_color_profile(), + s::PreviewImage => self.get_preview_image(), + s::Frame => self.get_frame(), + s::JPEGReconstruction => { + let Some(size) = self.config.jpeg_reconstruction else { + return Err(DecodeError::InternalError( + "Subscribe to JPEG reconstruction event but without a config!", + )); + }; + + self.jpeg_buffer.resize(size, 0); + check_dec_status(unsafe { - JxlDecoderGetBasicInfo(self.dec.ptr, info.as_mut_ptr()) + d::JxlDecoderSetJPEGBuffer( + self.dec.ptr, + self.jpeg_buffer.as_mut_ptr(), + self.jpeg_buffer.len(), + ) })?; - if let Some(pr) = self.dec.parallel_runner { - pr.callback_basic_info(unsafe { &*info.as_ptr() }); - } - - self.basic_info = Some(Arc::new(unsafe { info.assume_init() })); - Ok(State::BasicInfo(unsafe { - self.basic_info.as_ref().unwrap_unchecked().clone() - })) + Ok(State::Continue) } - s::ColorEncoding => { - let Some(config) = self.config.color_profile else { - return Err(DecodeError::InvalidUsage( - "Subscribe to color encoding event but without a color profile", - )); - }; + s::JPEGNeedMoreOutput => { + let remaining = unsafe { d::JxlDecoderReleaseJPEGBuffer(self.dec.ptr) }; + + self.jpeg_buffer + .resize(self.jpeg_buffer.len() + remaining, 0); + + check_dec_status(unsafe { + d::JxlDecoderSetJPEGBuffer( + self.dec.ptr, + self.jpeg_buffer.as_mut_ptr(), + self.jpeg_buffer.len(), + ) + })?; - if config.icc_profile { - let mut icc_size = 0; - let mut icc_profile = Vec::new(); - - check_dec_status(unsafe { - JxlDecoderGetICCProfileSize(self.dec.ptr, config.target, &mut icc_size) - })?; - icc_profile.resize(icc_size, 0); - - check_dec_status(unsafe { - JxlDecoderGetColorAsICCProfile( - self.dec.ptr, - config.target, - icc_profile.as_mut_ptr(), - icc_size, - ) - })?; - - Ok(State::IccProfile(icc_profile)) - } else { - let mut color_encoding = MaybeUninit::uninit(); - - check_dec_status(unsafe { - JxlDecoderGetColorAsEncodedProfile( - self.dec.ptr, - config.target, - color_encoding.as_mut_ptr(), - ) - })?; - Ok(State::ColorProfile(unsafe { color_encoding.assume_init() })) - } + Ok(State::Continue) } _ => unimplemented!(), } } + + fn get_basic_info(&mut self) -> Result { + let mut info = MaybeUninit::uninit(); + check_dec_status(unsafe { d::JxlDecoderGetBasicInfo(self.dec.ptr, info.as_mut_ptr()) })?; + + if let Some(pr) = self.dec.parallel_runner { + pr.callback_basic_info(unsafe { &*info.as_ptr() }); + } + + unsafe { + self.basic_info = Some(Arc::new(info.assume_init())); + Ok(State::BasicInfo( + self.basic_info.as_ref().unwrap_unchecked().clone(), + )) + } + } + + fn get_color_profile(&mut self) -> Result { + let Some(config) = self.config.color_profile else { + return Err(DecodeError::InternalError( + "Subscribe to color encoding event but without a color profile config!", + )); + }; + + if config.icc_profile { + let mut icc_size = 0; + let mut icc_profile = Vec::new(); + + check_dec_status(unsafe { + d::JxlDecoderGetICCProfileSize(self.dec.ptr, config.target, &mut icc_size) + })?; + icc_profile.resize(icc_size, 0); + + check_dec_status(unsafe { + d::JxlDecoderGetColorAsICCProfile( + self.dec.ptr, + config.target, + icc_profile.as_mut_ptr(), + icc_size, + ) + })?; + + Ok(State::IccProfile(icc_profile)) + } else { + let mut color_encoding = MaybeUninit::uninit(); + + check_dec_status(unsafe { + d::JxlDecoderGetColorAsEncodedProfile( + self.dec.ptr, + config.target, + color_encoding.as_mut_ptr(), + ) + })?; + Ok(State::ColorProfile(Box::new(unsafe { + color_encoding.assume_init() + }))) + } + } + + fn get_preview_image(&mut self) -> Result { + let Some(pixel_format) = self.config.preview else { + return Err(DecodeError::InternalError( + "Subscribe to preview image event but without a pixel format!", + )); + }; + + let mut size = 0; + + check_dec_status(unsafe { + d::JxlDecoderPreviewOutBufferSize(self.dec.ptr, &pixel_format, &mut size) + })?; + + let mut buffer = vec![0; size]; + check_dec_status(unsafe { + d::JxlDecoderSetPreviewOutBuffer( + self.dec.ptr, + &pixel_format, + buffer.as_mut_ptr().cast(), + buffer.len(), + ) + })?; + + Ok(State::PreviewImage(Pixels::new(buffer, &pixel_format))) + } + + fn get_frame(&mut self) -> Result { + let mut header = MaybeUninit::uninit(); + check_dec_status(unsafe { + d::JxlDecoderGetFrameHeader(self.dec.ptr, header.as_mut_ptr()) + })?; + let header = unsafe { header.assume_init() }; + + let mut buffer = vec![0; header.name_length as usize + 1]; + check_dec_status(unsafe { + d::JxlDecoderGetFrameName(self.dec.ptr, buffer.as_mut_ptr().cast(), buffer.len()) + })?; + let name = CString::from_vec_with_nul(buffer) + .map_err(|_| DecodeError::InternalError("Invalid frame name"))?; + + Ok(State::Frame) + } } impl<'dec, 'pr, 'mm> Iterator for Session<'dec, 'pr, 'mm> {