1use super::tables::{CAT_PROBS, ENERGY_CLASS};
2use super::{Quant, TokenCosts};
3use std::sync::LazyLock;
4
5pub const MAX_LEVEL: i32 = 16450;
6
7struct Costs {
8 bit: [[u16; 2]; 256],
9 extra: [u32; 67],
10 cat6: [[u32; 2]; 14],
11}
12
13static COSTS: LazyLock<Costs> = LazyLock::new(|| {
14 let mut bit = [[0u16; 2]; 256];
15 for (p, b) in bit.iter_mut().enumerate().skip(1) {
16 let p0 = p as f64 / 256.0;
17 b[0] = (-p0.log2() * 256.0).round() as u16;
18 b[1] = (-(1.0 - p0).log2() * 256.0).round() as u16;
19 }
20 let mut extra = [0u32; 67];
21 for (l, e) in extra.iter_mut().enumerate() {
22 let (cat, base) = match l {
23 0..=4 => (0usize, 0usize),
24 5..=6 => (1, 5),
25 7..=10 => (2, 7),
26 11..=18 => (3, 11),
27 19..=34 => (4, 19),
28 _ => (5, 35),
29 };
30 let probs = CAT_PROBS[cat];
31 let v = l.saturating_sub(base);
32 *e = probs
33 .iter()
34 .enumerate()
35 .map(|(k, &p)| bit[p as usize][(v >> (probs.len() - 1 - k)) & 1] as u32)
36 .sum();
37 }
38 let mut cat6 = [[0u32; 2]; 14];
39 for (k, c) in cat6.iter_mut().enumerate() {
40 let p = CAT_PROBS[6][k] as usize;
41 *c = [bit[p][0] as u32, bit[p][1] as u32];
42 }
43 Costs { bit, extra, cat6 }
44});
45
46#[inline(always)]
47pub fn token(l: u32) -> usize {
48 match l {
49 0..=4 => l as usize,
50 5..=6 => 5,
51 7..=10 => 6,
52 11..=18 => 7,
53 19..=34 => 8,
54 35..=66 => 9,
55 _ => 10,
56 }
57}
58
59#[inline(always)]
60fn level_cost(c: &Costs, tok_cost: &[u32; 11], l: u32) -> u32 {
61 if l == 0 {
62 return tok_cost[0];
63 }
64 let t = token(l);
65 let extra = if t < 10 {
66 c.extra[l as usize]
67 } else {
68 let v = l - 67;
69 (0..14)
70 .map(|k| c.cat6[k][((v >> (13 - k)) & 1) as usize])
71 .sum()
72 };
73 tok_cost[t] + extra + 256
74}
75
76pub fn bit_cost(p: u8, bit: bool) -> u32 {
77 COSTS.bit[p as usize][bit as usize] as u32
78}
79
80pub struct Scan<'a> {
81 pub scan: &'a [u16],
82 pub nb: &'a [(u16, u16)],
83 pub bands: &'a [u8],
84 pub iscan_t: &'a [u16],
85 pub log2n: usize,
86}
87
88impl Scan<'_> {
89 #[inline(always)]
90 pub fn tpos(&self, pos: usize) -> usize {
91 let m = (1 << self.log2n) - 1;
92 ((pos & m) << self.log2n) | (pos >> self.log2n)
93 }
94}
95
96#[inline(always)]
97fn recip(dq: i32) -> (u32, u32) {
98 let shift = 14 + (dq as u32).ilog2();
99 (((1u64 << shift) / dq as u64) as u32 + 1, shift)
100}
101
102#[inline(always)]
103fn ctx_at(tc: &[u8], s: &Scan, c: usize, ctx0: usize) -> usize {
104 if c == 0 {
105 ctx0
106 } else {
107 let (a, b) = s.nb[c];
108 ((1 + tc[a as usize] + tc[b as usize]) >> 1) as usize
109 }
110}
111
112pub fn rate<const N2: usize>(
113 levels: &[i16],
114 eob: usize,
115 s: &Scan,
116 ctx0: usize,
117 costs: &TokenCosts,
118) -> u32 {
119 let c = &*COSTS;
120 let n2 = s.scan.len();
121 let mut tc = [0u8; N2];
122 let mut r = 0u32;
123 let mut check = true;
124 for (k, &lv) in levels[..eob].iter().enumerate() {
125 let band = s.bands[k] as usize;
126 let ctx = ctx_at(&tc, s, k, ctx0);
127 if check {
128 r += costs.more[band][ctx][1];
129 }
130 let l = lv.unsigned_abs() as u32;
131 r += level_cost(c, &costs.token[band][ctx], l);
132 tc[s.scan[k] as usize] = ENERGY_CLASS[token(l)];
133 check = l != 0;
134 }
135 if eob < n2 {
136 let band = s.bands[eob] as usize;
137 r += costs.more[band][ctx_at(&tc, s, eob, ctx0)][0];
138 }
139 r
140}
141
142#[allow(clippy::too_many_arguments)]
143pub fn quantise<const N2: usize>(
144 coef: &[i16],
145 dqbuf: &mut [i16],
146 s: &Scan,
147 q: Quant,
148 deadzone: i32,
149 levels: &mut [i16; 1024],
150) -> (usize, usize) {
151 let n = 1usize << s.log2n;
152 let k32 = (s.log2n == 5) as u32;
153 let lmax_ac = MAX_LEVEL.min((32767 << k32) / q.ac) as u32;
154 let lmax_dc = MAX_LEVEL.min((32767 << k32) / q.dc) as u32;
155 let inv_ac = recip(q.ac);
156 let bias_ac = ((q.ac * deadzone) >> 6) as u32;
157 let coef = &coef[..N2];
158 let dqbuf = &mut dqbuf[..N2];
159 let iscan = &s.iscan_t[..N2];
160 let mask = n as u32 - 1;
161 let mut eob = 0u32;
162 let mut rows = 0u32;
163 let mut lev = [0i16; N2];
164 for i in 0..N2 {
165 let x = coef[i] as i32;
166 let a = (x.unsigned_abs() << k32) + bias_ac;
167 let l = ((a * inv_ac.0) >> inv_ac.1).min(lmax_ac);
168 let d = ((l * q.ac as u32) >> k32) as i32;
169 let sgn = x >> 31;
170 dqbuf[i] = ((d ^ sgn) - sgn) as i16;
171 lev[i] = ((l as i32 ^ sgn) - sgn) as i16;
172 let nz = if l != 0 { u32::MAX } else { 0 };
173 eob = eob.max((iscan[i] as u32 + 1) & nz);
174 rows = rows.max(i as u32 & mask & nz);
175 }
176 let mut eob = eob as usize;
177 let rows = rows as usize;
178 {
179 let x = coef[0] as i32;
180 let a = (x.unsigned_abs() << k32) + ((q.dc * deadzone) >> 6) as u32;
181 let r = recip(q.dc);
182 let l = ((a * r.0) >> r.1).min(lmax_dc);
183 let d = ((l * q.dc as u32) >> k32) as i32;
184 dqbuf[0] = (if x < 0 { -d } else { d }) as i16;
185 lev[0] = (if x < 0 { -(l as i32) } else { l as i32 }) as i16;
186 if l != 0 {
187 eob = eob.max(1);
188 }
189 }
190 for (c, l) in levels[..eob].iter_mut().enumerate() {
191 *l = lev[s.tpos(s.scan[c] as usize)];
192 }
193 (eob, if eob == 0 { 0 } else { rows + 1 })
194}
195
196#[allow(clippy::too_many_arguments)]
197pub fn trellis<const N2: usize>(
198 coef: &[i16],
199 s: &Scan,
200 q: Quant,
201 ctx0: usize,
202 costs: &TokenCosts,
203 lambda: u32,
204 levels: &mut [i16; 1024],
205) -> (usize, u32) {
206 let c = &*COSTS;
207 let n2 = N2;
208 let k32 = (s.log2n == 5) as u32;
209 let w: u64 = if k32 == 1 { 16 } else { 4 };
210 let lam = lambda as u64;
211 let lmax_ac = MAX_LEVEL.min((32767 << k32) / q.ac) as u32;
212 let lmax_dc = MAX_LEVEL.min((32767 << k32) / q.dc) as u32;
213 let inv_ac = recip(q.ac);
214 let inv_dc = recip(q.dc);
215 let coef = &coef[..N2];
216 let iscan = &s.iscan_t[..N2];
217 let half = q.ac as u32 >> 1;
218 let mut lr = [0u16; N2];
219 let mut eob0 = 0u32;
220 for i in 0..N2 {
221 let a = coef[i].unsigned_abs() as u32;
222 let l = ((((a << k32) + half) * inv_ac.0) >> inv_ac.1).min(lmax_ac);
223 lr[i] = l as u16;
224 let nz = if l != 0 { u32::MAX } else { 0 };
225 eob0 = eob0.max((iscan[i] as u32 + 1) & nz);
226 }
227 {
228 let a = coef[0].unsigned_abs() as u32;
229 let l = ((((a << k32) + (q.dc as u32 >> 1)) * inv_dc.0) >> inv_dc.1).min(lmax_dc);
230 lr[0] = l as u16;
231 if l != 0 {
232 eob0 = eob0.max(1);
233 }
234 }
235 let eob0 = eob0 as usize;
236 let mut abs = [0u16; N2];
237 let mut l0 = [0u16; N2];
238 let mut total = 0u64;
239 for k in 0..eob0 {
240 let i = s.tpos(s.scan[k] as usize);
241 let a = coef[i].unsigned_abs();
242 abs[k] = a;
243 l0[k] = lr[i];
244 total += (a as u64 * a as u64) * w;
245 }
246 let mut zeroed = 0u64;
247 let mut tc = [0u8; N2];
248 let band0 = s.bands[0] as usize;
249 let mut best_j = total + lam * costs.more[band0][ctx0][0] as u64;
250 let mut best_e = 0usize;
251 let mut best_r = costs.more[band0][ctx0][0];
252 let mut cum_d = 0u64;
253 let mut cum_r = 0u32;
254 let mut check = true;
255 for k in 0..eob0 {
256 let band = s.bands[k] as usize;
257 let ctx = ctx_at(&tc, s, k, ctx0);
258 if check {
259 cum_r += costs.more[band][ctx][1];
260 }
261 let dq = if k == 0 { q.dc } else { q.ac } as i64;
262 let a = abs[k] as i64;
263 let tok = &costs.token[band][ctx];
264 let next_more = if k + 1 < n2 {
265 costs.more[s.bands[k + 1] as usize][ctx][1]
266 } else {
267 0
268 };
269 let lv = l0[k] as u32;
270 zeroed += (a * a) as u64 * w;
271 let mut pick = lv;
272 let mut pick_r = level_cost(c, tok, lv);
273 let err = a - ((lv as i64 * dq) >> k32);
274 let mut pick_d = (err * err) as u64 * w;
275 if lv > 0 {
276 let base_j = pick_d + lam * (pick_r + next_more) as u64;
277 let l1 = lv - 1;
278 let r1 = level_cost(c, tok, l1);
279 let e1 = a - ((l1 as i64 * dq) >> k32);
280 let d1 = (e1 * e1) as u64 * w;
281 let j1 = d1 + lam * (r1 + if l1 > 0 { next_more } else { 0 }) as u64;
282 if j1 < base_j {
283 pick = l1;
284 pick_r = r1;
285 pick_d = d1;
286 }
287 }
288 cum_r += pick_r;
289 cum_d += pick_d;
290 l0[k] = pick as u16;
291 tc[s.scan[k] as usize] = ENERGY_CLASS[token(pick)];
292 check = pick != 0;
293 if pick != 0 {
294 let r_end = if k + 1 < n2 {
295 costs.more[s.bands[k + 1] as usize][ctx_at(&tc, s, k + 1, ctx0)][0]
296 } else {
297 0
298 };
299 let j = cum_d + (total - zeroed) + lam * (cum_r + r_end) as u64;
300 if j < best_j {
301 best_j = j;
302 best_e = k + 1;
303 best_r = cum_r + r_end;
304 }
305 }
306 }
307 for k in 0..best_e {
308 let i = s.tpos(s.scan[k] as usize);
309 let l = l0[k] as i16;
310 levels[k] = if coef[i] < 0 { -l } else { l };
311 }
312 (best_e, best_r)
313}
314
315pub fn dequantise(levels: &[i16], eob: usize, s: &Scan, q: Quant, dqbuf: &mut [i16]) -> usize {
316 let n = 1usize << s.log2n;
317 let k32 = (s.log2n == 5) as u32;
318 dqbuf[..n * n].fill(0);
319 let mut rows = 0;
320 for (k, &lv) in levels[..eob].iter().enumerate() {
321 let i = s.tpos(s.scan[k] as usize);
322 let dq = if i == 0 { q.dc } else { q.ac };
323 let d = ((lv.unsigned_abs() as u32 * dq as u32) >> k32) as i32;
324 dqbuf[i] = (if lv < 0 { -d } else { d }) as i16;
325 if lv != 0 {
326 rows = rows.max(i & (n - 1));
327 }
328 }
329 rows + 1
330}