1use crate::Result;
2use jxl_bitstream::{Bitstream, U};
3use jxl_image::{BitDepth, Extensions, ImageHeader, SizeHeader};
4use jxl_oxide_common::{define_bundle, Bundle, Name};
5
6define_bundle! {
7 #[derive(Debug)]
9 pub struct FrameHeader ctx(headers: &ImageHeader) error(crate::Error) {
10 all_default: ty(Bool) default(true),
11 pub frame_type: ty(Bundle(FrameType)) cond(!all_default) default(FrameType::RegularFrame),
12 pub encoding: ty(Bundle(Encoding)) cond(!all_default) default(Encoding::VarDct),
13 pub flags: ty(Bundle(FrameFlags)) cond(!all_default),
14 pub do_ycbcr: ty(Bool) cond(!all_default && !headers.metadata.xyb_encoded),
15 encoded_color_channels:
16 ty(u(0))
17 cond(false)
18 default({
19 let acutally_grayscale = encoding == Encoding::Modular
20 && !do_ycbcr
21 && !headers.metadata.xyb_encoded
22 && headers.metadata.grayscale();
23 if acutally_grayscale { 1 } else { 3 }
24 }),
25 pub jpeg_upsampling: ty(Array[u(2)]; 3) cond(do_ycbcr && !flags.use_lf_frame()),
26 pub upsampling: ty(U32(1, 2, 4, 8)) cond(!all_default && !flags.use_lf_frame()) default(1),
27 pub ec_upsampling:
28 ty(Vec[U32(1, 2, 4, 8)]; headers.metadata.ec_info.len())
29 cond(!all_default && !flags.use_lf_frame())
30 default(vec![1; headers.metadata.ec_info.len()]),
31 pub group_size_shift: ty(u(2)) cond(encoding == Encoding::Modular) default(1),
32 pub x_qm_scale:
33 ty(u(3))
34 cond(!all_default && headers.metadata.xyb_encoded && encoding == Encoding::VarDct)
35 default(Self::compute_default_xqms(encoding, headers.metadata.xyb_encoded)),
36 pub b_qm_scale:
37 ty(u(3))
38 cond(!all_default && headers.metadata.xyb_encoded && encoding == Encoding::VarDct)
39 default(2),
40 pub passes:
41 ty(Bundle(Passes))
42 cond(!all_default && frame_type != FrameType::ReferenceOnly),
43 pub lf_level: ty(1 + u(2)) cond(frame_type == FrameType::LfFrame) default(0),
44 pub have_crop: ty(Bool) cond(!all_default && frame_type != FrameType::LfFrame) default(false),
45 pub x0:
46 ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)); UnpackSigned)
47 cond(have_crop && frame_type != FrameType::ReferenceOnly),
48 pub y0:
49 ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)); UnpackSigned)
50 cond(have_crop && frame_type != FrameType::ReferenceOnly),
51 pub width:
52 ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)))
53 cond(have_crop)
54 default(headers.size.width),
55 pub height:
56 ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)))
57 cond(have_crop)
58 default(headers.size.height),
59 pub blending_info:
60 ty(Bundle(BlendingInfo))
61 ctx((
62 !headers.metadata.ec_info.is_empty(),
63 None,
64 CanvasSizeParams {
65 have_crop,
66 x0,
67 y0,
68 width,
69 height,
70 size: &headers.size,
71 },
72 ))
73 cond(!all_default && frame_type.is_normal_frame()),
74 pub ec_blending_info:
75 ty(Vec[Bundle(BlendingInfo)]; headers.metadata.ec_info.len())
76 ctx((
77 !headers.metadata.ec_info.is_empty(),
78 Some(blending_info.mode),
79 CanvasSizeParams {
80 have_crop,
81 x0,
82 y0,
83 width,
84 height,
85 size: &headers.size,
86 },
87 ))
88 cond(!all_default && frame_type.is_normal_frame()),
89 pub duration:
90 ty(U32(0, 1, u(8), u(32)))
91 cond(!all_default && frame_type.is_normal_frame() && headers.metadata.animation.is_some())
92 default(0),
93 pub timecode:
94 ty(u(32))
95 cond(!all_default && frame_type.is_normal_frame() && headers.metadata.animation.as_ref().map(|a| a.have_timecodes).unwrap_or(false))
96 default(0),
97 pub is_last:
98 ty(Bool)
99 cond(!all_default && frame_type.is_normal_frame())
100 default(frame_type == FrameType::RegularFrame),
101 pub save_as_reference:
102 ty(u(2))
103 cond(!all_default && frame_type != FrameType::LfFrame && !is_last)
104 default(0),
105 pub resets_canvas:
106 ty(Bool)
107 cond(false)
108 default(Self::resets_canvas(
109 blending_info.mode,
110 CanvasSizeParams {
111 have_crop,
112 x0,
113 y0,
114 width,
115 height,
116 size: &headers.size,
117 },
118 )),
119 pub save_before_ct:
120 ty(Bool)
121 cond(
122 !all_default && (
123 frame_type == FrameType::ReferenceOnly || (
124 resets_canvas &&
125 (!is_last && (duration == 0 || save_as_reference != 0) && frame_type != FrameType::LfFrame)
126 )
127 )
128 )
129 default(!frame_type.is_normal_frame()),
130 pub name: ty(Bundle(Name)) cond(!all_default),
131 pub restoration_filter: ty(Bundle(RestorationFilter)) ctx(encoding) cond(!all_default),
132 pub extensions: ty(Bundle(Extensions)) cond(!all_default),
133 pub bit_depth: ty(Bundle(BitDepth)) cond(false) default(headers.metadata.bit_depth),
134 }
135
136 #[derive(Debug)]
137 pub struct Passes error(crate::Error) {
138 pub num_passes: ty(U32(1, 2, 3, 4 + u(3))) default(1),
139 pub num_ds: ty(U32(0, 1, 2, 3 + u(1))) cond(num_passes != 1) default(0),
140 pub shift: ty(Vec[u(2)]; num_passes - 1) cond(num_passes != 1) default(vec![0; num_passes as usize - 1]),
141 pub downsample: ty(Vec[U32(1, 2, 4, 8)]; num_ds) cond(num_passes != 1) default(vec![1; num_ds as usize]),
142 pub last_pass: ty(Vec[U32(0, 1, 2, u(3))]; num_ds) cond(num_passes != 1) default(vec![0; num_ds as usize]),
143 }
144
145 #[derive(Debug)]
146 pub struct BlendingInfo ctx(context: (bool, Option<BlendMode>, CanvasSizeParams<'_>)) error(crate::Error) {
147 pub mode: ty(Bundle(BlendMode)),
148 pub alpha_channel:
149 ty(U32(0, 1, 2, 3 + u(3)))
150 cond(context.0 && (mode == BlendMode::Blend || mode == BlendMode::MulAdd))
151 default(0),
152 pub clamp:
153 ty(Bool)
154 cond((context.0 && (mode == BlendMode::Blend || mode == BlendMode::MulAdd)) || mode == BlendMode::Mul)
155 default(false),
156 pub source:
157 ty(u(2))
158 cond(!FrameHeader::resets_canvas(context.1.unwrap_or(mode), context.2))
159 default(0),
160 }
161
162 #[derive(Debug)]
163 pub struct RestorationFilter ctx(encoding: Encoding) error(crate::Error) {
164 all_default: ty(Bool) default(true),
165 pub gab: ty(Bundle(crate::filter::Gabor)) cond(!all_default),
166 pub epf: ty(Bundle(crate::filter::EdgePreservingFilter)) cond(!all_default),
167 pub extensions: ty(Bundle(Extensions)) cond(!all_default),
168 }
169}
170
171#[derive(Copy, Clone)]
172struct CanvasSizeParams<'a> {
173 have_crop: bool,
174 x0: i32,
175 y0: i32,
176 width: u32,
177 height: u32,
178 size: &'a SizeHeader,
179}
180
181impl FrameHeader {
182 fn test_full_image(canvas_size: CanvasSizeParams) -> bool {
183 let CanvasSizeParams {
184 x0,
185 y0,
186 width,
187 height,
188 size,
189 ..
190 } = canvas_size;
191
192 if x0 > 0 || y0 > 0 {
193 return false;
194 }
195
196 let right = x0 as i64 + (width as i64);
197 let bottom = y0 as i64 + (height as i64);
198 (right >= size.width as i64) && (bottom >= size.height as i64)
199 }
200
201 fn resets_canvas(blending_mode: BlendMode, canvas_size: CanvasSizeParams) -> bool {
202 blending_mode == BlendMode::Replace
203 && (!canvas_size.have_crop || Self::test_full_image(canvas_size))
204 }
205
206 fn compute_default_xqms(encoding: Encoding, xyb_encoded: bool) -> u32 {
207 if xyb_encoded && encoding == Encoding::VarDct {
208 3
209 } else {
210 2
211 }
212 }
213
214 #[inline]
216 pub fn is_keyframe(&self) -> bool {
217 self.frame_type.is_normal_frame() && (self.is_last || self.duration != 0)
218 }
219
220 #[inline]
221 pub fn can_reference(&self) -> bool {
222 !self.is_last
223 && (self.duration == 0 || self.save_as_reference != 0)
224 && self.frame_type != FrameType::LfFrame
225 }
226
227 pub fn sample_width(&self, upsampling: u32) -> u32 {
228 let &Self {
229 mut width,
230 lf_level,
231 ..
232 } = self;
233
234 if upsampling > 1 {
235 width = width.div_ceil(upsampling);
236 }
237 if lf_level > 0 {
238 let div = 1u32 << (3 * lf_level);
239 width = (width + div - 1) >> (3 * lf_level);
240 }
241
242 width
243 }
244
245 pub fn sample_height(&self, upsampling: u32) -> u32 {
246 let &Self {
247 mut height,
248 lf_level,
249 ..
250 } = self;
251
252 if upsampling > 1 {
253 height = height.div_ceil(upsampling);
254 }
255 if lf_level > 0 {
256 let div = 1u32 << (3 * lf_level);
257 height = (height + div - 1) >> (3 * lf_level);
258 }
259
260 height
261 }
262
263 pub fn color_sample_width(&self) -> u32 {
264 self.sample_width(self.upsampling)
265 }
266
267 pub fn color_sample_height(&self) -> u32 {
268 self.sample_height(self.upsampling)
269 }
270
271 #[inline]
273 pub fn encoded_color_channels(&self) -> usize {
274 self.encoded_color_channels as usize
275 }
276
277 pub fn num_groups(&self) -> u32 {
278 let width = self.color_sample_width();
279 let height = self.color_sample_height();
280 let group_dim = self.group_dim();
281
282 let hgroups = width.div_ceil(group_dim);
283 let vgroups = height.div_ceil(group_dim);
284
285 hgroups * vgroups
286 }
287
288 pub fn num_lf_groups(&self) -> u32 {
289 let width = self.color_sample_width();
290 let height = self.color_sample_height();
291 let lf_group_dim = self.lf_group_dim();
292
293 let hgroups = width.div_ceil(lf_group_dim);
294 let vgroups = height.div_ceil(lf_group_dim);
295
296 hgroups * vgroups
297 }
298
299 pub fn group_dim(&self) -> u32 {
300 128 << self.group_size_shift
301 }
302
303 pub fn groups_per_row(&self) -> u32 {
304 let group_dim = self.group_dim();
305 self.color_sample_width().div_ceil(group_dim)
306 }
307
308 pub fn lf_group_dim(&self) -> u32 {
309 self.group_dim() * 8
310 }
311
312 pub fn lf_groups_per_row(&self) -> u32 {
313 let lf_group_dim = self.lf_group_dim();
314 self.color_sample_width().div_ceil(lf_group_dim)
315 }
316
317 pub fn group_size_for(&self, group_idx: u32) -> (u32, u32) {
318 self.size_for(self.group_dim(), group_idx)
319 }
320
321 pub fn lf_group_size_for(&self, lf_group_idx: u32) -> (u32, u32) {
322 self.size_for(self.lf_group_dim(), lf_group_idx)
323 }
324
325 fn size_for(&self, group_dim: u32, group_idx: u32) -> (u32, u32) {
326 let width = self.color_sample_width();
327 let height = self.color_sample_height();
328 let full_rows = height / group_dim;
329 let rows_remainder = height % group_dim;
330 let full_cols = width / group_dim;
331 let cols_remainder = width % group_dim;
332
333 let stride = full_cols + (cols_remainder > 0) as u32;
334 let row = group_idx / stride;
335 let col = group_idx % stride;
336
337 let group_width = if col >= full_cols {
338 cols_remainder
339 } else {
340 group_dim
341 };
342 let group_height = if row >= full_rows {
343 rows_remainder
344 } else {
345 group_dim
346 };
347 (group_width, group_height)
348 }
349
350 pub fn lf_group_idx_from_group_idx(&self, group_idx: u32) -> u32 {
351 let groups_per_row = self.groups_per_row();
352 let lf_group_col = (group_idx % groups_per_row) / 8;
353 let lf_group_row = (group_idx / groups_per_row) / 8;
354 lf_group_col + lf_group_row * self.lf_groups_per_row()
355 }
356
357 pub fn group_idx_from_coord(&self, x: u32, y: u32) -> Option<u32> {
358 let shift = 7 + self.group_size_shift;
359 let group_x = x >> shift;
360 let group_y = y >> shift;
361 let group_idx = group_y * self.groups_per_row() + group_x;
362 if group_x >= self.groups_per_row() || group_idx >= self.num_groups() {
363 None
364 } else {
365 Some(group_idx)
366 }
367 }
368
369 pub fn is_group_collides_region(&self, group_idx: u32, region: (u32, u32, u32, u32)) -> bool {
370 let group_dim = self.group_dim();
371 let group_per_row = self.groups_per_row();
372 let group_left = (group_idx % group_per_row) * group_dim;
373 let group_top = (group_idx / group_per_row) * group_dim;
374 is_aabb_collides(region, (group_left, group_top, group_dim, group_dim))
375 }
376
377 pub fn is_lf_group_collides_region(
378 &self,
379 lf_group_idx: u32,
380 region: (u32, u32, u32, u32),
381 ) -> bool {
382 let lf_group_dim = self.lf_group_dim();
383 let lf_group_per_row = self.lf_groups_per_row();
384 let group_left = (lf_group_idx % lf_group_per_row) * lf_group_dim;
385 let group_top = (lf_group_idx / lf_group_per_row) * lf_group_dim;
386 is_aabb_collides(region, (group_left, group_top, lf_group_dim, lf_group_dim))
387 }
388}
389
390#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
391#[repr(u8)]
392pub enum FrameType {
393 #[default]
394 RegularFrame = 0,
395 LfFrame,
396 ReferenceOnly,
397 SkipProgressive,
398}
399
400impl FrameType {
401 pub fn is_normal_frame(&self) -> bool {
402 matches!(self, Self::RegularFrame | Self::SkipProgressive)
403 }
404
405 pub fn is_progressive_frame(&self) -> bool {
406 matches!(self, Self::RegularFrame | Self::LfFrame)
407 }
408}
409
410impl<Ctx> Bundle<Ctx> for FrameType {
411 type Error = crate::Error;
412
413 fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
414 Ok(match bitstream.read_bits(2)? {
415 0 => Self::RegularFrame,
416 1 => Self::LfFrame,
417 2 => Self::ReferenceOnly,
418 3 => Self::SkipProgressive,
419 _ => unreachable!(),
420 })
421 }
422}
423
424#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
425#[repr(u8)]
426pub enum Encoding {
427 #[default]
428 VarDct = 0,
429 Modular,
430}
431
432impl<Ctx> Bundle<Ctx> for Encoding {
433 type Error = crate::Error;
434
435 fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
436 Ok(match bitstream.read_bits(1)? {
437 0 => Self::VarDct,
438 1 => Self::Modular,
439 _ => unreachable!(),
440 })
441 }
442}
443
444#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
445pub struct FrameFlags(u64);
446
447impl FrameFlags {
448 const NOISE: u64 = 0x1;
449 const PATCHES: u64 = 0x2;
450 const SPLINES: u64 = 0x10;
451 const USE_LF_FRAME: u64 = 0x20;
452 const SKIP_ADAPTIVE_LF_SMOOTHING: u64 = 0x80;
453
454 pub fn noise(&self) -> bool {
455 self.0 & Self::NOISE != 0
456 }
457
458 pub fn patches(&self) -> bool {
459 self.0 & Self::PATCHES != 0
460 }
461
462 pub fn splines(&self) -> bool {
463 self.0 & Self::SPLINES != 0
464 }
465
466 pub fn use_lf_frame(&self) -> bool {
467 self.0 & Self::USE_LF_FRAME != 0
468 }
469
470 pub fn skip_adaptive_lf_smoothing(&self) -> bool {
471 self.0 & Self::SKIP_ADAPTIVE_LF_SMOOTHING != 0
472 }
473}
474
475impl<Ctx> Bundle<Ctx> for FrameFlags {
476 type Error = crate::Error;
477
478 fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
479 Ok(Self(bitstream.read_u64()?))
480 }
481}
482
483#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
484#[repr(u8)]
485pub enum BlendMode {
486 #[default]
487 Replace = 0,
488 Add = 1,
489 Blend = 2,
490 MulAdd = 3,
491 Mul = 4,
492}
493
494impl<Ctx> Bundle<Ctx> for BlendMode {
495 type Error = crate::Error;
496
497 fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
498 Ok(match bitstream.read_u32(0, 1, 2, 3 + U(2))? {
499 0 => Self::Replace,
500 1 => Self::Add,
501 2 => Self::Blend,
502 3 => Self::MulAdd,
503 4 => Self::Mul,
504 value => {
505 return Err(jxl_bitstream::Error::InvalidEnum {
506 name: "BlendMode",
507 value,
508 }
509 .into())
510 }
511 })
512 }
513}
514
515impl BlendMode {
516 #[inline]
517 pub fn use_alpha(self) -> bool {
518 matches!(self, Self::Blend | Self::MulAdd)
519 }
520}
521
522fn is_aabb_collides(rect0: (u32, u32, u32, u32), rect1: (u32, u32, u32, u32)) -> bool {
523 let (x0, y0, w0, h0) = rect0;
524 let (x1, y1, w1, h1) = rect1;
525 (x0 < x1 + w1) && (x0 + w0 > x1) && (y0 < y1 + h1) && (y0 + h0 > y1)
526}