1use std::sync::Arc;
2
3use jxl_frame::{FrameHeader, data::GlobalModular};
4use jxl_grid::{AlignedGrid, AllocTracker, MutableSubgrid};
5use jxl_image::{BitDepth, ImageHeader};
6use jxl_modular::{ChannelShift, Sample};
7use jxl_threadpool::JxlThreadPool;
8use jxl_vardct::LfChannelDequantization;
9
10use crate::{FrameRender, FrameRenderHandle, IndexedFrame, Region, Result, util};
11
12#[derive(Debug)]
13pub enum ImageBuffer {
14 F32(AlignedGrid<f32>),
15 I32(AlignedGrid<i32>),
16 I16(AlignedGrid<i16>),
17}
18
19impl ImageBuffer {
20 #[inline]
21 pub fn from_modular_channel<S: Sample>(g: AlignedGrid<S>) -> Self {
22 let g = match S::try_into_grid_i16(g) {
23 Ok(g) => return Self::I16(g),
24 Err(g) => g,
25 };
26 match S::try_into_grid_i32(g) {
27 Ok(g) => Self::I32(g),
28 Err(_) => panic!(),
29 }
30 }
31
32 #[inline]
33 pub fn zeroed_f32(width: usize, height: usize, tracker: Option<&AllocTracker>) -> Result<Self> {
34 let grid = AlignedGrid::with_alloc_tracker(width, height, tracker)?;
35 Ok(Self::F32(grid))
36 }
37
38 #[inline]
39 pub fn try_clone(&self) -> Result<Self> {
40 match self {
41 Self::F32(g) => g.try_clone().map(Self::F32),
42 Self::I32(g) => g.try_clone().map(Self::I32),
43 Self::I16(g) => g.try_clone().map(Self::I16),
44 }
45 .map_err(From::from)
46 }
47
48 #[inline]
49 pub fn width(&self) -> usize {
50 match self {
51 Self::F32(g) => g.width(),
52 Self::I32(g) => g.width(),
53 Self::I16(g) => g.width(),
54 }
55 }
56
57 #[inline]
58 pub fn height(&self) -> usize {
59 match self {
60 Self::F32(g) => g.height(),
61 Self::I32(g) => g.height(),
62 Self::I16(g) => g.height(),
63 }
64 }
65
66 #[inline]
67 pub fn tracker(&self) -> Option<AllocTracker> {
68 match self {
69 Self::F32(g) => g.tracker(),
70 Self::I32(g) => g.tracker(),
71 Self::I16(g) => g.tracker(),
72 }
73 }
74
75 #[inline]
76 pub fn as_float(&self) -> Option<&AlignedGrid<f32>> {
77 if let Self::F32(g) = self {
78 Some(g)
79 } else {
80 None
81 }
82 }
83
84 #[inline]
85 pub fn as_float_mut(&mut self) -> Option<&mut AlignedGrid<f32>> {
86 if let Self::F32(g) = self {
87 Some(g)
88 } else {
89 None
90 }
91 }
92
93 pub fn convert_to_float_modular(
94 &mut self,
95 bit_depth: BitDepth,
96 ) -> Result<&mut AlignedGrid<f32>> {
97 Ok(match self {
98 Self::F32(g) => g,
99 Self::I32(g) => {
100 let mut out =
101 AlignedGrid::with_alloc_tracker(g.width(), g.height(), g.tracker().as_ref())?;
102 for (o, &i) in out.buf_mut().iter_mut().zip(g.buf()) {
103 *o = bit_depth.parse_integer_sample(i);
104 }
105
106 *self = Self::F32(out);
107 self.as_float_mut().unwrap()
108 }
109 Self::I16(g) => {
110 let mut out =
111 AlignedGrid::with_alloc_tracker(g.width(), g.height(), g.tracker().as_ref())?;
112 for (o, &i) in out.buf_mut().iter_mut().zip(g.buf()) {
113 *o = bit_depth.parse_integer_sample(i as i32);
114 }
115
116 *self = Self::F32(out);
117 self.as_float_mut().unwrap()
118 }
119 })
120 }
121
122 pub fn cast_to_float(&mut self) -> Result<&mut AlignedGrid<f32>> {
123 Ok(match self {
124 Self::F32(g) => g,
125 Self::I32(g) => {
126 let mut out =
127 AlignedGrid::with_alloc_tracker(g.width(), g.height(), g.tracker().as_ref())?;
128 for (o, &i) in out.buf_mut().iter_mut().zip(g.buf()) {
129 *o = i as f32;
130 }
131
132 *self = Self::F32(out);
133 self.as_float_mut().unwrap()
134 }
135 Self::I16(g) => {
136 let mut out =
137 AlignedGrid::with_alloc_tracker(g.width(), g.height(), g.tracker().as_ref())?;
138 for (o, &i) in out.buf_mut().iter_mut().zip(g.buf()) {
139 *o = i as f32;
140 }
141
142 *self = Self::F32(out);
143 self.as_float_mut().unwrap()
144 }
145 })
146 }
147
148 pub fn convert_to_float_modular_xyb<'g>(
149 yxb: [&'g mut Self; 3],
150 lf_dequant: &LfChannelDequantization,
151 ) -> Result<[&'g mut AlignedGrid<f32>; 3]> {
152 match yxb {
153 [Self::F32(_), Self::F32(_), Self::F32(_)] => {
154 panic!("channels are already converted");
155 }
156 [Self::I32(y), Self::I32(_), Self::I32(b)] => {
157 for (b, &y) in b.buf_mut().iter_mut().zip(y.buf()) {
158 *b = b.saturating_add(y);
159 }
160 }
161 [Self::I16(y), Self::I16(_), Self::I16(b)] => {
162 for (b, &y) in b.buf_mut().iter_mut().zip(y.buf()) {
163 *b = b.saturating_add(y);
164 }
165 }
166 _ => panic!(),
167 }
168
169 let [y, x, b] = yxb;
170 let y = y.cast_to_float()?;
171 let x = x.cast_to_float()?;
172 let b = b.cast_to_float()?;
173 let buf_y = y.buf_mut();
174 let buf_x = x.buf_mut();
175 let buf_b = b.buf_mut();
176 let m_x_lf = lf_dequant.m_x_lf_unscaled();
177 let m_y_lf = lf_dequant.m_y_lf_unscaled();
178 let m_b_lf = lf_dequant.m_b_lf_unscaled();
179
180 for ((y, x), b) in buf_y.iter_mut().zip(buf_x).zip(buf_b) {
181 let py = *y;
182 let px = *x;
183 *y = px * m_x_lf;
184 *x = py * m_y_lf;
185 *b *= m_b_lf;
186 }
187
188 Ok([y, x, b])
189 }
190
191 pub(crate) fn upsample_nn(&self, factor: u32) -> Result<ImageBuffer> {
192 #[inline]
193 fn inner<S: Copy>(
194 original: &[S],
195 target: &mut [S],
196 width: usize,
197 height: usize,
198 factor: u32,
199 ) {
200 assert_eq!(original.len(), width * height);
201 assert_eq!(target.len(), original.len() << (factor * 2));
202 let step = 1usize << factor;
203 let stride = width << factor;
204 for y in 0..height {
205 let original = &original[y * width..];
206 let target = &mut target[y * step * stride..];
207 for (x, &value) in original[..width].iter().enumerate() {
208 target[x * step..][..step].fill(value);
209 }
210 for row in 1..step {
211 target.copy_within(..stride, stride * row);
212 }
213 }
214 }
215
216 if factor == 0 {
217 return self.try_clone();
218 }
219
220 let tracker = self.tracker();
221 let width = self.width();
222 let height = self.height();
223 Ok(match self {
224 Self::F32(g) => {
225 let up_width = width << factor;
226 let up_height = height << factor;
227 let mut out =
228 AlignedGrid::with_alloc_tracker(up_width, up_height, tracker.as_ref())?;
229
230 let original = g.buf();
231 let target = out.buf_mut();
232 inner(original, target, width, height, factor);
233 Self::F32(out)
234 }
235 Self::I32(g) => {
236 let up_width = width << factor;
237 let up_height = height << factor;
238 let mut out =
239 AlignedGrid::with_alloc_tracker(up_width, up_height, tracker.as_ref())?;
240
241 let original = g.buf();
242 let target = out.buf_mut();
243 inner(original, target, width, height, factor);
244 Self::I32(out)
245 }
246 Self::I16(g) => {
247 let up_width = width << factor;
248 let up_height = height << factor;
249 let mut out =
250 AlignedGrid::with_alloc_tracker(up_width, up_height, tracker.as_ref())?;
251
252 let original = g.buf();
253 let target = out.buf_mut();
254 inner(original, target, width, height, factor);
255 Self::I16(out)
256 }
257 })
258 }
259}
260
261#[derive(Debug)]
262pub struct ImageWithRegion {
263 buffer: Vec<ImageBuffer>,
264 regions: Vec<(Region, ChannelShift)>,
265 color_channels: usize,
266 ct_done: bool,
267 blend_done: bool,
268 tracker: Option<AllocTracker>,
269}
270
271impl ImageWithRegion {
272 pub(crate) fn new(color_channels: usize, tracker: Option<&AllocTracker>) -> Self {
273 Self {
274 buffer: Vec::new(),
275 regions: Vec::new(),
276 color_channels,
277 ct_done: false,
278 blend_done: false,
279 tracker: tracker.cloned(),
280 }
281 }
282
283 pub(crate) fn try_clone(&self) -> Result<Self> {
284 Ok(Self {
285 buffer: self
286 .buffer
287 .iter()
288 .map(|x| x.try_clone())
289 .collect::<std::result::Result<_, _>>()?,
290 regions: self.regions.clone(),
291 color_channels: self.color_channels,
292 ct_done: self.ct_done,
293 blend_done: false,
294 tracker: self.tracker.clone(),
295 })
296 }
297
298 #[inline]
299 pub(crate) fn alloc_tracker(&self) -> Option<&AllocTracker> {
300 self.tracker.as_ref()
301 }
302
303 #[inline]
304 pub fn channels(&self) -> usize {
305 self.buffer.len()
306 }
307
308 #[inline]
309 pub fn buffer(&self) -> &[ImageBuffer] {
310 &self.buffer
311 }
312
313 #[inline]
314 pub fn buffer_mut(&mut self) -> &mut [ImageBuffer] {
315 &mut self.buffer
316 }
317
318 #[inline]
319 pub fn take_buffer(&mut self) -> Vec<ImageBuffer> {
320 std::mem::take(&mut self.buffer)
321 }
322
323 #[inline]
324 pub fn regions_and_shifts(&self) -> &[(Region, ChannelShift)] {
325 &self.regions
326 }
327
328 #[inline]
329 pub fn append_channel(&mut self, buffer: ImageBuffer, region: Region) {
330 assert_eq!(buffer.width(), region.width as usize);
331 assert_eq!(buffer.height(), region.height as usize);
332 self.buffer.push(buffer);
333 self.regions.push((region, ChannelShift::from_shift(0)));
334 }
335
336 #[inline]
337 pub fn append_channel_shifted(
338 &mut self,
339 buffer: ImageBuffer,
340 original_region: Region,
341 shift: ChannelShift,
342 ) {
343 let (width, height) = shift.shift_size((original_region.width, original_region.height));
344 assert_eq!(buffer.width(), width as usize);
345 assert_eq!(buffer.height(), height as usize);
346 self.buffer.push(buffer);
347 self.regions.push((original_region, shift));
348 }
349
350 #[inline]
351 pub fn replace_channel(&mut self, index: usize, buffer: ImageBuffer, region: Region) {
352 assert_eq!(buffer.width(), region.width as usize);
353 assert_eq!(buffer.height(), region.height as usize);
354 self.buffer[index] = buffer;
355 self.regions[index] = (region, ChannelShift::from_shift(0));
356 }
357
358 #[inline]
359 pub fn replace_channel_shifted(
360 &mut self,
361 index: usize,
362 buffer: ImageBuffer,
363 original_region: Region,
364 shift: ChannelShift,
365 ) {
366 let (width, height) = shift.shift_size((original_region.width, original_region.height));
367 assert_eq!(buffer.width(), width as usize);
368 assert_eq!(buffer.height(), height as usize);
369 self.buffer[index] = buffer;
370 self.regions[index] = (original_region, shift);
371 }
372
373 #[inline]
374 pub(crate) fn swap_channel_f32(
375 &mut self,
376 index: usize,
377 buffer: &mut AlignedGrid<f32>,
378 region: Region,
379 ) {
380 assert_eq!(buffer.width(), region.width as usize);
381 assert_eq!(buffer.height(), region.height as usize);
382 let ImageBuffer::F32(original_buffer) = &mut self.buffer[index] else {
383 panic!("original buffer is not F32");
384 };
385 std::mem::swap(original_buffer, buffer);
386 self.regions[index] = (region, ChannelShift::from_shift(0));
387 }
388
389 pub fn extend_from_gmodular<S: Sample>(&mut self, gmodular: GlobalModular<S>) {
390 let Some(image) = gmodular.modular.into_image() else {
391 return;
392 };
393 for (g, info) in image.into_image_channels_with_info() {
394 let (width, height) = info.original_size();
395 let shift = info.shift();
396
397 let original_region = Region::with_size(width, height);
398 let buffer = ImageBuffer::from_modular_channel(g);
399 self.append_channel_shifted(buffer, original_region, shift);
400 }
401 }
402
403 pub(crate) fn clone_gray(&mut self) -> Result<()> {
404 assert_eq!(self.color_channels, 1);
405
406 let gray = self.buffer[0].try_clone()?;
407 let region = self.regions[0];
408 self.buffer.insert(1, gray.try_clone()?);
409 self.regions.insert(1, region);
410 self.buffer.insert(2, gray);
411 self.regions.insert(2, region);
412
413 self.color_channels = 3;
414 Ok(())
415 }
416
417 pub(crate) fn convert_modular_color(&mut self, bit_depth: BitDepth) -> Result<()> {
418 assert!(self.buffer.len() >= self.color_channels);
419 for g in self.buffer.iter_mut().take(self.color_channels) {
420 g.convert_to_float_modular(bit_depth)?;
421 }
422 Ok(())
423 }
424
425 pub(crate) fn convert_modular_xyb(
426 &mut self,
427 lf_dequant: &LfChannelDequantization,
428 ) -> Result<()> {
429 assert_eq!(self.color_channels, 3);
430 let [y, x, b, ..] = &mut *self.buffer else {
431 panic!()
432 };
433 ImageBuffer::convert_to_float_modular_xyb([y, x, b], lf_dequant)?;
434 Ok(())
435 }
436
437 pub(crate) fn upsample_lf(&self, lf_level: u32) -> Result<Self> {
438 let factor = lf_level * 3;
439 let mut out = Self::new(self.color_channels, self.tracker.as_ref());
440 for (&(region, shift), buffer) in self.regions.iter().zip(&self.buffer) {
441 let up_region = region.upsample(factor);
442 let buffer = buffer.upsample_nn(factor)?;
443 out.append_channel_shifted(buffer, up_region, shift);
444 }
445 Ok(out)
446 }
447
448 pub(crate) fn upsample_jpeg(
449 &mut self,
450 valid_region: Region,
451 bit_depth: BitDepth,
452 ) -> Result<()> {
453 assert_eq!(self.color_channels, 3);
454 self.convert_modular_color(bit_depth)?;
455
456 for (g, (region, shift)) in self.buffer.iter_mut().zip(&mut self.regions).take(3) {
457 let downsampled_image_region = region.downsample_with_shift(*shift);
458 let downsampled_valid_region = valid_region.downsample_with_shift(*shift);
459 let left = downsampled_valid_region
460 .left
461 .abs_diff(downsampled_image_region.left);
462 let top = downsampled_valid_region
463 .top
464 .abs_diff(downsampled_image_region.top);
465 let width = downsampled_valid_region.width;
466 let height = downsampled_valid_region.height;
467
468 let subgrid = g.as_float().unwrap().as_subgrid().subgrid(
469 left as usize..(left + width) as usize,
470 top as usize..(top + height) as usize,
471 );
472 let out = crate::filter::apply_jpeg_upsampling_single(
473 subgrid,
474 *shift,
475 valid_region,
476 self.tracker.as_ref(),
477 )?;
478
479 *g = ImageBuffer::F32(out);
480 *region = valid_region;
481 *shift = ChannelShift::from_shift(0);
482 }
483
484 Ok(())
485 }
486
487 pub(crate) fn upsample_nonseparable(
488 &mut self,
489 image_header: &ImageHeader,
490 frame_header: &FrameHeader,
491 upsampled_valid_region: Region,
492 ec_to_color_only: bool,
493 ) -> Result<()> {
494 if frame_header.upsampling != 1 && self.buffer[0].as_float().is_none() {
495 debug_assert!(!image_header.metadata.xyb_encoded);
496 }
497
498 let color_channels = self.color_channels;
499 let color_shift = frame_header.upsampling.trailing_zeros();
500
501 for (idx, ((region, shift), g)) in self.regions.iter_mut().zip(&mut self.buffer).enumerate()
502 {
503 let tracker = g.tracker();
504 let ChannelShift::Shifts(upsampling_factor) = *shift else {
505 panic!("invalid channel shift for upsampling: {shift:?}");
506 };
507 let bit_depth = if let Some(ec_idx) = idx.checked_sub(color_channels) {
508 image_header.metadata.ec_info[ec_idx].bit_depth
509 } else {
510 image_header.metadata.bit_depth
511 };
512
513 let target_factor = if ec_to_color_only { color_shift } else { 0 };
514 if upsampling_factor == target_factor {
515 continue;
516 }
517 let grid = g.convert_to_float_modular(bit_depth)?;
518
519 let downsampled_image_region = region.downsample(upsampling_factor);
520 let downsampled_valid_region = upsampled_valid_region.downsample(upsampling_factor);
521 let left = downsampled_valid_region
522 .left
523 .abs_diff(downsampled_image_region.left);
524 let top = downsampled_valid_region
525 .top
526 .abs_diff(downsampled_image_region.top);
527 let width = downsampled_valid_region.width;
528 let height = downsampled_valid_region.height;
529
530 let subgrid = grid.as_subgrid().subgrid(
531 left as usize..(left + width) as usize,
532 top as usize..(top + height) as usize,
533 );
534 *region = downsampled_valid_region;
535 let out = tracing::trace_span!(
536 "Non-separable upsampling",
537 channel_idx = idx,
538 upsampling_factor,
539 target_factor
540 )
541 .in_scope(|| {
542 crate::features::upsample(
543 subgrid,
544 region,
545 image_header,
546 upsampling_factor - target_factor,
547 tracker.as_ref(),
548 )
549 })?;
550 if let Some(out) = out {
551 *g = ImageBuffer::F32(out);
552 }
553 *shift = ChannelShift::from_shift(target_factor);
554 }
555
556 Ok(())
557 }
558
559 pub(crate) fn prepare_color_upsampling(&mut self, frame_header: &FrameHeader) {
560 let upsampling_factor = frame_header.upsampling.trailing_zeros();
561 for (region, shift) in &mut self.regions {
562 match shift {
563 ChannelShift::Raw(..=-1, _) | ChannelShift::Raw(_, ..=-1) => continue,
564 ChannelShift::Raw(h, v) => {
565 *h = h.wrapping_add_unsigned(upsampling_factor);
566 *v = v.wrapping_add_unsigned(upsampling_factor);
567 }
568 ChannelShift::Shifts(shift) => {
569 *shift += upsampling_factor;
570 }
571 ChannelShift::JpegUpsampling {
572 has_h_subsample: false,
573 h_subsample: false,
574 has_v_subsample: false,
575 v_subsample: false,
576 } => {
577 *shift = ChannelShift::Shifts(upsampling_factor);
578 }
579 ChannelShift::JpegUpsampling { .. } => {
580 panic!("unexpected chroma subsampling {shift:?}");
581 }
582 }
583 *region = region.upsample(upsampling_factor);
584 }
585 }
586
587 #[inline]
588 pub(crate) fn remove_color_channels(&mut self, count: usize) {
589 assert!(self.color_channels >= count);
590 self.buffer.drain(count..self.color_channels);
591 self.regions.drain(count..self.color_channels);
592 self.color_channels = count;
593 }
594
595 pub(crate) fn fill_opaque_alpha(&mut self, ec_info: &[jxl_image::ExtraChannelInfo]) {
596 for (ec_idx, ec_info) in ec_info.iter().enumerate() {
597 if !matches!(ec_info.ty, jxl_image::ExtraChannelType::Alpha { .. }) {
598 continue;
599 }
600
601 let channel = &mut self.buffer[self.color_channels + ec_idx];
602 let opaque_int = match ec_info.bit_depth {
603 BitDepth::IntegerSample { bits_per_sample } => (1u32 << bits_per_sample) - 1,
604 BitDepth::FloatSample {
605 bits_per_sample,
606 exp_bits,
607 } => {
608 let mantissa_bits = bits_per_sample - exp_bits - 1;
609 ((1u32 << (exp_bits - 1)) - 1) << mantissa_bits
610 }
611 };
612 match channel {
613 ImageBuffer::I16(g) => {
614 g.buf_mut().fill(opaque_int as i16);
615 }
616 ImageBuffer::I32(g) => {
617 g.buf_mut().fill(opaque_int as i32);
618 }
619 ImageBuffer::F32(g) => {
620 g.buf_mut().fill(1.0);
621 }
622 }
623 }
624 }
625
626 #[inline]
627 pub fn color_channels(&self) -> usize {
628 self.color_channels
629 }
630
631 #[inline]
632 pub(crate) fn ct_done(&self) -> bool {
633 self.ct_done
634 }
635
636 #[inline]
637 pub(crate) fn set_ct_done(&mut self, ct_done: bool) {
638 self.ct_done = ct_done;
639 }
640
641 #[inline]
642 pub(crate) fn set_blend_done(&mut self, blend_done: bool) {
643 self.blend_done = blend_done;
644 }
645}
646
647impl ImageWithRegion {
648 pub(crate) fn as_color_floats(&self) -> [&AlignedGrid<f32>; 3] {
649 assert_eq!(self.color_channels, 3);
650 let [a, b, c, ..] = &*self.buffer else {
651 panic!()
652 };
653 let a = a.as_float().unwrap();
654 let b = b.as_float().unwrap();
655 let c = c.as_float().unwrap();
656 [a, b, c]
657 }
658
659 pub(crate) fn as_color_floats_mut(&mut self) -> [&mut AlignedGrid<f32>; 3] {
660 assert_eq!(self.color_channels, 3);
661 let [a, b, c, ..] = &mut *self.buffer else {
662 panic!()
663 };
664 let a = a.as_float_mut().unwrap();
665 let b = b.as_float_mut().unwrap();
666 let c = c.as_float_mut().unwrap();
667 [a, b, c]
668 }
669
670 pub(crate) fn color_groups_with_group_id(
671 &mut self,
672 frame_header: &jxl_frame::FrameHeader,
673 ) -> Vec<(u32, [MutableSubgrid<'_, f32>; 3])> {
674 assert_eq!(self.color_channels, 3);
675 let [fb_x, fb_y, fb_b, ..] = &mut *self.buffer else {
676 panic!();
677 };
678
679 let group_dim = frame_header.group_dim();
680 let base_group_x = self.regions[0].0.left as u32 / group_dim;
681 let base_group_y = self.regions[0].0.top as u32 / group_dim;
682 let width = self.regions[0].0.width;
683 let frame_groups_per_row = frame_header.groups_per_row();
684 let groups_per_row = width.div_ceil(group_dim);
685
686 let [fb_x, fb_y, fb_b] = [(0usize, fb_x), (1, fb_y), (2, fb_b)].map(|(idx, fb)| {
687 let fb = fb.as_float_mut().unwrap().as_subgrid_mut();
688 let hshift = self.regions[idx].1.hshift();
689 let vshift = self.regions[idx].1.vshift();
690 let group_dim = group_dim as usize;
691 fb.into_groups(group_dim >> hshift, group_dim >> vshift)
692 });
693
694 fb_x.into_iter()
695 .zip(fb_y)
696 .zip(fb_b)
697 .enumerate()
698 .map(|(idx, ((fb_x, fb_y), fb_b))| {
699 let idx = idx as u32;
700 let group_x = base_group_x + (idx % groups_per_row);
701 let group_y = base_group_y + (idx / groups_per_row);
702 let group_idx = group_y * frame_groups_per_row + group_x;
703 (group_idx, [fb_x, fb_y, fb_b])
704 })
705 .collect()
706 }
707}
708
709pub struct RenderedImage<S: Sample> {
710 image: Arc<FrameRenderHandle<S>>,
711}
712
713impl<S: Sample> RenderedImage<S> {
714 pub(crate) fn new(image: Arc<FrameRenderHandle<S>>) -> Self {
715 Self { image }
716 }
717}
718
719impl<S: Sample> RenderedImage<S> {
720 pub(crate) fn blend(
721 &self,
722 oriented_image_region: Option<Region>,
723 pool: &JxlThreadPool,
724 ) -> Result<Arc<ImageWithRegion>> {
725 let image_header = self.image.frame.image_header();
726 let image_region = self.image.image_region;
727 let oriented_image_region = oriented_image_region
728 .unwrap_or_else(|| util::apply_orientation_to_image_region(image_header, image_region));
729
730 let mut grid_lock = self.image.wait_until_render()?;
731 if let FrameRender::Blended(image) = &*grid_lock {
732 return Ok(Arc::clone(image));
733 }
734
735 let render = std::mem::replace(&mut *grid_lock, FrameRender::ErrTaken);
736 let FrameRender::Done(mut grid) = render else {
737 panic!();
738 };
739
740 let skip_composition = composite_preprocess(&self.image.frame, &mut grid, pool)?;
741 if skip_composition {
742 let image = Arc::new(grid);
743 *grid_lock = FrameRender::Blended(Arc::clone(&image));
744 return Ok(image);
745 }
746
747 *grid_lock = FrameRender::Rendering;
748 drop(grid_lock);
749
750 composite(
751 &self.image.frame,
752 &mut grid,
753 self.image.refs.clone(),
754 oriented_image_region,
755 pool,
756 )?;
757
758 let image = Arc::new(grid);
759 drop(
760 self.image
761 .done_render(FrameRender::Blended(Arc::clone(&image))),
762 );
763 Ok(image)
764 }
765
766 pub(crate) fn try_take_blended(&self) -> Option<ImageWithRegion> {
767 let mut grid_lock = self.image.render.lock().unwrap();
768 match std::mem::take(&mut *grid_lock) {
769 FrameRender::Blended(image) => {
770 let cloned_image = Arc::clone(&image);
771 let image_inner = Arc::into_inner(cloned_image);
772 if image_inner.is_none() {
773 *grid_lock = FrameRender::Blended(image);
774 }
775 image_inner
776 }
777 render => {
778 *grid_lock = render;
779 None
780 }
781 }
782 }
783}
784
785pub(crate) fn composite_preprocess(
789 frame: &IndexedFrame,
790 grid: &mut ImageWithRegion,
791 pool: &JxlThreadPool,
792) -> Result<bool> {
793 let image_header = frame.image_header();
794 let frame_header = frame.header();
795
796 if frame_header.can_reference() {
797 let bit_depth_it =
798 std::iter::repeat_n(image_header.metadata.bit_depth, grid.color_channels)
799 .chain(image_header.metadata.ec_info.iter().map(|ec| ec.bit_depth));
800 for (buffer, bit_depth) in grid.buffer.iter_mut().zip(bit_depth_it) {
801 buffer.convert_to_float_modular(bit_depth)?;
802 }
803 }
804
805 let skip_blending = !frame_header.frame_type.is_normal_frame() || frame_header.resets_canvas;
806
807 if !(grid.ct_done() || frame_header.save_before_ct || skip_blending && frame_header.is_last) {
808 util::convert_color_for_record(image_header, frame_header.do_ycbcr, grid, pool)?;
809 }
810
811 if skip_blending {
812 grid.blend_done = true;
813 }
814
815 Ok(skip_blending)
816}
817
818pub(crate) fn composite<S: Sample>(
819 frame: &IndexedFrame,
820 grid: &mut ImageWithRegion,
821 refs: [Option<crate::Reference<S>>; 4],
822 oriented_image_region: Region,
823 pool: &JxlThreadPool,
824) -> Result<()> {
825 let image_header = frame.image_header();
826 let frame_header = frame.header();
827 let frame_region = oriented_image_region
828 .translate(-frame_header.x0, -frame_header.y0)
829 .downsample(frame_header.lf_level * 3);
830 let frame_region = util::pad_lf_region(frame_header, frame_region);
831 let frame_region = util::pad_color_region(image_header, frame_header, frame_region);
832 let frame_region = frame_region.upsample(frame_header.upsampling.ilog2());
833 let frame_region = if frame_header.frame_type.is_normal_frame() {
834 let full_image_region_in_frame =
835 Region::with_size(image_header.size.width, image_header.size.height)
836 .translate(-frame_header.x0, -frame_header.y0);
837 frame_region.intersection(full_image_region_in_frame)
838 } else {
839 frame_region
840 };
841
842 let image = crate::blend::blend(image_header, refs, frame, grid, frame_region, pool)?;
843 *grid = image;
844 Ok(())
845}