nihav-duck/src/codecs/vp7enc/blocks.rs

   1 use nihav_core::frame::*;
   2 use nihav_codec_support::codecs::{MV, ZERO_MV};
   3 use nihav_codec_support::data::GenericCache;
   4 use super::super::vp78::{PredMode, MVSplitMode, SubMVRef};
   5 use super::super::vp78data::*;
   6 use super::super::vp78dsp::*;
   7 use super::super::vp7data::*;
   8 use super::super::vp7dsp::*;
   9
  10 #[derive(Clone,Copy)]
  11 pub enum MBType {
  12     Intra(PredMode, PredMode),
  13     Intra4x4([PredMode; 16], [u8; 16], PredMode),
  14     InterNoMV(bool, [u8; 4]),
  15     InterNearest(bool, [u8; 4]),
  16     InterNear(bool, [u8; 4]),
  17     InterMV(bool, [u8; 4], MV),
  18     InterSplitMV(bool, [u8; 4], MVSplitMode, [SubMVRef; 16], [MV; 16]),
  19 }
  20
  21 impl MBType {
  22     pub fn is_intra(&self) -> bool {
  23         match *self {
  24             MBType::Intra(_, _) |
  25             MBType::Intra4x4(_, _, _) => true,
  26             _ => false,
  27         }
  28     }
  29     pub fn get_last(&self) -> bool {
  30         match *self {
  31             MBType::InterNoMV(last, _) |
  32             MBType::InterNearest(last, _) |
  33             MBType::InterNear(last, _) |
  34             MBType::InterMV(last, _, _) |
  35             MBType::InterSplitMV(last, _, _, _, _) => last,
  36             _ => false,
  37         }
  38     }
  39 }
  40
  41 impl Default for MBType {
  42     fn default() -> Self { MBType::Intra(PredMode::DCPred, PredMode::DCPred) }
  43 }
  44
  45 pub fn get_block_difference(dst: &mut [i16; 16], src1: &[u8; 16], src2: &[u8; 16]) {
  46     for (dst, (&src1, &src2)) in dst.iter_mut().zip(src1.iter().zip(src2.iter())) {
  47         *dst = i16::from(src1) - i16::from(src2);
  48     }
  49 }
  50 pub fn get_difference_dist(old: &[u8; 16], new: &[u8; 16], diff: &[i16; 16]) -> u32 {
  51     let mut dist = 0;
  52     for ((&old, &new), &diff) in old.iter().zip(new.iter()).zip(diff.iter()) {
  53         let nval = (i16::from(new) + diff).max(0).min(255);
  54         let oval = i16::from(old);
  55         dist += (i32::from(nval - oval) * i32::from(nval - oval)) as u32;
  56     }
  57     dist
  58 }
  59
  60 pub fn requant_y2_dc(val: &mut i16, q: usize) {
  61     *val = *val / Y2_DC_QUANTS[q] * Y2_DC_QUANTS[q];
  62 }
  63
  64 pub trait DCTBlock {
  65     fn has_nz(&self) -> bool;
  66     fn fdct(&mut self);
  67     fn idct(&mut self);
  68     fn requant_y(&mut self, q: usize);
  69     fn quant(&mut self, q: usize, ctype: usize);
  70     fn dequant(&mut self, q: usize, ctype: usize);
  71 }
  72
  73 impl DCTBlock for [i16; 16] {
  74     fn has_nz(&self) -> bool {
  75         for &el in self.iter() {
  76             if el != 0 {
  77                 return true;
  78             }
  79         }
  80         false
  81     }
  82     #[allow(clippy::erasing_op)]
  83     #[allow(clippy::identity_op)]
  84     fn fdct(&mut self) {
  85         let mut tmp = [0i16; 16];
  86         for i in 0..4 {
  87             let s0 = i32::from(self[i + 4 * 0]);
  88             let s1 = i32::from(self[i + 4 * 1]);
  89             let s2 = i32::from(self[i + 4 * 2]);
  90             let s3 = i32::from(self[i + 4 * 3]);
  91
  92             let t0 = (s0 + s3).wrapping_mul(23170) + 0x2000;
  93             let t1 = (s1 + s2).wrapping_mul(23170);
  94             let t2 = s0 - s3;
  95             let t3 = s1 - s2;
  96             let t4 = t2.wrapping_mul(30274) + t3.wrapping_mul(12540) + 0x2000;
  97             let t5 = t2.wrapping_mul(12540) - t3.wrapping_mul(30274) + 0x2000;
  98
  99             tmp[i + 0 * 4] = ((t0 + t1) >> 14) as i16;
 100             tmp[i + 1 * 4] = ( t4       >> 14) as i16;
 101             tmp[i + 2 * 4] = ((t0 - t1) >> 14) as i16;
 102             tmp[i + 3 * 4] = ( t5       >> 14) as i16;
 103         }
 104         for (src, dst) in tmp.chunks(4).zip(self.chunks_mut(4)) {
 105             let s0 = i32::from(src[0]);
 106             let s1 = i32::from(src[1]);
 107             let s2 = i32::from(src[2]);
 108             let s3 = i32::from(src[3]);
 109
 110             let t0 = (s0 + s3).wrapping_mul(23170) + 0x8000;
 111             let t1 = (s1 + s2).wrapping_mul(23170);
 112             let t2 = s0 - s3;
 113             let t3 = s1 - s2;
 114             let t4 = t2.wrapping_mul(30274) + t3.wrapping_mul(12540) + 0x8000;
 115             let t5 = t2.wrapping_mul(12540) - t3.wrapping_mul(30274) + 0x8000;
 116
 117             dst[0] = ((t0 + t1) >> 16) as i16;
 118             dst[1] = ( t4       >> 16) as i16;
 119             dst[2] = ((t0 - t1) >> 16) as i16;
 120             dst[3] = ( t5       >> 16) as i16;
 121         }
 122     }
 123     fn idct(&mut self) { idct4x4(self) }
 124     fn requant_y(&mut self, q: usize) {
 125         self[0] = self[0] / Y_DC_QUANTS[q] * Y_DC_QUANTS[q];
 126         for el in self[1..].iter_mut() {
 127             *el = *el / Y_AC_QUANTS[q] * Y_AC_QUANTS[q];
 128         }
 129     }
 130     fn quant(&mut self, q: usize, ctype: usize) {
 131         let (q_dc, q_ac) = match ctype {
 132                 0 | 3 => (Y_DC_QUANTS[q],  Y_AC_QUANTS[q]),
 133                 2     => (UV_DC_QUANTS[q], UV_AC_QUANTS[q]),
 134                 _     => (Y2_DC_QUANTS[q], Y2_AC_QUANTS[q]),
 135             };
 136         self[0] /= q_dc;
 137         for el in self[1..].iter_mut() {
 138             *el /= q_ac;
 139         }
 140     }
 141     fn dequant(&mut self, q: usize, ctype: usize) {
 142         let (q_dc, q_ac) = match ctype {
 143                 0 | 3 => (Y_DC_QUANTS[q],  Y_AC_QUANTS[q]),
 144                 2     => (UV_DC_QUANTS[q], UV_AC_QUANTS[q]),
 145                 _     => (Y2_DC_QUANTS[q], Y2_AC_QUANTS[q]),
 146             };
 147         self[0] *= q_dc;
 148         for el in self[1..].iter_mut() {
 149             *el *= q_ac;
 150         }
 151     }
 152 }
 153
 154 pub trait IPredBlock16 {
 155     fn ipred16(&mut self, stride: usize, mode: PredMode, ipred: &IPredContext);
 156 }
 157 pub trait IPredBlock8 {
 158     fn ipred8 (&mut self, stride: usize, mode: PredMode, ipred: &IPredContext);
 159 }
 160 pub trait IPredBlock4 {
 161     fn ipred4 (&mut self, stride: usize, mode: PredMode, ipred: &IPredContext);
 162 }
 163
 164 impl IPredBlock16 for [u8; 256] {
 165     fn ipred16(&mut self, stride: usize, mode: PredMode, ipred: &IPredContext) {
 166         match mode {
 167             PredMode::DCPred => IPred16x16::ipred_dc(self, 0, stride, ipred),
 168             PredMode::HPred  => IPred16x16::ipred_h (self, 0, stride, ipred),
 169             PredMode::VPred  => IPred16x16::ipred_v (self, 0, stride, ipred),
 170             PredMode::TMPred => IPred16x16::ipred_tm(self, 0, stride, ipred),
 171             _ => {},
 172         }
 173     }
 174 }
 175 impl IPredBlock8 for [u8; 64] {
 176     fn ipred8(&mut self, stride: usize, mode: PredMode, ipred: &IPredContext) {
 177         match mode {
 178             PredMode::DCPred => IPred8x8::ipred_dc(self, 0, stride, ipred),
 179             PredMode::HPred  => IPred8x8::ipred_h (self, 0, stride, ipred),
 180             PredMode::VPred  => IPred8x8::ipred_v (self, 0, stride, ipred),
 181             PredMode::TMPred => IPred8x8::ipred_tm(self, 0, stride, ipred),
 182             _ => {},
 183         }
 184     }
 185 }
 186 impl IPredBlock4 for &mut [u8] {
 187     fn ipred4(&mut self, stride: usize, mode: PredMode, ipred: &IPredContext) {
 188         match mode {
 189             PredMode::DCPred => IPred4x4::ipred_dc(self, 0, stride, ipred),
 190             PredMode::HPred  => IPred4x4::ipred_he(self, 0, stride, ipred),
 191             PredMode::VPred  => IPred4x4::ipred_ve(self, 0, stride, ipred),
 192             PredMode::TMPred => IPred4x4::ipred_tm(self, 0, stride, ipred),
 193             PredMode::LDPred => IPred4x4::ipred_ld(self, 0, stride, ipred),
 194             PredMode::RDPred => IPred4x4::ipred_rd(self, 0, stride, ipred),
 195             PredMode::VRPred => IPred4x4::ipred_vr(self, 0, stride, ipred),
 196             PredMode::VLPred => IPred4x4::ipred_vl(self, 0, stride, ipred),
 197             PredMode::HDPred => IPred4x4::ipred_hd(self, 0, stride, ipred),
 198             PredMode::HUPred => IPred4x4::ipred_hu(self, 0, stride, ipred),
 199             _ => {},
 200         }
 201     }
 202 }
 203 impl IPredBlock4 for [u8; 16] {
 204     fn ipred4(&mut self, stride: usize, mode: PredMode, ipred: &IPredContext) {
 205         (self as &mut [u8]).ipred4(stride, mode, ipred);
 206     }
 207 }
 208
 209 pub struct LumaIterator<'a> {
 210     luma:   &'a [u8; 256],
 211     blkno:  usize,
 212 }
 213
 214 impl<'a> Iterator for LumaIterator<'a> {
 215     type Item = [u8; 16];
 216     fn next(&mut self) -> Option<Self::Item> {
 217         if self.blkno < 16 {
 218             let mut blk = [0; 16];
 219             let off = (self.blkno & 3) * 4 + (self.blkno >> 2) * 16 * 4;
 220             for (dst, src) in blk.chunks_exact_mut(4).zip(self.luma[off..].chunks(16)) {
 221                 dst.copy_from_slice(&src[..4]);
 222             }
 223             self.blkno += 1;
 224             Some(blk)
 225         } else {
 226             None
 227         }
 228     }
 229 }
 230
 231 pub struct ChromaIterator<'a> {
 232     chroma: &'a [u8; 64],
 233     blkno:  usize,
 234 }
 235
 236 impl<'a> Iterator for ChromaIterator<'a> {
 237     type Item = [u8; 16];
 238     fn next(&mut self) -> Option<Self::Item> {
 239         if self.blkno < 4 {
 240             let mut blk = [0; 16];
 241             let off = (self.blkno & 1) * 4 + (self.blkno >> 1) * 8 * 4;
 242             for (dst, src) in blk.chunks_exact_mut(4).zip(self.chroma[off..].chunks(8)) {
 243                 dst.copy_from_slice(&src[..4]);
 244             }
 245             self.blkno += 1;
 246             Some(blk)
 247         } else {
 248             None
 249         }
 250     }
 251 }
 252
 253 pub struct SrcBlock {
 254     pub luma:   [u8; 256],
 255     pub chroma: [[u8; 64]; 2],
 256 }
 257
 258 impl Default for SrcBlock {
 259     fn default() -> Self {
 260         unsafe { std::mem::zeroed() }
 261     }
 262 }
 263
 264 impl SrcBlock {
 265     pub fn new() -> Self { Self::default() }
 266     pub fn is_flat(&self) -> bool {
 267         let y0 = self.luma[0];
 268         for &el in self.luma[1..].iter() {
 269             if el != y0 {
 270                 return false;
 271             }
 272         }
 273         true
 274     }
 275     pub fn apply_ipred_luma(&self, mode: PredMode, ipred: &IPredContext, dst: &mut Residue) {
 276         let mut tmp = [0; 256];
 277         (&mut tmp).ipred16(16, mode, ipred);
 278         dst.set_luma_from_diff(&self.luma, &tmp);
 279     }
 280     pub fn fill_ipred_luma(&mut self, mode: PredMode, ipred: &IPredContext) {
 281         self.luma.ipred16(16, mode, ipred);
 282     }
 283     pub fn apply_ipred_chroma(&self, mode: PredMode, ipred_u: &IPredContext, ipred_v: &IPredContext, dst: &mut Residue) {
 284         let mut tmp = [[0u8; 64]; 2];
 285         tmp[0].ipred8(8, mode, ipred_u);
 286         tmp[1].ipred8(8, mode, ipred_v);
 287         dst.set_chroma_from_diff(&self.chroma, &tmp);
 288     }
 289     pub fn fill_ipred_chroma(&mut self, mode: PredMode, ipred_u: &IPredContext, ipred_v: &IPredContext) {
 290         self.chroma[0].ipred8(8, mode, ipred_u);
 291         self.chroma[1].ipred8(8, mode, ipred_v);
 292     }
 293
 294     pub fn luma_blocks(&self) -> LumaIterator {
 295         LumaIterator{ luma: &self.luma, blkno: 0 }
 296     }
 297     pub fn chroma_blocks(&self, plane: usize) -> ChromaIterator {
 298         ChromaIterator{ chroma: &self.chroma[plane], blkno: 0 }
 299     }
 300 }
 301
 302 #[derive(Clone)]
 303 pub struct Residue {
 304     pub luma:   [[i16; 16]; 16],
 305     pub dcs:    [i16; 16],
 306     pub chroma: [[[i16; 16]; 4]; 2],
 307     pub has_dc: bool,
 308     pub q:      u8,
 309 }
 310
 311 impl Default for Residue {
 312     fn default() -> Self {
 313         unsafe { std::mem::zeroed() }
 314     }
 315 }
 316
 317 impl Residue {
 318     pub fn new() -> Self { Self::default() }
 319     pub fn reset(&mut self) {
 320         self.has_dc = false;
 321         self.q = 242;
 322     }
 323     pub fn add_residue(&mut self, dst: &mut SrcBlock) {
 324         self.dequant();
 325         self.idct();
 326
 327         for (dst, src) in dst.luma.chunks_mut(16 * 4).zip(self.luma.chunks(4)) {
 328             for (x, blk) in src.iter().enumerate() {
 329                 for (drow, srow) in dst[x * 4..].chunks_mut(16).zip(blk.chunks(4)) {
 330                     for (del, &sel) in drow.iter_mut().zip(srow.iter()) {
 331                         *del = (i16::from(*del) + sel).max(0).min(255) as u8;
 332                     }
 333                 }
 334             }
 335         }
 336         for (dchroma, schroma) in dst.chroma.iter_mut().zip(self.chroma.iter()) {
 337             for (dst, src) in dchroma.chunks_mut(8 * 4).zip(schroma.chunks(2)) {
 338                 for (x, blk) in src.iter().enumerate() {
 339                     for (drow, srow) in dst[x * 4..].chunks_mut(8).zip(blk.chunks(4)) {
 340                         for (del, &sel) in drow.iter_mut().zip(srow.iter()) {
 341                             *del = (i16::from(*del) + sel).max(0).min(255) as u8;
 342                         }
 343                     }
 344                 }
 345             }
 346         }
 347     }
 348     pub fn add_residue_chroma(&mut self, dst: &mut SrcBlock) {
 349         let q = self.q as usize;
 350         for (dchroma, schroma) in dst.chroma.iter_mut().zip(self.chroma.iter_mut()) {
 351             for (dst, src) in dchroma.chunks_mut(8 * 4).zip(schroma.chunks_mut(2)) {
 352                 for (x, blk) in src.iter_mut().enumerate() {
 353                     blk[0] *= UV_DC_QUANTS[q];
 354                     for el in blk[1..].iter_mut() {
 355                         if *el != 0 {
 356                             *el *= UV_AC_QUANTS[q];
 357                         }
 358                     }
 359                     blk.idct();
 360                     for (drow, srow) in dst[x * 4..].chunks_mut(8).zip(blk.chunks(4)) {
 361                         for (del, &sel) in drow.iter_mut().zip(srow.iter()) {
 362                             *del = (i16::from(*del) + sel).max(0).min(255) as u8;
 363                         }
 364                     }
 365                 }
 366             }
 367         }
 368     }
 369     pub fn set_luma_from_diff(&mut self, blk1: &[u8; 256], blk2: &[u8; 256]) {
 370         for (dst, (src1, src2)) in self.luma.chunks_mut(4).zip(blk1.chunks(16 * 4).zip(blk2.chunks(16 * 4))) {
 371             for (x, blk) in dst.iter_mut().enumerate() {
 372                 for (dst, (row1, row2)) in blk.chunks_mut(4).zip(src1[x * 4..].chunks(16).zip(src2[x * 4..].chunks(16))) {
 373                     for (dst, (&a, &b)) in dst.iter_mut().zip(row1.iter().zip(row2.iter())) {
 374                         *dst = i16::from(a) - i16::from(b);
 375                     }
 376                 }
 377             }
 378         }
 379     }
 380     pub fn set_chroma_from_diff(&mut self, blk1: &[[u8; 64]; 2], blk2: &[[u8; 64]; 2]) {
 381         for (chroma, (src1, src2)) in self.chroma.iter_mut().zip(blk1.iter().zip(blk2.iter())) {
 382             for (dst, (src1, src2)) in chroma.chunks_mut(2).zip(src1.chunks(8 * 4).zip(src2.chunks(8 * 4))) {
 383                 for (x, blk) in dst.iter_mut().enumerate() {
 384                     for (dst, (row1, row2)) in blk.chunks_mut(4).zip(src1[x * 4..].chunks(8).zip(src2[x * 4..].chunks(8))) {
 385                         for (dst, (&a, &b)) in dst.iter_mut().zip(row1.iter().zip(row2.iter())) {
 386                             *dst = i16::from(a) - i16::from(b);
 387                         }
 388                     }
 389                 }
 390             }
 391         }
 392     }
 393     pub fn fdct(&mut self) {
 394         self.fdct_luma();
 395         self.fdct_chroma();
 396     }
 397     pub fn fdct_luma(&mut self) {
 398         for blk in self.luma.iter_mut() {
 399             blk.fdct();
 400         }
 401     }
 402     pub fn fdct_chroma(&mut self) {
 403         for chroma in self.chroma.iter_mut() {
 404             for blk in chroma.iter_mut() {
 405                 blk.fdct();
 406             }
 407         }
 408     }
 409     pub fn fdct_dc_block(&mut self) {
 410         for (dc, blk) in self.dcs.iter_mut().zip(self.luma.iter_mut()) {
 411             *dc = blk[0];
 412             blk[0] = 0;
 413         }
 414         self.dcs.fdct();
 415         self.has_dc = true;
 416     }
 417     pub fn idct(&mut self) {
 418         self.idct_luma();
 419         self.idct_chroma();
 420     }
 421     pub fn idct_luma(&mut self) {
 422         if self.has_dc {
 423             self.dcs.idct();
 424             for (&dc, blk) in self.dcs.iter().zip(self.luma.iter_mut()) {
 425                 blk[0] = dc;
 426             }
 427         }
 428         for blk in self.luma.iter_mut() {
 429             blk.idct();
 430         }
 431     }
 432     pub fn idct_chroma(&mut self) {
 433         for chroma in self.chroma.iter_mut() {
 434             for blk in chroma.iter_mut() {
 435                 blk.idct();
 436             }
 437         }
 438     }
 439     pub fn quant(&mut self, q: usize) {
 440         self.quant_luma(q);
 441         self.quant_chroma(q);
 442         self.q = q as u8;
 443     }
 444     pub fn quant_luma(&mut self, q: usize) {
 445         if self.has_dc {
 446             self.dcs[0] /= Y2_DC_QUANTS[q];
 447             for el in self.dcs[1..].iter_mut() {
 448                 if *el != 0 {
 449                     *el /= Y2_AC_QUANTS[q];
 450                 }
 451             }
 452         }
 453         for blk in self.luma.iter_mut() {
 454             blk[0] /= Y_DC_QUANTS[q];
 455             for el in blk[1..].iter_mut() {
 456                 if *el != 0 {
 457                     *el /= Y_AC_QUANTS[q];
 458                 }
 459             }
 460         }
 461         self.q = q as u8;
 462     }
 463     pub fn quant_chroma(&mut self, q: usize) {
 464         for chroma in self.chroma.iter_mut() {
 465             for blk in chroma.iter_mut() {
 466                 blk[0] /= UV_DC_QUANTS[q];
 467                 for el in blk[1..].iter_mut() {
 468                     if *el != 0 {
 469                         *el /= UV_AC_QUANTS[q];
 470                     }
 471                 }
 472             }
 473         }
 474         self.q = q as u8;
 475     }
 476     pub fn dequant(&mut self) {
 477         self.dequant_luma();
 478         self.dequant_chroma();
 479     }
 480     pub fn dequant_luma(&mut self) {
 481         let q = self.q as usize;
 482         if self.has_dc {
 483             self.dcs[0] *= Y2_DC_QUANTS[q];
 484             for el in self.dcs[1..].iter_mut() {
 485                 if *el != 0 {
 486                     *el *= Y2_AC_QUANTS[q];
 487                 }
 488             }
 489         }
 490         for blk in self.luma.iter_mut() {
 491             blk[0] *= Y_DC_QUANTS[q];
 492             for el in blk[1..].iter_mut() {
 493                 if *el != 0 {
 494                     *el *= Y_AC_QUANTS[q];
 495                 }
 496             }
 497         }
 498     }
 499     pub fn dequant_chroma(&mut self) {
 500         let q = self.q as usize;
 501         for chroma in self.chroma.iter_mut() {
 502             for blk in chroma.iter_mut() {
 503                 blk[0] *= UV_DC_QUANTS[q];
 504                 for el in blk[1..].iter_mut() {
 505                     if *el != 0 {
 506                         *el *= UV_AC_QUANTS[q];
 507                     }
 508                 }
 509             }
 510         }
 511     }
 512 }
 513
 514 pub fn load_blocks(src: &NAVideoBuffer<u8>, sblocks: &mut Vec<SrcBlock>) {
 515     let data = src.get_data();
 516     let y = &data[src.get_offset(0)..];
 517     let u = &data[src.get_offset(1)..];
 518     let v = &data[src.get_offset(2)..];
 519     let ystride = src.get_stride(0);
 520     let ustride = src.get_stride(1);
 521     let vstride = src.get_stride(2);
 522     let (width, height) = src.get_dimensions(0);
 523
 524     sblocks.clear();
 525     for (ystrip, (ustrip, vstrip)) in y.chunks(ystride * 16).take((height + 15) / 16).zip(u.chunks(ustride * 8).zip(v.chunks(vstride * 8))) {
 526         for x in (0..width).step_by(16) {
 527             let mut sblk = SrcBlock::default();
 528
 529             for (dst, src) in sblk.luma.chunks_mut(16).zip(ystrip[x..].chunks(ystride)) {
 530                 dst.copy_from_slice(&src[..16]);
 531             }
 532             for (dst, src) in sblk.chroma[0].chunks_mut(8).zip(ustrip[x / 2..].chunks(ustride)) {
 533                 dst.copy_from_slice(&src[..8]);
 534             }
 535             for (dst, src) in sblk.chroma[1].chunks_mut(8).zip(vstrip[x / 2..].chunks(vstride)) {
 536                 dst.copy_from_slice(&src[..8]);
 537             }
 538             sblocks.push(sblk);
 539         }
 540     }
 541 }
 542
 543 pub struct YModePred {
 544     pub cache:  GenericCache<PredMode>,
 545 }
 546
 547 impl YModePred {
 548     fn resize(&mut self, mb_w: usize) {
 549         self.cache = GenericCache::new(4, mb_w * 4 + 1, PredMode::DCPred);
 550     }
 551     pub fn set_mode(&mut self, mb_x: usize, mode: PredMode) {
 552         for row in self.cache.data[self.cache.xpos + mb_x * 4..].chunks_mut(self.cache.stride).take(4) {
 553             for el in row[..4].iter_mut() {
 554                 *el = mode.to_b_mode();
 555             }
 556         }
 557     }
 558     pub fn set_modes4x4(&mut self, mb_x: usize, imodes: &[PredMode; 16], ctx: &mut [u8; 16]) {
 559         let mut off = self.cache.xpos + mb_x * 4;
 560         for y in 0..4 {
 561             for x in 0..4 {
 562                 let top_idx = self.cache.data[off + x - self.cache.stride].to_b_index();
 563                 let left_idx = self.cache.data[off + x - 1].to_b_index();
 564                 self.cache.data[off + x] = imodes[x + y * 4];
 565                 ctx[x + y * 4] = ((top_idx * 10) + left_idx) as u8;
 566             }
 567             off += self.cache.stride;
 568         }
 569     }
 570 }
 571
 572 impl Default for YModePred {
 573     fn default() -> Self {
 574         Self {
 575             cache: GenericCache::new(0, 0, PredMode::DCPred)
 576         }
 577     }
 578 }
 579
 580 #[derive(Default)]
 581 pub struct BlockPCtx {
 582     pub nz_y2:      u8,
 583     pub nz_y_top:   [bool; 4],
 584     pub nz_y_left:  [bool; 4],
 585     pub nz_c_top:   [[bool; 2]; 2],
 586     pub nz_c_left:  [[bool; 2]; 2],
 587 }
 588
 589 #[derive(Default)]
 590 pub struct PredContext {
 591     pub mb_w:       usize,
 592     pub mb_h:       usize,
 593
 594     pub top_line_y: Vec<u8>,
 595     pub top_line_u: Vec<u8>,
 596     pub top_line_v: Vec<u8>,
 597     pub tl_y:       u8,
 598     pub tl_u:       u8,
 599     pub tl_v:       u8,
 600
 601     pub left_y:     [u8; 16],
 602     pub left_u:     [u8; 16],
 603     pub left_v:     [u8; 16],
 604
 605     pub dc_last:    [i16; 2],
 606     pub dc_count:   [usize; 2],
 607     dc_last_saved:  [i16; 2],
 608     dc_count_saved: [usize; 2],
 609     pub nz_y2_top:  Vec<bool>,
 610     pub nz_y2_left: bool,
 611     pub nz_y_top:   Vec<bool>,
 612     pub nz_y_left:  [bool; 4],
 613     pub nz_c_top:   [Vec<bool>; 2],
 614     pub nz_c_left:  [[bool; 2]; 2],
 615
 616     pub ymodes:     YModePred,
 617
 618     pub mvs:        Vec<MV>,
 619     pub mv_stride:  usize,
 620     pub version:    u8,
 621 }
 622
 623 impl PredContext {
 624     pub fn new() -> Self { Self::default() }
 625     pub fn resize(&mut self, mb_w: usize, mb_h: usize) {
 626         self.mb_w = mb_w;
 627         self.mb_h = mb_h;
 628
 629         self.top_line_y.resize(mb_w * 16 + 1, 0);
 630         self.top_line_u.resize(mb_w * 8 + 1, 0);
 631         self.top_line_v.resize(mb_w * 8 + 1, 0);
 632
 633         self.nz_y2_top.resize(mb_w, false);
 634         self.nz_y_top.resize(mb_w * 4, false);
 635         self.nz_c_top[0].resize(mb_w * 2, false);
 636         self.nz_c_top[1].resize(mb_w * 2, false);
 637
 638         self.ymodes.resize(mb_w);
 639
 640         self.mv_stride = mb_w * 4;
 641         self.mvs.resize(self.mv_stride * mb_h * 4, ZERO_MV);
 642     }
 643
 644     pub fn reset(&mut self) {
 645         for el in self.top_line_y.iter_mut() { *el = 0x80; }
 646         for el in self.top_line_u.iter_mut() { *el = 0x80; }
 647         for el in self.top_line_v.iter_mut() { *el = 0x80; }
 648         self.left_y = [0x80; 16];
 649         self.left_u = [0x80; 16];
 650         self.left_v = [0x80; 16];
 651         self.tl_y = 0x80;
 652         self.tl_u = 0x80;
 653         self.tl_v = 0x80;
 654
 655         for el in self.nz_y_top.iter_mut() { *el = false; }
 656         self.nz_y_left = [false; 4];
 657         for el in self.nz_y2_top.iter_mut() { *el = false; }
 658         self.nz_y2_left = false;
 659         for el in self.nz_c_top[0].iter_mut() { *el = false; }
 660         for el in self.nz_c_top[1].iter_mut() { *el = false; }
 661         self.nz_c_left = [[false; 2]; 2];
 662
 663         self.ymodes.cache.reset();
 664
 665         for mv in self.mvs.iter_mut() { *mv = ZERO_MV; }
 666     }
 667     pub fn reset_intra(&mut self) {
 668         self.dc_last  = [0; 2];
 669         self.dc_count = [0; 2];
 670         self.dc_last_saved  = [0; 2];
 671         self.dc_count_saved = [0; 2];
 672     }
 673     pub fn save_dc_pred(&mut self) {
 674         self.dc_last_saved  = self.dc_last;
 675         self.dc_count_saved = self.dc_count;
 676     }
 677     #[allow(dead_code)]
 678     pub fn restore_dc_pred(&mut self) {
 679         self.dc_last  = self.dc_last_saved;
 680         self.dc_count = self.dc_count_saved;
 681     }
 682     pub fn update_mb_row(&mut self) {
 683         self.left_y = [0x80; 16];
 684         self.left_u = [0x80; 16];
 685         self.left_v = [0x80; 16];
 686         self.tl_y = 0x80;
 687         self.tl_u = 0x80;
 688         self.tl_v = 0x80;
 689         self.ymodes.cache.update_row();
 690     }
 691     pub fn update_mb(&mut self, sblk: &SrcBlock, mb_x: usize) {
 692         for (dst, src) in self.left_y.iter_mut().zip(sblk.luma.chunks_exact(16)) {
 693             *dst = src[15];
 694         }
 695         self.tl_y = self.top_line_y[mb_x * 16 + 16];
 696         self.top_line_y[mb_x * 16 + 1..][..16].copy_from_slice(&sblk.luma[15 * 16..]);
 697
 698         for (dst, src) in self.left_u.iter_mut().zip(sblk.chroma[0].chunks_exact(8)) {
 699             *dst = src[7];
 700         }
 701         self.tl_u = self.top_line_u[mb_x * 8 + 8];
 702         self.top_line_u[mb_x * 8 + 1..][..8].copy_from_slice(&sblk.chroma[0][7 * 8..]);
 703
 704         for (dst, src) in self.left_v.iter_mut().zip(sblk.chroma[1].chunks_exact(8)) {
 705             *dst = src[7];
 706         }
 707         self.tl_v = self.top_line_v[mb_x * 8 + 8];
 708         self.top_line_v[mb_x * 8 + 1..][..8].copy_from_slice(&sblk.chroma[1][7 * 8..]);
 709     }
 710     pub fn fill_ipred(&mut self, plane: usize, mb_x: usize, ipred: &mut IPredContext) {
 711         match plane {
 712             0 => {
 713                 if ipred.has_top {
 714                     ipred.top.copy_from_slice(&self.top_line_y[mb_x * 16 + 1..][..16]);
 715                     ipred.tl = self.tl_y;
 716                 }
 717                 ipred.left.copy_from_slice(&self.left_y);
 718                 ipred.has_left = mb_x > 0;
 719             },
 720             1 => {
 721                 if ipred.has_top {
 722                     ipred.top[..8].copy_from_slice(&self.top_line_u[mb_x * 8 + 1..][..8]);
 723                     ipred.tl = self.tl_u;
 724                 }
 725                 ipred.left.copy_from_slice(&self.left_u);
 726                 ipred.has_left = mb_x > 0;
 727             },
 728             _ => {
 729                 if ipred.has_top {
 730                     ipred.top[..8].copy_from_slice(&self.top_line_v[mb_x * 8 + 1..][..8]);
 731                     ipred.tl = self.tl_v;
 732                 }
 733                 ipred.left.copy_from_slice(&self.left_v);
 734                 ipred.has_left = mb_x > 0;
 735             },
 736         }
 737     }
 738     pub fn get_ipred_tr(&self, mb_x: usize) -> [u8; 4] {
 739         if mb_x < self.mb_w - 1 {
 740             let mut tr = [0; 4];
 741             tr.copy_from_slice(&self.top_line_y[mb_x * 16 + 1 + 16..][..4]);
 742             tr
 743         } else {
 744             [0x80; 4]
 745         }
 746     }
 747     pub fn fill_pctx(&self, mb_x: usize, pctx: &mut BlockPCtx) {
 748         pctx.nz_y2     = (self.nz_y2_left as u8) + (self.nz_y2_top[mb_x] as u8);
 749         pctx.nz_y_left = self.nz_y_left;
 750         pctx.nz_y_top.copy_from_slice(&self.nz_y_top[mb_x * 4..][..4]);
 751         pctx.nz_c_left = self.nz_c_left;
 752         pctx.nz_c_top  = [[self.nz_c_top[0][mb_x * 2], self.nz_c_top[0][mb_x * 2 + 1]],
 753                           [self.nz_c_top[1][mb_x * 2], self.nz_c_top[1][mb_x * 2 + 1]]];
 754     }
 755     pub fn set_nz(&mut self, mb_x: usize, blk: &Residue) {
 756         if blk.has_dc {
 757             let has_nz = blk.dcs.has_nz();
 758             self.nz_y2_left = has_nz;
 759             self.nz_y2_top[mb_x] = has_nz;
 760         }
 761         for (y, blk_row) in blk.luma.chunks(4).enumerate() {
 762             for (x, blk) in blk_row.iter().enumerate() {
 763                 let has_nz = blk.has_nz();
 764                 self.nz_y_left[y] = has_nz;
 765                 self.nz_y_top[mb_x * 4 + x] = has_nz;
 766             }
 767         }
 768         for (c, chroma) in blk.chroma.iter().enumerate() {
 769             for (y, blk_row) in chroma.chunks(2).enumerate() {
 770                 for (x, blk) in blk_row.iter().enumerate() {
 771                     let has_nz = blk.has_nz();
 772                     self.nz_c_left[c][y] = has_nz;
 773                     self.nz_c_top[c][mb_x * 2 + x] = has_nz;
 774                 }
 775             }
 776         }
 777     }
 778
 779     pub fn get_y2_dc_pred(&self, last: bool) -> i16 {
 780         let ref_id = !last as usize;
 781         if self.dc_count[ref_id] > 3 {
 782             self.dc_last[ref_id]
 783         } else {
 784             0
 785         }
 786     }
 787     pub fn predict_y2_dc(&mut self, dc: &mut i16, last: bool) {
 788         let ref_id = !last as usize;
 789         let pdc = self.dc_last[ref_id];
 790         let orig_dc = *dc;
 791
 792         if self.dc_count[ref_id] > 3 {
 793             *dc -= pdc;
 794         }
 795
 796         if (pdc == 0) || (orig_dc == 0) || ((pdc ^ orig_dc) < 0) {
 797             self.dc_count[ref_id] = 0;
 798         } else if pdc == orig_dc {
 799             self.dc_count[ref_id] += 1;
 800         }
 801         self.dc_last[ref_id] = orig_dc;
 802     }
 803
 804     pub fn fill_mv(&mut self, mb_x: usize, mb_y: usize, mv: MV) {
 805         let mut iidx = mb_x * 4 + mb_y * 4 * self.mv_stride;
 806         for _ in 0..4 {
 807             for x in 0..4 {
 808                 self.mvs[iidx + x] = mv;
 809             }
 810             iidx += self.mb_w * 4;
 811         }
 812     }
 813     pub fn find_mv_pred(&self, mb_x: usize, mb_y: usize) -> ([u8; 4], MV, MV, MV) {
 814         let mut nearest_mv = ZERO_MV;
 815         let mut near_mv = ZERO_MV;
 816
 817         let mut ct: [u8; 4] = [0; 4];
 818
 819         let start = if self.version == 0 { 1 } else { 0 };
 820         let mvwrap = (self.mb_w as isize) + 1;
 821         for (yoff, xoff, weight, blk_no) in CAND_POS.iter() {
 822             let cx = (mb_x as isize) + (*xoff as isize);
 823             let cy = (mb_y as isize) + (*yoff as isize);
 824             let mvpos = cx + cy * mvwrap;
 825             if (mvpos < start) || ((mvpos % mvwrap) == (mvwrap - 1)) {
 826                 ct[0] += weight;
 827                 continue;
 828             }
 829             let cx = (mvpos % mvwrap) as usize;
 830             let cy = (mvpos / mvwrap) as usize;
 831             let bx = (*blk_no as usize) & 3;
 832             let by = (*blk_no as usize) >> 2;
 833             let blk_pos = cx * 4 + bx + (cy * 4 + by) * self.mv_stride;
 834             let mv = self.mvs[blk_pos];
 835             if mv == ZERO_MV {
 836                 ct[0] += weight;
 837                 continue;
 838             }
 839             let idx;
 840             if (nearest_mv == ZERO_MV) || (nearest_mv == mv) {
 841                 nearest_mv = mv;
 842                 idx = 1;
 843             } else if near_mv == ZERO_MV {
 844                 near_mv = mv;
 845                 idx = 2;
 846             } else {
 847                 idx = if mv == near_mv { 2 } else { 3 };
 848             }
 849             ct[idx] += weight;
 850         }
 851         let pred_mv = if ct[1] > ct[2] {
 852                 if ct[1] >= ct[0] { nearest_mv } else { ZERO_MV }
 853             } else {
 854                 if ct[2] >= ct[0] { near_mv } else { ZERO_MV }
 855             };
 856
 857         let mvprobs = [INTER_MODE_PROBS[ct[0] as usize][0],
 858                        INTER_MODE_PROBS[ct[1] as usize][1],
 859                        INTER_MODE_PROBS[ct[2] as usize][2],
 860                        INTER_MODE_PROBS[ct[2] as usize][3]];
 861
 862         (mvprobs, nearest_mv, near_mv, pred_mv)
 863     }
 864 }