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     huffman:    bool,
 176
 177     version:    u8,
 178     profile:    u8,
 179
 180     models:     VP56Models,
 181     stats:      VP56ModelsStat,
 182     pmodels:    VP56Models,
 183
 184     last_frame: NABufferType,
 185     gold_frame: NABufferType,
 186     last_gold:  bool,
 187     me_mode:    MVSearchMode,
 188     me_range:   i16,
 189
 190     force_q:    Option<usize>,
 191 }
 192
 193 impl VP6Encoder {
 194     fn new() -> Self {
 195         let vt = alloc_video_buffer(NAVideoInfo::new(24, 24, false, VP_YUVA420_FORMAT), 4).unwrap();
 196         let mc_buf = vt.get_vbuf().unwrap();
 197         Self {
 198             stream:     None,
 199             pkt:        None,
 200             key_int:    10,
 201             frmcount:   0,
 202             mb_w:       0,
 203             mb_h:       0,
 204             dc_pred:    VP56DCPred::new(),
 205             top_ctx:    [Vec::new(), Vec::new(), Vec::new(), Vec::new()],
 206             fenc:       FrameEncoder::new(),
 207             ratectl:    RateControl::new(),
 208             mc_buf,
 209
 210             huffman:    false,
 211
 212             version:    VERSION_VP60,
 213             profile:    VP6_SIMPLE_PROFILE,
 214
 215             models:     VP56Models::new(),
 216             pmodels:    VP56Models::new(),
 217             stats:      VP56ModelsStat::new(),
 218
 219             last_frame: NABufferType::None,
 220             gold_frame: NABufferType::None,
 221             last_gold:  false,
 222             me_mode:    MVSearchMode::default(),
 223             me_range:   16,
 224
 225             force_q:    None,
 226         }
 227     }
 228     fn decide_encoding(&mut self) -> bool {
 229         false
 230     }
 231     fn estimate_blocks(&mut self, is_intra: bool) {
 232         for top_ctx in self.top_ctx.iter_mut() {
 233             for el in top_ctx.iter_mut() {
 234                 *el = false;
 235             }
 236         }
 237         let mut last_mbt = VPMBType::InterNoMV;
 238         let mut mb_idx = 0;
 239         for _mb_y in 0..self.mb_h {
 240             let mut left_dc = [false; 4];
 241             for mb_x in 0..self.mb_w {
 242                 let mb_type = self.fenc.mb_types[mb_idx];
 243                 if !is_intra {
 244                     estimate_mb_type(mb_type, last_mbt, ((self.fenc.num_mv[mb_idx] + 1) % 3) as usize, &mut self.stats);
 245                     last_mbt = mb_type;
 246                     match mb_type {
 247                         VPMBType::InterMV | VPMBType::GoldenMV => {
 248                             estimate_mv(self.fenc.coded_mv[mb_idx][3], &mut self.stats);
 249                         },
 250                         VPMBType::InterFourMV => {
 251                             for (&sub_type, &mv) in self.fenc.fmv_sub[mb_idx].iter().zip(self.fenc.coded_mv[mb_idx].iter()) {
 252                                 if sub_type == VPMBType::InterMV {
 253                                     estimate_mv(mv, &mut self.stats);
 254                                 }
 255                             }
 256                         },
 257                         _ => {},
 258                     };
 259                 }
 260                 let mb = self.fenc.get_mb(mb_idx);
 261                 for i in 0..4 {
 262                     let cur_idx = mb_x * 2 + (i & 1);
 263                     let mut dc_mode = 0;
 264                     if self.top_ctx[0][cur_idx] {
 265                         dc_mode += 1;
 266                     }
 267                     if left_dc[i >> 1] {
 268                         dc_mode += 1;
 269                     }
 270                     self.top_ctx[0][cur_idx] = mb.coeffs[i][0] != 0;
 271                     left_dc[i >> 1] = mb.coeffs[i][0] != 0;
 272                     estimate_block(&mb.coeffs[i], dc_mode, &mut self.stats.coeff_models[0], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 273                 }
 274
 275                 let mut dc_mode = 0;
 276                 if self.top_ctx[1][mb_x] {
 277                     dc_mode += 1;
 278                 }
 279                 if left_dc[2] {
 280                     dc_mode += 1;
 281                 }
 282                 self.top_ctx[1][mb_x] = mb.coeffs[4][0] != 0;
 283                 left_dc[2] = mb.coeffs[4][0] != 0;
 284                 estimate_block(&mb.coeffs[4], dc_mode, &mut self.stats.coeff_models[1], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 285
 286                 let mut dc_mode = 0;
 287                 if self.top_ctx[2][mb_x] {
 288                     dc_mode += 1;
 289                 }
 290                 if left_dc[3] {
 291                     dc_mode += 1;
 292                 }
 293                 self.top_ctx[2][mb_x] = mb.coeffs[5][0] != 0;
 294                 left_dc[3] = mb.coeffs[5][0] != 0;
 295                 estimate_block(&mb.coeffs[5], dc_mode, &mut self.stats.coeff_models[1], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 296
 297                 mb_idx += 1;
 298             }
 299         }
 300     }
 301     fn prepare_huff_models(&mut self) {
 302         for i in 0..2 {
 303             self.models.vp6huff.dc_token_tree[i].build_codes(&self.models.coeff_models[i].dc_value_probs);
 304         }
 305         for i in 0..2 {
 306             for mode in 0..3 {
 307                 for band in 0..6 {
 308                     self.models.vp6huff.ac_token_tree[i][mode][band].build_codes(&self.models.coeff_models[i].ac_val_probs[mode][band]);
 309                 }
 310             }
 311         }
 312         for i in 0..2 {
 313             self.models.vp6huff.zero_run_tree[i].build_codes_zero_run(&self.models.vp6models.zero_run_probs[i]);
 314         }
 315     }
 316     fn determine_coeff_runs_luma(&self, hstate: &mut HuffState, mb_pos: usize, blk: usize) {
 317         let mb = self.fenc.get_mb(mb_pos);
 318
 319         if !hstate.dc_zr_coded[0] {
 320             if mb.coeffs[blk].no_dc() {
 321                 hstate.dc_zero_run[0] = 0;
 322                 let mut blk_no = (blk + 1) & 3;
 323                 let mut mb_no = mb_pos + ((blk + 1) >> 2);
 324                 let mut cmb = mb;
 325                 let mut last_mb_no = mb_pos;
 326                 while (hstate.dc_zero_run[0] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 327                     if mb_no != last_mb_no {
 328                         cmb = self.fenc.get_mb(mb_no);
 329                         last_mb_no = mb_no;
 330                     }
 331                     if !cmb.coeffs[blk_no].no_dc() {
 332                         break;
 333                     }
 334                     hstate.dc_zero_run[0] += 1;
 335                     blk_no += 1;
 336                     if blk_no == 4 {
 337                         blk_no = 0;
 338                         mb_no += 1;
 339                     }
 340                 }
 341             }
 342         }
 343         if !hstate.ac_zr_coded[0] {
 344             if mb.coeffs[blk].no_ac() {
 345                 hstate.ac_zero_run[0] = 0;
 346                 let mut blk_no = (blk + 1) & 3;
 347                 let mut mb_no = mb_pos + ((blk + 1) >> 2);
 348                 let mut cmb = mb;
 349                 let mut last_mb_no = mb_pos;
 350                 while (hstate.ac_zero_run[0] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 351                     if mb_no != last_mb_no {
 352                         cmb = self.fenc.get_mb(mb_no);
 353                         last_mb_no = mb_no;
 354                     }
 355                     if !cmb.coeffs[blk_no].no_ac() {
 356                         break;
 357                     }
 358                     hstate.ac_zero_run[0] += 1;
 359                     blk_no += 1;
 360                     if blk_no == 4 {
 361                         blk_no = 0;
 362                         mb_no += 1;
 363                     }
 364                 }
 365             }
 366         }
 367     }
 368     fn determine_coeff_runs_chroma(&self, hstate: &mut HuffState, mb_pos: usize, plane: usize) {
 369         let mb = self.fenc.get_mb(mb_pos);
 370         let blk = plane + 3;
 371
 372         if !hstate.dc_zr_coded[1] {
 373             if mb.coeffs[blk].no_dc() {
 374                 hstate.dc_zero_run[1] = 0;
 375                 let mut blk_no = if blk == 4 { 5 } else { 4 };
 376                 let mut mb_no = mb_pos + if blk == 4 { 0 } else { 1 };
 377                 while (hstate.dc_zero_run[1] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 378                     let mb = self.fenc.get_mb(mb_no);
 379                     if !mb.coeffs[blk_no].no_dc() {
 380                         break;
 381                     }
 382                     hstate.dc_zero_run[1] += 1;
 383                     blk_no += 1;
 384                     if blk_no == 6 {
 385                         blk_no = 4;
 386                         mb_no += 1;
 387                     }
 388                 }
 389             }
 390         }
 391         if !hstate.ac_zr_coded[1] {
 392             if 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     }
 411     fn encode_intra(&mut self, bw: &mut ByteWriter, quant: usize) -> EncoderResult<bool> {
 412         self.models.reset(false);
 413         self.models.reset_mbtype_models();
 414         self.stats.reset();
 415
 416         self.pmodels.reset(false);
 417         self.pmodels.reset_mbtype_models();
 418
 419         let multistream = self.huffman || self.version != VERSION_VP60;
 420
 421         self.fenc.prepare_intra_blocks();
 422         self.fenc.apply_dc_prediction(&mut self.dc_pred);
 423         self.estimate_blocks(true);
 424         self.stats.generate(&mut self.models, true);
 425
 426         // header
 427         bw.write_byte(((quant as u8) << 1) | (multistream as u8))?;
 428         bw.write_byte((self.version << 3) | (self.profile << 1))?;
 429         bw.write_u16be(0)?; // part 2 offset placeholder
 430
 431         let mut bc = BoolEncoder::new(bw);
 432
 433         bc.put_bits(self.mb_h as u32, 8)?;
 434         bc.put_bits(self.mb_w as u32, 8)?;
 435         bc.put_bits(self.mb_h as u32, 8)?; // display MB height
 436         bc.put_bits(self.mb_w as u32, 8)?; // display MB width
 437         bc.put_bits(0, 2)?; // scaline mode
 438         // todo other advanced profile bits
 439         bc.put_bits(self.huffman as u32, 1)?; // Huffman mode
 440
 441         encode_coeff_models(&mut bc, &mut self.models, &self.pmodels, true, false)?;
 442         self.pmodels = self.models.clone();
 443
 444         if multistream || (self.profile == VP6_SIMPLE_PROFILE) {
 445             bc.flush()?;
 446
 447             // patch coefficient offset
 448             let offset = bw.tell();
 449             if offset >= 65535 {
 450                 return Err(EncoderError::Bug);
 451             }
 452             bw.seek(SeekFrom::Start(2))?;
 453             bw.write_u16be(offset as u16)?;
 454             bw.seek(SeekFrom::End(0))?;
 455
 456             bc = BoolEncoder::new(bw);
 457         }
 458         let writer = if !self.huffman {
 459                 VP6Writer::BoolCoder(bc)
 460             } else {
 461                 VP6Writer::Huffman(HuffEncoder::new(bw))
 462             };
 463         self.encode_coeffs(writer)?;
 464         Ok(true)
 465     }
 466     fn encode_inter(&mut self, bw: &mut ByteWriter, quant: usize) -> EncoderResult<bool> {
 467         self.stats.reset();
 468
 469         let multistream = self.huffman || self.version != VERSION_VP60;
 470         let loop_filter = false;
 471
 472         self.fenc.prepare_intra_blocks();
 473         self.fenc.prepare_inter_blocks(false);
 474         if !self.last_gold {
 475             self.fenc.prepare_inter_blocks(true);
 476         }
 477         let lambda = if self.force_q.is_some() { 1.0 } else { self.ratectl.lambda };
 478         self.fenc.select_inter_blocks(self.last_frame.get_vbuf().unwrap(), self.mc_buf.clone(), !self.last_gold, lambda);
 479         // todo implement forced intra
 480         let (_force_intra, golden_frame) = self.fenc.decide_frame_type();
 481         self.fenc.apply_dc_prediction(&mut self.dc_pred);
 482         self.fenc.predict_mvs();
 483         self.estimate_blocks(false);
 484
 485         self.stats.generate(&mut self.models, false);
 486
 487         // header
 488         bw.write_byte(0x80 | ((quant as u8) << 1) | (multistream as u8))?;
 489         bw.write_u16be(0)?; // part 2 offset placeholder
 490
 491         let mut bc = BoolEncoder::new(bw);
 492
 493         bc.put_bits(golden_frame as u32, 1)?; // refresh golden frame
 494         if self.profile == VP6_ADVANCED_PROFILE {
 495             bc.put_bits(loop_filter as u32, 1)?; // use loop filter
 496             if loop_filter {
 497                 bc.put_bits(0, 1)?; // loop filter selector
 498             }
 499             if self.version == VERSION_VP62 {
 500                 bc.put_bits(0, 1)?; // auto select PM
 501             }
 502         }
 503         // todo other advanced profile bits
 504         bc.put_bits(self.huffman as u32, 1)?;
 505
 506         encode_mode_prob_models(&mut bc, &mut self.models, &self.pmodels, &self.stats.mbtype_models)?;
 507         encode_mv_models(&mut bc, &self.models.mv_models, &self.pmodels.mv_models)?;
 508         encode_coeff_models(&mut bc, &mut self.models, &self.pmodels, false, false)?;
 509         self.pmodels = self.models.clone();
 510
 511         let mut last_mbt = VPMBType::InterNoMV;
 512         for mb_idx in 0..self.mb_w * self.mb_h {
 513             let mb_type = self.fenc.mb_types[mb_idx];
 514             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)?;
 515             last_mbt = mb_type;
 516             match mb_type {
 517                 VPMBType::InterMV | VPMBType::GoldenMV => {
 518                     encode_mv(&mut bc, self.fenc.coded_mv[mb_idx][3], &self.models)?;
 519                 },
 520                 VPMBType::InterFourMV => {
 521                     for &sub_type in self.fenc.fmv_sub[mb_idx].iter() {
 522                         let id = match sub_type {
 523                                 VPMBType::InterNoMV     => 0,
 524                                 VPMBType::InterMV       => 1,
 525                                 VPMBType::InterNearest  => 2,
 526                                 VPMBType::InterNear     => 3,
 527                                 _ => unreachable!(),
 528                             };
 529                         bc.put_bits(id, 2)?;
 530                     }
 531                     for (&sub_type, &mv) in self.fenc.fmv_sub[mb_idx].iter().zip(self.fenc.coded_mv[mb_idx].iter()) {
 532                         if sub_type == VPMBType::InterMV {
 533                             encode_mv(&mut bc, mv, &self.models)?;
 534                         }
 535                     }
 536                 },
 537                 _ => {},
 538             };
 539         }
 540
 541         if multistream || (self.profile == VP6_SIMPLE_PROFILE) {
 542             bc.flush()?;
 543
 544             // patch coefficient offset
 545             let offset = bw.tell();
 546             if offset >= 65535 {
 547                 return Err(EncoderError::Bug);
 548             }
 549             bw.seek(SeekFrom::Start(1))?;
 550             bw.write_u16be(offset as u16)?;
 551             bw.seek(SeekFrom::End(0))?;
 552
 553             bc = BoolEncoder::new(bw);
 554         }
 555         let writer = if !self.huffman {
 556                 VP6Writer::BoolCoder(bc)
 557             } else {
 558                 VP6Writer::Huffman(HuffEncoder::new(bw))
 559             };
 560         self.encode_coeffs(writer)?;
 561         Ok(golden_frame)
 562     }
 563     fn encode_coeffs(&mut self, mut writer: VP6Writer) -> EncoderResult<()> {
 564         if self.huffman {
 565             self.prepare_huff_models();
 566         }
 567
 568         let mut hstate = HuffState::new();
 569         for top_ctx in self.top_ctx.iter_mut() {
 570             for el in top_ctx.iter_mut() {
 571                 *el = false;
 572             }
 573         }
 574         let mut mb_pos = 0;
 575         for _mb_y in 0..self.mb_h {
 576             let mut left_dc = [false; 4];
 577             for mb_x in 0..self.mb_w {
 578                 let mb = self.fenc.get_mb(mb_pos);
 579                 for i in 0..4 {
 580                     let cur_idx = mb_x * 2 + (i & 1);
 581                     let mut dc_mode = 0;
 582                     if self.top_ctx[0][cur_idx] {
 583                         dc_mode += 1;
 584                     }
 585                     if left_dc[i >> 1] {
 586                         dc_mode += 1;
 587                     }
 588                     self.top_ctx[0][cur_idx] = mb.coeffs[i][0] != 0;
 589                     left_dc[i >> 1] = mb.coeffs[i][0] != 0;
 590                     if self.huffman {
 591                         self.determine_coeff_runs_luma(&mut hstate, mb_pos, i);
 592                     }
 593                     match writer {
 594                         VP6Writer::BoolCoder(ref mut bc) => encode_block(bc, &mb.coeffs[i], dc_mode, &self.models.coeff_models[0], &self.models.vp6models)?,
 595                         VP6Writer::Huffman(ref mut huff) => encode_block_huff(huff, &self.models.vp6models.zigzag, &mb.coeffs[i], 0, &mut hstate, &self.models.vp6huff)?,
 596                     };
 597                 }
 598
 599                 for plane in 1..3 {
 600                     let mut dc_mode = 0;
 601                     if self.top_ctx[plane][mb_x] {
 602                         dc_mode += 1;
 603                     }
 604                     if left_dc[plane + 1] {
 605                         dc_mode += 1;
 606                     }
 607                     self.top_ctx[plane][mb_x] = mb.coeffs[plane + 3][0] != 0;
 608                     left_dc[plane + 1] = mb.coeffs[plane + 3][0] != 0;
 609                     if self.huffman {
 610                         self.determine_coeff_runs_chroma(&mut hstate, mb_pos, plane);
 611                     }
 612                     match writer {
 613                         VP6Writer::BoolCoder(ref mut bc) => encode_block(bc, &mb.coeffs[plane + 3], dc_mode, &self.models.coeff_models[1], &self.models.vp6models)?,
 614                         VP6Writer::Huffman(ref mut huff) => encode_block_huff(huff, &self.models.vp6models.zigzag, &mb.coeffs[plane + 3], 1, &mut hstate, &self.models.vp6huff)?,
 615                     };
 616                 }
 617
 618                 mb_pos += 1;
 619             }
 620         }
 621
 622         match writer {
 623             VP6Writer::BoolCoder(bc) => bc.flush()?,
 624             VP6Writer::Huffman(huff) => huff.flush()?,
 625         };
 626
 627         Ok(())
 628     }
 629 }
 630
 631 impl NAEncoder for VP6Encoder {
 632     fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
 633         match encinfo.format {
 634             NACodecTypeInfo::None => {
 635                 let mut ofmt = EncodeParameters::default();
 636                 ofmt.format = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, true, YUV420_FORMAT));
 637                 Ok(ofmt)
 638             },
 639             NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
 640             NACodecTypeInfo::Video(vinfo) => {
 641                 let outinfo = NAVideoInfo::new((vinfo.width + 3) & !3, (vinfo.height + 3) & !3, true, YUV420_FORMAT);
 642                 let mut ofmt = *encinfo;
 643                 ofmt.format = NACodecTypeInfo::Video(outinfo);
 644                 Ok(ofmt)
 645             }
 646         }
 647     }
 648     fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
 649         match encinfo.format {
 650             NACodecTypeInfo::None => Err(EncoderError::FormatError),
 651             NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
 652             NACodecTypeInfo::Video(vinfo) => {
 653                 if vinfo.format != YUV420_FORMAT {
 654                     return Err(EncoderError::FormatError);
 655                 }
 656                 if ((vinfo.width | vinfo.height) & 15) != 0 {
 657                     return Err(EncoderError::FormatError);
 658                 }
 659                 if (vinfo.width | vinfo.height) >= (1 << 12) {
 660                     return Err(EncoderError::FormatError);
 661                 }
 662
 663                 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, true, vinfo.format);
 664                 let info = NACodecInfo::new("vp6", NACodecTypeInfo::Video(out_info), None);
 665                 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
 666                 stream.set_num(stream_id as usize);
 667                 let stream = stream.into_ref();
 668
 669                 self.last_frame = alloc_video_buffer(out_info, 4)?;
 670                 self.gold_frame = alloc_video_buffer(out_info, 4)?;
 671
 672                 self.stream = Some(stream.clone());
 673
 674                 self.mb_w = (vinfo.width  + 15) >> 4;
 675                 self.mb_h = (vinfo.height + 15) >> 4;
 676                 self.fenc.resize(self.mb_w, self.mb_h);
 677                 self.ratectl.init(self.mb_w, self.mb_h, encinfo.bitrate, encinfo.tb_num, encinfo.tb_den);
 678
 679                 self.dc_pred.resize(self.mb_w);
 680                 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]];
 681
 682                 self.version = VERSION_VP60;
 683                 self.profile = VP6_SIMPLE_PROFILE;
 684
 685                 Ok(stream)
 686             },
 687         }
 688     }
 689     fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
 690         let buf = frm.get_buffer();
 691         if let Some(ref vbuf) = buf.get_vbuf() {
 692             let mut dbuf = Vec::with_capacity(4);
 693             let mut gw   = GrowableMemoryWriter::new_write(&mut dbuf);
 694             let mut bw   = ByteWriter::new(&mut gw);
 695
 696             // todo integrate with rate control
 697             let is_intra = (self.frmcount == 0) || self.decide_encoding();
 698             let quant = if let Some(q) = self.force_q {
 699                     q
 700                 } else {
 701                     self.ratectl.guess_quant(is_intra, self.huffman)
 702                 };
 703
 704             self.fenc.read_mbs(vbuf);
 705             self.fenc.set_quant(quant);
 706             self.fenc.me_mode = self.me_mode;
 707             self.fenc.me_range = self.me_range;
 708             let golden_frame = if is_intra {
 709                     self.encode_intra(&mut bw, quant)?
 710                 } else {
 711                     self.fenc.estimate_mvs(self.last_frame.get_vbuf().unwrap(), self.mc_buf.clone(), false);
 712                     if !self.last_gold {
 713                         self.fenc.estimate_mvs(self.gold_frame.get_vbuf().unwrap(), self.mc_buf.clone(), true);
 714                     }
 715                     self.encode_inter(&mut bw, quant)?
 716                 };
 717             self.fenc.reconstruct_frame(&mut self.dc_pred, self.last_frame.get_vbuf().unwrap());
 718             self.last_gold = golden_frame;
 719             if golden_frame {
 720                 let mut dfrm = self.gold_frame.get_vbuf().unwrap();
 721                 let src = self.last_frame.get_vbuf().unwrap();
 722
 723                 let dst = dfrm.get_data_mut().unwrap();
 724                 dst.copy_from_slice(src.get_data());
 725             }
 726
 727             if self.force_q.is_none() {
 728                 self.ratectl.update(dbuf.len() * 8);
 729             }
 730
 731             self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf));
 732             if self.key_int > 0 {
 733                 self.frmcount += 1;
 734             }
 735             if self.frmcount == self.key_int {
 736                 self.frmcount = 0;
 737             }
 738             Ok(())
 739         } else {
 740             Err(EncoderError::InvalidParameters)
 741         }
 742     }
 743     fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
 744         let mut npkt = None;
 745         std::mem::swap(&mut self.pkt, &mut npkt);
 746         Ok(npkt)
 747     }
 748     fn flush(&mut self) -> EncoderResult<()> {
 749         self.frmcount = 0;
 750         Ok(())
 751     }
 752 }
 753
 754 const HUFFMAN_OPTION: &str = "huffman";
 755 const QUANT_OPTION: &str = "quant";
 756 const VERSION_OPTION: &str = "version";
 757 const MV_SEARCH_OPTION: &str = "mv_mode";
 758 const MV_RANGE_OPTION: &str = "mv_range";
 759
 760 const ENCODER_OPTS: &[NAOptionDefinition] = &[
 761     NAOptionDefinition {
 762         name: KEYFRAME_OPTION, description: KEYFRAME_OPTION_DESC,
 763         opt_type: NAOptionDefinitionType::Int(Some(0), Some(128)) },
 764     NAOptionDefinition {
 765         name: HUFFMAN_OPTION, description: "use Huffman encoding",
 766         opt_type: NAOptionDefinitionType::Bool },
 767     NAOptionDefinition {
 768         name: VERSION_OPTION, description: "codec minor version",
 769         opt_type: NAOptionDefinitionType::String(Some(&["vp60", "vp61", "vp62"])) },
 770     NAOptionDefinition {
 771         name: QUANT_OPTION, description: "force fixed quantiser for encoding",
 772         opt_type: NAOptionDefinitionType::Int(Some(-1), Some(63)) },
 773     NAOptionDefinition {
 774         name: MV_SEARCH_OPTION, description: "motion search mode",
 775         opt_type: NAOptionDefinitionType::String(Some(&["full", "dia", "hex"])) },
 776     NAOptionDefinition {
 777         name: MV_RANGE_OPTION, description: "motion search range (in pixels)",
 778         opt_type: NAOptionDefinitionType::Int(Some(0), Some(30)) },
 779 ];
 780
 781 impl NAOptionHandler for VP6Encoder {
 782     fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
 783     fn set_options(&mut self, options: &[NAOption]) {
 784         for option in options.iter() {
 785             for opt_def in ENCODER_OPTS.iter() {
 786                 if opt_def.check(option).is_ok() {
 787                     match option.name {
 788                         KEYFRAME_OPTION => {
 789                             if let NAValue::Int(intval) = option.value {
 790                                 self.key_int = intval as u8;
 791                             }
 792                         },
 793                         HUFFMAN_OPTION => {
 794                             if let NAValue::Bool(bval) = option.value {
 795                                 self.huffman = bval;
 796                             }
 797                         },
 798                         VERSION_OPTION => {
 799                             if let NAValue::String(ref string) = option.value {
 800                                 self.version = match string.as_str() {
 801                                         "vp60" => VERSION_VP60,
 802                                         "vp61" => VERSION_VP61,
 803                                         "vp62" => VERSION_VP62,
 804                                         _ => unreachable!(),
 805                                     };
 806                             }
 807                         },
 808                         QUANT_OPTION => {
 809                             if let NAValue::Int(intval) = option.value {
 810                                 self.force_q = if intval < 0 { None } else { Some(intval as usize) };
 811                             }
 812                         },
 813                         MV_SEARCH_OPTION => {
 814                             if let NAValue::String(ref string) = option.value {
 815                                 if let Ok(mv_mode) = string.parse::<MVSearchMode>() {
 816                                     self.me_mode = mv_mode;
 817                                 }
 818                             }
 819                         },
 820                         MV_RANGE_OPTION => {
 821                             if let NAValue::Int(intval) = option.value {
 822                                 self.me_range = intval as i16;
 823                             }
 824                         },
 825                         _ => {},
 826                     };
 827                 }
 828             }
 829         }
 830     }
 831     fn query_option_value(&self, name: &str) -> Option<NAValue> {
 832         match name {
 833             KEYFRAME_OPTION => Some(NAValue::Int(i64::from(self.key_int))),
 834             HUFFMAN_OPTION => Some(NAValue::Bool(self.huffman)),
 835             VERSION_OPTION => {
 836                     let ver = match self.version {
 837                             VERSION_VP60 => "vp60",
 838                             VERSION_VP61 => "vp61",
 839                             VERSION_VP62 => "vp62",
 840                             _ => unreachable!(),
 841                         };
 842                     Some(NAValue::String(ver.to_string()))
 843                 },
 844             QUANT_OPTION => if let Some(q) = self.force_q {
 845                     Some(NAValue::Int(q as i64))
 846                 } else {
 847                     Some(NAValue::Int(-1))
 848                 },
 849             MV_SEARCH_OPTION => Some(NAValue::String(self.me_mode.to_string())),
 850             MV_RANGE_OPTION => Some(NAValue::Int(i64::from(self.me_range))),
 851             _ => None,
 852         }
 853     }
 854 }
 855
 856 pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
 857     Box::new(VP6Encoder::new())
 858 }
 859
 860 #[cfg(test)]
 861 mod test {
 862     use nihav_core::codecs::*;
 863     use nihav_core::demuxers::*;
 864     use nihav_core::muxers::*;
 865     use crate::*;
 866     use nihav_commonfmt::*;
 867     use nihav_codec_support::test::enc_video::*;
 868
 869     fn encode_test(out_name: &'static str, enc_options: &[NAOption], hash: &[u32; 4]) {
 870         let mut dmx_reg = RegisteredDemuxers::new();
 871         generic_register_all_demuxers(&mut dmx_reg);
 872         let mut dec_reg = RegisteredDecoders::new();
 873         duck_register_all_decoders(&mut dec_reg);
 874         let mut mux_reg = RegisteredMuxers::new();
 875         generic_register_all_muxers(&mut mux_reg);
 876         let mut enc_reg = RegisteredEncoders::new();
 877         duck_register_all_encoders(&mut enc_reg);
 878
 879         let dec_config = DecoderTestParams {
 880                 demuxer:        "avi",
 881                 in_name:        "assets/Duck/ot171_vp40.avi",
 882                 stream_type:    StreamType::Video,
 883                 limit:          Some(1),
 884                 dmx_reg, dec_reg,
 885             };
 886         let enc_config = EncoderTestParams {
 887                 muxer:          "avi",
 888                 enc_name:       "vp6",
 889                 out_name,
 890                 mux_reg, enc_reg,
 891             };
 892         let dst_vinfo = NAVideoInfo {
 893                 width:   0,
 894                 height:  0,
 895                 format:  YUV420_FORMAT,
 896                 flipped: true,
 897                 bits:    12,
 898             };
 899         let enc_params = EncodeParameters {
 900                 format:  NACodecTypeInfo::Video(dst_vinfo),
 901                 quality: 0,
 902                 bitrate: 25000,
 903                 tb_num:  0,
 904                 tb_den:  0,
 905                 flags:   0,
 906             };
 907         //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
 908         test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
 909                           hash);
 910     }
 911     #[test]
 912     fn test_vp6_encoder_bc() {
 913         let enc_options = &[
 914                 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(42) },
 915             ];
 916         encode_test("vp6-bool.avi", enc_options, &[0xb57f49e5, 0x6b48accd, 0xc28fadb3, 0xc89a30d2]);
 917     }
 918     #[test]
 919     fn test_vp6_encoder_rc() {
 920         let enc_options = &[
 921             ];
 922         encode_test("vp6-rc.avi", enc_options, &[0x790baca9, 0x663eafcf, 0x36d1bed8, 0xddf882de]);
 923     }
 924     #[test]
 925     fn test_vp6_encoder_huff() {
 926         let enc_options = &[
 927                 NAOption { name: super::HUFFMAN_OPTION, value: NAValue::Bool(true) },
 928                 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(42) },
 929             ];
 930         encode_test("vp6-huff.avi", enc_options, &[0x6e9bb23d, 0xde296d92, 0x4c225bae, 0x3651e31f]);
 931     }
 932 }