1use super::dist;
2use super::halfpel::HalfPel;
3use super::inter::{self, BlockPos, MvLimits};
4use super::intra::{EdgeFlags, Edges};
5use super::motion::{Search, Subpel, sad_per_bit};
6use super::mvcost::MvCost;
7use super::*;
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub struct Decision {
11 pub mode: u8,
12 pub sub_modes: [u8; 4],
13 pub uv_mode: u8,
14 pub ref_frame: [i8; 2],
15 pub mv: [Mv; 2],
16 pub interp_filter: u8,
17 pub dist: u32,
18 pub rate: u32,
19 pub cost: u64,
20}
21
22impl Decision {
23 pub fn is_inter(&self) -> bool {
24 self.ref_frame[0] > INTRA_FRAME
25 }
26}
27
28#[inline(always)]
29pub fn cost(dist: u32, rate: u32, sad_per_bit: u32) -> u64 {
30 ((dist as u64) << 8) + rate as u64 * sad_per_bit as u64
31}
32
33pub struct IntraBlock<'a> {
34 pub src: &'a [u8],
35 pub src_stride: usize,
36 pub recon: &'a [u8],
37 pub recon_stride: usize,
38 pub x: usize,
39 pub y: usize,
40 pub w: usize,
41 pub h: usize,
42 pub tx_size: u8,
43 pub have_left: bool,
44 pub have_above: bool,
45 pub max_x: usize,
46 pub max_y: usize,
47}
48
49#[derive(Default)]
50pub struct IntraScratch {
51 canvas: Vec<u8>,
52 pred: Vec<u8>,
53 edges: Vec<(usize, usize, Edges)>,
54}
55
56const CANVAS: usize = 65;
57
58impl IntraScratch {
59 fn fill(&mut self, b: &IntraBlock) {
60 if self.canvas.is_empty() {
61 self.canvas = vec![0; 65 * 65];
62 self.pred = vec![0; 64 * 64];
63 self.edges = Vec::with_capacity(256);
64 }
65 let c = &mut self.canvas;
66 if b.have_above {
67 let row = (b.y - 1) * b.recon_stride;
68 for i in 0..=b.w {
69 let x = (b.x + i).saturating_sub(1).min(b.max_x);
70 c[i] = b.recon[row + x];
71 }
72 }
73 for i in 0..b.h {
74 let r = (i + 1) * CANVAS;
75 if b.have_left {
76 c[r] = b.recon[(b.y + i).min(b.max_y) * b.recon_stride + b.x - 1];
77 }
78 c[r + 1..r + 1 + b.w].copy_from_slice(&b.src[i * b.src_stride..i * b.src_stride + b.w]);
79 }
80 self.edges.clear();
81 let step = 4usize << b.tx_size;
82 let n4w = b.w / 4;
83 for ty in (0..b.h).step_by(step) {
84 for tx in (0..b.w).step_by(step) {
85 let f = EdgeFlags {
86 have_left: b.have_left || tx > 0,
87 have_above: b.have_above || ty > 0,
88 not_on_right: tx / 4 + (step / 4) < n4w,
89 max_x: (b.max_x + 1 - b.x).min(b.w),
90 max_y: (b.max_y + 1 - b.y).min(b.h),
91 };
92 let e = Edges::build(&self.canvas, CANVAS, tx + 1, ty + 1, b.tx_size, f);
93 self.edges.push((tx, ty, e));
94 }
95 }
96 }
97
98 fn predict(&mut self, mode: u8) {
99 for (tx, ty, e) in &self.edges {
100 e.predict(mode, &mut self.pred[ty * 64 + tx..], 64);
101 }
102 }
103
104 pub fn evaluate(&mut self, b: &IntraBlock, mode: u8) -> u32 {
105 self.fill(b);
106 self.predict(mode);
107 dist::satd(b.src, b.src_stride, &self.pred, 64, b.w, b.h)
108 }
109
110 pub fn best(
111 &mut self,
112 b: &IntraBlock,
113 mode_rates: &[u32; 10],
114 sad_per_bit: u32,
115 modes: &[u8],
116 ) -> (u8, u32, u32, u64) {
117 self.fill(b);
118 let mut out = (DC_PRED, 0, 0, u64::MAX);
119 for &m in modes {
120 self.predict(m);
121 let d = dist::satd(b.src, b.src_stride, &self.pred, 64, b.w, b.h);
122 let j = cost(d, mode_rates[m as usize], sad_per_bit);
123 if j < out.3 {
124 out = (m, d, mode_rates[m as usize], j);
125 }
126 }
127 out
128 }
129}
130
131pub const ALL_INTRA: [u8; 10] = [
132 DC_PRED, V_PRED, H_PRED, TM_PRED, D45_PRED, D135_PRED, D117_PRED, D153_PRED, D207_PRED,
133 D63_PRED,
134];
135pub const FAST_INTRA: [u8; 4] = [DC_PRED, V_PRED, H_PRED, TM_PRED];
136
137pub struct IntraChoice<'a> {
138 pub luma: IntraBlock<'a>,
139 pub chroma: [IntraBlock<'a>; 2],
140 pub y_mode_rates: &'a [u32; 10],
141 pub uv_mode_rates: &'a [[u32; 10]; 10],
142 pub lambda: u32,
143 pub modes: &'a [u8],
144}
145
146fn best_uv(
147 c: &IntraChoice,
148 y_mode: u8,
149 spb: u32,
150 scratch: &mut [IntraScratch; 2],
151) -> (u8, u32, u32) {
152 let rates = &c.uv_mode_rates[y_mode as usize];
153 let mut uv = (DC_PRED, 0u32, 0u32, u64::MAX);
154 scratch[0].fill(&c.chroma[0]);
155 scratch[1].fill(&c.chroma[1]);
156 for &m in c.modes {
157 let mut d = 0;
158 for (s, b) in scratch.iter_mut().zip(&c.chroma) {
159 s.predict(m);
160 d += dist::satd(b.src, b.src_stride, &s.pred, 64, b.w, b.h);
161 }
162 let j = cost(d, rates[m as usize], spb);
163 if j < uv.3 {
164 uv = (m, d, rates[m as usize], j);
165 }
166 }
167 (uv.0, uv.1, uv.2)
168}
169
170fn intra_decision(
171 mode: u8,
172 sub_modes: [u8; 4],
173 uv: (u8, u32, u32),
174 yd: u32,
175 yr: u32,
176 spb: u32,
177) -> Decision {
178 let (dist, rate) = (yd + uv.1, yr + uv.2);
179 Decision {
180 mode,
181 sub_modes,
182 uv_mode: uv.0,
183 ref_frame: [INTRA_FRAME, NONE_FRAME],
184 mv: [Mv::ZERO; 2],
185 interp_filter: EIGHTTAP,
186 dist,
187 rate,
188 cost: cost(dist, rate, spb),
189 }
190}
191
192pub fn best_intra(c: &IntraChoice, scratch: &mut [IntraScratch; 2]) -> Decision {
193 let spb = sad_per_bit(c.lambda);
194 let (mode, yd, yr, _) = scratch[0].best(&c.luma, c.y_mode_rates, spb, c.modes);
195 let uv = best_uv(c, mode, spb, scratch);
196 intra_decision(mode, [mode; 4], uv, yd, yr, spb)
197}
198
199pub fn best_intra_4x4(
200 c: &IntraChoice,
201 above: [u8; 2],
202 left: [u8; 2],
203 sub_rate: &dyn Fn(u8, u8, u8) -> u32,
204 scratch: &mut [IntraScratch; 2],
205) -> Decision {
206 let spb = sad_per_bit(c.lambda);
207 let b = &c.luma;
208 assert!(b.w == 8 && b.h == 8 && b.tx_size == 0);
209 let s = &mut scratch[0];
210 s.fill(b);
211 let mut modes = [DC_PRED; 4];
212 let (mut yd, mut yr) = (0, 0);
213 for k in 0..4 {
214 let (tx, ty) = (4 * (k & 1), 4 * (k >> 1));
215 let a = if k >= 2 { modes[k - 2] } else { above[k] };
216 let l = if k & 1 == 1 {
217 modes[k - 1]
218 } else {
219 left[k >> 1]
220 };
221 let src = &b.src[ty * b.src_stride + tx..];
222 let mut best = (DC_PRED, 0, 0, u64::MAX);
223 for &m in c.modes {
224 s.edges[k].2.predict(m, &mut s.pred[ty * 64 + tx..], 64);
225 let d = dist::satd(src, b.src_stride, &s.pred[ty * 64 + tx..], 64, 4, 4);
226 let r = sub_rate(a, l, m);
227 let j = cost(d, r, spb);
228 if j < best.3 {
229 best = (m, d, r, j);
230 }
231 }
232 modes[k] = best.0;
233 yd += best.1;
234 yr += best.2;
235 }
236 let uv = best_uv(c, modes[3], spb, scratch);
237 intra_decision(modes[3], modes, uv, yd, yr, spb)
238}
239
240#[derive(Clone, Copy)]
241pub struct InterRef<'a> {
242 pub frame: i8,
243 pub luma: Plane<'a>,
244 pub nearest: Mv,
245 pub near: Mv,
246 pub best: Mv,
247 pub hints: &'a [Mv],
248 pub ref_rate: u32,
249 pub half: Option<&'a HalfPel>,
250}
251
252#[derive(Clone, Copy)]
253pub struct InterChoice<'a> {
254 pub src: &'a [u8],
255 pub src_stride: usize,
256 pub pos: BlockPos,
257 pub w: usize,
258 pub h: usize,
259 pub refs: &'a [InterRef<'a>],
260 pub mode_rates: &'a [u32; 4],
261 pub is_inter_rate: u32,
262 pub mv_cost: &'a MvCost,
263 pub allow_hp: bool,
264 pub filter: u8,
265 pub lambda: u32,
266 pub range: i32,
267 pub subpel: Subpel,
268}
269
270#[derive(Clone, Copy, Debug)]
271pub struct CompoundPair {
272 pub refs: [usize; 2],
273 pub rate: u32,
274}
275
276#[derive(Default)]
277pub struct InterScratch {
278 pred: Vec<u8>,
279 pred2: Vec<u8>,
280 avg: Vec<u8>,
281}
282
283fn predict_satd(
284 c: &InterChoice,
285 limits: &MvLimits,
286 refs: &[&InterRef],
287 mvs: &[Mv],
288 scratch: &mut InterScratch,
289) -> u32 {
290 let (x, y) = ((c.pos.mi_col * 8) as isize, (c.pos.mi_row * 8) as isize);
291 let at = |mv: Mv| {
292 let m = limits.clamp(mv);
293 ((x << 4) + 2 * m.col as isize, (y << 4) + 2 * m.row as isize)
294 };
295 let (x0, y0) = at(mvs[0]);
296 inter::mc(
297 &refs[0].luma,
298 x0,
299 y0,
300 c.w,
301 c.h,
302 c.filter,
303 &mut scratch.pred,
304 64,
305 );
306 if refs.len() > 1 {
307 let (x1, y1) = at(mvs[1]);
308 inter::mc(
309 &refs[1].luma,
310 x1,
311 y1,
312 c.w,
313 c.h,
314 c.filter,
315 &mut scratch.pred2,
316 64,
317 );
318 for i in 0..c.h {
319 let r = i * 64..i * 64 + c.w;
320 inter::average(
321 &scratch.pred[r.clone()],
322 &scratch.pred2[r.clone()],
323 c.w,
324 1,
325 &mut scratch.avg[r],
326 c.w,
327 );
328 }
329 return dist::satd(c.src, c.src_stride, &scratch.avg, 64, c.w, c.h);
330 }
331 dist::satd(c.src, c.src_stride, &scratch.pred, 64, c.w, c.h)
332}
333
334pub fn best_inter(c: &InterChoice, scratch: &mut InterScratch) -> Decision {
335 best_inter_with(c, &[], scratch)
336}
337
338pub fn best_inter_with(
339 c: &InterChoice,
340 compound: &[CompoundPair],
341 scratch: &mut InterScratch,
342) -> Decision {
343 if scratch.pred.is_empty() {
344 scratch.pred = vec![0; 64 * 64];
345 scratch.pred2 = vec![0; 64 * 64];
346 scratch.avg = vec![0; 64 * 64];
347 }
348 let spb = sad_per_bit(c.lambda);
349 let p = &c.pos;
350 let limits = MvLimits::new(p.mi_row, p.mi_col, p.bsize, p.mi_rows, p.mi_cols);
351 let (x, y) = (p.mi_col * 8, p.mi_row * 8);
352 let mut out = Decision {
353 mode: ZEROMV,
354 sub_modes: [ZEROMV; 4],
355 uv_mode: DC_PRED,
356 ref_frame: [LAST_FRAME, NONE_FRAME],
357 mv: [Mv::ZERO; 2],
358 interp_filter: c.filter,
359 dist: u32::MAX,
360 rate: 0,
361 cost: u64::MAX,
362 };
363 let mut offer = |mode: u8, frames: [i8; 2], mv: [Mv; 2], d: u32, rate: u32| {
364 let j = cost(d, rate, spb);
365 if j < out.cost {
366 out = Decision {
367 mode,
368 sub_modes: [mode; 4],
369 uv_mode: DC_PRED,
370 ref_frame: frames,
371 mv,
372 interp_filter: c.filter,
373 dist: d,
374 rate,
375 cost: j,
376 };
377 }
378 };
379 let mut found = [(Mv::ZERO, 0u32); 4];
380 for (i, r) in c.refs.iter().enumerate() {
381 let base = c.is_inter_rate + r.ref_rate;
382 let frames = [r.frame, NONE_FRAME];
383 let mut tried = [Mv::new(i16::MIN, i16::MIN); 3];
384 for (k, (mode, mv)) in [(ZEROMV, Mv::ZERO), (NEARESTMV, r.nearest), (NEARMV, r.near)]
385 .into_iter()
386 .enumerate()
387 {
388 if tried[..k].contains(&mv) {
389 continue;
390 }
391 tried[k] = mv;
392 let d = predict_satd(c, &limits, &[r], &[mv], scratch);
393 offer(
394 mode,
395 frames,
396 [mv, Mv::ZERO],
397 d,
398 base + c.mode_rates[(mode - NEARESTMV) as usize],
399 );
400 }
401 let mut cands = [Mv::ZERO; 12];
402 let mut n = 0;
403 for &m in [r.best, r.nearest, r.near].iter().chain(r.hints.iter()) {
404 if n < cands.len() {
405 cands[n] = m;
406 n += 1;
407 }
408 }
409 let s = Search {
410 src: c.src,
411 src_stride: c.src_stride,
412 reference: r.luma,
413 x,
414 y,
415 w: c.w,
416 h: c.h,
417 limits,
418 mv_cost: c.mv_cost,
419 ref_mv: r.best,
420 allow_hp: c.allow_hp,
421 sad_per_bit: spb,
422 range: c.range,
423 subpel: c.subpel,
424 filter: c.filter,
425 half: r.half,
426 };
427 let m = s.run_with(
428 &cands[..n],
429 (&mut scratch.pred[..64 * 64]).try_into().unwrap(),
430 );
431 if let Some(f) = found.get_mut(i) {
432 *f = (m.mv, m.rate);
433 }
434 if !tried.contains(&m.mv) {
435 let rate = base + c.mode_rates[(NEWMV - NEARESTMV) as usize] + m.rate;
436 offer(NEWMV, frames, [m.mv, Mv::ZERO], m.dist, rate);
437 }
438 }
439 for pair in compound {
440 let [i, k] = pair.refs;
441 if i.max(k) >= c.refs.len().min(found.len()) {
442 continue;
443 }
444 let (a, b) = (&c.refs[i], &c.refs[k]);
445 let base = c.is_inter_rate + pair.rate;
446 let frames = [a.frame, b.frame];
447 for (mode, mvs, extra) in [
448 (ZEROMV, [Mv::ZERO; 2], 0),
449 (NEARESTMV, [a.nearest, b.nearest], 0),
450 (NEARMV, [a.near, b.near], 0),
451 (NEWMV, [found[i].0, found[k].0], found[i].1 + found[k].1),
452 ] {
453 let d = predict_satd(c, &limits, &[a, b], &mvs, scratch);
454 offer(
455 mode,
456 frames,
457 mvs,
458 d,
459 base + c.mode_rates[(mode - NEARESTMV) as usize] + extra,
460 );
461 }
462 }
463 out
464}
465
466pub fn best_filter(
467 c: &InterChoice,
468 d: &Decision,
469 filter_rates: &[u32; 3],
470 scratch: &mut InterScratch,
471) -> (u8, u32, u64) {
472 if scratch.pred.is_empty() {
473 *scratch = InterScratch {
474 pred: vec![0; 64 * 64],
475 pred2: vec![0; 64 * 64],
476 avg: vec![0; 64 * 64],
477 };
478 }
479 let spb = sad_per_bit(c.lambda);
480 let p = &c.pos;
481 let limits = MvLimits::new(p.mi_row, p.mi_col, p.bsize, p.mi_rows, p.mi_cols);
482 let find = |f: i8| c.refs.iter().find(|r| r.frame == f);
483 let Some(a) = find(d.ref_frame[0]) else {
484 return (d.interp_filter, d.dist, u64::MAX);
485 };
486 let b = if d.ref_frame[1] > INTRA_FRAME {
487 find(d.ref_frame[1])
488 } else {
489 None
490 };
491 let mut best = (d.interp_filter, d.dist, u64::MAX);
492 for f in [EIGHTTAP, EIGHTTAP_SMOOTH, EIGHTTAP_SHARP] {
493 let cf = InterChoice { filter: f, ..*c };
494 let dist = match b {
495 Some(b) => predict_satd(&cf, &limits, &[a, b], &d.mv, scratch),
496 None => predict_satd(&cf, &limits, &[a], &d.mv[..1], scratch),
497 };
498 let j = cost(dist, filter_rates[f as usize], spb);
499 if j < best.2 {
500 best = (f, dist, j);
501 }
502 }
503 best
504}