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

   1 use nihav_core::codecs::EncoderResult;
   2 use nihav_codec_support::codecs::{MV, ZERO_MV};
   3 use crate::codecs::vpenc::coder::*;
   4 use super::super::vpcommon::*;
   5 use super::super::vp78::*;
   6 use super::super::vp78data::*;
   7 use super::super::vp7data::*;
   8 use super::blocks::MBType;
   9 use super::models::*;
  10 pub use crate::codecs::vpenc::coder::{BoolEncoder, Estimator};
  11
  12 const KF_Y_MODE_TREE: &[TokenSeq<PredMode>] = &[
  13     bit_seq!(PredMode::BPred;         F;       0),
  14     bit_seq!(PredMode::DCPred;        T, F, F; 0, 1, 2),
  15     bit_seq!(PredMode::VPred;         T, F, T; 0, 1, 2),
  16     bit_seq!(PredMode::HPred;         T, T, F; 0, 1, 3),
  17     bit_seq!(PredMode::TMPred;        T, T, T; 0, 1, 3),
  18 ];
  19
  20 const Y_MODE_TREE: &[TokenSeq<PredMode>] = &[
  21     bit_seq!(PredMode::DCPred;        F;       0),
  22     bit_seq!(PredMode::VPred;         T, F, F; 0, 1, 2),
  23     bit_seq!(PredMode::HPred;         T, F, T; 0, 1, 2),
  24     bit_seq!(PredMode::TMPred;        T, T, F; 0, 1, 3),
  25     bit_seq!(PredMode::BPred;         T, T, T; 0, 1, 3),
  26 ];
  27
  28 const UV_MODE_TREE: &[TokenSeq<PredMode>] = &[
  29     bit_seq!(PredMode::DCPred;        F;       0),
  30     bit_seq!(PredMode::VPred;         T, F;    0, 1),
  31     bit_seq!(PredMode::HPred;         T, T, F; 0, 1, 2),
  32     bit_seq!(PredMode::TMPred;        T, T, T; 0, 1, 2),
  33 ];
  34
  35 const B_MODE_TREE: &[TokenSeq<PredMode>] = &[
  36     bit_seq!(PredMode::DCPred;        F;                    0),
  37     bit_seq!(PredMode::TMPred;        T, F;                 0, 1),
  38     bit_seq!(PredMode::VPred;         T, T, F;              0, 1, 2),
  39     bit_seq!(PredMode::HPred;         T, T, T, F, F;        0, 1, 2, 3, 4),
  40     bit_seq!(PredMode::RDPred;        T, T, T, F, T, F;     0, 1, 2, 3, 4, 5),
  41     bit_seq!(PredMode::VRPred;        T, T, T, F, T, T;     0, 1, 2, 3, 4, 5),
  42     bit_seq!(PredMode::LDPred;        T, T, T, T, F;        0, 1, 2, 3, 6),
  43     bit_seq!(PredMode::VLPred;        T, T, T, T, T, F;     0, 1, 2, 3, 6, 7),
  44     bit_seq!(PredMode::HDPred;        T, T, T, T, T, T, F;  0, 1, 2, 3, 6, 7, 8),
  45     bit_seq!(PredMode::HUPred;        T, T, T, T, T, T, T;  0, 1, 2, 3, 6, 7, 8),
  46 ];
  47
  48 const MV_REF_TREE: &[TokenSeq<VPMBType>] = &[
  49     bit_seq!(VPMBType::InterNoMV;       F;              0),
  50     bit_seq!(VPMBType::InterNearest;    T, F;           0, 1),
  51     bit_seq!(VPMBType::InterNear;       T, T, F;        0, 1, 2),
  52     bit_seq!(VPMBType::InterMV;         T, T, T, F;     0, 1, 2, 3),
  53     bit_seq!(VPMBType::InterFourMV;     T, T, T, T;     0, 1, 2, 3),
  54 ];
  55
  56 const COEF_TREE: &[TokenSeq<DCTToken>] = &[
  57     bit_seq!(DCTToken::EOB;         F;                   0),
  58     bit_seq!(DCTToken::Zero;        T, F;                0, 1),
  59     bit_seq!(DCTToken::One;         T, T, F;             0, 1, 2),
  60     bit_seq!(DCTToken::Two;         T, T, T, F, F;       0, 1, 2, 3, 4),
  61     bit_seq!(DCTToken::Three;       T, T, T, F, T, F;    0, 1, 2, 3, 4, 5),
  62     bit_seq!(DCTToken::Four;        T, T, T, F, T, T;    0, 1, 2, 3, 4, 5),
  63     bit_seq!(DCTToken::Cat1;        T, T, T, T, F, F;    0, 1, 2, 3, 6, 7),
  64     bit_seq!(DCTToken::Cat2;        T, T, T, T, F, T;    0, 1, 2, 3, 6, 7),
  65     bit_seq!(DCTToken::Cat3;        T, T, T, T, T, F, F; 0, 1, 2, 3, 6, 8, 9),
  66     bit_seq!(DCTToken::Cat4;        T, T, T, T, T, F, T; 0, 1, 2, 3, 6, 8, 9),
  67     bit_seq!(DCTToken::Cat5;        T, T, T, T, T, T, F; 0, 1, 2, 3, 6, 8, 10),
  68     bit_seq!(DCTToken::Cat6;        T, T, T, T, T, T, T; 0, 1, 2, 3, 6, 8, 10),
  69 ];
  70
  71 const MV_TREE: &[TokenSeq<i16>] = &[
  72     bit_seq!(0; F, F, F, F; 0, 2, 3, 4),
  73     bit_seq!(1; F, F, F, T; 0, 2, 3, 4),
  74     bit_seq!(2; F, F, T, F; 0, 2, 3, 5),
  75     bit_seq!(3; F, F, T, T; 0, 2, 3, 5),
  76     bit_seq!(4; F, T, F, F; 0, 2, 6, 7),
  77     bit_seq!(5; F, T, F, T; 0, 2, 6, 7),
  78     bit_seq!(6; F, T, T, F; 0, 2, 6, 8),
  79     bit_seq!(7; F, T, T, T; 0, 2, 6, 8),
  80     bit_seq!(8; T;          0),
  81 ];
  82
  83 const MV_SPLIT_MODE_TREE: &[TokenSeq<MVSplitMode>] = &[
  84     bit_seq!(MVSplitMode::Sixteenths;   F;          0),
  85     bit_seq!(MVSplitMode::Quarters;     T, F;       0, 1),
  86     bit_seq!(MVSplitMode::TopBottom;    T, T, F;    0, 1, 2),
  87     bit_seq!(MVSplitMode::LeftRight;    T, T, T;    0, 1, 2),
  88 ];
  89
  90 const SUB_MV_REF_TREE: &[TokenSeq<SubMVRef>] = &[
  91     bit_seq!(SubMVRef::Left;    F;          0),
  92     bit_seq!(SubMVRef::Above;   T, F;       0, 1),
  93     bit_seq!(SubMVRef::Zero;    T, T, F;    0, 1, 2),
  94     bit_seq!(SubMVRef::New;     T, T, T;    0, 1, 2),
  95 ];
  96
  97 const FEATURE_TREE: &[TokenSeq<u8>] = &[
  98     bit_seq!(0; F, F; 0, 1),
  99     bit_seq!(1; F, T; 0, 1),
 100     bit_seq!(2; T, F; 0, 2),
 101     bit_seq!(3; T, T; 0, 2)
 102 ];
 103
 104 pub trait VP7BoolEncoder {
 105     fn put_byte(&mut self, val: u8) -> EncoderResult<()>;
 106     fn write_large_coef(&mut self, val: i16, cat: usize) -> EncoderResult<()>;
 107     fn encode_subblock(&mut self, blk: &[i16; 16], ctype: usize, pctx: u8, models: &VP7Models) -> EncoderResult<()>;
 108     fn encode_mv_component(&mut self, val: i16, probs: &[u8; 17]) -> EncoderResult<()>;
 109     fn encode_mv(&mut self, mv: MV, models: &VP7Models) -> EncoderResult<()>;
 110     fn encode_sub_mv(&mut self, stype: SubMVRef, mv: MV, models: &VP7Models) -> EncoderResult<()>;
 111     fn encode_feature(&mut self, id: usize, feat: Option<u8>, models: &VP7Models) -> EncoderResult<()>;
 112     fn encode_mb_type(&mut self, is_intra: bool, mb_type: &MBType, models: &VP7Models) -> EncoderResult<()>;
 113 }
 114
 115 impl<'a, 'b> VP7BoolEncoder for BoolEncoder<'a, 'b> {
 116     fn put_byte(&mut self, val: u8) -> EncoderResult<()> {
 117         self.put_bits(u32::from(val), 8)
 118     }
 119     fn write_large_coef(&mut self, val: i16, cat: usize) -> EncoderResult<()> {
 120         let base = VP56_COEF_BASE[cat];
 121         let mut probs = VP56_COEF_ADD_PROBS[cat].iter();
 122         let add = val.abs() - base;
 123         let mut mask = 1 << (VP6_COEF_ADD_BITS[cat] - 1);
 124         while mask != 0 {
 125             self.put_bool((add & mask) != 0, *probs.next().unwrap())?;
 126             mask >>= 1;
 127         }
 128         self.put_bool(val < 0, 128)?;
 129
 130         Ok(())
 131     }
 132     fn encode_subblock(&mut self, blk: &[i16; 16], ctype: usize, pctx: u8, models: &VP7Models) -> EncoderResult<()> {
 133         let probs = &models.coef_probs[ctype];
 134
 135         let start = if ctype != 0 { 0 } else { 1 };
 136         let mut cval = pctx as usize;
 137
 138         let mut last = 16;
 139         for &idx in DEFAULT_SCAN_ORDER.iter().skip(start) {
 140             if blk[idx] != 0 {
 141                 last = idx;
 142             }
 143         }
 144
 145         if last == 16 {
 146             self.write_el(DCTToken::EOB, COEF_TREE, &probs[COEF_BANDS[start]][cval])?;
 147             return Ok(());
 148         }
 149
 150         for i in start..16 {
 151             let val = blk[DEFAULT_SCAN_ORDER[i]];
 152             let token = match val.abs() {
 153                     0 => DCTToken::Zero,
 154                     1 => DCTToken::One,
 155                     2 => DCTToken::Two,
 156                     3 => DCTToken::Three,
 157                     4 => DCTToken::Four,
 158                     5..=6 => DCTToken::Cat1,
 159                     7 ..=10 => DCTToken::Cat2,
 160                     11..=18 => DCTToken::Cat3,
 161                     19..=34 => DCTToken::Cat4,
 162                     35..=66 => DCTToken::Cat5,
 163                     _       => DCTToken::Cat6,
 164                 };
 165             self.write_el(token, COEF_TREE, &probs[COEF_BANDS[i]][cval])?;
 166             match token {
 167                 DCTToken::Zero => {},
 168                 DCTToken::One |
 169                 DCTToken::Two |
 170                 DCTToken::Three |
 171                 DCTToken::Four => self.put_bool(val < 0, 128)?,
 172                 DCTToken::Cat1 => self.write_large_coef(val, 0)?,
 173                 DCTToken::Cat2 => self.write_large_coef(val, 1)?,
 174                 DCTToken::Cat3 => self.write_large_coef(val, 2)?,
 175                 DCTToken::Cat4 => self.write_large_coef(val, 3)?,
 176                 DCTToken::Cat5 => self.write_large_coef(val, 4)?,
 177                 DCTToken::Cat6 => self.write_large_coef(val, 5)?,
 178                 DCTToken::EOB => {},
 179             };
 180             cval = val.abs().min(2) as usize;
 181
 182             if DEFAULT_SCAN_ORDER[i] == last {
 183                 if DEFAULT_SCAN_ORDER[i] != 15 {
 184                     self.write_el(DCTToken::EOB, COEF_TREE, &probs[COEF_BANDS[i + 1]][cval])?;
 185                 }
 186                 break;
 187             }
 188         }
 189         Ok(())
 190     }
 191     fn encode_mv_component(&mut self, val: i16, probs: &[u8; 17]) -> EncoderResult<()> {
 192         let aval = val.abs();
 193         self.write_el(aval.min(8), MV_TREE, probs)?;
 194         if aval >= 8 {
 195             for &ord in LONG_VECTOR_ORDER.iter() {
 196                 self.put_bool(((aval >> ord) & 1) != 0, probs[ord + 9])?;
 197             }
 198             if (aval & 0xF0) != 0 {
 199                 self.put_bool((aval & (1 << 3)) != 0, probs[3 + 9])?;
 200             }
 201         }
 202         if val != 0 {
 203             self.put_bool(val < 0, probs[1])?;
 204         }
 205         Ok(())
 206     }
 207     fn encode_mv(&mut self, mv: MV, models: &VP7Models) -> EncoderResult<()> {
 208         self.encode_mv_component(mv.y, &models.mv_probs[0])?;
 209         self.encode_mv_component(mv.x, &models.mv_probs[1])?;
 210         Ok(())
 211     }
 212     fn encode_sub_mv(&mut self, stype: SubMVRef, mv: MV, models: &VP7Models) -> EncoderResult<()> {
 213         self.write_el(stype, SUB_MV_REF_TREE, &SUB_MV_REF_PROBS)?;
 214         if stype == SubMVRef::New {
 215             self.encode_mv_component(mv.y, &models.mv_probs[0])?;
 216             self.encode_mv_component(mv.x, &models.mv_probs[1])?;
 217         }
 218         Ok(())
 219     }
 220     fn encode_feature(&mut self, id: usize, feat: Option<u8>, models: &VP7Models) -> EncoderResult<()> {
 221         self.put_bool(feat.is_some(), models.feature_present[id])?;
 222         if let Some(num) = feat {
 223             self.write_el(num, FEATURE_TREE, &models.feature_tree_probs[id])?;
 224         }
 225         Ok(())
 226     }
 227     fn encode_mb_type(&mut self, is_intra: bool, mb_type: &MBType, models: &VP7Models) -> EncoderResult<()> {
 228         if !is_intra {
 229             self.put_bool(!mb_type.is_intra(), models.prob_intra_pred)?;
 230             if !mb_type.is_intra() {
 231                 let last = mb_type.get_last();
 232                 self.put_bool(!last, models.prob_last_pred)?;
 233             }
 234         }
 235         match *mb_type {
 236             MBType::Intra(ymode, uvmode) => {
 237                 if is_intra {
 238                     self.write_el(ymode, KF_Y_MODE_TREE, KF_Y_MODE_TREE_PROBS)?;
 239                     self.write_el(uvmode, UV_MODE_TREE, KF_UV_MODE_TREE_PROBS)?;
 240                 } else {
 241                     self.write_el(ymode, Y_MODE_TREE, &models.kf_ymode_prob)?;
 242                     self.write_el(uvmode, UV_MODE_TREE, &models.kf_uvmode_prob)?;
 243                 }
 244             },
 245             MBType::Intra4x4(ymodes, yctx, uvmode) => {
 246                 if is_intra {
 247                     self.write_el(PredMode::BPred, KF_Y_MODE_TREE, KF_Y_MODE_TREE_PROBS)?;
 248                     for (&ypred, &yctx) in ymodes.iter().zip(yctx.iter()) {
 249                         let top_idx  = (yctx / 10) as usize;
 250                         let left_idx = (yctx % 10) as usize;
 251                         self.write_el(ypred, B_MODE_TREE, &KF_B_MODE_TREE_PROBS[top_idx][left_idx])?;
 252                     }
 253                     self.write_el(uvmode, UV_MODE_TREE, KF_UV_MODE_TREE_PROBS)?;
 254                 } else {
 255                     self.write_el(PredMode::BPred, Y_MODE_TREE, &models.kf_ymode_prob)?;
 256                     for &ypred in ymodes.iter() {
 257                         self.write_el(ypred, B_MODE_TREE, B_MODE_TREE_PROBS)?;
 258                     }
 259                     self.write_el(uvmode, UV_MODE_TREE, &models.kf_uvmode_prob)?;
 260                 }
 261             },
 262             MBType::InterNoMV(_last, ref mv_probs) => {
 263                 self.write_el(VPMBType::InterNoMV, MV_REF_TREE, mv_probs)?;
 264             },
 265             MBType::InterNearest(_last, ref mv_probs) => {
 266                 self.write_el(VPMBType::InterNearest, MV_REF_TREE, mv_probs)?;
 267             },
 268             MBType::InterNear(_last, ref mv_probs) => {
 269                 self.write_el(VPMBType::InterNear, MV_REF_TREE, mv_probs)?;
 270             },
 271             MBType::InterMV(_last, ref mv_probs, mv) => {
 272                 self.write_el(VPMBType::InterMV, MV_REF_TREE, mv_probs)?;
 273                 self.encode_mv(mv, models)?;
 274             },
 275             MBType::InterSplitMV(_last, ref mv_probs, split_mode, stypes, mvs) => {
 276                 self.write_el(VPMBType::InterFourMV, MV_REF_TREE, mv_probs)?;
 277                 self.write_el(split_mode, MV_SPLIT_MODE_TREE, &MV_SPLIT_MODE_PROBS)?;
 278                 match split_mode {
 279                     MVSplitMode::TopBottom | MVSplitMode::LeftRight => {
 280                         for (&stype, &mv) in stypes.iter().zip(mvs.iter()).take(2) {
 281                             self.encode_sub_mv(stype, mv, models)?;
 282                         }
 283                     },
 284                     MVSplitMode::Quarters => {
 285                         for (&stype, &mv) in stypes.iter().zip(mvs.iter()).take(4) {
 286                             self.encode_sub_mv(stype, mv, models)?;
 287                         }
 288                     },
 289                     MVSplitMode::Sixteenths => {
 290                         for (&stype, &mv) in stypes.iter().zip(mvs.iter()) {
 291                             self.encode_sub_mv(stype, mv, models)?;
 292                         }
 293                     },
 294                 };
 295             },
 296         };
 297         Ok(())
 298     }
 299 }
 300
 301 pub fn encode_dct_coef_prob_upd(bc: &mut BoolEncoder, coef_probs: &[[[[u8; 11]; 3]; 8]; 4], prev: &[[[[u8; 11]; 3]; 8]; 4]) -> EncoderResult<()> {
 302     for ((new, old), upd) in coef_probs.iter().zip(prev.iter()).zip(DCT_UPDATE_PROBS.iter()) {
 303         for ((new, old), upd) in new.iter().zip(old.iter()).zip(upd.iter()) {
 304             for ((new, old), upd) in new.iter().zip(old.iter()).zip(upd.iter()) {
 305                 for ((&new, &old), &upd) in new.iter().zip(old.iter()).zip(upd.iter()) {
 306                     bc.put_bool(new != old, upd)?;
 307                     if new != old {
 308                         bc.put_byte(new)?;
 309                     }
 310                 }
 311             }
 312         }
 313     }
 314     Ok(())
 315 }
 316
 317 pub fn encode_mv_prob_upd(bc: &mut BoolEncoder, mv_probs: &[[u8; 17]; 2], prev: &[[u8; 17]; 2]) -> EncoderResult<()> {
 318     for ((new, old), upd) in mv_probs.iter().zip(prev.iter()).zip(MV_UPDATE_PROBS.iter()) {
 319         for ((&new, &old), &upd) in new.iter().zip(old.iter()).zip(upd.iter()) {
 320             bc.put_bool(new != old, upd)?;
 321             if new != old {
 322                 bc.put_bits(u32::from(new) >> 1, 7)?;
 323             }
 324         }
 325     }
 326     Ok(())
 327 }
 328
 329 pub trait VP7Estimator {
 330     fn estimate_subblock(&self, blk: &[i16; 16], ctype: usize, pctx: u8, models: &mut VP7ModelsStat);
 331     fn estimate_mv_component(&self, val: i16, probs: &mut [ProbCounter; 17]);
 332     fn estimate_mv(&self, mv: MV, models: &mut VP7ModelsStat);
 333     fn estimate_sub_mv(&self, stype: SubMVRef, mv: MV, models: &mut VP7ModelsStat);
 334     fn estimate_mb_type(&self, is_intra: bool, mb_type: &MBType, models: &mut VP7ModelsStat);
 335     fn estimate_feature(&self, id: usize, feat: Option<u8>, models: &mut VP7ModelsStat);
 336 }
 337
 338 impl VP7Estimator for Estimator {
 339     fn estimate_subblock(&self, blk: &[i16; 16], ctype: usize, pctx: u8, models: &mut VP7ModelsStat) {
 340         let probs = &mut models.coef_probs[ctype];
 341
 342         let start = if ctype != 0 { 0 } else { 1 };
 343         let mut cval = pctx as usize;
 344
 345         let mut last = 16;
 346         for &idx in DEFAULT_SCAN_ORDER.iter().skip(start) {
 347             if blk[idx] != 0 {
 348                 last = idx;
 349             }
 350         }
 351
 352         if last == 16 {
 353             self.write_el(DCTToken::EOB, COEF_TREE, &mut probs[COEF_BANDS[start]][cval]);
 354             return;
 355         }
 356
 357         for i in start..16 {
 358             let val = blk[DEFAULT_SCAN_ORDER[i]];
 359             let token = match val.abs() {
 360                     0 => DCTToken::Zero,
 361                     1 => DCTToken::One,
 362                     2 => DCTToken::Two,
 363                     3 => DCTToken::Three,
 364                     4 => DCTToken::Four,
 365                     5..=6 => DCTToken::Cat1,
 366                     7 ..=10 => DCTToken::Cat2,
 367                     11..=18 => DCTToken::Cat3,
 368                     19..=34 => DCTToken::Cat4,
 369                     35..=66 => DCTToken::Cat5,
 370                     _       => DCTToken::Cat6,
 371                 };
 372             self.write_el(token, COEF_TREE, &mut probs[COEF_BANDS[i]][cval]);
 373             cval = val.abs().min(2) as usize;
 374
 375             if DEFAULT_SCAN_ORDER[i] == last {
 376                 if DEFAULT_SCAN_ORDER[i] != 15 {
 377                     self.write_el(DCTToken::EOB, COEF_TREE, &mut probs[COEF_BANDS[i + 1]][cval]);
 378                 }
 379                 break;
 380             }
 381         }
 382     }
 383     fn estimate_mv_component(&self, val: i16, probs: &mut [ProbCounter; 17]) {
 384         let aval = val.abs();
 385         self.write_el(aval.min(8), MV_TREE, probs);
 386         if aval >= 8 {
 387             for &ord in LONG_VECTOR_ORDER.iter() {
 388                 probs[ord + 9].add(((aval >> ord) & 1) != 0);
 389             }
 390             if (aval & 0xF0) != 0 {
 391                 probs[3 + 9].add((aval & (1 << 3)) != 0);
 392             }
 393         }
 394         if val != 0 {
 395             probs[1].add(val < 0);
 396         }
 397     }
 398     fn estimate_mv(&self, mv: MV, models: &mut VP7ModelsStat) {
 399         self.estimate_mv_component(mv.y, &mut models.mv_probs[0]);
 400         self.estimate_mv_component(mv.x, &mut models.mv_probs[1]);
 401     }
 402     fn estimate_sub_mv(&self, stype: SubMVRef, mv: MV, models: &mut VP7ModelsStat) {
 403         if stype == SubMVRef::New {
 404             self.estimate_mv_component(mv.y, &mut models.mv_probs[0]);
 405             self.estimate_mv_component(mv.x, &mut models.mv_probs[1]);
 406         }
 407     }
 408     fn estimate_mb_type(&self, is_intra: bool, mb_type: &MBType, models: &mut VP7ModelsStat) {
 409         if !is_intra {
 410             models.prob_intra_pred.add(!mb_type.is_intra());
 411             if !mb_type.is_intra() {
 412                 let last = mb_type.get_last();
 413                 models.prob_last_pred.add(!last);
 414             }
 415         }
 416         match *mb_type {
 417             MBType::Intra(ymode, cmode) => {
 418                 if !is_intra {
 419                     self.write_el(ymode, Y_MODE_TREE, &mut models.kf_ymode_prob);
 420                     self.write_el(cmode, UV_MODE_TREE, &mut models.kf_uvmode_prob);
 421                 }
 422             },
 423             MBType::Intra4x4(_, _, cmode) => {
 424                 if !is_intra {
 425                     self.write_el(PredMode::BPred, Y_MODE_TREE, &mut models.kf_ymode_prob);
 426                     self.write_el(cmode, UV_MODE_TREE, &mut models.kf_uvmode_prob);
 427                 }
 428             },
 429             MBType::InterMV(_last, _, mv) => {
 430                 self.estimate_mv(mv, models);
 431             },
 432             MBType::InterSplitMV(_last, _, split_mode, stypes, mvs) => {
 433                 match split_mode {
 434                     MVSplitMode::TopBottom | MVSplitMode::LeftRight => {
 435                         for (&stype, &mv) in stypes.iter().zip(mvs.iter()).take(2) {
 436                             self.estimate_sub_mv(stype, mv, models);
 437                         }
 438                     },
 439                     MVSplitMode::Quarters => {
 440                         for (&stype, &mv) in stypes.iter().zip(mvs.iter()).take(4) {
 441                             self.estimate_sub_mv(stype, mv, models);
 442                         }
 443                     },
 444                     MVSplitMode::Sixteenths => {
 445                         for (&stype, &mv) in stypes.iter().zip(mvs.iter()) {
 446                             self.estimate_sub_mv(stype, mv, models);
 447                         }
 448                     },
 449                 };
 450             },
 451             _ => {},
 452         };
 453     }
 454     fn estimate_feature(&self, id: usize, feat: Option<u8>, models: &mut VP7ModelsStat) {
 455         models.feature_present[id].add(feat.is_some());
 456         if let Some(num) = feat {
 457             self.write_el(num, FEATURE_TREE, &mut models.feature_tree_probs[id]);
 458         }
 459     }
 460 }
 461
 462 fn code_nits<T: PartialEq>(el: T, tree: &[TokenSeq<T>], probs: &[u8]) -> u32 {
 463     let mut nits = 0;
 464     for entry in tree.iter() {
 465         if entry.val == el {
 466             for seq in entry.seq.iter() {
 467                 nits += Estimator::est_nits(seq.bit, probs[seq.idx as usize]);
 468             }
 469             return nits;
 470         }
 471     }
 472     0
 473 }
 474 pub fn b_mode_nits(mode: PredMode) -> u32 {
 475     code_nits(mode, B_MODE_TREE, &KF_B_MODE_TREE_PROBS[2][2]) // todo find better context
 476 }
 477 fn mv_component_nits(val: i16, probs: &[u8; 17]) -> u32 {
 478     let aval = val.abs();
 479     let mut nits = code_nits(aval.min(8), MV_TREE, probs);
 480     if aval >= 8 {
 481         for &ord in LONG_VECTOR_ORDER.iter() {
 482             nits += Estimator::est_nits(((aval >> ord) & 1) != 0, probs[ord + 9]);
 483         }
 484         if (aval & 0xF0) != 0 {
 485             nits += Estimator::est_nits((aval & (1 << 3)) != 0, probs[3 + 9]);
 486         }
 487     }
 488     if val != 0 {
 489         nits += u32::from(PROB_BITS[128]);
 490     }
 491     nits
 492 }
 493 pub fn inter_mv_nits(mv: MV, mvprobs: &[u8; 4], nearest_mv: MV, near_mv: MV, pred_mv: MV, models: &VP7Models) -> u32 {
 494     if mv == ZERO_MV {
 495         code_nits(VPMBType::InterNoMV, MV_REF_TREE, mvprobs)
 496     } else if mv == nearest_mv {
 497         code_nits(VPMBType::InterNearest, MV_REF_TREE, mvprobs)
 498     } else if mv == near_mv {
 499         code_nits(VPMBType::InterNear, MV_REF_TREE, mvprobs)
 500     } else {
 501         let dmv = mv - pred_mv;
 502         let mut nits = code_nits(VPMBType::InterMV, MV_REF_TREE, mvprobs);
 503         nits += mv_component_nits(dmv.y, &models.mv_probs[0]);
 504         nits += mv_component_nits(dmv.x, &models.mv_probs[1]);
 505         nits
 506     }
 507 }
 508 pub fn sub_mv_mode_nits(mode: MVSplitMode) -> u32 {
 509     code_nits(mode, MV_SPLIT_MODE_TREE, &MV_SPLIT_MODE_PROBS)
 510 }
 511 pub fn sub_mv_nits(mv: MV, left_mv: MV, top_mv: MV, pred_mv: MV, models: &VP7Models) -> (SubMVRef, u32) {
 512     if mv == ZERO_MV {
 513         (SubMVRef::Zero, code_nits(SubMVRef::Zero, SUB_MV_REF_TREE, &SUB_MV_REF_PROBS))
 514     } else if mv == left_mv {
 515         (SubMVRef::Left, code_nits(SubMVRef::Left, SUB_MV_REF_TREE, &SUB_MV_REF_PROBS))
 516     } else if mv == top_mv {
 517         (SubMVRef::Above, code_nits(SubMVRef::Above, SUB_MV_REF_TREE, &SUB_MV_REF_PROBS))
 518     } else {
 519         let dmv = mv - pred_mv;
 520         let mut nits = code_nits(SubMVRef::New, SUB_MV_REF_TREE, &SUB_MV_REF_PROBS);
 521         nits += mv_component_nits(dmv.y, &models.mv_probs[0]);
 522         nits += mv_component_nits(dmv.x, &models.mv_probs[1]);
 523         (SubMVRef::New, nits)
 524     }
 525 }
 526 fn est_large_coef(val: i16, cat: usize) -> u32 {
 527     let base = VP56_COEF_BASE[cat];
 528     let mut probs = VP56_COEF_ADD_PROBS[cat].iter();
 529     let add = val.abs() - base;
 530     let mut mask = 1 << (VP6_COEF_ADD_BITS[cat] - 1);
 531     let mut nits = 0;
 532     while mask != 0 {
 533         nits += Estimator::est_nits((add & mask) != 0, *probs.next().unwrap());
 534         mask >>= 1;
 535     }
 536     nits += u32::from(PROB_BITS[128]);
 537
 538     nits
 539 }
 540
 541 pub fn estimate_subblock_nits(blk: &[i16; 16], ctype: usize, pctx: u8, probs: &[[[u8; 11]; 3]; 8]) -> u32 {
 542     let start = if ctype != 0 { 0 } else { 1 };
 543     let mut cval = pctx as usize;
 544
 545     let mut last = 16;
 546     for &idx in DEFAULT_SCAN_ORDER.iter().skip(start) {
 547         if blk[idx] != 0 {
 548             last = idx;
 549         }
 550     }
 551
 552     if last == 16 {
 553         return code_nits(DCTToken::EOB, COEF_TREE, &probs[COEF_BANDS[start]][cval]);
 554     }
 555
 556     let mut nits = 0;
 557     for i in start..16 {
 558         let val = blk[DEFAULT_SCAN_ORDER[i]];
 559         let token = match val.abs() {
 560                 0 => DCTToken::Zero,
 561                 1 => DCTToken::One,
 562                 2 => DCTToken::Two,
 563                 3 => DCTToken::Three,
 564                 4 => DCTToken::Four,
 565                 5..=6 => DCTToken::Cat1,
 566                 7 ..=10 => DCTToken::Cat2,
 567                 11..=18 => DCTToken::Cat3,
 568                 19..=34 => DCTToken::Cat4,
 569                 35..=66 => DCTToken::Cat5,
 570                 _       => DCTToken::Cat6,
 571             };
 572         nits += code_nits(token, COEF_TREE, &probs[COEF_BANDS[i]][cval]);
 573         nits += match token {
 574                 DCTToken::Zero => 0,
 575                 DCTToken::One |
 576                 DCTToken::Two |
 577                 DCTToken::Three |
 578                 DCTToken::Four => u32::from(PROB_BITS[128]),
 579                 DCTToken::Cat1 => est_large_coef(val, 0),
 580                 DCTToken::Cat2 => est_large_coef(val, 1),
 581                 DCTToken::Cat3 => est_large_coef(val, 2),
 582                 DCTToken::Cat4 => est_large_coef(val, 3),
 583                 DCTToken::Cat5 => est_large_coef(val, 4),
 584                 DCTToken::Cat6 => est_large_coef(val, 5),
 585                 DCTToken::EOB => 0,
 586             };
 587         cval = val.abs().min(2) as usize;
 588
 589         if DEFAULT_SCAN_ORDER[i] == last {
 590             if DEFAULT_SCAN_ORDER[i] != 15 {
 591                 nits += code_nits(DCTToken::EOB, COEF_TREE, &probs[COEF_BANDS[i + 1]][cval]);
 592             }
 593             break;
 594         }
 595     }
 596     nits
 597 }
 598
 599 const VP6_COEF_ADD_BITS: [u8; 6] = [ 1, 2, 3, 4, 5, 11 ];