use std::{
collections::HashMap,
sync::{Arc, Condvar, Mutex, MutexGuard},
};
use jxl_frame::data::{HfGlobal, LfGlobal, LfGroup};
use jxl_modular::{ChannelShift, Sample};
use crate::{
Error, ImageWithRegion, IndexedFrame, Reference, Region, Result, image::RenderedImage,
};
pub type RenderOp<S> =
Arc<dyn Fn(FrameRender<S>, Region) -> FrameRender<S> + Send + Sync + 'static>;
#[derive(Debug)]
pub struct RenderCache<S: Sample> {
pub(crate) lf_global: Option<LfGlobal<S>>,
pub(crate) hf_global: Option<HfGlobal>,
pub(crate) lf_groups: HashMap<u32, LfGroup<S>>,
}
impl<S: Sample> RenderCache<S> {
pub fn new(frame: &crate::IndexedFrame) -> Self {
let frame_header = frame.header();
let jpeg_upsampling = frame_header.jpeg_upsampling;
let shifts_cbycr: [_; 3] =
std::array::from_fn(|idx| ChannelShift::from_jpeg_upsampling(jpeg_upsampling, idx));
let lf_width = frame_header.color_sample_width().div_ceil(8);
let lf_height = frame_header.color_sample_height().div_ceil(8);
let mut whd = [(lf_width, lf_height); 3];
for ((w, h), shift) in whd.iter_mut().zip(shifts_cbycr) {
let (shift_w, shift_h) = shift.shift_size((lf_width, lf_height));
*w = shift_w;
*h = shift_h;
}
Self {
lf_global: None,
hf_global: None,
lf_groups: HashMap::new(),
}
}
}
#[derive(Default)]
pub enum FrameRender<S: Sample> {
#[default]
None,
Rendering,
InProgress(Box<RenderCache<S>>),
Done(ImageWithRegion),
Blended(Arc<ImageWithRegion>),
Err(crate::Error),
ErrTaken,
}
impl<S: Sample> std::fmt::Debug for FrameRender<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::None => write!(f, "None"),
Self::Rendering => write!(f, "Rendering"),
Self::InProgress(_) => write!(f, "InProgress(_)"),
Self::Done(_) => write!(f, "Done(_)"),
Self::Blended(_) => write!(f, "Blended(_)"),
Self::Err(e) => f.debug_tuple("Err").field(e).finish(),
Self::ErrTaken => write!(f, "ErrTaken"),
}
}
}
pub struct FrameRenderHandle<S: Sample> {
pub(crate) frame: Arc<IndexedFrame>,
pub(crate) image_region: Region,
pub(crate) render: Mutex<FrameRender<S>>,
pub(crate) condvar: Condvar,
pub(crate) render_op: RenderOp<S>,
pub(crate) refs: [Option<Reference<S>>; 4],
}
impl<S: Sample> std::fmt::Debug for FrameRenderHandle<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FrameRenderHandle")
.field("render", &self.render)
.field("condvar", &self.condvar)
.finish_non_exhaustive()
}
}
impl<S: Sample> FrameRenderHandle<S> {
#[inline]
pub fn new(
frame: Arc<IndexedFrame>,
image_region: Region,
render_op: RenderOp<S>,
refs: [Option<Reference<S>>; 4],
) -> Self {
Self {
frame,
image_region,
render: Mutex::new(FrameRender::None),
condvar: Condvar::new(),
render_op,
refs,
}
}
#[inline]
pub fn from_cache(
frame: Arc<IndexedFrame>,
image_region: Region,
cache: RenderCache<S>,
render_op: RenderOp<S>,
refs: [Option<Reference<S>>; 4],
) -> Self {
let render = FrameRender::InProgress(Box::new(cache));
Self {
frame,
image_region,
render: Mutex::new(render),
condvar: Condvar::new(),
render_op,
refs,
}
}
pub fn run_with_image(self: Arc<Self>) -> Result<RenderedImage<S>> {
let render = if let Some(state) = self.start_render()? {
let _guard = tracing::trace_span!("Run with image", index = self.frame.idx).entered();
let render_result = (self.render_op)(state, self.image_region);
match render_result {
FrameRender::InProgress(_) => {
drop(self.done_render(render_result));
return Err(Error::IncompleteFrame);
}
FrameRender::Err(e) => {
drop(self.done_render(FrameRender::ErrTaken));
return Err(e);
}
_ => {}
}
self.done_render(render_result)
} else {
self.wait_until_render()?
};
drop(render);
Ok(RenderedImage::new(self))
}
pub fn run(&self, image_region: Region) {
if let Some(state) = self.start_render_silent() {
let _guard = tracing::trace_span!("Run", index = self.frame.idx).entered();
let render_result = (self.render_op)(state, image_region);
drop(self.done_render(render_result));
}
}
pub fn reset(&self) -> FrameRender<S> {
let mut render_ref = self.render.lock().unwrap();
std::mem::replace(&mut *render_ref, FrameRender::None)
}
fn start_render(&self) -> Result<Option<FrameRender<S>>> {
let mut render_ref = self.render.lock().unwrap();
let render = std::mem::replace(&mut *render_ref, FrameRender::Rendering);
match render {
FrameRender::None | FrameRender::InProgress(_) => Ok(Some(render)),
FrameRender::Err(e) => {
*render_ref = FrameRender::ErrTaken;
Err(e)
}
FrameRender::ErrTaken => {
*render_ref = FrameRender::ErrTaken;
Err(Error::FailedReference)
}
render => {
*render_ref = render;
Ok(None)
}
}
}
fn start_render_silent(&self) -> Option<FrameRender<S>> {
let mut render_ref = self.render.lock().unwrap();
let render = std::mem::replace(&mut *render_ref, FrameRender::Rendering);
match render {
FrameRender::None | FrameRender::InProgress(_) => Some(render),
render => {
*render_ref = render;
None
}
}
}
pub(crate) fn wait_until_render(&self) -> Result<MutexGuard<'_, FrameRender<S>>> {
let mut render_ref = self.render.lock().unwrap();
loop {
let render = std::mem::replace(&mut *render_ref, FrameRender::None);
match render {
FrameRender::Rendering => {
tracing::trace!(index = self.frame.idx, "Waiting...");
*render_ref = render;
render_ref = self.condvar.wait(render_ref).unwrap();
}
FrameRender::Done(_) | FrameRender::Blended(_) => {
*render_ref = render;
return Ok(render_ref);
}
FrameRender::None | FrameRender::InProgress(_) => {
return Err(Error::IncompleteFrame);
}
FrameRender::Err(e) => return Err(e),
FrameRender::ErrTaken => return Err(Error::FailedReference),
}
}
}
pub(crate) fn done_render(&self, render: FrameRender<S>) -> MutexGuard<'_, FrameRender<S>> {
assert!(!matches!(render, FrameRender::Rendering));
let mut guard = self.render.lock().unwrap();
*guard = render;
self.condvar.notify_all();
guard
}
}