nihav-duck/src/codecs/vp6enc/mod.rs

   1 use nihav_core::codecs::*;
   2 use nihav_core::io::byteio::*;
   3 use super::vp6data::*;
   4 use super::vpcommon::*;
   5
   6 mod coder;
   7 use coder::*;
   8 #[allow(clippy::identity_op)]
   9 #[allow(clippy::needless_range_loop)]
  10 mod dsp;
  11 use dsp::MVSearchMode;
  12 mod huff;
  13 use huff::*;
  14 #[allow(clippy::needless_range_loop)]
  15 mod mb;
  16 use mb::*;
  17 mod models;
  18 use models::*;
  19 mod ratectl;
  20 use ratectl::*;
  21 mod rdo;
  22
  23 const VERSION_VP61: u8 = VERSION_VP60 + 1;
  24
  25 enum VP6Writer<'a, 'b> {
  26     BoolCoder(BoolEncoder<'a, 'b>),
  27     Huffman(HuffEncoder<'a, 'b>),
  28 }
  29
  30 #[derive(Default)]
  31 pub struct VP56DCPred {
  32     dc_y:       Vec<i16>,
  33     dc_u:       Vec<i16>,
  34     dc_v:       Vec<i16>,
  35     ldc_y:      [i16; 2],
  36     ldc_u:      i16,
  37     ldc_v:      i16,
  38     ref_y:      Vec<u8>,
  39     ref_c:      Vec<u8>,
  40     ref_left:   u8,
  41     y_idx:      usize,
  42     c_idx:      usize,
  43
  44     last_dc:    [[i16; 4]; 3],
  45 }
  46
  47 const INVALID_REF: u8 = 42;
  48
  49 impl VP56DCPred {
  50     fn new() -> Self { Self::default() }
  51     fn resize(&mut self, mb_w: usize) {
  52         self.dc_y.resize(mb_w * 2 + 2, 0);
  53         self.dc_u.resize(mb_w     + 2, 0);
  54         self.dc_v.resize(mb_w     + 2, 0);
  55         self.ref_y.resize(mb_w * 2 + 2, INVALID_REF);
  56         self.ref_c.resize(mb_w     + 2, INVALID_REF);
  57         self.ref_c[0] = 0;
  58     }
  59     fn reset(&mut self) {
  60         self.update_row();
  61         for el in self.ref_y.iter_mut().skip(1) { *el = INVALID_REF; }
  62         for el in self.ref_c.iter_mut().skip(1) { *el = INVALID_REF; }
  63
  64         self.last_dc = [[0; 4]; 3];
  65         self.last_dc[0][1] = 0x80;
  66         self.last_dc[0][2] = 0x80;
  67     }
  68     fn update_row(&mut self) {
  69         self.y_idx = 1;
  70         self.c_idx = 1;
  71         self.ldc_y = [0; 2];
  72         self.ldc_u = 0;
  73         self.ldc_v = 0;
  74         self.ref_left = INVALID_REF;
  75     }
  76     fn next_mb(&mut self) {
  77         self.y_idx += 2;
  78         self.c_idx += 1;
  79     }
  80     fn predict_dc(&mut self, mb_type: VPMBType, blk_no: usize, coeffs: &mut [i16; 64], fwd: bool) {
  81         let is_luma = blk_no < 4;
  82         let (plane, dcs) = match blk_no {
  83                 4 => (1, &mut self.dc_u),
  84                 5 => (2, &mut self.dc_v),
  85                 _ => (0, &mut self.dc_y),
  86              };
  87         let (dc_ref, dc_idx) = if is_luma {
  88                 (&mut self.ref_y, self.y_idx + (blk_no & 1))
  89             } else {
  90                 (&mut self.ref_c, self.c_idx)
  91             };
  92         let ref_id = mb_type.get_ref_id();
  93         let mut dc_pred = 0;
  94         let mut count = 0;
  95         let has_left_blk = is_luma && ((blk_no & 1) == 1);
  96         if has_left_blk || self.ref_left == ref_id {
  97             dc_pred += match blk_no {
  98                     0 | 1 => self.ldc_y[0],
  99                     2 | 3 => self.ldc_y[1],
 100                     4     => self.ldc_u,
 101                     _     => self.ldc_v,
 102                 };
 103             count += 1;
 104         }
 105         if dc_ref[dc_idx] == ref_id {
 106             dc_pred += dcs[dc_idx];
 107             count += 1;
 108         }
 109         if count == 0 {
 110             dc_pred = self.last_dc[ref_id as usize][plane];
 111         } else if count == 2 {
 112             dc_pred /= 2;
 113         }
 114         if !fwd {
 115             coeffs[0] += dc_pred;
 116         }
 117
 118         let dc = coeffs[0];
 119         if blk_no != 4 { // update top block reference only for the second chroma component
 120             dc_ref[dc_idx] = ref_id;
 121         }
 122         match blk_no {
 123             0 | 1 => {
 124                 self.ldc_y[0] = dc;
 125             },
 126             2 | 3 => {
 127                 self.ldc_y[1] = dc;
 128             },
 129             4 => {
 130                 self.ldc_u = dc;
 131             },
 132             _ => {
 133                 self.ldc_v = dc;
 134                 self.ref_left = ref_id;
 135             },
 136         };
 137         dcs[dc_idx] = dc;
 138
 139         self.last_dc[ref_id as usize][plane] = dc;
 140         if fwd {
 141             coeffs[0] -= dc_pred;
 142         }
 143     }
 144 }
 145
 146 trait ZeroBlock {
 147     fn no_dc(&self) -> bool;
 148     fn no_ac(&self) -> bool;
 149 }
 150
 151 impl ZeroBlock for [i16; 64] {
 152     fn no_dc(&self) -> bool {
 153         self[0] == 0
 154     }
 155     fn no_ac(&self) -> bool {
 156         for &el in self[1..].iter() {
 157             if el != 0 {
 158                 return false;
 159             }
 160         }
 161         true
 162     }
 163 }
 164
 165 struct VP6Encoder {
 166     stream:     Option<NAStreamRef>,
 167     pkt:        Option<NAPacket>,
 168     key_int:    u8,
 169     frmcount:   u8,
 170     mb_w:       usize,
 171     mb_h:       usize,
 172     dc_pred:    VP56DCPred,
 173     top_ctx:    [Vec<bool>; 4],
 174     mc_buf:     NAVideoBufferRef<u8>,
 175     fenc:       FrameEncoder,
 176     ratectl:    RateControl,
 177
 178     flipped:    bool,
 179     huffman:    bool,
 180
 181     version:    u8,
 182     profile:    u8,
 183
 184     models:     VP56Models,
 185     stats:      VP56ModelsStat,
 186     pmodels:    VP56Models,
 187
 188     last_frame: NABufferType,
 189     gold_frame: NABufferType,
 190     last_gold:  bool,
 191     me_mode:    MVSearchMode,
 192     me_range:   i16,
 193
 194     force_q:    Option<usize>,
 195     fast:       bool,
 196 }
 197
 198 impl VP6Encoder {
 199     fn new(flipped: bool) -> Self {
 200         let vt = alloc_video_buffer(NAVideoInfo::new(24, 24, false, VP_YUVA420_FORMAT), 4).unwrap();
 201         let mc_buf = vt.get_vbuf().unwrap();
 202         Self {
 203             stream:     None,
 204             pkt:        None,
 205             key_int:    10,
 206             frmcount:   0,
 207             mb_w:       0,
 208             mb_h:       0,
 209             dc_pred:    VP56DCPred::new(),
 210             top_ctx:    [Vec::new(), Vec::new(), Vec::new(), Vec::new()],
 211             fenc:       FrameEncoder::new(),
 212             ratectl:    RateControl::new(),
 213             mc_buf,
 214
 215             flipped,
 216             huffman:    false,
 217
 218             version:    VERSION_VP60,
 219             profile:    VP6_SIMPLE_PROFILE,
 220
 221             models:     VP56Models::new(),
 222             pmodels:    VP56Models::new(),
 223             stats:      VP56ModelsStat::new(),
 224
 225             last_frame: NABufferType::None,
 226             gold_frame: NABufferType::None,
 227             last_gold:  false,
 228             me_mode:    MVSearchMode::default(),
 229             me_range:   16,
 230
 231             force_q:    None,
 232             fast:       false,
 233         }
 234     }
 235     fn decide_encoding(&mut self) -> bool {
 236         false
 237     }
 238     fn estimate_blocks(&mut self, is_intra: bool) {
 239         for top_ctx in self.top_ctx.iter_mut() {
 240             for el in top_ctx.iter_mut() {
 241                 *el = false;
 242             }
 243         }
 244         let mut last_mbt = VPMBType::InterNoMV;
 245         let mut mb_idx = 0;
 246         for _mb_y in 0..self.mb_h {
 247             let mut left_dc = [false; 4];
 248             for mb_x in 0..self.mb_w {
 249                 let mb_type = self.fenc.mb_types[mb_idx];
 250                 if !is_intra {
 251                     estimate_mb_type(mb_type, last_mbt, ((self.fenc.num_mv[mb_idx] + 1) % 3) as usize, &mut self.stats);
 252                     last_mbt = mb_type;
 253                     match mb_type {
 254                         VPMBType::InterMV | VPMBType::GoldenMV => {
 255                             estimate_mv(self.fenc.coded_mv[mb_idx][3], &mut self.stats);
 256                         },
 257                         VPMBType::InterFourMV => {
 258                             for (&sub_type, &mv) in self.fenc.fmv_sub[mb_idx].iter().zip(self.fenc.coded_mv[mb_idx].iter()) {
 259                                 if sub_type == VPMBType::InterMV {
 260                                     estimate_mv(mv, &mut self.stats);
 261                                 }
 262                             }
 263                         },
 264                         _ => {},
 265                     };
 266                 }
 267                 let mb = self.fenc.get_mb(mb_idx);
 268                 for i in 0..4 {
 269                     let cur_idx = mb_x * 2 + (i & 1);
 270                     let mut dc_mode = 0;
 271                     if self.top_ctx[0][cur_idx] {
 272                         dc_mode += 1;
 273                     }
 274                     if left_dc[i >> 1] {
 275                         dc_mode += 1;
 276                     }
 277                     self.top_ctx[0][cur_idx] = mb.coeffs[i][0] != 0;
 278                     left_dc[i >> 1] = mb.coeffs[i][0] != 0;
 279                     estimate_block(&mb.coeffs[i], dc_mode, &mut self.stats.coeff_models[0], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 280                 }
 281
 282                 let mut dc_mode = 0;
 283                 if self.top_ctx[1][mb_x] {
 284                     dc_mode += 1;
 285                 }
 286                 if left_dc[2] {
 287                     dc_mode += 1;
 288                 }
 289                 self.top_ctx[1][mb_x] = mb.coeffs[4][0] != 0;
 290                 left_dc[2] = mb.coeffs[4][0] != 0;
 291                 estimate_block(&mb.coeffs[4], dc_mode, &mut self.stats.coeff_models[1], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 292
 293                 let mut dc_mode = 0;
 294                 if self.top_ctx[2][mb_x] {
 295                     dc_mode += 1;
 296                 }
 297                 if left_dc[3] {
 298                     dc_mode += 1;
 299                 }
 300                 self.top_ctx[2][mb_x] = mb.coeffs[5][0] != 0;
 301                 left_dc[3] = mb.coeffs[5][0] != 0;
 302                 estimate_block(&mb.coeffs[5], dc_mode, &mut self.stats.coeff_models[1], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 303
 304                 mb_idx += 1;
 305             }
 306         }
 307     }
 308     fn prepare_huff_models(&mut self) {
 309         for i in 0..2 {
 310             self.models.vp6huff.dc_token_tree[i].build_codes(&self.models.coeff_models[i].dc_value_probs);
 311         }
 312         for i in 0..2 {
 313             for mode in 0..3 {
 314                 for band in 0..6 {
 315                     self.models.vp6huff.ac_token_tree[i][mode][band].build_codes(&self.models.coeff_models[i].ac_val_probs[mode][band]);
 316                 }
 317             }
 318         }
 319         for i in 0..2 {
 320             self.models.vp6huff.zero_run_tree[i].build_codes_zero_run(&self.models.vp6models.zero_run_probs[i]);
 321         }
 322     }
 323     fn determine_coeff_runs_luma(&self, hstate: &mut HuffState, mb_pos: usize, blk: usize) {
 324         let mb = self.fenc.get_mb(mb_pos);
 325
 326         if !hstate.dc_zr_coded[0] && mb.coeffs[blk].no_dc() {
 327             hstate.dc_zero_run[0] = 0;
 328             let mut blk_no = (blk + 1) & 3;
 329             let mut mb_no = mb_pos + ((blk + 1) >> 2);
 330             let mut cmb = mb;
 331             let mut last_mb_no = mb_pos;
 332             while (hstate.dc_zero_run[0] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 333                 if mb_no != last_mb_no {
 334                     cmb = self.fenc.get_mb(mb_no);
 335                     last_mb_no = mb_no;
 336                 }
 337                 if !cmb.coeffs[blk_no].no_dc() {
 338                     break;
 339                 }
 340                 hstate.dc_zero_run[0] += 1;
 341                 blk_no += 1;
 342                 if blk_no == 4 {
 343                     blk_no = 0;
 344                     mb_no += 1;
 345                 }
 346             }
 347         }
 348         if !hstate.ac_zr_coded[0] && mb.coeffs[blk].no_ac() {
 349             hstate.ac_zero_run[0] = 0;
 350             let mut blk_no = (blk + 1) & 3;
 351             let mut mb_no = mb_pos + ((blk + 1) >> 2);
 352             let mut cmb = mb;
 353             let mut last_mb_no = mb_pos;
 354             while (hstate.ac_zero_run[0] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 355                 if mb_no != last_mb_no {
 356                     cmb = self.fenc.get_mb(mb_no);
 357                     last_mb_no = mb_no;
 358                 }
 359                 if !cmb.coeffs[blk_no].no_ac() {
 360                     break;
 361                 }
 362                 hstate.ac_zero_run[0] += 1;
 363                 blk_no += 1;
 364                 if blk_no == 4 {
 365                     blk_no = 0;
 366                     mb_no += 1;
 367                 }
 368             }
 369         }
 370     }
 371     fn determine_coeff_runs_chroma(&self, hstate: &mut HuffState, mb_pos: usize, plane: usize) {
 372         let mb = self.fenc.get_mb(mb_pos);
 373         let blk = plane + 3;
 374
 375         if !hstate.dc_zr_coded[1] && mb.coeffs[blk].no_dc() {
 376             hstate.dc_zero_run[1] = 0;
 377             let mut blk_no = if blk == 4 { 5 } else { 4 };
 378             let mut mb_no = mb_pos + if blk == 4 { 0 } else { 1 };
 379             while (hstate.dc_zero_run[1] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 380                 let mb = self.fenc.get_mb(mb_no);
 381                 if !mb.coeffs[blk_no].no_dc() {
 382                     break;
 383                 }
 384                 hstate.dc_zero_run[1] += 1;
 385                 blk_no += 1;
 386                 if blk_no == 6 {
 387                     blk_no = 4;
 388                     mb_no += 1;
 389                 }
 390             }
 391         }
 392         if !hstate.ac_zr_coded[1] && mb.coeffs[blk].no_ac() {
 393             hstate.ac_zero_run[1] = 0;
 394             let mut blk_no = if blk == 4 { 5 } else { 4 };
 395             let mut mb_no = mb_pos + if blk == 4 { 0 } else { 1 };
 396             while (hstate.ac_zero_run[1] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 397                 let mb = self.fenc.get_mb(mb_no);
 398                 if !mb.coeffs[blk_no].no_ac() {
 399                     break;
 400                 }
 401                 hstate.ac_zero_run[1] += 1;
 402                 blk_no += 1;
 403                 if blk_no == 6 {
 404                     blk_no = 4;
 405                     mb_no += 1;
 406                 }
 407             }
 408         }
 409     }
 410     fn encode_intra(&mut self, bw: &mut ByteWriter, quant: usize) -> EncoderResult<bool> {
 411         self.models.reset(false);
 412         self.models.reset_mbtype_models();
 413         self.stats.reset();
 414
 415         self.pmodels.reset(false);
 416         self.pmodels.reset_mbtype_models();
 417
 418         let multistream = self.huffman || self.version != VERSION_VP60;
 419
 420         self.fenc.prepare_intra_blocks();
 421         self.fenc.apply_dc_prediction(&mut self.dc_pred);
 422         self.estimate_blocks(true);
 423         self.stats.generate(&mut self.models, true);
 424
 425         // header
 426         bw.write_byte(((quant as u8) << 1) | (multistream as u8))?;
 427         bw.write_byte((self.version << 3) | (self.profile << 1))?;
 428         bw.write_u16be(0)?; // part 2 offset placeholder
 429
 430         let mut bc = BoolEncoder::new(bw);
 431
 432         bc.put_bits(self.mb_h as u32, 8)?;
 433         bc.put_bits(self.mb_w as u32, 8)?;
 434         bc.put_bits(self.mb_h as u32, 8)?; // display MB height
 435         bc.put_bits(self.mb_w as u32, 8)?; // display MB width
 436         bc.put_bits(0, 2)?; // scaline mode
 437         // todo other advanced profile bits
 438         bc.put_bits(self.huffman as u32, 1)?; // Huffman mode
 439
 440         encode_coeff_models(&mut bc, &mut self.models, &self.pmodels, true, false)?;
 441         self.pmodels = self.models.clone();
 442
 443         if multistream || (self.profile == VP6_SIMPLE_PROFILE) {
 444             bc.flush()?;
 445
 446             // patch coefficient offset
 447             let offset = bw.tell();
 448             if offset >= 65535 {
 449                 return Err(EncoderError::Bug);
 450             }
 451             bw.seek(SeekFrom::Start(2))?;
 452             bw.write_u16be(offset as u16)?;
 453             bw.seek(SeekFrom::End(0))?;
 454
 455             bc = BoolEncoder::new(bw);
 456         }
 457         let writer = if !self.huffman {
 458                 VP6Writer::BoolCoder(bc)
 459             } else {
 460                 VP6Writer::Huffman(HuffEncoder::new(bw))
 461             };
 462         self.encode_coeffs(writer)?;
 463         Ok(true)
 464     }
 465     fn encode_inter(&mut self, bw: &mut ByteWriter, quant: usize) -> EncoderResult<bool> {
 466         self.stats.reset();
 467
 468         let multistream = self.huffman || self.version != VERSION_VP60;
 469         let loop_filter = false;
 470
 471         self.fenc.prepare_intra_blocks();
 472         self.fenc.prepare_inter_blocks(false);
 473         if !self.last_gold {
 474             self.fenc.prepare_inter_blocks(true);
 475         }
 476         let lambda = if self.force_q.is_some() { 1.0 } else { self.ratectl.lambda };
 477         self.fenc.select_inter_blocks(self.last_frame.get_vbuf().unwrap(), self.mc_buf.clone(), !self.last_gold, lambda);
 478         // todo implement forced intra
 479         let (_force_intra, golden_frame) = self.fenc.decide_frame_type();
 480         self.fenc.apply_dc_prediction(&mut self.dc_pred);
 481         self.fenc.predict_mvs();
 482
 483         self.write_inter_frame(bw, quant, multistream, loop_filter, golden_frame)?;
 484
 485         Ok(golden_frame)
 486     }
 487     fn write_inter_frame(&mut self, bw: &mut ByteWriter, quant: usize, multistream: bool, loop_filter: bool, golden_frame: bool) -> EncoderResult<()> {
 488         self.estimate_blocks(false);
 489
 490         self.stats.generate(&mut self.models, false);
 491
 492         // header
 493         bw.write_byte(0x80 | ((quant as u8) << 1) | (multistream as u8))?;
 494         bw.write_u16be(0)?; // part 2 offset placeholder
 495
 496         let mut bc = BoolEncoder::new(bw);
 497
 498         bc.put_bits(golden_frame as u32, 1)?; // refresh golden frame
 499         if self.profile == VP6_ADVANCED_PROFILE {
 500             bc.put_bits(loop_filter as u32, 1)?; // use loop filter
 501             if loop_filter {
 502                 bc.put_bits(0, 1)?; // loop filter selector
 503             }
 504             if self.version == VERSION_VP62 {
 505                 bc.put_bits(0, 1)?; // auto select PM
 506             }
 507         }
 508         // todo other advanced profile bits
 509         bc.put_bits(self.huffman as u32, 1)?;
 510
 511         encode_mode_prob_models(&mut bc, &mut self.models, &self.pmodels, &self.stats.mbtype_models)?;
 512         encode_mv_models(&mut bc, &self.models.mv_models, &self.pmodels.mv_models)?;
 513         encode_coeff_models(&mut bc, &mut self.models, &self.pmodels, false, false)?;
 514         self.pmodels = self.models.clone();
 515
 516         let mut last_mbt = VPMBType::InterNoMV;
 517         for mb_idx in 0..self.mb_w * self.mb_h {
 518             let mb_type = self.fenc.mb_types[mb_idx];
 519             encode_mb_type(&mut bc, self.fenc.mb_types[mb_idx], last_mbt, ((self.fenc.num_mv[mb_idx] + 1) % 3) as usize, &self.models)?;
 520             last_mbt = mb_type;
 521             match mb_type {
 522                 VPMBType::InterMV | VPMBType::GoldenMV => {
 523                     encode_mv(&mut bc, self.fenc.coded_mv[mb_idx][3], &self.models)?;
 524                 },
 525                 VPMBType::InterFourMV => {
 526                     for &sub_type in self.fenc.fmv_sub[mb_idx].iter() {
 527                         let id = match sub_type {
 528                                 VPMBType::InterNoMV     => 0,
 529                                 VPMBType::InterMV       => 1,
 530                                 VPMBType::InterNearest  => 2,
 531                                 VPMBType::InterNear     => 3,
 532                                 _ => unreachable!(),
 533                             };
 534                         bc.put_bits(id, 2)?;
 535                     }
 536                     for (&sub_type, &mv) in self.fenc.fmv_sub[mb_idx].iter().zip(self.fenc.coded_mv[mb_idx].iter()) {
 537                         if sub_type == VPMBType::InterMV {
 538                             encode_mv(&mut bc, mv, &self.models)?;
 539                         }
 540                     }
 541                 },
 542                 _ => {},
 543             };
 544         }
 545
 546         if multistream || (self.profile == VP6_SIMPLE_PROFILE) {
 547             bc.flush()?;
 548
 549             // patch coefficient offset
 550             let offset = bw.tell();
 551             if offset >= 65535 {
 552                 return Err(EncoderError::Bug);
 553             }
 554             bw.seek(SeekFrom::Start(1))?;
 555             bw.write_u16be(offset as u16)?;
 556             bw.seek(SeekFrom::End(0))?;
 557
 558             bc = BoolEncoder::new(bw);
 559         }
 560         let writer = if !self.huffman {
 561                 VP6Writer::BoolCoder(bc)
 562             } else {
 563                 VP6Writer::Huffman(HuffEncoder::new(bw))
 564             };
 565         self.encode_coeffs(writer)?;
 566
 567         Ok(())
 568     }
 569     fn encode_inter_fast(&mut self, bw: &mut ByteWriter, quant: usize) -> EncoderResult<bool> {
 570         self.stats.reset();
 571
 572         let multistream = self.huffman || self.version != VERSION_VP60;
 573         let loop_filter = false;
 574
 575         let last_frm = self.last_frame.get_vbuf().unwrap();
 576         let gold_frm = if !self.last_gold {
 577                 Some(self.gold_frame.get_vbuf().unwrap())
 578             } else {
 579                 None
 580             };
 581         let lambda = if self.force_q.is_some() { 1.0 } else { self.ratectl.lambda };
 582         self.fenc.select_inter_blocks_fast(last_frm, gold_frm, self.mc_buf.clone(), lambda);
 583         let golden_frame = false;
 584         self.fenc.apply_dc_prediction(&mut self.dc_pred);
 585         self.fenc.predict_mvs();
 586         self.estimate_blocks(false);
 587
 588         self.write_inter_frame(bw, quant, multistream, loop_filter, golden_frame)?;
 589
 590         Ok(false)
 591     }
 592     fn encode_coeffs(&mut self, mut writer: VP6Writer) -> EncoderResult<()> {
 593         if self.huffman {
 594             self.prepare_huff_models();
 595         }
 596
 597         let mut hstate = HuffState::new();
 598         for top_ctx in self.top_ctx.iter_mut() {
 599             for el in top_ctx.iter_mut() {
 600                 *el = false;
 601             }
 602         }
 603         let mut mb_pos = 0;
 604         for _mb_y in 0..self.mb_h {
 605             let mut left_dc = [false; 4];
 606             for mb_x in 0..self.mb_w {
 607                 let mb = self.fenc.get_mb(mb_pos);
 608                 for i in 0..4 {
 609                     let cur_idx = mb_x * 2 + (i & 1);
 610                     let mut dc_mode = 0;
 611                     if self.top_ctx[0][cur_idx] {
 612                         dc_mode += 1;
 613                     }
 614                     if left_dc[i >> 1] {
 615                         dc_mode += 1;
 616                     }
 617                     self.top_ctx[0][cur_idx] = mb.coeffs[i][0] != 0;
 618                     left_dc[i >> 1] = mb.coeffs[i][0] != 0;
 619                     if self.huffman {
 620                         self.determine_coeff_runs_luma(&mut hstate, mb_pos, i);
 621                     }
 622                     match writer {
 623                         VP6Writer::BoolCoder(ref mut bc) => encode_block(bc, &mb.coeffs[i], dc_mode, &self.models.coeff_models[0], &self.models.vp6models)?,
 624                         VP6Writer::Huffman(ref mut huff) => encode_block_huff(huff, &self.models.vp6models.zigzag, &mb.coeffs[i], 0, &mut hstate, &self.models.vp6huff)?,
 625                     };
 626                 }
 627
 628                 for plane in 1..3 {
 629                     let mut dc_mode = 0;
 630                     if self.top_ctx[plane][mb_x] {
 631                         dc_mode += 1;
 632                     }
 633                     if left_dc[plane + 1] {
 634                         dc_mode += 1;
 635                     }
 636                     self.top_ctx[plane][mb_x] = mb.coeffs[plane + 3][0] != 0;
 637                     left_dc[plane + 1] = mb.coeffs[plane + 3][0] != 0;
 638                     if self.huffman {
 639                         self.determine_coeff_runs_chroma(&mut hstate, mb_pos, plane);
 640                     }
 641                     match writer {
 642                         VP6Writer::BoolCoder(ref mut bc) => encode_block(bc, &mb.coeffs[plane + 3], dc_mode, &self.models.coeff_models[1], &self.models.vp6models)?,
 643                         VP6Writer::Huffman(ref mut huff) => encode_block_huff(huff, &self.models.vp6models.zigzag, &mb.coeffs[plane + 3], 1, &mut hstate, &self.models.vp6huff)?,
 644                     };
 645                 }
 646
 647                 mb_pos += 1;
 648             }
 649         }
 650
 651         match writer {
 652             VP6Writer::BoolCoder(bc) => bc.flush()?,
 653             VP6Writer::Huffman(huff) => huff.flush()?,
 654         };
 655
 656         Ok(())
 657     }
 658 }
 659
 660 impl NAEncoder for VP6Encoder {
 661     fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
 662         match encinfo.format {
 663             NACodecTypeInfo::None => {
 664                 Ok(EncodeParameters {
 665                     format: NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, self.flipped, YUV420_FORMAT)),
 666                     ..Default::default()
 667                 })
 668             },
 669             NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
 670             NACodecTypeInfo::Video(vinfo) => {
 671                 let outinfo = NAVideoInfo::new((vinfo.width + 3) & !3, (vinfo.height + 3) & !3, self.flipped, YUV420_FORMAT);
 672                 let mut ofmt = *encinfo;
 673                 ofmt.format = NACodecTypeInfo::Video(outinfo);
 674                 Ok(ofmt)
 675             }
 676         }
 677     }
 678     fn get_capabilities(&self) -> u64 { 0 }
 679     fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
 680         match encinfo.format {
 681             NACodecTypeInfo::None => Err(EncoderError::FormatError),
 682             NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
 683             NACodecTypeInfo::Video(vinfo) => {
 684                 if vinfo.format != YUV420_FORMAT {
 685                     return Err(EncoderError::FormatError);
 686                 }
 687                 if ((vinfo.width | vinfo.height) & 15) != 0 {
 688                     return Err(EncoderError::FormatError);
 689                 }
 690                 if (vinfo.width | vinfo.height) >= (1 << 12) {
 691                     return Err(EncoderError::FormatError);
 692                 }
 693
 694                 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, self.flipped, vinfo.format);
 695                 let info = NACodecInfo::new(if self.flipped { "vp6" } else { "vp6f" }, NACodecTypeInfo::Video(out_info), None);
 696                 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
 697                 stream.set_num(stream_id as usize);
 698                 let stream = stream.into_ref();
 699
 700                 self.last_frame = alloc_video_buffer(out_info, 4)?;
 701                 self.gold_frame = alloc_video_buffer(out_info, 4)?;
 702
 703                 self.stream = Some(stream.clone());
 704
 705                 self.mb_w = (vinfo.width  + 15) >> 4;
 706                 self.mb_h = (vinfo.height + 15) >> 4;
 707                 self.fenc.resize(self.mb_w, self.mb_h);
 708                 self.ratectl.init(self.mb_w, self.mb_h, encinfo.bitrate, encinfo.tb_num, encinfo.tb_den);
 709
 710                 self.dc_pred.resize(self.mb_w);
 711                 self.top_ctx = [vec![false; self.mb_w * 2], vec![false; self.mb_w], vec![false; self.mb_w], vec![false; self.mb_w * 2]];
 712
 713                 self.version = VERSION_VP60;
 714                 self.profile = VP6_SIMPLE_PROFILE;
 715
 716                 Ok(stream)
 717             },
 718         }
 719     }
 720     fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
 721         let buf = frm.get_buffer();
 722         if let Some(ref vbuf) = buf.get_vbuf() {
 723             let mut dbuf = Vec::with_capacity(4);
 724             let mut gw   = GrowableMemoryWriter::new_write(&mut dbuf);
 725             let mut bw   = ByteWriter::new(&mut gw);
 726
 727             // todo integrate with rate control
 728             let is_intra = (self.frmcount == 0) || self.decide_encoding();
 729             let quant = if let Some(q) = self.force_q {
 730                     q
 731                 } else {
 732                     self.ratectl.guess_quant(is_intra, self.huffman)
 733                 };
 734
 735             self.fenc.read_mbs(vbuf);
 736             self.fenc.set_quant(quant);
 737             self.fenc.me_mode = self.me_mode;
 738             self.fenc.me_range = self.me_range;
 739             let golden_frame = if is_intra {
 740                     self.encode_intra(&mut bw, quant)?
 741                 } else if !self.fast {
 742                     self.fenc.estimate_mvs(self.last_frame.get_vbuf().unwrap(), self.mc_buf.clone(), false);
 743                     if !self.last_gold {
 744                         self.fenc.estimate_mvs(self.gold_frame.get_vbuf().unwrap(), self.mc_buf.clone(), true);
 745                     }
 746                     self.encode_inter(&mut bw, quant)?
 747                 } else {
 748                     self.encode_inter_fast(&mut bw, quant)?
 749                 };
 750             self.fenc.reconstruct_frame(&mut self.dc_pred, self.last_frame.get_vbuf().unwrap());
 751             self.last_gold = golden_frame;
 752             if golden_frame {
 753                 let mut dfrm = self.gold_frame.get_vbuf().unwrap();
 754                 let src = self.last_frame.get_vbuf().unwrap();
 755
 756                 let dst = dfrm.get_data_mut().unwrap();
 757                 dst.copy_from_slice(src.get_data());
 758             }
 759
 760             if self.force_q.is_none() {
 761                 self.ratectl.update(dbuf.len() * 8);
 762             }
 763
 764             self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf));
 765             if self.key_int > 0 {
 766                 self.frmcount += 1;
 767             }
 768             if self.frmcount == self.key_int {
 769                 self.frmcount = 0;
 770             }
 771             Ok(())
 772         } else {
 773             Err(EncoderError::InvalidParameters)
 774         }
 775     }
 776     fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
 777         let mut npkt = None;
 778         std::mem::swap(&mut self.pkt, &mut npkt);
 779         Ok(npkt)
 780     }
 781     fn flush(&mut self) -> EncoderResult<()> {
 782         self.frmcount = 0;
 783         Ok(())
 784     }
 785 }
 786
 787 const HUFFMAN_OPTION: &str = "huffman";
 788 const QUANT_OPTION: &str = "quant";
 789 const VERSION_OPTION: &str = "version";
 790 const MV_SEARCH_OPTION: &str = "mv_mode";
 791 const MV_RANGE_OPTION: &str = "mv_range";
 792 const FAST_OPTION: &str = "fast";
 793
 794 const ENCODER_OPTS: &[NAOptionDefinition] = &[
 795     NAOptionDefinition {
 796         name: KEYFRAME_OPTION, description: KEYFRAME_OPTION_DESC,
 797         opt_type: NAOptionDefinitionType::Int(Some(0), Some(128)) },
 798     NAOptionDefinition {
 799         name: HUFFMAN_OPTION, description: "use Huffman encoding",
 800         opt_type: NAOptionDefinitionType::Bool },
 801     NAOptionDefinition {
 802         name: VERSION_OPTION, description: "codec minor version",
 803         opt_type: NAOptionDefinitionType::String(Some(&["vp60", "vp61", "vp62"])) },
 804     NAOptionDefinition {
 805         name: QUANT_OPTION, description: "force fixed quantiser for encoding",
 806         opt_type: NAOptionDefinitionType::Int(Some(-1), Some(63)) },
 807     NAOptionDefinition {
 808         name: MV_SEARCH_OPTION, description: "motion search mode",
 809         opt_type: NAOptionDefinitionType::String(Some(&["full", "dia", "hex"])) },
 810     NAOptionDefinition {
 811         name: MV_RANGE_OPTION, description: "motion search range (in pixels)",
 812         opt_type: NAOptionDefinitionType::Int(Some(0), Some(30)) },
 813     NAOptionDefinition {
 814         name: FAST_OPTION, description: "faster (but worse) encoding",
 815         opt_type: NAOptionDefinitionType::Bool },
 816 ];
 817
 818 impl NAOptionHandler for VP6Encoder {
 819     fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
 820     fn set_options(&mut self, options: &[NAOption]) {
 821         for option in options.iter() {
 822             for opt_def in ENCODER_OPTS.iter() {
 823                 if opt_def.check(option).is_ok() {
 824                     match option.name {
 825                         KEYFRAME_OPTION => {
 826                             if let NAValue::Int(intval) = option.value {
 827                                 self.key_int = intval as u8;
 828                             }
 829                         },
 830                         HUFFMAN_OPTION => {
 831                             if let NAValue::Bool(bval) = option.value {
 832                                 self.huffman = bval;
 833                             }
 834                         },
 835                         VERSION_OPTION => {
 836                             if let NAValue::String(ref string) = option.value {
 837                                 self.version = match string.as_str() {
 838                                         "vp60" => VERSION_VP60,
 839                                         "vp61" => VERSION_VP61,
 840                                         "vp62" => VERSION_VP62,
 841                                         _ => unreachable!(),
 842                                     };
 843                             }
 844                         },
 845                         QUANT_OPTION => {
 846                             if let NAValue::Int(intval) = option.value {
 847                                 self.force_q = if intval < 0 { None } else { Some(intval as usize) };
 848                             }
 849                         },
 850                         MV_SEARCH_OPTION => {
 851                             if let NAValue::String(ref string) = option.value {
 852                                 if let Ok(mv_mode) = string.parse::<MVSearchMode>() {
 853                                     self.me_mode = mv_mode;
 854                                 }
 855                             }
 856                         },
 857                         MV_RANGE_OPTION => {
 858                             if let NAValue::Int(intval) = option.value {
 859                                 self.me_range = intval as i16;
 860                             }
 861                         },
 862                         FAST_OPTION => {
 863                             if let NAValue::Bool(bval) = option.value {
 864                                 self.fast = bval;
 865                             }
 866                         },
 867                         _ => {},
 868                     };
 869                 }
 870             }
 871         }
 872     }
 873     fn query_option_value(&self, name: &str) -> Option<NAValue> {
 874         match name {
 875             KEYFRAME_OPTION => Some(NAValue::Int(i64::from(self.key_int))),
 876             HUFFMAN_OPTION => Some(NAValue::Bool(self.huffman)),
 877             VERSION_OPTION => {
 878                     let ver = match self.version {
 879                             VERSION_VP60 => "vp60",
 880                             VERSION_VP61 => "vp61",
 881                             VERSION_VP62 => "vp62",
 882                             _ => unreachable!(),
 883                         };
 884                     Some(NAValue::String(ver.to_string()))
 885                 },
 886             QUANT_OPTION => if let Some(q) = self.force_q {
 887                     Some(NAValue::Int(q as i64))
 888                 } else {
 889                     Some(NAValue::Int(-1))
 890                 },
 891             MV_SEARCH_OPTION => Some(NAValue::String(self.me_mode.to_string())),
 892             MV_RANGE_OPTION => Some(NAValue::Int(i64::from(self.me_range))),
 893             FAST_OPTION => Some(NAValue::Bool(self.fast)),
 894             _ => None,
 895         }
 896     }
 897 }
 898
 899 pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
 900     Box::new(VP6Encoder::new(true))
 901 }
 902
 903 pub fn get_encoder_flv() -> Box<dyn NAEncoder + Send> {
 904     Box::new(VP6Encoder::new(false))
 905 }
 906
 907 #[cfg(test)]
 908 mod test {
 909     use nihav_core::codecs::*;
 910     use nihav_core::demuxers::*;
 911     use nihav_core::muxers::*;
 912     use crate::*;
 913     use nihav_commonfmt::*;
 914     use nihav_codec_support::test::enc_video::*;
 915
 916     fn encode_test(out_name: &'static str, enc_options: &[NAOption], hash: &[u32; 4]) {
 917         let mut dmx_reg = RegisteredDemuxers::new();
 918         generic_register_all_demuxers(&mut dmx_reg);
 919         let mut dec_reg = RegisteredDecoders::new();
 920         duck_register_all_decoders(&mut dec_reg);
 921         let mut mux_reg = RegisteredMuxers::new();
 922         generic_register_all_muxers(&mut mux_reg);
 923         let mut enc_reg = RegisteredEncoders::new();
 924         duck_register_all_encoders(&mut enc_reg);
 925
 926         // sample: https://samples.mplayerhq.hu/V-codecs/VP4/ot171_vp40.avi
 927         let dec_config = DecoderTestParams {
 928                 demuxer:        "avi",
 929                 in_name:        "assets/Duck/ot171_vp40.avi",
 930                 stream_type:    StreamType::Video,
 931                 limit:          Some(1),
 932                 dmx_reg, dec_reg,
 933             };
 934         let enc_config = EncoderTestParams {
 935                 muxer:          "avi",
 936                 enc_name:       "vp6",
 937                 out_name,
 938                 mux_reg, enc_reg,
 939             };
 940         let dst_vinfo = NAVideoInfo {
 941                 width:   0,
 942                 height:  0,
 943                 format:  YUV420_FORMAT,
 944                 flipped: true,
 945                 bits:    12,
 946             };
 947         let enc_params = EncodeParameters {
 948                 format:  NACodecTypeInfo::Video(dst_vinfo),
 949                 quality: 0,
 950                 bitrate: 25000,
 951                 tb_num:  0,
 952                 tb_den:  0,
 953                 flags:   0,
 954             };
 955         //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
 956         test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
 957                           hash);
 958     }
 959     #[test]
 960     fn test_vp6_encoder_bc() {
 961         let enc_options = &[
 962                 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(42) },
 963             ];
 964         encode_test("vp6-bool.avi", enc_options, &[0xb57f49e5, 0x6b48accd, 0xc28fadb3, 0xc89a30d2]);
 965     }
 966     #[test]
 967     fn test_vp6_encoder_fast() {
 968         let enc_options = &[
 969                 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(42) },
 970                 NAOption { name: super::FAST_OPTION, value: NAValue::Bool(true) },
 971             ];
 972         encode_test("vp6-fast.avi", enc_options, &[0xb8037ce1, 0xc00ade72, 0x3c0b73c2, 0xbfc4113d]);
 973     }
 974     #[test]
 975     fn test_vp6_encoder_rc() {
 976         let enc_options = &[
 977             ];
 978         encode_test("vp6-rc.avi", enc_options, &[0x790baca9, 0x663eafcf, 0x36d1bed8, 0xddf882de]);
 979     }
 980     #[test]
 981     fn test_vp6_encoder_huff() {
 982         let enc_options = &[
 983                 NAOption { name: super::HUFFMAN_OPTION, value: NAValue::Bool(true) },
 984                 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(42) },
 985             ];
 986         encode_test("vp6-huff.avi", enc_options, &[0x6e9bb23d, 0xde296d92, 0x4c225bae, 0x3651e31f]);
 987     }
 988 }