jxl_modular/
param.rs

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}