codecs/vp6enc/coder.rs

   1 use nihav_core::codecs::{EncoderResult, EncoderError};
   2 use nihav_codec_support::codecs::MV;
   3 use crate::codecs::vpenc::coder::*;
   4 use crate::codecs::vpenc::models::*;
   5 use super::super::vpcommon::*;
   6 use super::super::vp6data::*;
   7 use super::models::*;
   8
   9 pub use crate::codecs::vpenc::coder::{BoolEncoder, Estimator};
  10
  11 pub const MODE_TREE: &[TokenSeq<VPMBType>] = &[
  12     bit_seq!(VPMBType::Intra;         T, F, F; 0, 2, 5),
  13     bit_seq!(VPMBType::InterFourMV;   T, F, T; 0, 2, 5),
  14     bit_seq!(VPMBType::InterNoMV;     F, F, F; 0, 1, 3),
  15     bit_seq!(VPMBType::InterMV;       F, F, T; 0, 1, 3),
  16     bit_seq!(VPMBType::InterNearest;  F, T, F; 0, 1, 4),
  17     bit_seq!(VPMBType::InterNear;     F, T, T; 0, 1, 4),
  18     bit_seq!(VPMBType::GoldenNoMV;    T, T, F, F; 0, 2, 6, 7),
  19     bit_seq!(VPMBType::GoldenMV;      T, T, F, T; 0, 2, 6, 7),
  20     bit_seq!(VPMBType::GoldenNearest; T, T, T, F; 0, 2, 6, 8),
  21     bit_seq!(VPMBType::GoldenNear;    T, T, T, T; 0, 2, 6, 8),
  22 ];
  23
  24 const MODE_TREE_DIFF: &[TokenSeq<u8>] = &[
  25     bit_seq!(1; F, T; 0, 1),
  26     bit_seq!(2; F, F; 0, 1),
  27     bit_seq!(3; T, F, T; 0, 2, 3),
  28     bit_seq!(4; T, F, F, T; 0, 2, 3, 4),
  29     bit_seq!(5; T, F, F, F, T; 0, 2, 3, 4, 5),
  30     bit_seq!(6; T, F, F, F, F; 0, 2, 3, 4, 5),
  31     bit_seq!(7; T, T; 0, 2),
  32 ];
  33
  34 const MODE_TREE_DIFF_PROBS: &[u8; 6] = &[171, 83, 199, 140, 125, 104];
  35
  36 const SHORT_MV_TREE: &[TokenSeq<u8>] = &[
  37     bit_seq!(0; F, F, F; 0, 1, 2),
  38     bit_seq!(1; F, F, T; 0, 1, 2),
  39     bit_seq!(2; F, T, F; 0, 1, 3),
  40     bit_seq!(3; F, T, T; 0, 1, 3),
  41     bit_seq!(4; T, F, F; 0, 4, 5),
  42     bit_seq!(5; T, F, T; 0, 4, 5),
  43     bit_seq!(6; T, T, F; 0, 4, 6),
  44     bit_seq!(7; T, T, T; 0, 4, 6),
  45 ];
  46
  47 const EOB: i8 = 42;
  48
  49 const DC_TREE: &[TokenSeq<i8>] = &[
  50     bit_seq!(  0; F; 0),
  51     bit_seq!(  1; T, F; 0, 2),
  52     bit_seq!(  2; T, T, F, F; 0, 2, 3, 4),
  53     bit_seq!(  3; T, T, F, T, F; 0, 2, 3, 4, 5),
  54     bit_seq!(  4; T, T, F, T, T; 0, 2, 3, 4, 5),
  55     bit_seq!( -1; T, T, T, F, F; 0, 2, 3, 6, 7),
  56     bit_seq!( -2; T, T, T, F, T; 0, 2, 3, 6, 7),
  57     bit_seq!( -3; T, T, T, T, F, F; 0, 2, 3, 6, 8, 9),
  58     bit_seq!( -4; T, T, T, T, F, T; 0, 2, 3, 6, 8, 9),
  59     bit_seq!( -5; T, T, T, T, T, F; 0, 2, 3, 6, 8, 10),
  60     bit_seq!( -6; T, T, T, T, T, T; 0, 2, 3, 6, 8, 10),
  61 ];
  62
  63 const NZ_COEF_TREE: &[TokenSeq<i8>] = &[
  64     bit_seq!(  1; F; 2),
  65     bit_seq!(  2; T, F, F; 2, 3, 4),
  66     bit_seq!(  3; T, F, T, F; 2, 3, 4, 5),
  67     bit_seq!(  4; T, F, T, T; 2, 3, 4, 5),
  68     bit_seq!( -1; T, T, F, F; 2, 3, 6, 7),
  69     bit_seq!( -2; T, T, F, T; 2, 3, 6, 7),
  70     bit_seq!( -3; T, T, T, F, F; 2, 3, 6, 8, 9),
  71     bit_seq!( -4; T, T, T, F, T; 2, 3, 6, 8, 9),
  72     bit_seq!( -5; T, T, T, T, F; 2, 3, 6, 8, 10),
  73     bit_seq!( -6; T, T, T, T, T; 2, 3, 6, 8, 10),
  74 ];
  75
  76 const COEF_TREE: &[TokenSeq<i8>] = &[
  77     bit_seq!(  0; F, T; 0, 1),
  78     bit_seq!(EOB; F, F; 0, 1),
  79     bit_seq!(  1; T, F; 0, 2),
  80     bit_seq!(  2; T, T, F, F; 0, 2, 3, 4),
  81     bit_seq!(  3; T, T, F, T, F; 0, 2, 3, 4, 5),
  82     bit_seq!(  4; T, T, F, T, T; 0, 2, 3, 4, 5),
  83     bit_seq!( -1; T, T, T, F, F; 0, 2, 3, 6, 7),
  84     bit_seq!( -2; T, T, T, F, T; 0, 2, 3, 6, 7),
  85     bit_seq!( -3; T, T, T, T, F, F; 0, 2, 3, 6, 8, 9),
  86     bit_seq!( -4; T, T, T, T, F, T; 0, 2, 3, 6, 8, 9),
  87     bit_seq!( -5; T, T, T, T, T, F; 0, 2, 3, 6, 8, 10),
  88     bit_seq!( -6; T, T, T, T, T, T; 0, 2, 3, 6, 8, 10),
  89 ];
  90
  91 fn coef_to_cat(coef: i16) -> i8 {
  92     match coef.abs() {
  93         0 ..=4  => coef.abs() as i8,
  94         5 ..=6  => -1,
  95         7 ..=10 => -2,
  96         11..=18 => -3,
  97         19..=34 => -4,
  98         35..=66 => -5,
  99         _       => -6,
 100     }
 101 }
 102
 103 const ZERO_RUN_TREE: &[TokenSeq<u8>] = &[
 104     bit_seq!(1; F, F, F; 0, 1, 2),
 105     bit_seq!(2; F, F, T; 0, 1, 2),
 106     bit_seq!(3; F, T, F; 0, 1, 3),
 107     bit_seq!(4; F, T, T; 0, 1, 3),
 108     bit_seq!(5; T, F, F, F; 0, 4, 5, 6),
 109     bit_seq!(6; T, F, F, T; 0, 4, 5, 6),
 110     bit_seq!(7; T, F, T, F; 0, 4, 5, 7),
 111     bit_seq!(8; T, F, T, T; 0, 4, 5, 7),
 112     bit_seq!(9; T, T; 0, 4),
 113 ];
 114
 115 pub trait EncoderTrait {
 116     fn write_cat(&mut self, cat: i8, tree: &[TokenSeq<i8>], tok_probs: &[u8], val_probs: &[u8; 11]) -> EncoderResult<()>;
 117     fn write_large_coef(&mut self, val: i16, cat: usize) -> EncoderResult<()>;
 118     fn write_dc(&mut self, val: i16, tok_probs: &[u8; 5], val_probs: &[u8; 11]) -> EncoderResult<()>;
 119     fn write_ac(&mut self, val: i16, tree: &[TokenSeq<i8>], probs: &[u8; 11]) -> EncoderResult<()>;
 120     fn write_zero_run(&mut self, val: usize, probs: &[u8; 14]) -> EncoderResult<()>;
 121 }
 122
 123 impl<'a, 'b> EncoderTrait for BoolEncoder<'a, 'b> {
 124     fn write_cat(&mut self, cat: i8, tree: &[TokenSeq<i8>], tok_probs: &[u8], val_probs: &[u8; 11]) -> EncoderResult<()> {
 125         for entry in tree.iter() {
 126             if entry.val == cat {
 127                 for seq in entry.seq.iter() {
 128                     let prob = if seq.idx < 5 {
 129                             tok_probs[seq.idx as usize]
 130                         } else {
 131                             val_probs[seq.idx as usize]
 132                         };
 133                     self.put_bool(seq.bit, prob)?;
 134                 }
 135                 return Ok(());
 136             }
 137         }
 138         Err(EncoderError::Bug)
 139     }
 140     fn write_large_coef(&mut self, val: i16, cat: usize) -> EncoderResult<()> {
 141         let base = VP56_COEF_BASE[cat];
 142         let mut probs = VP56_COEF_ADD_PROBS[cat].iter();
 143         let add = val.abs() - base;
 144         let mut mask = 1 << (VP6_COEF_ADD_BITS[cat] - 1);
 145         while mask != 0 {
 146             self.put_bool((add & mask) != 0, *probs.next().unwrap())?;
 147             mask >>= 1;
 148         }
 149         self.put_bool(val < 0, 128)?;
 150
 151         Ok(())
 152     }
 153     fn write_dc(&mut self, val: i16, tok_probs: &[u8; 5], val_probs: &[u8; 11]) -> EncoderResult<()> {
 154         let cat = coef_to_cat(val);
 155         self.write_cat(cat, DC_TREE, tok_probs, val_probs)?;
 156         if cat < 0 {
 157             self.write_large_coef(val, (-cat - 1) as usize)?;
 158         } else if val != 0 {
 159             self.put_bool(val < 0, 128)?;
 160         }
 161         Ok(())
 162     }
 163     fn write_ac(&mut self, val: i16, tree: &[TokenSeq<i8>], probs: &[u8; 11]) -> EncoderResult<()> {
 164         let cat = coef_to_cat(val);
 165         self.write_cat(cat, tree, probs, probs)?;
 166         if cat < 0 {
 167             self.write_large_coef(val, (-cat - 1) as usize)?;
 168         } else if val != 0 {
 169             self.put_bool(val < 0, 128)?;
 170         }
 171         Ok(())
 172     }
 173     fn write_zero_run(&mut self, val: usize, probs: &[u8; 14]) -> EncoderResult<()> {
 174         self.write_el(val.min(9) as u8, ZERO_RUN_TREE, probs)?;
 175         if val >= 9 {
 176             let add = val - 9;
 177             for i in 0..6 {
 178                 self.put_bool(((add >> i) & 1) != 0, probs[i + 8])?;
 179             }
 180         }
 181         Ok(())
 182     }
 183 }
 184
 185 fn rescale_mb_mode_prob(prob: u32, total: u32) -> u8 {
 186     (255 * prob / (1 + total)) as u8
 187 }
 188
 189 fn calc_mb_model_probs(prob_xmitted: &[u8; 20], mbtype_models: &mut [VP56MBTypeModel; 10]) {
 190     for mode in 0..10 {
 191         let mdl = &mut mbtype_models[mode];
 192         let mut cnt = [0u32; 10];
 193         let mut total = 0;
 194         for i in 0..10 {
 195             if i == mode { continue; }
 196             cnt[i] = 100 * u32::from(prob_xmitted[i * 2]);
 197             total += cnt[i];
 198         }
 199         let sum = u32::from(prob_xmitted[mode * 2]) + u32::from(prob_xmitted[mode * 2 + 1]);
 200         mdl.probs[9] = 255 - rescale_mb_mode_prob(u32::from(prob_xmitted[mode * 2 + 1]), sum);
 201
 202         let inter_mv0_weight = (cnt[0] as u32) + (cnt[2] as u32);
 203         let inter_mv1_weight = (cnt[3] as u32) + (cnt[4] as u32);
 204         let gold_mv0_weight = (cnt[5] as u32) + (cnt[6] as u32);
 205         let gold_mv1_weight = (cnt[8] as u32) + (cnt[9] as u32);
 206         let mix_weight = (cnt[1] as u32) + (cnt[7] as u32);
 207         mdl.probs[0] = 1 + rescale_mb_mode_prob(inter_mv0_weight + inter_mv1_weight, total);
 208         mdl.probs[1] = 1 + rescale_mb_mode_prob(inter_mv0_weight, inter_mv0_weight + inter_mv1_weight);
 209         mdl.probs[2] = 1 + rescale_mb_mode_prob(mix_weight, mix_weight + gold_mv0_weight + gold_mv1_weight);
 210         mdl.probs[3] = 1 + rescale_mb_mode_prob(cnt[0] as u32, inter_mv0_weight);
 211         mdl.probs[4] = 1 + rescale_mb_mode_prob(cnt[3] as u32, inter_mv1_weight);
 212         mdl.probs[5] = 1 + rescale_mb_mode_prob(cnt[1], mix_weight);
 213         mdl.probs[6] = 1 + rescale_mb_mode_prob(gold_mv0_weight, gold_mv0_weight + gold_mv1_weight);
 214         mdl.probs[7] = 1 + rescale_mb_mode_prob(cnt[5], gold_mv0_weight);
 215         mdl.probs[8] = 1 + rescale_mb_mode_prob(cnt[8], gold_mv1_weight);
 216     }
 217 }
 218
 219 fn calc_mbtype_bits(prob_xmitted: &[u8; 20], stats: &[[usize; 10]; 10], mdl: &mut [VP56MBTypeModel; 10]) -> u32 {
 220     const MB_TYPES: [VPMBType; 10] = [
 221         VPMBType::InterNoMV,
 222         VPMBType::Intra,
 223         VPMBType::InterMV,
 224         VPMBType::InterNearest,
 225         VPMBType::InterNear,
 226         VPMBType::GoldenNoMV,
 227         VPMBType::GoldenMV,
 228         VPMBType::InterFourMV,
 229         VPMBType::GoldenNearest,
 230         VPMBType::GoldenNear
 231     ];
 232
 233     calc_mb_model_probs(prob_xmitted, mdl);
 234     let mut nits = 0;
 235     for (last, (srow, mdl)) in stats.iter().zip(mdl.iter()).enumerate() {
 236         for (cur, &ccount) in srow.iter().enumerate() {
 237             let ccount = ccount as u32;
 238             nits += Estimator::est_nits(cur == last, mdl.probs[9]) * ccount;
 239             if cur != last {
 240                 for entry in MODE_TREE.iter() {
 241                     if entry.val == MB_TYPES[cur] {
 242                         for seq in entry.seq.iter() {
 243                             nits += Estimator::est_nits(seq.bit, mdl.probs[seq.idx as usize]) * ccount;
 244                         }
 245                         break;
 246                     }
 247                 }
 248             }
 249         }
 250     }
 251
 252     Estimator::nits_to_bits(nits)
 253 }
 254
 255 fn find_model_vq(prob_xmitted: &[u8; 20], vq: &[[u8; 20]; 16]) -> usize {
 256     let mut best_idx = 0;
 257     let mut best_dist = i16::MAX;
 258
 259     for (idx, row) in vq.iter().enumerate() {
 260         let mut dist = 0;
 261         for i in 0..20 {
 262             let a = prob_xmitted[i ^ 1];
 263             let b = row[i];
 264             dist += (i16::from(a) - i16::from(b)).abs();
 265         }
 266         if dist == 0 {
 267             return idx;
 268         }
 269         if dist < best_dist {
 270             best_dist = dist;
 271             best_idx  = idx;
 272         }
 273     }
 274
 275     best_idx
 276 }
 277
 278 // todo per-delta decision, incremental updates and such
 279 fn deltas_bits(probs: &[u8; 20], base: &[u8; 20], stats: &[[usize; 10]; 10], tmp: &mut [VP56MBTypeModel; 10], deltas: &mut [i16; 20]) -> u32 {
 280     const DELTA_PROBS: [u8; 8] = [
 281         PROB_BITS[205],
 282         PROB_BITS[256 - 205] + PROB_BITS[171] + PROB_BITS[256 - 83] + PROB_BITS[128],
 283         PROB_BITS[256 - 205] + PROB_BITS[171] + PROB_BITS[83] + PROB_BITS[128],
 284         PROB_BITS[256 - 205] + PROB_BITS[256 - 171] + PROB_BITS[199] + PROB_BITS[256 - 140] + PROB_BITS[128],
 285         PROB_BITS[256 - 205] + PROB_BITS[256 - 171] + PROB_BITS[199] + PROB_BITS[140] + PROB_BITS[256 - 125] + PROB_BITS[128],
 286         PROB_BITS[256 - 205] + PROB_BITS[256 - 171] + PROB_BITS[199] + PROB_BITS[140] + PROB_BITS[125] + PROB_BITS[256 - 104] + PROB_BITS[128],
 287         PROB_BITS[256 - 205] + PROB_BITS[256 - 171] + PROB_BITS[199] + PROB_BITS[140] + PROB_BITS[125] + PROB_BITS[104] + PROB_BITS[128],
 288         PROB_BITS[256 - 205] + PROB_BITS[256 - 171] + PROB_BITS[256 - 199] + 8 * PROB_BITS[128],
 289     ];
 290
 291     let mut nits = 0;
 292     let mut tprobs = [0u8; 20];
 293
 294     for i in 0..20 {
 295         let old = i16::from(base[i]);
 296         let new = i16::from(probs[i]);
 297         let mut diff = (new - old) & !3;
 298         if old + diff > 255 {
 299             diff -= 4;
 300         } else if old + diff < 0 || (old + diff == 0 && new != 0) {
 301             diff += 4;
 302         }
 303         tprobs[i] = (old + diff) as u8;
 304         deltas[i] = diff;
 305         nits += u32::from(DELTA_PROBS[(diff.abs() >> 2).min(7) as usize]);
 306     }
 307
 308     Estimator::nits_to_bits(nits) + calc_mbtype_bits(&tprobs, stats, tmp) + 5
 309 }
 310
 311 pub fn encode_mode_prob_models(bc: &mut BoolEncoder, models: &mut VP56Models, pmodels: &VP56Models, stats: &[[[usize; 10]; 10]; 3]) -> EncoderResult<()> {
 312     let mut tmp = [VP56MBTypeModel::default(); 10];
 313     let mut tprob = [0; 20];
 314     for ctx in 0..3 {
 315         let mut models_changed = models.prob_xmitted[ctx] != pmodels.prob_xmitted[ctx];
 316         if models_changed {
 317             let old_bits = calc_mbtype_bits(&pmodels.prob_xmitted[ctx], &stats[ctx], &mut tmp);
 318             let new_bits = calc_mbtype_bits(&models.prob_xmitted[ctx], &stats[ctx], &mut tmp) + 4;
 319             if new_bits < old_bits {
 320                 let idx = find_model_vq(&models.prob_xmitted[ctx], &VP56_MODE_VQ[ctx]);
 321                 for i in 0..20 {
 322                     tprob[i ^ 1] = VP56_MODE_VQ[ctx][idx][i];
 323                 }
 324                 let vq_bits = calc_mbtype_bits(&tprob, &stats[ctx], &mut tmp) + 4;
 325                 if vq_bits < old_bits {
 326                     bc.put_bool(true, 174)?;
 327                     bc.put_bits(idx as u32, 4)?;
 328                     let mut diffs_present = tprob != models.prob_xmitted[ctx];
 329                     let mut deltas = [0; 20];
 330                     let delta_cost = deltas_bits(&models.prob_xmitted[ctx], &tprob, &stats[ctx], &mut tmp, &mut deltas);
 331                     if delta_cost + 1 >= new_bits {
 332                         diffs_present = false;
 333                     }
 334                     if diffs_present {
 335                         bc.put_bool(true, 254)?;
 336                         for i in 0..20 {
 337                             let diff = deltas[i ^ 1] >> 2;
 338                             bc.put_bool(diff != 0, 205)?;
 339                             if diff != 0 {
 340                                 let d0 = diff.abs().min(7) as u8;
 341                                 bc.put_bool(diff < 0, 128)?;
 342                                 bc.write_el(d0, MODE_TREE_DIFF, MODE_TREE_DIFF_PROBS)?;
 343                                 if d0 == 7 {
 344                                     bc.put_bits(diff.abs() as u32, 7)?;
 345                                 }
 346                                 tprob[i ^ 1] = (i16::from(tprob[i ^ 1]) + deltas[i ^ 1]) as u8;
 347                             }
 348                         }
 349                     }
 350                     if !diffs_present {
 351                         bc.put_bool(false, 254)?;
 352                     }
 353                 } else {
 354                     models_changed = false;
 355                 }
 356             } else {
 357                 models_changed = false;
 358             }
 359         }
 360         if !models_changed {
 361             bc.put_bool(false, 174)?;
 362             bc.put_bool(false, 254)?;
 363             models.prob_xmitted[ctx] = pmodels.prob_xmitted[ctx];
 364         } else {
 365             models.prob_xmitted[ctx] = tprob;
 366         }
 367     }
 368     for ctx in 0..3 {
 369         let prob_xmitted = &models.prob_xmitted[ctx];
 370         calc_mb_model_probs(prob_xmitted, &mut models.mbtype_models[ctx]);
 371     }
 372     Ok(())
 373 }
 374
 375 pub fn encode_mv_models(bc: &mut BoolEncoder, models: &[VP56MVModel; 2], pmodels: &[VP56MVModel; 2]) -> EncoderResult<()> {
 376     for (i, (mdl, pmdl)) in models.iter().zip(pmodels.iter()).enumerate() {
 377         bc.encode_probability(mdl.nz_prob, pmdl.nz_prob, HAS_NZ_PROB[i])?;
 378         bc.encode_probability(mdl.sign_prob, pmdl.sign_prob, HAS_SIGN_PROB[i])?;
 379     }
 380     for (i, (mdl, pmdl)) in models.iter().zip(pmodels.iter()).enumerate() {
 381         for (&coded_prob, (&prob, &pprob)) in HAS_TREE_PROB[i].iter().zip(mdl.tree_probs.iter().zip(pmdl.tree_probs.iter())) {
 382             bc.encode_probability(prob, pprob, coded_prob)?;
 383         }
 384     }
 385     for (i, (mdl, pmdl)) in models.iter().zip(pmodels.iter()).enumerate() {
 386         for (&coded_prob, (&prob, &pprob)) in HAS_RAW_PROB[i].iter().zip(mdl.raw_probs.iter().zip(pmdl.raw_probs.iter())) {
 387             bc.encode_probability(prob, pprob, coded_prob)?;
 388         }
 389     }
 390     Ok(())
 391 }
 392
 393 pub fn encode_coeff_models(bc: &mut BoolEncoder, models: &mut VP56Models, pmodels: &VP56Models, is_intra: bool, interlaced: bool) -> EncoderResult<()> {
 394     let mut def_prob = [128u8; 11];
 395     for plane in 0..2 {
 396         for i in 0..11 {
 397             let pprob = pmodels.coeff_models[plane].dc_value_probs[i];
 398             let prob = models.coeff_models[plane].dc_value_probs[i];
 399             let changed = (is_intra && prob != def_prob[i]) || (!is_intra && prob != pprob);
 400             bc.put_bool(changed, HAS_COEF_PROBS[plane][i])?;
 401             if changed {
 402                 bc.put_probability(prob)?;
 403                 def_prob[i] = prob;
 404             }
 405         }
 406     }
 407
 408     bc.put_bool(false, 128)?;
 409     reset_scan(&mut models.vp6models, interlaced);
 410     /* for scan
 411         for i in 1..64 {
 412             if bc.read_prob(HAS_SCAN_UPD_PROBS[i]) {
 413                 models.vp6models.scan_order[i]  = bc.read_bits(4) as usize;
 414             }
 415         }
 416         update_scan(&mut models.vp6models);
 417     */
 418
 419     for comp in 0..2 {
 420         for i in 0..14 {
 421             bc.encode_probability(models.vp6models.zero_run_probs[comp][i], pmodels.vp6models.zero_run_probs[comp][i], HAS_ZERO_RUN_PROBS[comp][i])?;
 422         }
 423     }
 424
 425     for ctype in 0..3 {
 426         for plane in 0..2 {
 427             for group in 0..6 {
 428                 for i in 0..11 {
 429                     let pprob = pmodels.coeff_models[plane].ac_val_probs[ctype][group][i];
 430                     let prob = models.coeff_models[plane].ac_val_probs[ctype][group][i];
 431                     let changed = (is_intra && prob != def_prob[i]) || (!is_intra && prob != pprob);
 432                     bc.put_bool(changed, VP6_AC_PROBS[ctype][plane][group][i])?;
 433                     if changed {
 434                         bc.put_probability(prob)?;
 435                         def_prob[i] = prob;
 436                     }
 437                 }
 438             }
 439         }
 440     }
 441
 442     for plane in 0..2 {
 443         let mdl = &mut models.coeff_models[plane];
 444         for i in 0..3 {
 445             for k in 0..5 {
 446                 mdl.dc_token_probs[0][i][k] = rescale_prob(mdl.dc_value_probs[k], &VP6_DC_WEIGHTS[k][i], 255);
 447             }
 448         }
 449     }
 450     Ok(())
 451 }
 452
 453 pub fn encode_block(bc: &mut BoolEncoder, blk: &[i16; 64], dc_mode: usize, model: &VP56CoeffModel, vp6model: &VP6Models) -> EncoderResult<()> {
 454     let mut last = 64;
 455     for i in (0..64).rev() {
 456         if blk[vp6model.zigzag[i]] != 0 {
 457             last = i;
 458             break;
 459         }
 460     }
 461     if last < 64 {
 462         bc.write_dc(blk[0], &model.dc_token_probs[0][dc_mode], &model.dc_value_probs)?;
 463         let mut idx = 1;
 464         let mut last_idx = 0;
 465         let mut last_val = blk[0];
 466         while idx <= last {
 467             let val = blk[vp6model.zigzag[idx]];
 468             let has_nnz = (idx == 1) || (last_val != 0);
 469             if (val != 0) || has_nnz {
 470                 if last_val == 0 && idx != 1 {
 471                     let zrun = idx - last_idx;
 472                     bc.write_zero_run(zrun, &vp6model.zero_run_probs[if last_idx + 1 >= 7 { 1 } else { 0 }])?;
 473                 }
 474                 let ac_band = VP6_IDX_TO_AC_BAND[idx];
 475                 let ac_mode = last_val.abs().min(2) as usize;
 476                 let tree = if has_nnz { COEF_TREE } else { NZ_COEF_TREE };
 477                 bc.write_ac(val, tree, &model.ac_val_probs[ac_mode][ac_band])?;
 478                 last_val = val;
 479                 last_idx = idx;
 480             }
 481             idx += 1;
 482         }
 483         if idx < 64 {
 484             let ac_band = VP6_IDX_TO_AC_BAND[idx];
 485             let ac_mode = last_val.abs().min(2) as usize;
 486             bc.write_el(EOB, COEF_TREE, &model.ac_val_probs[ac_mode][ac_band])?;
 487         }
 488     } else {
 489         bc.write_cat(0, DC_TREE, &model.dc_token_probs[0][dc_mode], &model.dc_value_probs)?;
 490         let ac_band = VP6_IDX_TO_AC_BAND[1];
 491         bc.write_el(EOB, COEF_TREE, &model.ac_val_probs[0][ac_band])?;
 492     }
 493     Ok(())
 494 }
 495
 496 fn map_mb_type(mbtype: VPMBType) -> usize {
 497     match mbtype {
 498         VPMBType::InterNoMV     => 0,
 499         VPMBType::Intra         => 1,
 500         VPMBType::InterMV       => 2,
 501         VPMBType::InterNearest  => 3,
 502         VPMBType::InterNear     => 4,
 503         VPMBType::GoldenNoMV    => 5,
 504         VPMBType::GoldenMV      => 6,
 505         VPMBType::InterFourMV   => 7,
 506         VPMBType::GoldenNearest => 8,
 507         VPMBType::GoldenNear    => 9,
 508     }
 509 }
 510
 511 pub fn encode_mb_type(bc: &mut BoolEncoder, mb_type: VPMBType, last_mb_type: VPMBType, ctx: usize, model: &VP56Models) -> EncoderResult<()> {
 512     let probs = &model.mbtype_models[ctx][map_mb_type(last_mb_type)].probs;
 513     bc.put_bool(mb_type == last_mb_type, probs[9])?;
 514     if mb_type != last_mb_type {
 515         bc.write_el(mb_type, MODE_TREE, probs)?;
 516     }
 517     Ok(())
 518 }
 519
 520 fn encode_mv_component(bc: &mut BoolEncoder, mv: i16, model: &VP56MVModel) -> EncoderResult<()> {
 521     let aval = mv.abs();
 522     bc.put_bool(aval >= 8, model.nz_prob)?;
 523     if aval < 8 {
 524         bc.write_el(aval as u8, SHORT_MV_TREE, &model.tree_probs)?;
 525     } else {
 526         for &ord in LONG_VECTOR_ORDER.iter() {
 527             bc.put_bool(((aval >> ord) & 1) != 0, model.raw_probs[ord])?;
 528         }
 529         if (aval & 0xF0) != 0 {
 530             bc.put_bool((aval & (1 << 3)) != 0, model.raw_probs[3])?;
 531         }
 532     }
 533     if aval != 0 {
 534         bc.put_bool(mv < 0, model.sign_prob)?;
 535     }
 536     Ok(())
 537 }
 538
 539 pub fn encode_mv(bc: &mut BoolEncoder, mv: MV, model: &VP56Models) -> EncoderResult<()> {
 540     encode_mv_component(bc, mv.x, &model.mv_models[0])?;
 541     encode_mv_component(bc, mv.y, &model.mv_models[1])?;
 542     Ok(())
 543 }
 544
 545 pub trait VP6EstimatorTrait {
 546     fn write_cat(&self, cat: i8, tree: &[TokenSeq<i8>], probs: &mut [ProbCounter; 11]);
 547     fn write_dc(&self, val: i16, probs: &mut [ProbCounter; 11]);
 548     fn write_ac(&self, val: i16, tree: &[TokenSeq<i8>], probs: &mut [ProbCounter; 11]);
 549     fn write_zero_run(&self, val: usize, probs: &mut [ProbCounter; 14]);
 550 }
 551
 552 impl VP6EstimatorTrait for Estimator {
 553     fn write_cat(&self, cat: i8, tree: &[TokenSeq<i8>], probs: &mut [ProbCounter; 11]) {
 554         for entry in tree.iter() {
 555             if entry.val == cat {
 556                 for seq in entry.seq.iter() {
 557                     probs[seq.idx as usize].add(seq.bit);
 558                 }
 559                 return;
 560             }
 561         }
 562     }
 563     fn write_dc(&self, val: i16, probs: &mut [ProbCounter; 11]) {
 564         self.write_cat(coef_to_cat(val), DC_TREE, probs);
 565     }
 566     fn write_ac(&self, val: i16, tree: &[TokenSeq<i8>], probs: &mut [ProbCounter; 11]) {
 567         self.write_cat(coef_to_cat(val), tree, probs);
 568     }
 569     fn write_zero_run(&self, val: usize, probs: &mut [ProbCounter; 14]) {
 570         self.write_el(val.min(9) as u8, ZERO_RUN_TREE, probs);
 571         if val >= 9 {
 572             let add = val - 9;
 573             for i in 0..6 {
 574                 probs[i + 8].add(((add >> i) & 1) != 0);
 575             }
 576         }
 577     }
 578 }
 579
 580 pub fn estimate_block(blk: &[i16; 64], _dc_mode: usize, model: &mut VP56CoeffModelStat, vp6model: &mut VP6ModelsStat, scan: &[usize; 64]) {
 581     let bc = Estimator::new();
 582
 583     let mut last = 64;
 584     for i in (0..64).rev() {
 585         if blk[scan[i]] != 0 {
 586             last = i;
 587             break;
 588         }
 589     }
 590     if last < 64 {
 591         bc.write_dc(blk[0], &mut model.dc_value_probs);
 592         let mut idx = 1;
 593         let mut last_idx = 0;
 594         let mut last_val = blk[0];
 595         while idx <= last {
 596             let val = blk[scan[idx]];
 597             let has_nnz = (idx == 1) || (last_val != 0);
 598             if (val != 0) || has_nnz {
 599                 if last_val == 0 && idx != 1 {
 600                     let zrun = idx - last_idx;
 601                     bc.write_zero_run(zrun, &mut vp6model.zero_run_probs[if last_idx + 1 >= 7 { 1 } else { 0 }]);
 602                 }
 603                 let ac_band = VP6_IDX_TO_AC_BAND[idx];
 604                 let ac_mode = last_val.abs().min(2) as usize;
 605                 let tree = if has_nnz { COEF_TREE } else { NZ_COEF_TREE };
 606                 bc.write_ac(val, tree, &mut model.ac_val_probs[ac_mode][ac_band]);
 607                 last_val = val;
 608                 last_idx = idx;
 609             }
 610             idx += 1;
 611         }
 612         if idx < 64 {
 613             let ac_band = VP6_IDX_TO_AC_BAND[idx];
 614             let ac_mode = last_val.abs().min(2) as usize;
 615             bc.write_el(EOB, COEF_TREE, &mut model.ac_val_probs[ac_mode][ac_band]);
 616         }
 617     } else {
 618         bc.write_cat(0, DC_TREE, &mut model.dc_value_probs);
 619         let ac_band = VP6_IDX_TO_AC_BAND[1];
 620         bc.write_el(EOB, COEF_TREE, &mut model.ac_val_probs[0][ac_band]);
 621     }
 622 }
 623
 624 pub fn estimate_mb_type(mb_type: VPMBType, last_mb_type: VPMBType, ctx: usize, model: &mut VP56ModelsStat) {
 625     model.mbtype_models[ctx][map_mb_type(last_mb_type)][map_mb_type(mb_type)] += 1;
 626 }
 627
 628 fn estimate_mv_component(mv: i16, model: &mut VP56MVModelStat) {
 629     let bc = Estimator::new();
 630     let aval = mv.abs();
 631     model.nz_prob.add(aval >= 8);
 632     if aval < 8 {
 633         bc.write_el(aval as u8, SHORT_MV_TREE, &mut model.tree_probs);
 634     } else {
 635         for &ord in LONG_VECTOR_ORDER.iter() {
 636             model.raw_probs[ord].add(((aval >> ord) & 1) != 0);
 637         }
 638         if (aval & 0xF0) != 0 {
 639             model.raw_probs[3].add((aval & (1 << 3)) != 0);
 640         }
 641     }
 642     if aval != 0 {
 643         model.sign_prob.add(mv < 0);
 644     }
 645 }
 646
 647 pub fn estimate_mv(mv: MV, model: &mut VP56ModelsStat) {
 648     estimate_mv_component(mv.x, &mut model.mv_models[0]);
 649     estimate_mv_component(mv.y, &mut model.mv_models[1]);
 650 }
 651
 652 const VP56_MODE_VQ: [[[u8; 20]; 16]; 3] = [
 653   [
 654     [   9,  15,  32,  25,   7,  19,   9,  21,   1,  12,  14,  12,   3,  18,  14,  23,   3,  10,   0,   4 ],
 655     [  48,  39,   1,   2,  11,  27,  29,  44,   7,  27,   1,   4,   0,   3,   1,   6,   1,   2,   0,   0 ],
 656     [  21,  32,   1,   2,   4,  10,  32,  43,   6,  23,   2,   3,   1,  19,   1,   6,  12,  21,   0,   7 ],
 657     [  69,  83,   0,   0,   0,   2,  10,  29,   3,  12,   0,   1,   0,   3,   0,   3,   2,   2,   0,   0 ],
 658     [  11,  20,   1,   4,  18,  36,  43,  48,  13,  35,   0,   2,   0,   5,   3,  12,   1,   2,   0,   0 ],
 659     [  70,  44,   0,   1,   2,  10,  37,  46,   8,  26,   0,   2,   0,   2,   0,   2,   0,   1,   0,   0 ],
 660     [   8,  15,   0,   1,   8,  21,  74,  53,  22,  42,   0,   1,   0,   2,   0,   3,   1,   2,   0,   0 ],
 661     [ 141,  42,   0,   0,   1,   4,  11,  24,   1,  11,   0,   1,   0,   1,   0,   2,   0,   0,   0,   0 ],
 662     [   8,  19,   4,  10,  24,  45,  21,  37,   9,  29,   0,   3,   1,   7,  11,  25,   0,   2,   0,   1 ],
 663     [  46,  42,   0,   1,   2,  10,  54,  51,  10,  30,   0,   2,   0,   2,   0,   1,   0,   1,   0,   0 ],
 664     [  28,  32,   0,   0,   3,  10,  75,  51,  14,  33,   0,   1,   0,   2,   0,   1,   1,   2,   0,   0 ],
 665     [ 100,  46,   0,   1,   3,   9,  21,  37,   5,  20,   0,   1,   0,   2,   1,   2,   0,   1,   0,   0 ],
 666     [  27,  29,   0,   1,   9,  25,  53,  51,  12,  34,   0,   1,   0,   3,   1,   5,   0,   2,   0,   0 ],
 667     [  80,  38,   0,   0,   1,   4,  69,  33,   5,  16,   0,   1,   0,   1,   0,   0,   0,   1,   0,   0 ],
 668     [  16,  20,   0,   0,   2,   8, 104,  49,  15,  33,   0,   1,   0,   1,   0,   1,   1,   1,   0,   0 ],
 669     [ 194,  16,   0,   0,   1,   1,   1,   9,   1,   3,   0,   0,   0,   1,   0,   1,   0,   0,   0,   0 ],
 670   ], [
 671     [  41,  22,   1,   0,   1,  31,   0,   0,   0,   0,   0,   1,   1,   7,   0,   1,  98,  25,   4,  10 ],
 672     [ 123,  37,   6,   4,   1,  27,   0,   0,   0,   0,   5,   8,   1,   7,   0,   1,  12,  10,   0,   2 ],
 673     [  26,  14,  14,  12,   0,  24,   0,   0,   0,   0,  55,  17,   1,   9,   0,  36,   5,   7,   1,   3 ],
 674     [ 209,   5,   0,   0,   0,  27,   0,   0,   0,   0,   0,   1,   0,   1,   0,   1,   0,   0,   0,   0 ],
 675     [   2,   5,   4,   5,   0, 121,   0,   0,   0,   0,   0,   3,   2,   4,   1,   4,   2,   2,   0,   1 ],
 676     [ 175,   5,   0,   1,   0,  48,   0,   0,   0,   0,   0,   2,   0,   1,   0,   2,   0,   1,   0,   0 ],
 677     [  83,   5,   2,   3,   0, 102,   0,   0,   0,   0,   1,   3,   0,   2,   0,   1,   0,   0,   0,   0 ],
 678     [ 233,   6,   0,   0,   0,   8,   0,   0,   0,   0,   0,   1,   0,   1,   0,   0,   0,   1,   0,   0 ],
 679     [  34,  16, 112,  21,   1,  28,   0,   0,   0,   0,   6,   8,   1,   7,   0,   3,   2,   5,   0,   2 ],
 680     [ 159,  35,   2,   2,   0,  25,   0,   0,   0,   0,   3,   6,   0,   5,   0,   1,   4,   4,   0,   1 ],
 681     [  75,  39,   5,   7,   2,  48,   0,   0,   0,   0,   3,  11,   2,  16,   1,   4,   7,  10,   0,   2 ],
 682     [ 212,  21,   0,   1,   0,   9,   0,   0,   0,   0,   1,   2,   0,   2,   0,   0,   2,   2,   0,   0 ],
 683     [   4,   2,   0,   0,   0, 172,   0,   0,   0,   0,   0,   1,   0,   2,   0,   0,   2,   0,   0,   0 ],
 684     [ 187,  22,   1,   1,   0,  17,   0,   0,   0,   0,   3,   6,   0,   4,   0,   1,   4,   4,   0,   1 ],
 685     [ 133,   6,   1,   2,   1,  70,   0,   0,   0,   0,   0,   2,   0,   4,   0,   3,   1,   1,   0,   0 ],
 686     [ 251,   1,   0,   0,   0,   2,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0 ],
 687   ], [
 688     [   2,   3,   2,   3,   0,   2,   0,   2,   0,   0,  11,   4,   1,   4,   0,   2,   3,   2,   0,   4 ],
 689     [  49,  46,   3,   4,   7,  31,  42,  41,   0,   0,   2,   6,   1,   7,   1,   4,   2,   4,   0,   1 ],
 690     [  26,  25,   1,   1,   2,  10,  67,  39,   0,   0,   1,   1,   0,  14,   0,   2,  31,  26,   1,   6 ],
 691     [ 103,  46,   1,   2,   2,  10,  33,  42,   0,   0,   1,   4,   0,   3,   0,   1,   1,   3,   0,   0 ],
 692     [  14,  31,   9,  13,  14,  54,  22,  29,   0,   0,   2,   6,   4,  18,   6,  13,   1,   5,   0,   1 ],
 693     [  85,  39,   0,   0,   1,   9,  69,  40,   0,   0,   0,   1,   0,   3,   0,   1,   2,   3,   0,   0 ],
 694     [  31,  28,   0,   0,   3,  14, 130,  34,   0,   0,   0,   1,   0,   3,   0,   1,   3,   3,   0,   1 ],
 695     [ 171,  25,   0,   0,   1,   5,  25,  21,   0,   0,   0,   1,   0,   1,   0,   0,   0,   0,   0,   0 ],
 696     [  17,  21,  68,  29,   6,  15,  13,  22,   0,   0,   6,  12,   3,  14,   4,  10,   1,   7,   0,   3 ],
 697     [  51,  39,   0,   1,   2,  12,  91,  44,   0,   0,   0,   2,   0,   3,   0,   1,   2,   3,   0,   1 ],
 698     [  81,  25,   0,   0,   2,   9, 106,  26,   0,   0,   0,   1,   0,   1,   0,   1,   1,   1,   0,   0 ],
 699     [ 140,  37,   0,   1,   1,   8,  24,  33,   0,   0,   1,   2,   0,   2,   0,   1,   1,   2,   0,   0 ],
 700     [  14,  23,   1,   3,  11,  53,  90,  31,   0,   0,   0,   3,   1,   5,   2,   6,   1,   2,   0,   0 ],
 701     [ 123,  29,   0,   0,   1,   7,  57,  30,   0,   0,   0,   1,   0,   1,   0,   1,   0,   1,   0,   0 ],
 702     [  13,  14,   0,   0,   4,  20, 175,  20,   0,   0,   0,   1,   0,   1,   0,   1,   1,   1,   0,   0 ],
 703     [ 202,  23,   0,   0,   1,   3,   2,   9,   0,   0,   0,   1,   0,   1,   0,   1,   0,   0,   0,   0 ],
 704   ]
 705 ];