pito_codecs/vp9/enc/pred/
intra.rs1use super::*;
2
3#[derive(Clone)]
4pub struct Edges {
5 above: [u8; 65],
6 left: [u8; 32],
7 have_left: bool,
8 have_above: bool,
9 log2: usize,
10}
11
12#[derive(Clone, Copy, Debug)]
13pub struct EdgeFlags {
14 pub have_left: bool,
15 pub have_above: bool,
16 pub not_on_right: bool,
17 pub max_x: usize,
18 pub max_y: usize,
19}
20
21impl Edges {
22 pub fn build(
23 plane: &[u8],
24 stride: usize,
25 x: usize,
26 y: usize,
27 tx_size: u8,
28 f: EdgeFlags,
29 ) -> Edges {
30 let log2 = tx_size as usize + 2;
31 let size = 1usize << log2;
32 let mut above = [0u8; 65];
33 let mut left = [0u8; 32];
34 if !f.have_above {
35 above[1..=2 * size].fill(127);
36 } else {
37 let row = &plane[(y - 1) * stride..];
38 if x + size <= f.max_x + 1 {
39 above[1..=size].copy_from_slice(&row[x..x + size]);
40 } else {
41 for i in 0..size {
42 above[1 + i] = row[(x + i).min(f.max_x)];
43 }
44 }
45 if f.not_on_right && tx_size == 0 {
46 for i in size..2 * size {
47 above[1 + i] = row[(x + i).min(f.max_x)];
48 }
49 } else {
50 let v = above[size];
51 above[1 + size..=2 * size].fill(v);
52 }
53 }
54 above[0] = if f.have_above && f.have_left {
55 plane[(y - 1) * stride + (x - 1).min(f.max_x)]
56 } else if f.have_above {
57 129
58 } else {
59 127
60 };
61 if f.have_left {
62 for i in 0..size {
63 left[i] = plane[(y + i).min(f.max_y) * stride + x - 1];
64 }
65 } else {
66 left[..size].fill(129);
67 }
68 Edges {
69 above,
70 left,
71 have_left: f.have_left,
72 have_above: f.have_above,
73 log2,
74 }
75 }
76
77 pub fn size(&self) -> usize {
78 1 << self.log2
79 }
80
81 pub fn predict(&self, mode: u8, dst: &mut [u8], dst_stride: usize) {
82 match self.log2 {
83 2 => self.predict_n::<4>(mode, dst, dst_stride),
84 3 => self.predict_n::<8>(mode, dst, dst_stride),
85 4 => self.predict_n::<16>(mode, dst, dst_stride),
86 _ => self.predict_n::<32>(mode, dst, dst_stride),
87 }
88 }
89
90 fn predict_n<const N: usize>(&self, mode: u8, dst: &mut [u8], s: usize) {
91 let above = &self.above;
92 let left = &self.left;
93 let a = |i: isize| above[(i + 1) as usize] as i32;
94 let l = |i: usize| left[i] as i32;
95 let r1 = |v: i32| ((v + 1) >> 1) as u8;
96 let r2 = |v: i32| ((v + 2) >> 2) as u8;
97 match mode {
98 V_PRED => {
99 let row: &[u8; N] = above[1..=N].try_into().unwrap();
100 for i in 0..N {
101 dst[i * s..i * s + N].copy_from_slice(row);
102 }
103 }
104 H_PRED => {
105 for i in 0..N {
106 dst[i * s..i * s + N].fill(left[i]);
107 }
108 }
109 TM_PRED => {
110 let tl = a(-1);
111 for i in 0..N {
112 let d = l(i) - tl;
113 let row = &mut dst[i * s..i * s + N];
114 for j in 0..N {
115 row[j] = (above[1 + j] as i32 + d).clamp(0, 255) as u8;
116 }
117 }
118 }
119 D45_PRED => {
120 let mut e = [0u8; 64];
121 for k in 0..2 * N - 1 {
122 e[k] = if k + 2 < 2 * N {
123 r2(a(k as isize) + 2 * a(k as isize + 1) + a(k as isize + 2))
124 } else {
125 above[2 * N]
126 };
127 }
128 e[2 * N - 1] = above[2 * N];
129 for i in 0..N {
130 dst[i * s..i * s + N].copy_from_slice(&e[i..i + N]);
131 }
132 }
133 D63_PRED => {
134 for i in 0..N {
135 let row = &mut dst[i * s..i * s + N];
136 for j in 0..N {
137 let k = (i / 2 + j) as isize;
138 row[j] = if i & 1 == 1 {
139 r2(a(k) + 2 * a(k + 1) + a(k + 2))
140 } else {
141 r1(a(k) + a(k + 1))
142 };
143 }
144 }
145 }
146 D207_PRED => {
147 let mut p = [[0u8; N]; N];
148 for j in 0..N {
149 p[N - 1][j] = left[N - 1];
150 }
151 for i in 0..N - 1 {
152 p[i][0] = r1(l(i) + l(i + 1));
153 }
154 for i in 0..N.saturating_sub(2) {
155 p[i][1] = r2(l(i) + 2 * l(i + 1) + l(i + 2));
156 }
157 p[N - 2][1] = r2(l(N - 2) + 3 * l(N - 1));
158 for j in 2..N {
159 for i in (0..N - 1).rev() {
160 p[i][j] = p[i + 1][j - 2];
161 }
162 }
163 put(&p, dst, s);
164 }
165 D117_PRED => {
166 let mut p = [[0u8; N]; N];
167 for j in 0..N {
168 p[0][j] = r1(a(j as isize - 1) + a(j as isize));
169 }
170 p[1][0] = r2(l(0) + 2 * a(-1) + a(0));
171 for j in 1..N {
172 p[1][j] = r2(a(j as isize - 2) + 2 * a(j as isize - 1) + a(j as isize));
173 }
174 p[2][0] = r2(a(-1) + 2 * l(0) + l(1));
175 for i in 3..N {
176 p[i][0] = r2(l(i - 3) + 2 * l(i - 2) + l(i - 1));
177 }
178 for i in 2..N {
179 for j in 1..N {
180 p[i][j] = p[i - 2][j - 1];
181 }
182 }
183 put(&p, dst, s);
184 }
185 D135_PRED => {
186 let mut p = [[0u8; N]; N];
187 p[0][0] = r2(l(0) + 2 * a(-1) + a(0));
188 for j in 1..N {
189 p[0][j] = r2(a(j as isize - 2) + 2 * a(j as isize - 1) + a(j as isize));
190 }
191 p[1][0] = r2(a(-1) + 2 * l(0) + l(1));
192 for i in 2..N {
193 p[i][0] = r2(l(i - 2) + 2 * l(i - 1) + l(i));
194 }
195 for i in 1..N {
196 for j in 1..N {
197 p[i][j] = p[i - 1][j - 1];
198 }
199 }
200 put(&p, dst, s);
201 }
202 D153_PRED => {
203 let mut p = [[0u8; N]; N];
204 p[0][0] = r1(l(0) + a(-1));
205 for i in 1..N {
206 p[i][0] = r1(l(i - 1) + l(i));
207 }
208 p[0][1] = r2(l(0) + 2 * a(-1) + a(0));
209 p[1][1] = r2(a(-1) + 2 * l(0) + l(1));
210 for i in 2..N {
211 p[i][1] = r2(l(i - 2) + 2 * l(i - 1) + l(i));
212 }
213 for j in 2..N {
214 p[0][j] = r2(a(j as isize - 3) + 2 * a(j as isize - 2) + a(j as isize - 1));
215 }
216 for i in 1..N {
217 for j in 2..N {
218 p[i][j] = p[i - 1][j - 2];
219 }
220 }
221 put(&p, dst, s);
222 }
223 _ => {
224 let log2 = N.trailing_zeros();
225 let sl = || left[..N].iter().map(|&v| v as u32).sum::<u32>();
226 let sa = || above[1..=N].iter().map(|&v| v as u32).sum::<u32>();
227 let v = match (self.have_left, self.have_above) {
228 (true, true) => (sl() + sa() + N as u32) >> (log2 + 1),
229 (true, false) => (sl() + (N as u32 >> 1)) >> log2,
230 (false, true) => (sa() + (N as u32 >> 1)) >> log2,
231 (false, false) => 128,
232 } as u8;
233 for i in 0..N {
234 dst[i * s..i * s + N].fill(v);
235 }
236 }
237 }
238 }
239}
240
241#[inline(always)]
242fn put<const N: usize>(p: &[[u8; N]; N], dst: &mut [u8], s: usize) {
243 for i in 0..N {
244 dst[i * s..i * s + N].copy_from_slice(&p[i]);
245 }
246}
247
248pub fn predict_in_place(
249 plane: &mut [u8],
250 stride: usize,
251 x: usize,
252 y: usize,
253 mode: u8,
254 tx_size: u8,
255 f: EdgeFlags,
256) {
257 let e = Edges::build(plane, stride, x, y, tx_size, f);
258 e.predict(mode, &mut plane[y * stride + x..], stride);
259}