1use super::dist;
2use super::halfpel::HalfPel;
3use super::inter::{self, MvLimits};
4use super::mvcost::MvCost;
5use super::*;
6
7pub fn sad_per_bit(lambda: u32) -> u32 {
8 (0.247 * (lambda as f64).sqrt() + 2.41).round() as u32
9}
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub enum Subpel {
13 None,
14 Half,
15 Quarter,
16 Full,
17}
18
19#[derive(Clone, Copy)]
20pub struct Search<'a> {
21 pub src: &'a [u8],
22 pub src_stride: usize,
23 pub reference: Plane<'a>,
24 pub x: usize,
25 pub y: usize,
26 pub w: usize,
27 pub h: usize,
28 pub limits: MvLimits,
29 pub mv_cost: &'a MvCost,
30 pub ref_mv: Mv,
31 pub allow_hp: bool,
32 pub sad_per_bit: u32,
33 pub range: i32,
34 pub subpel: Subpel,
35 pub filter: u8,
36 pub half: Option<&'a HalfPel>,
37}
38
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub struct MotionResult {
41 pub mv: Mv,
42 pub dist: u32,
43 pub rate: u32,
44}
45
46const HEX: [(i32, i32); 6] = [(-2, 0), (-1, -2), (1, -2), (2, 0), (1, 2), (-1, 2)];
47const DIAMOND: [(i32, i32); 4] = [(0, -1), (-1, 0), (1, 0), (0, 1)];
48const SQUARE: [(i32, i32); 8] = [
49 (-1, -1),
50 (0, -1),
51 (1, -1),
52 (-1, 0),
53 (1, 0),
54 (-1, 1),
55 (0, 1),
56 (1, 1),
57];
58
59impl Search<'_> {
60 fn usehp(&self) -> bool {
61 self.allow_hp && self.ref_mv.use_hp()
62 }
63
64 #[inline(always)]
65 fn mv_bits(&self, mv: Mv) -> u32 {
66 self.mv_cost.rate(mv, self.ref_mv, self.usehp())
67 }
68
69 #[inline(always)]
70 fn penalty(&self, mv: Mv) -> u32 {
71 (self.mv_bits(mv) * self.sad_per_bit + 128) >> 8
72 }
73
74 #[inline(always)]
75 fn full_sad(&self, row: i32, col: i32) -> u32 {
76 let r = &self.reference;
77 let o = r.at(
78 self.x as isize + col as isize,
79 self.y as isize + row as isize,
80 );
81 dist::sad(
82 self.src,
83 self.src_stride,
84 &r.data[o..],
85 r.stride,
86 self.w,
87 self.h,
88 )
89 }
90
91 fn full_limits(&self) -> (i32, i32, i32, i32) {
92 let l = &self.limits;
93 let b = self.reference.border as i32 - 8;
94 let (x, y) = (self.x as i32, self.y as i32);
95 let (w, h) = (self.reference.width as i32, self.reference.height as i32);
96 let rmin = ((l.row_min + 7) >> 3).max(-y - b);
97 let rmax = (l.row_max >> 3).min(h + b - y - self.h as i32);
98 let cmin = ((l.col_min + 7) >> 3).max(-x - b);
99 let cmax = (l.col_max >> 3).min(w + b - x - self.w as i32);
100 (rmin, rmax.max(rmin), cmin, cmax.max(cmin))
101 }
102
103 pub fn full_pel(&self, candidates: &[Mv]) -> (i32, i32, u32) {
104 let (rmin, rmax, cmin, cmax) = self.full_limits();
105 let fit = |r: i32, c: i32| (r.clamp(rmin, rmax), c.clamp(cmin, cmax));
106 let cost = |r: i32, c: i32| {
107 self.full_sad(r, c) + self.penalty(Mv::new((r * 8) as i16, (c * 8) as i16))
108 };
109 let (mut br, mut bc) = fit(0, 0);
110 let mut best = cost(br, bc);
111 let mut seen = [(i32::MIN, i32::MIN); 8];
112 for (k, m) in candidates.iter().enumerate() {
113 let (r, c) = fit(
114 (m.row as i32 + 4).div_euclid(8),
115 (m.col as i32 + 4).div_euclid(8),
116 );
117 if (r, c) == (br, bc) || seen[..k.min(8)].contains(&(r, c)) {
118 continue;
119 }
120 if k < 8 {
121 seen[k] = (r, c);
122 }
123 let v = cost(r, c);
124 if v < best {
125 (best, br, bc) = (v, r, c);
126 }
127 }
128 let (cr, cc) = (br, bc);
129 let inside = |r: i32, c: i32| {
130 (rmin..=rmax).contains(&r)
131 && (cmin..=cmax).contains(&c)
132 && (r - cr).abs() <= self.range
133 && (c - cc).abs() <= self.range
134 };
135 let mut dir = usize::MAX;
136 for step in 0..=self.range {
137 let (pr, pc) = (br, bc);
138 let mut moved = usize::MAX;
139 let ks: &[usize] = if dir == usize::MAX {
140 &[0, 1, 2, 3, 4, 5]
141 } else {
142 &[(dir + 5) % 6, dir, (dir + 1) % 6]
143 };
144 for &k in ks {
145 let (dc, dr) = HEX[k];
146 let (r, c) = (pr + dr, pc + dc);
147 if inside(r, c) {
148 let v = cost(r, c);
149 if v < best {
150 (best, br, bc, moved) = (v, r, c, k);
151 }
152 }
153 }
154 if moved == usize::MAX || step == self.range {
155 break;
156 }
157 dir = moved;
158 }
159 loop {
160 let (pr, pc) = (br, bc);
161 for &(dc, dr) in &DIAMOND {
162 let (r, c) = (pr + dr, pc + dc);
163 if inside(r, c) {
164 let v = cost(r, c);
165 if v < best {
166 (best, br, bc) = (v, r, c);
167 }
168 }
169 }
170 if (br, bc) == (pr, pc) {
171 break;
172 }
173 }
174 let (pr, pc) = (br, bc);
175 for &(dc, dr) in &SQUARE {
176 let (r, c) = (pr + dr, pc + dc);
177 if inside(r, c) && dr != 0 && dc != 0 {
178 let v = cost(r, c);
179 if v < best {
180 (best, br, bc) = (v, r, c);
181 }
182 }
183 }
184 (br, bc, best)
185 }
186
187 fn subpel_cost(&self, mv: Mv, pred: &mut [u8; 64 * 64]) -> u32 {
188 let x16 = ((self.x as isize) << 4) + 2 * mv.col as isize;
189 let y16 = ((self.y as isize) << 4) + 2 * mv.row as isize;
190 if let Some(hp) = self.half
191 && hp.filter() == self.filter
192 && hp.covers(x16 >> 4, y16 >> 4, self.w, self.h)
193 && let Some((p, ps)) = hp.at(x16, y16)
194 {
195 return dist::satd(self.src, self.src_stride, p, ps, self.w, self.h) + self.penalty(mv);
196 }
197 inter::mc(
198 &self.reference,
199 x16,
200 y16,
201 self.w,
202 self.h,
203 self.filter,
204 pred,
205 self.w,
206 );
207 dist::satd(self.src, self.src_stride, pred, self.w, self.w, self.h) + self.penalty(mv)
208 }
209
210 pub fn run(&self, candidates: &[Mv]) -> MotionResult {
211 self.run_with(candidates, &mut [0u8; 64 * 64])
212 }
213
214 pub fn run_with(&self, candidates: &[Mv], pred: &mut [u8; 64 * 64]) -> MotionResult {
215 let (r, c, _) = self.full_pel(candidates);
216 let mut best_mv = Mv::new((r * 8) as i16, (c * 8) as i16);
217 let mut best = self.subpel_cost(best_mv, pred);
218 let mut steps: &[i16] = match self.subpel {
219 Subpel::None => &[],
220 Subpel::Half => &[4],
221 Subpel::Quarter => &[4, 2],
222 Subpel::Full => &[4, 2, 1],
223 };
224 if !self.usehp() && steps.last() == Some(&1) {
225 steps = &steps[..steps.len() - 1];
226 }
227 for &s in steps {
228 let center = best_mv;
229 let mut cross = [u32::MAX; 4];
230 for (k, &(dc, dr)) in DIAMOND.iter().enumerate() {
231 let m = Mv::new(center.row + dr as i16 * s, center.col + dc as i16 * s);
232 if !self.limits.contains(m) {
233 continue;
234 }
235 let v = self.subpel_cost(m, pred);
236 cross[k] = v;
237 if v < best {
238 (best, best_mv) = (v, m);
239 }
240 }
241 let dr = if cross[0] < cross[3] { -s } else { s };
242 let dc = if cross[1] < cross[2] { -s } else { s };
243 let m = Mv::new(center.row + dr, center.col + dc);
244 if self.limits.contains(m) {
245 let v = self.subpel_cost(m, pred);
246 if v < best {
247 (best, best_mv) = (v, m);
248 }
249 }
250 }
251 let rate = self.mv_bits(best_mv);
252 MotionResult {
253 mv: best_mv,
254 dist: best - self.penalty(best_mv),
255 rate,
256 }
257 }
258}