1use jxl_grid::AllocTracker;
2
3use super::MaConfig;
4
5#[derive(Debug, Clone)]
6pub struct ModularParams<'a, 'b> {
7 pub group_dim: u32,
8 pub bit_depth: u32,
9 pub channels: Vec<ModularChannelParams>,
10 pub ma_config: Option<&'a MaConfig>,
11 pub tracker: Option<&'b AllocTracker>,
12 pub narrow_buffer: bool,
13}
14
15impl<'a, 'b> ModularParams<'a, 'b> {
16 pub fn new(
17 width: u32,
18 height: u32,
19 group_dim: u32,
20 bit_depth: u32,
21 channel_shifts: Vec<ChannelShift>,
22 ma_config: Option<&'a MaConfig>,
23 tracker: Option<&'b AllocTracker>,
24 ) -> Self {
25 let channels = channel_shifts
26 .into_iter()
27 .map(|shift| ModularChannelParams {
28 width,
29 height,
30 shift,
31 })
32 .collect();
33 Self::with_channels(group_dim, bit_depth, channels, ma_config, tracker)
34 }
35
36 pub fn with_channels(
37 group_dim: u32,
38 bit_depth: u32,
39 channels: Vec<ModularChannelParams>,
40 ma_config: Option<&'a MaConfig>,
41 tracker: Option<&'b AllocTracker>,
42 ) -> Self {
43 Self {
44 group_dim,
45 bit_depth,
46 channels,
47 ma_config,
48 tracker,
49 narrow_buffer: false,
50 }
51 }
52}
53
54#[derive(Debug, Clone)]
55pub struct ModularChannelParams {
56 pub width: u32,
57 pub height: u32,
58 pub shift: ChannelShift,
59}
60
61impl ModularChannelParams {
62 pub fn new(width: u32, height: u32) -> Self {
63 Self {
64 width,
65 height,
66 shift: ChannelShift::from_shift(0),
67 }
68 }
69
70 pub fn jpeg(width: u32, height: u32, jpeg_upsampling: [u32; 3], idx: usize) -> Self {
71 Self {
72 width,
73 height,
74 shift: ChannelShift::from_jpeg_upsampling(jpeg_upsampling, idx),
75 }
76 }
77
78 pub fn with_shift(width: u32, height: u32, shift: ChannelShift) -> Self {
79 Self {
80 width,
81 height,
82 shift,
83 }
84 }
85}
86
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub enum ChannelShift {
89 JpegUpsampling {
90 has_h_subsample: bool,
91 h_subsample: bool,
92 has_v_subsample: bool,
93 v_subsample: bool,
94 },
95 Shifts(u32),
96 Raw(i32, i32),
97}
98
99impl ChannelShift {
100 #[inline]
101 pub fn from_shift(shift: u32) -> ChannelShift {
102 Self::Shifts(shift)
103 }
104
105 pub fn from_jpeg_upsampling(jpeg_upsampling: [u32; 3], idx: usize) -> Self {
106 let upsampling = jpeg_upsampling[idx];
107 let hscale = jpeg_upsampling.into_iter().any(|v| v == 1 || v == 2);
108 let vscale = jpeg_upsampling.into_iter().any(|v| v == 1 || v == 3);
109 let (h, v) = match upsampling {
110 0 => (hscale, vscale),
111 1 => (false, false),
112 2 => (false, vscale),
113 3 => (hscale, false),
114 _ => panic!("Invalid jpeg_upsampling value of {}", upsampling),
115 };
116 Self::JpegUpsampling {
117 has_h_subsample: hscale,
118 h_subsample: h,
119 has_v_subsample: vscale,
120 v_subsample: v,
121 }
122 }
123
124 #[inline]
125 pub fn hshift(&self) -> i32 {
126 match self {
127 Self::JpegUpsampling { h_subsample, .. } => *h_subsample as i32,
128 Self::Shifts(s) => *s as i32,
129 Self::Raw(h, _) => *h,
130 }
131 }
132
133 #[inline]
134 pub fn vshift(&self) -> i32 {
135 match self {
136 Self::JpegUpsampling { v_subsample, .. } => *v_subsample as i32,
137 Self::Shifts(s) => *s as i32,
138 Self::Raw(_, v) => *v,
139 }
140 }
141
142 pub fn shift_size(&self, (width, height): (u32, u32)) -> (u32, u32) {
143 match *self {
144 Self::JpegUpsampling {
145 has_h_subsample,
146 has_v_subsample,
147 h_subsample,
148 v_subsample,
149 } => {
150 let width = if has_h_subsample {
151 let size = (width + 1) / 2;
152 if h_subsample {
153 size
154 } else {
155 size * 2
156 }
157 } else {
158 width
159 };
160 let height = if has_v_subsample {
161 let size = (height + 1) / 2;
162 if v_subsample {
163 size
164 } else {
165 size * 2
166 }
167 } else {
168 height
169 };
170 (width, height)
171 }
172 Self::Shifts(s) => {
173 let add = (1 << s) - 1;
174 ((width + add) >> s, (height + add) >> s)
175 }
176 Self::Raw(h, v) => {
177 let add_h = (1 << h) - 1;
178 let add_v = (1 << v) - 1;
179 ((width + add_h) >> h, (height + add_v) >> v)
180 }
181 }
182 }
183}