]> git.nihav.org Git - nihav.git/blame - nihav-codec-support/src/codecs/blockdsp.rs
codec-support: add division support for FFTComplex
[nihav.git] / nihav-codec-support / src / codecs / blockdsp.rs
CommitLineData
c6004662 1//! Various pixel block manipulation functions.
b4d5b851 2use nihav_core::frame::*;
cf64af13 3
c6004662 4/// Puts YUV420 16x16 macroblock data onto picture in the requested place.
cf64af13
KS
5pub fn put_blocks(buf: &mut NAVideoBuffer<u8>, xpos: usize, ypos: usize, blk: &[[i16;64]; 6]) {
6 let stridey = buf.get_stride(0);
7 let strideu = buf.get_stride(1);
8 let stridev = buf.get_stride(2);
9 let mut idxy = buf.get_offset(0) + xpos * 16 + ypos * 16 * stridey;
10 let mut idxu = buf.get_offset(1) + xpos * 8 + ypos * 8 * strideu;
11 let mut idxv = buf.get_offset(2) + xpos * 8 + ypos * 8 * stridev;
12
1a967e6b 13 let data = buf.get_data_mut().unwrap();
9037cf6b 14 let framebuf: &mut [u8] = data.as_mut_slice();
cf64af13
KS
15
16 for j in 0..8 {
17 for k in 0..8 {
18 let mut v = blk[0][k + j * 8];
e243ceb4 19 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
20 framebuf[idxy + k] = v as u8;
21 }
22 for k in 0..8 {
23 let mut v = blk[1][k + j * 8];
e243ceb4 24 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
25 framebuf[idxy + k + 8] = v as u8;
26 }
27 idxy += stridey;
28 }
29 for j in 0..8 {
30 for k in 0..8 {
31 let mut v = blk[2][k + j * 8];
e243ceb4 32 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
33 framebuf[idxy + k] = v as u8;
34 }
35 for k in 0..8 {
36 let mut v = blk[3][k + j * 8];
e243ceb4 37 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
38 framebuf[idxy + k + 8] = v as u8;
39 }
40 idxy += stridey;
41 }
42
43 for j in 0..8 {
44 for k in 0..8 {
45 let mut v = blk[4][k + j * 8];
e243ceb4 46 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
47 framebuf[idxu + k] = v as u8;
48 }
49 for k in 0..8 {
50 let mut v = blk[5][k + j * 8];
e243ceb4 51 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
52 framebuf[idxv + k] = v as u8;
53 }
54 idxu += strideu;
55 idxv += stridev;
56 }
57}
58
c6004662 59/// Adds YUV420 16x16 macroblock coefficients to the picture in the requested place.
cf64af13
KS
60pub fn add_blocks(buf: &mut NAVideoBuffer<u8>, xpos: usize, ypos: usize, blk: &[[i16;64]; 6]) {
61 let stridey = buf.get_stride(0);
62 let strideu = buf.get_stride(1);
63 let stridev = buf.get_stride(2);
64 let mut idxy = buf.get_offset(0) + xpos * 16 + ypos * 16 * stridey;
65 let mut idxu = buf.get_offset(1) + xpos * 8 + ypos * 8 * strideu;
66 let mut idxv = buf.get_offset(2) + xpos * 8 + ypos * 8 * stridev;
67
1a967e6b 68 let data = buf.get_data_mut().unwrap();
9037cf6b 69 let framebuf: &mut [u8] = data.as_mut_slice();
cf64af13
KS
70
71 for j in 0..8 {
72 for k in 0..8 {
e243ceb4
KS
73 let mut v = blk[0][k + j * 8] + i16::from(framebuf[idxy + k]);
74 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
75 framebuf[idxy + k] = v as u8;
76 }
77 for k in 0..8 {
e243ceb4
KS
78 let mut v = blk[1][k + j * 8] + i16::from(framebuf[idxy + k + 8]);
79 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
80 framebuf[idxy + k + 8] = v as u8;
81 }
82 idxy += stridey;
83 }
84 for j in 0..8 {
85 for k in 0..8 {
e243ceb4
KS
86 let mut v = blk[2][k + j * 8] + i16::from(framebuf[idxy + k]);
87 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
88 framebuf[idxy + k] = v as u8;
89 }
90 for k in 0..8 {
e243ceb4
KS
91 let mut v = blk[3][k + j * 8] + i16::from(framebuf[idxy + k + 8]);
92 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
93 framebuf[idxy + k + 8] = v as u8;
94 }
95 idxy += stridey;
96 }
97
98 for j in 0..8 {
99 for k in 0..8 {
e243ceb4
KS
100 let mut v = blk[4][k + j * 8] + i16::from(framebuf[idxu + k]);
101 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
102 framebuf[idxu + k] = v as u8;
103 }
104 for k in 0..8 {
e243ceb4
KS
105 let mut v = blk[5][k + j * 8] + i16::from(framebuf[idxv + k]);
106 if v < 0 { v = 0; } else if v > 255 { v = 255; }
cf64af13
KS
107 framebuf[idxv + k] = v as u8;
108 }
109 idxu += strideu;
110 idxv += stridev;
111 }
112}
113
c6004662 114/// Copies block from the picture with pixels beyond the picture borders being replaced with replicated edge pixels.
86081fed 115pub fn edge_emu(src: &NAVideoBuffer<u8>, xpos: isize, ypos: isize, bw: usize, bh: usize, dst: &mut [u8], dstride: usize, comp: usize, align: u8) {
cf64af13
KS
116 let stride = src.get_stride(comp);
117 let offs = src.get_offset(comp);
86081fed
KS
118 let (w_, h_) = src.get_dimensions(comp);
119 let (hss, vss) = src.get_info().get_format().get_chromaton(comp).unwrap().get_subsampling();
cf64af13
KS
120 let data = src.get_data();
121 let framebuf: &[u8] = data.as_slice();
122
86081fed
KS
123 let (w, h) = if align == 0 {
124 (w_, h_)
125 } else {
126 let wa = if align > hss { (1 << (align - hss)) - 1 } else { 0 };
127 let ha = if align > vss { (1 << (align - vss)) - 1 } else { 0 };
128 ((w_ + wa) & !wa, (h_ + ha) & !ha)
129 };
130
cf64af13
KS
131 for y in 0..bh {
132 let srcy;
133 if (y as isize) + ypos < 0 { srcy = 0; }
134 else if (y as isize) + ypos >= (h as isize) { srcy = h - 1; }
135 else { srcy = ((y as isize) + ypos) as usize; }
136
137 for x in 0..bw {
138 let srcx;
139 if (x as isize) + xpos < 0 { srcx = 0; }
140 else if (x as isize) + xpos >= (w as isize) { srcx = w - 1; }
141 else { srcx = ((x as isize) + xpos) as usize; }
142 dst[x + y * dstride] = framebuf[offs + srcx + srcy * stride];
143 }
144 }
145}
146
c6004662
KS
147/// A generic type for motion interpolation function used by [`copy_blocks`]
148///
149/// The function expects following parameters:
150/// * destination buffer
151/// * destination buffer stride
152/// * source buffer
153/// * source buffer stride
154/// * block width
155/// * block height
156///
157/// [`copy_blocks`]: ./fn.copy_blocks.html
fdb4b2fb
KS
158pub type BlkInterpFunc = fn(&mut [u8], usize, &[u8], usize, usize, usize);
159
c6004662
KS
160/// Performs motion compensation on YUV420 macroblock.
161///
162/// Arguments:
163/// * `dx` and `dy` - destination coordinates
164/// * `sx` and `sy` - source coordinates
165/// * `bw` and `bh` - block dimensions
166/// * `preborder` and `postborder` - number of pixels before and after interpolated one used by the interpolation filter.
167/// * `mode` - interpolation mode (essentially the index for the `interp` array)
cf64af13
KS
168pub fn copy_blocks(dst: &mut NAVideoBuffer<u8>, src: &NAVideoBuffer<u8>,
169 dx: usize, dy: usize, sx: isize, sy: isize, bw: usize, bh: usize,
170 preborder: usize, postborder: usize,
fdb4b2fb 171 mode: usize, interp: &[BlkInterpFunc])
cf64af13
KS
172{
173 let pre = if mode != 0 { preborder as isize } else { 0 };
174 let post = if mode != 0 { postborder as isize } else { 0 };
175 let (w, h) = src.get_dimensions(0);
176
177 if (sx - pre < 0) || ((sx >> 1) - pre < 0) || (sx + (bw as isize) + post > (w as isize)) ||
178 (sy - pre < 0) || ((sy >> 1) - pre < 0) || (sy + (bh as isize) + post > (h as isize)) {
179 let ebuf_stride: usize = 32;
e243ceb4 180 let mut ebuf: Vec<u8> = vec![0; ebuf_stride * (bh + ((pre + post) as usize))];
cf64af13
KS
181
182 for comp in 0..3 {
2bfdf329 183 let dstride = dst.get_stride(comp);
cf64af13 184 let doff = dst.get_offset(comp);
1a967e6b 185 let ddta = dst.get_data_mut().unwrap();
4a9d2671 186 let dbuf: &mut [u8] = ddta.as_mut_slice();
cf64af13
KS
187 let x = if comp > 0 { dx/2 } else { dx };
188 let y = if comp > 0 { dy/2 } else { dy };
189 let sx_ = (if comp > 0 { sx >> 1 } else { sx }) - pre;
190 let sy_ = (if comp > 0 { sy >> 1 } else { sy }) - pre;
191 let bw_ = (if comp > 0 { bw/2 } else { bw }) + ((pre + post) as usize);
192 let bh_ = (if comp > 0 { bh/2 } else { bh }) + ((pre + post) as usize);
193 edge_emu(src, sx_ - pre, sy_ - pre, bw_, bh_,
86081fed 194 ebuf.as_mut_slice(), ebuf_stride, comp, 0);
cf64af13
KS
195 let bw_ = if comp > 0 { bw/2 } else { bw };
196 let bh_ = if comp > 0 { bh/2 } else { bh };
197 (interp[mode])(&mut dbuf[doff + x + y * dstride..], dstride, ebuf.as_slice(), ebuf_stride, bw_, bh_);
198 }
cf64af13
KS
199 } else {
200 for comp in 0..3 {
201 let sstride = src.get_stride(comp);
202 let soff = src.get_offset(comp);
203 let sdta = src.get_data();
204 let sbuf: &[u8] = sdta.as_slice();
2bfdf329 205 let dstride = dst.get_stride(comp);
cf64af13 206 let doff = dst.get_offset(comp);
1a967e6b 207 let ddta = dst.get_data_mut().unwrap();
4a9d2671 208 let dbuf: &mut [u8] = ddta.as_mut_slice();
cf64af13
KS
209 let x = if comp > 0 { dx/2 } else { dx };
210 let y = if comp > 0 { dy/2 } else { dy };
211 let sx_ = ((if comp > 0 { sx >> 1 } else { sx }) - pre) as usize;
212 let sy_ = ((if comp > 0 { sy >> 1 } else { sy }) - pre) as usize;
213 let bw_ = if comp > 0 { bw/2 } else { bw };
214 let bh_ = if comp > 0 { bh/2 } else { bh };
215 (interp[mode])(&mut dbuf[doff + x + y * dstride..], dstride, &sbuf[(soff + sx_ + sy_ * sstride)..], sstride, bw_, bh_);
216 }
217 }
218}
c6670064 219
c6004662
KS
220/// Performs motion compensation on arbitrary block on some plane.
221///
222/// See [`copy_blocks`] for the arguments explanation.
223///
224/// [`copy_blocks`]: ./fn.copy_blocks.html
c6670064
KS
225pub fn copy_block(dst: &mut NASimpleVideoFrame<u8>, src: NAVideoBufferRef<u8>, comp: usize,
226 dx: usize, dy: usize, mv_x: i16, mv_y: i16, bw: usize, bh: usize,
227 preborder: usize, postborder: usize,
228 mode: usize, interp: &[BlkInterpFunc])
229{
230 let pre = if mode != 0 { preborder as isize } else { 0 };
231 let post = if mode != 0 { postborder as isize } else { 0 };
232 let (w, h) = src.get_dimensions(comp);
233 let sx = (dx as isize) + (mv_x as isize);
234 let sy = (dy as isize) + (mv_y as isize);
235
236 if (sx - pre < 0) || (sx + (bw as isize) + post > (w as isize)) ||
237 (sy - pre < 0) || (sy + (bh as isize) + post > (h as isize)) {
238 let ebuf_stride: usize = 32;
239 let mut ebuf: Vec<u8> = vec![0; ebuf_stride * (bh + ((pre + post) as usize))];
240
241 let dstride = dst.stride[comp];
242 let doff = dst.offset[comp];
243 let edge = (pre + post) as usize;
244 edge_emu(&src, sx - pre, sy - pre, bw + edge, bh + edge,
86081fed 245 ebuf.as_mut_slice(), ebuf_stride, comp, 0);
c6670064
KS
246 (interp[mode])(&mut dst.data[doff + dx + dy * dstride..], dstride,
247 ebuf.as_slice(), ebuf_stride, bw, bh);
248 } else {
249 let sstride = src.get_stride(comp);
250 let soff = src.get_offset(comp);
251 let sdta = src.get_data();
252 let sbuf: &[u8] = sdta.as_slice();
253 let dstride = dst.stride[comp];
254 let doff = dst.offset[comp];
255 let saddr = soff + ((sx - pre) as usize) + ((sy - pre) as usize) * sstride;
256 (interp[mode])(&mut dst.data[doff + dx + dy * dstride..], dstride,
257 &sbuf[saddr..], sstride, bw, bh);
258 }
259}
0faa067b
KS
260
261fn hpel_interp00(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize)
262{
263 let mut didx = 0;
264 let mut sidx = 0;
265 for _ in 0..bh {
266 dst[didx..][..bw].copy_from_slice(&src[sidx..][..bw]);
267 didx += dstride;
268 sidx += sstride;
269 }
270}
271
272fn hpel_interp01(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize)
273{
274 let mut didx = 0;
275 let mut sidx = 0;
276 for _ in 0..bh {
277 for x in 0..bw { dst[didx + x] = ((u16::from(src[sidx + x]) + u16::from(src[sidx + x + 1]) + 1) >> 1) as u8; }
278 didx += dstride;
279 sidx += sstride;
280 }
281}
282
283fn hpel_interp10(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize)
284{
285 let mut didx = 0;
286 let mut sidx = 0;
287 for _ in 0..bh {
288 for x in 0..bw { dst[didx + x] = ((u16::from(src[sidx + x]) + u16::from(src[sidx + x + sstride]) + 1) >> 1) as u8; }
289 didx += dstride;
290 sidx += sstride;
291 }
292}
293
294fn hpel_interp11(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize)
295{
296 let mut didx = 0;
297 let mut sidx = 0;
298 for _ in 0..bh {
299 for x in 0..bw {
300 dst[didx + x] = ((u16::from(src[sidx + x]) +
301 u16::from(src[sidx + x + 1]) +
302 u16::from(src[sidx + x + sstride]) +
303 u16::from(src[sidx + x + sstride + 1]) + 2) >> 2) as u8;
304 }
305 didx += dstride;
306 sidx += sstride;
307 }
308}
309
310/// Half-pixel interpolation functions.
311pub const HALFPEL_INTERP_FUNCS: &[BlkInterpFunc] = &[
312 hpel_interp00, hpel_interp01, hpel_interp10, hpel_interp11 ];
313