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