nihav-realmedia/src/codecs/rv40enc/bitstream.rs

   1 use nihav_core::frame::FrameType;
   2 use nihav_core::io::bitwriter::*;
   3 use nihav_core::io::intcode::*;
   4 use nihav_codec_support::codecs::MV;
   5 use super::types::*;
   6 use super::super::rv34codes::*;
   7 use super::super::rv40data::*;
   8
   9 pub fn write_slice_header(bw: &mut BitWriter, ftype: FrameType, q: usize, set_idx: usize, deblock: bool, pts: u32) {
  10     bw.write0();
  11     match ftype {
  12         FrameType::I => bw.write(0, 2),
  13         FrameType::P => bw.write(2, 2),
  14         FrameType::B => bw.write(3, 2),
  15         _ => unreachable!(),
  16     };
  17     bw.write(q as u32, 5);
  18     bw.write(0, 2); // unknown
  19     bw.write(set_idx as u32, 2);
  20     bw.write(!deblock as u32, 1);
  21     bw.write(pts, 13);
  22 }
  23
  24 pub fn write_slice_dimensions(bw: &mut BitWriter, width: usize, height: usize) {
  25     let wcode = match width {
  26             160 => 0,
  27             176 => 1,
  28             240 => 2,
  29             320 => 3,
  30             352 => 4,
  31             640 => 5,
  32             704 => 6,
  33             _   => 7,
  34         };
  35     bw.write(wcode, 3);
  36     if wcode == 7 {
  37         let mut w = width >> 2;
  38         while w >= 255 {
  39             bw.write(255, 8);
  40             w -= 255;
  41         }
  42         bw.write(w as u32, 8);
  43     }
  44
  45     let hcode = match height {
  46             120 => 0,
  47             132 => 1,
  48             144 => 2,
  49             240 => 3,
  50             288 => 4,
  51             480 => 5,
  52             180 => 6,
  53             360 => 7,
  54             576 => 8,
  55             _   => 9,
  56         };
  57     if hcode < 6 {
  58         bw.write(hcode, 3);
  59     } else {
  60         bw.write(hcode + 6, 4);
  61         if hcode == 9 {
  62             let mut h = height >> 2;
  63             while h >= 255 {
  64                 bw.write(255, 8);
  65                 h -= 255;
  66             }
  67             bw.write(h as u32, 8);
  68         }
  69     }
  70 }
  71
  72 pub fn write_slice_mb_idx(bw: &mut BitWriter, mb_idx: usize, num_mbs: usize) {
  73     let mba_bits = match num_mbs - 1 {
  74                0..=  47 =>  6,
  75               48..=  98 =>  7,
  76               99..= 395 =>  9,
  77              396..=1583 => 11,
  78             1584..=6335 => 13,
  79             6336..=9215 => 14,
  80             _ => unreachable!(),
  81         };
  82     bw.write(mb_idx as u32, mba_bits);
  83 }
  84
  85 pub fn write_skip_count(bw: &mut BitWriter, skip_count: u32) {
  86     bw.write_code(UintCodeType::Gamma, skip_count);
  87 }
  88
  89 fn write_mv(bw: &mut BitWriter, mv: MV) {
  90     let xcode = if mv.x > 0 { (mv.x - 1) * 2 + 1 } else { -mv.x * 2 } as u32;
  91     let ycode = if mv.y > 0 { (mv.y - 1) * 2 + 1 } else { -mv.y * 2 } as u32;
  92
  93     bw.write_code(UintCodeType::Gamma, xcode);
  94     bw.write_code(UintCodeType::Gamma, ycode);
  95 }
  96
  97 pub fn write_mb_header(bw: &mut BitWriter, ftype: FrameType, sstate: &SliceState, mbstate: &MBState) {
  98     let mb_idx = mbstate.get_mb_idx(sstate.mb_x, sstate.mb_y);
  99     let pred_mbt = mbstate.get_pred_mbtype(sstate, ftype == FrameType::B);
 100
 101     let set_id = pred_mbt.to_code();
 102
 103     if ftype != FrameType::I {
 104         let (codes, lens) = if ftype == FrameType::P {
 105                 (&RV40_PTYPE_CODES[set_id][..], &RV40_PTYPE_BITS[set_id][..])
 106             } else {
 107                 (&RV40_BTYPE_CODES[set_id][..], &RV40_BTYPE_BITS[set_id][..])
 108             };
 109         let idx = mbstate.mb_type[mb_idx].to_code();
 110         bw.write(codes[idx].into(), lens[idx]);
 111     }
 112     match mbstate.mb_type[mb_idx] {
 113         MBType::Intra16 => {
 114             if ftype == FrameType::I {
 115                 bw.write1();
 116             }
 117             bw.write(mbstate.ipred[mbstate.get_blk4_idx(sstate.mb_x, sstate.mb_y)] as u32, 2);
 118         },
 119         MBType::Intra => {
 120             if ftype == FrameType::I {
 121                 bw.write0();
 122                 bw.write1(); //dquant
 123             }
 124             let ystart = if sstate.has_t { 0 } else { 1 };
 125             let mut blk4_idx = mbstate.get_blk4_idx(sstate.mb_x, sstate.mb_y);
 126
 127             if !sstate.has_t {
 128                 let mut code = 0;
 129                 for &el in mbstate.ipred[blk4_idx..][..4].iter() {
 130                     code = code * 2 + if el == 0 { 0 } else { 1 };
 131                 }
 132                 bw.write(RV40_AIC_TOP_CODES[code].into(), RV40_AIC_TOP_BITS[code]);
 133                 blk4_idx += mbstate.blk4_stride;
 134             }
 135             for y in ystart..4 {
 136                 let mut x = 0;
 137                 while x < 4 {
 138                     let (lctx, tctx, trctx) = mbstate.get_ipred4x4_ctx(sstate.mb_x, sstate.mb_y, x, y);
 139                     let mode = mbstate.ipred[blk4_idx + x];
 140                     let ctx_word = if x < 3 {
 141                             ((trctx & 0xF) as u16) + (((tctx & 0xF) as u16) << 4) + (((lctx & 0xF) as u16) << 8)
 142                         } else { 0xFFF };
 143                     if let Some(idx) = RV40_AIC_PATTERNS.iter().position(|&x| x == ctx_word) {
 144                         let mode1 = mbstate.ipred[blk4_idx + x + 1];
 145                         let code = mode * 9 + mode1;
 146                         bw.write(RV40_AIC_MODE2_CODES[idx][code as usize].into(),
 147                                  RV40_AIC_MODE2_BITS[idx][code as usize]);
 148                         x += 2;
 149                     } else if tctx != -1 && lctx != -1 {
 150                         let idx = (tctx + lctx * 10) as usize;
 151                         let code = mode as usize;
 152                         bw.write(RV40_AIC_MODE1_CODES[idx][code].into(),
 153                                  RV40_AIC_MODE1_BITS[idx][code]);
 154                         x += 1;
 155                     } else {
 156                         match lctx {
 157                             -1 if tctx < 2 => {
 158                                 if mode == 0 {
 159                                     bw.write1();
 160                                 } else {
 161 assert_eq!(mode, 1);
 162                                     bw.write0();
 163                                 }
 164                             },
 165                             0 | 2 => {
 166                                 if mode == 0 {
 167                                     bw.write1();
 168                                 } else {
 169 assert_eq!(mode, 2);
 170                                     bw.write0();
 171                                 }
 172                             },
 173                             _ => {
 174 assert_eq!(mode, 0);
 175                             },
 176                         };
 177                         x += 1;
 178                     }
 179                 }
 180                 blk4_idx += mbstate.blk4_stride;
 181             }
 182         },
 183         MBType::P16x16 | MBType::P16x16Mix => {
 184             let diff_mv = mbstate.get_diff_mv(sstate, true, 0, 0);
 185             write_mv(bw, diff_mv);
 186         },
 187         MBType::P16x8 => {
 188             let diff_mv = mbstate.get_diff_mv(sstate, true, 0, 0);
 189             write_mv(bw, diff_mv);
 190             let diff_mv = mbstate.get_diff_mv(sstate, true, 0, 1);
 191             write_mv(bw, diff_mv);
 192         },
 193         MBType::P8x16 => {
 194             let diff_mv = mbstate.get_diff_mv(sstate, false, 0, 0);
 195             write_mv(bw, diff_mv);
 196             let diff_mv = mbstate.get_diff_mv(sstate, false, 1, 0);
 197             write_mv(bw, diff_mv);
 198         },
 199         MBType::P8x8 => {
 200             for i in 0..4 {
 201                 let diff_mv = mbstate.get_diff_mv(sstate, false, i & 1, i >> 1);
 202                 write_mv(bw, diff_mv);
 203             }
 204         },
 205         MBType::Forward => {
 206             let fwd_diff = mbstate.get_diff_mv_b(sstate, true);
 207             write_mv(bw, fwd_diff);
 208         },
 209         MBType::Backward => {
 210             let bwd_diff = mbstate.get_diff_mv_b(sstate, false);
 211             write_mv(bw, bwd_diff);
 212         },
 213         MBType::Bidir => {
 214             let fwd_diff = mbstate.get_diff_mv_b(sstate, true);
 215             let bwd_diff = mbstate.get_diff_mv_b(sstate, false);
 216             write_mv(bw, fwd_diff);
 217             write_mv(bw, bwd_diff);
 218         },
 219         MBType::Invalid => unreachable!(),
 220         _ => unimplemented!(),
 221     };
 222 }
 223
 224 trait CodeWriter {
 225     fn write(&self, bw: &mut BitWriter, code: u16);
 226 }
 227
 228 impl CodeWriter for RV34CodeReader {
 229     fn write(&self, bw: &mut BitWriter, to_write: u16) {
 230         for (&sym, (&code, &bits)) in self.syms.iter().zip(self.codes.iter().zip(self.lengths.iter())) {
 231             if sym == to_write {
 232                 bw.write(code, bits);
 233                 return;
 234             }
 235         }
 236 unreachable!();
 237     }
 238 }
 239
 240 impl CodeWriter for RV34CBPCodeReader {
 241     fn write(&self, bw: &mut BitWriter, to_write: u16) {
 242         for (&sym, (&code, &bits)) in self.syms.iter().zip(self.codes.iter().zip(self.lengths.iter())) {
 243             if u16::from(sym) == to_write {
 244                 bw.write(code, bits);
 245                 return;
 246             }
 247         }
 248 unreachable!();
 249     }
 250 }
 251
 252 struct CBPSet {
 253     cbp_pattern:    RV34CodeReader,
 254     cbp:            [RV34CBPCodeReader; 4]
 255 }
 256
 257 impl CBPSet {
 258     fn new(intra: bool, set: usize, subset: usize) -> Self {
 259         if intra {
 260             let cbp_pat = RV34CodeReader::new(&RV34_INTRA_CBPPAT[set][subset]);
 261             let cbp0 = RV34CBPCodeReader::new(&RV34_INTRA_CBP[set][subset]);
 262             let cbp1 = RV34CBPCodeReader::new(&RV34_INTRA_CBP[set][subset + 1*2]);
 263             let cbp2 = RV34CBPCodeReader::new(&RV34_INTRA_CBP[set][subset + 2*2]);
 264             let cbp3 = RV34CBPCodeReader::new(&RV34_INTRA_CBP[set][subset + 3*2]);
 265             CBPSet { cbp_pattern: cbp_pat, cbp: [cbp0, cbp1, cbp2, cbp3] }
 266         } else {
 267             let cbp_pat = RV34CodeReader::new(&RV34_INTER_CBPPAT[set]);
 268             let cbp0 = RV34CBPCodeReader::new(&RV34_INTER_CBP[set][0]);
 269             let cbp1 = RV34CBPCodeReader::new(&RV34_INTER_CBP[set][1]);
 270             let cbp2 = RV34CBPCodeReader::new(&RV34_INTER_CBP[set][2]);
 271             let cbp3 = RV34CBPCodeReader::new(&RV34_INTER_CBP[set][3]);
 272             CBPSet { cbp_pattern: cbp_pat, cbp: [cbp0, cbp1, cbp2, cbp3] }
 273         }
 274     }
 275 }
 276
 277 struct CoefSet {
 278     pat0:       Vec<RV34CodeReader>,
 279     pat1:       Vec<RV34CodeReader>,
 280     pat2:       Vec<RV34CodeReader>,
 281 }
 282
 283 impl CoefSet {
 284     fn new(intra: bool, set: usize) -> Self {
 285         if intra {
 286             let first0 = RV34CodeReader::new(&RV34_INTRA_FIRSTPAT[set][0]);
 287             let first1 = RV34CodeReader::new(&RV34_INTRA_FIRSTPAT[set][1]);
 288             let first2 = RV34CodeReader::new(&RV34_INTRA_FIRSTPAT[set][2]);
 289             let first3 = RV34CodeReader::new(&RV34_INTRA_FIRSTPAT[set][3]);
 290             let firsts = vec![first0, first1, first2, first3];
 291
 292             let second0 = RV34CodeReader::new(&RV34_INTRA_SECONDPAT[set][0]);
 293             let second1 = RV34CodeReader::new(&RV34_INTRA_SECONDPAT[set][1]);
 294             let seconds = vec![second0, second1];
 295
 296             let third0 = RV34CodeReader::new(&RV34_INTRA_THIRDPAT[set][0]);
 297             let third1 = RV34CodeReader::new(&RV34_INTRA_THIRDPAT[set][1]);
 298             let thirds = vec![third0, third1];
 299
 300             CoefSet { pat0: firsts, pat1: seconds, pat2: thirds }
 301         } else {
 302             let first0 = RV34CodeReader::new(&RV34_INTER_FIRSTPAT[set][0]);
 303             let first1 = RV34CodeReader::new(&RV34_INTER_FIRSTPAT[set][1]);
 304             let firsts = vec![first0, first1];
 305
 306             let second0 = RV34CodeReader::new(&RV34_INTER_SECONDPAT[set][0]);
 307             let second1 = RV34CodeReader::new(&RV34_INTER_SECONDPAT[set][1]);
 308             let seconds = vec![second0, second1];
 309
 310             let third0 = RV34CodeReader::new(&RV34_INTER_THIRDPAT[set][0]);
 311             let third1 = RV34CodeReader::new(&RV34_INTER_THIRDPAT[set][1]);
 312             let thirds = vec![third0, third1];
 313
 314             CoefSet { pat0: firsts, pat1: seconds, pat2: thirds }
 315         }
 316     }
 317 }
 318
 319 struct FullSet {
 320     cbp:        Vec<CBPSet>,
 321     cset:       CoefSet,
 322     coeffs:     RV34CodeReader,
 323 }
 324
 325 impl FullSet {
 326     fn new(intra: bool, set: usize) -> Self {
 327         if intra {
 328             let cbp0 = CBPSet::new(intra, set, 0);
 329             let cbp1 = CBPSet::new(intra, set, 1);
 330             let cbp: Vec<CBPSet> = vec![cbp0, cbp1];
 331             let cset = CoefSet::new(intra, set);
 332             let coeffs = RV34CodeReader::new(&RV34_INTRA_COEFFS[set]);
 333             FullSet { cbp, cset, coeffs }
 334         } else {
 335             let cbp0 = CBPSet::new(intra, set, 0);
 336             let cbp: Vec<CBPSet> = vec![cbp0];
 337             let cset = CoefSet::new(intra, set);
 338             let coeffs = RV34CodeReader::new(&RV34_INTER_COEFFS[set]);
 339             FullSet { cbp, cset, coeffs }
 340         }
 341     }
 342     fn write_block(&self, bw: &mut BitWriter, blk: &Block, subset_idx: usize, luma: bool) {
 343         let sblk0 = [blk.coeffs[0], blk.coeffs[1], blk.coeffs[4], blk.coeffs[5]];
 344         let sblk1 = [blk.coeffs[2], blk.coeffs[3], blk.coeffs[6], blk.coeffs[7]];
 345         let sblk2 = [blk.coeffs[8], blk.coeffs[12], blk.coeffs[9], blk.coeffs[13]]; // sub-block 2 has different order
 346         let sblk3 = [blk.coeffs[10], blk.coeffs[11], blk.coeffs[14], blk.coeffs[15]];
 347
 348         let idx0 = get_subblock_index(&sblk0);
 349         let idx1 = get_subblock_index(&sblk1);
 350         let idx2 = get_subblock_index(&sblk2);
 351         let idx3 = get_subblock_index(&sblk3);
 352
 353         let mut cflags = idx0;
 354         cflags = (cflags << 1) | ((idx1 != 0) as u16);
 355         cflags = (cflags << 1) | ((idx2 != 0) as u16);
 356         cflags = (cflags << 1) | ((idx3 != 0) as u16);
 357
 358         self.cset.pat0[subset_idx].write(bw, cflags);
 359
 360         if matches!(idx0, 0 | 27 | 54 | 81) { // only first coefficient is set
 361             write_single_coeff(bw, &self.coeffs, sblk0[0], 3);
 362         } else {
 363             write_coeffs(bw, &self.coeffs, &sblk0);
 364         }
 365         if idx1 != 0 {
 366             self.cset.pat1[!luma as usize].write(bw, idx1);
 367             write_coeffs(bw, &self.coeffs, &sblk1);
 368         }
 369         if idx2 != 0 {
 370             self.cset.pat1[!luma as usize].write(bw, idx2);
 371             write_coeffs(bw, &self.coeffs, &sblk2);
 372         }
 373         if idx3 != 0 {
 374             self.cset.pat2[!luma as usize].write(bw, idx3);
 375             write_coeffs(bw, &self.coeffs, &sblk3);
 376         }
 377     }
 378 }
 379
 380 fn write_coeffs(bw: &mut BitWriter, coeffs: &RV34CodeReader, blk: &[i16; 4]) {
 381     for (&val, &limit) in blk.iter().zip([3i16, 2, 2, 2].iter()) {
 382         write_single_coeff(bw, coeffs, val, limit);
 383     }
 384 }
 385
 386 fn write_single_coeff(bw: &mut BitWriter, coeffs: &RV34CodeReader, val: i16, limit: i16) {
 387     if val != 0 {
 388         if val.abs() >= limit {
 389             let mut val = (val.abs() - limit) as u16;
 390             if val > 23 {
 391                 val -= 22;
 392                 let bits = (15 - val.leading_zeros()) as u16;
 393                 coeffs.write(bw, bits + 23);
 394                 bw.write(u32::from(val - (1 << bits)), bits as u8);
 395             } else {
 396                 coeffs.write(bw, val);
 397             }
 398         }
 399         if val > 0 {
 400             bw.write0();
 401         } else {
 402             bw.write1();
 403         }
 404     }
 405 }
 406
 407 pub struct CodeSets {
 408     super_idx:  usize,
 409     set_idx:    usize,
 410     intra:      bool,
 411     is16:       bool,
 412     is_p16:     bool,
 413
 414     iset:       Vec<FullSet>,
 415     pset:       Vec<FullSet>,
 416 }
 417
 418 impl CodeSets {
 419     pub fn new() -> Self {
 420         let mut iset: Vec<FullSet> = Vec::with_capacity(5);
 421         for set in 0..5 { iset.push(FullSet::new(true, set)); }
 422         let mut pset: Vec<FullSet> = Vec::with_capacity(7);
 423         for set in 0..7 { pset.push(FullSet::new(false, set)); }
 424
 425         Self {
 426             iset, pset,
 427             super_idx:  0,
 428             set_idx:    0,
 429             intra:      false,
 430             is16:       false,
 431             is_p16:     false,
 432         }
 433     }
 434     pub fn init(&mut self, quant: usize, subset: usize) {
 435         let mut idx = quant;
 436         if (subset == 2) && (idx < 19) {
 437             idx += 10;
 438         } else if (subset != 0) && (idx < 26) {
 439             idx += 5;
 440         }
 441         if idx > 30 {
 442             idx = 30;
 443         }
 444         self.super_idx = idx;
 445     }
 446     pub fn set_params(&mut self, mbtype: &MacroblockType) {
 447         self.is_p16 = matches!(*mbtype, MacroblockType::InterMix(_));
 448         self.intra = mbtype.is_intra() || self.is_p16;
 449         self.is16 = mbtype.is_16();
 450         self.set_idx = if self.intra {
 451                 RV34_SET_IDX_INTRA[self.super_idx]
 452             } else {
 453                 RV34_SET_IDX_INTER[self.super_idx]
 454             };
 455     }
 456     fn write_cbp(&self, bw: &mut BitWriter, coded_pat: [bool; 24], cbp_code: &CBPSet) {
 457         let mut cbp_pat = 0u16;
 458         for i in 16..20 {
 459             cbp_pat = cbp_pat * 3 + (coded_pat[i] as u16) + (coded_pat[i + 4] as u16);
 460         }
 461         let mut nnz = 0usize;
 462         for blk4 in coded_pat[..16].chunks(4) {
 463             let cur_nz = blk4.contains(&true);
 464             if cur_nz {
 465                 nnz += 1;
 466             }
 467             cbp_pat = cbp_pat * 2 + (cur_nz as u16);
 468         }
 469         nnz = nnz.saturating_sub(1);
 470
 471         cbp_code.cbp_pattern.write(bw, cbp_pat);
 472         for blk4 in coded_pat[..16].chunks(4) {
 473             let pat = (blk4[3] as u16) * 32 + (blk4[2] as u16) * 16 + (blk4[1] as u16) * 2 + (blk4[0] as u16);
 474             if pat != 0 {
 475                 cbp_code.cbp[nnz].write(bw, pat);
 476             }
 477         }
 478         for i in 16..20 {
 479             if coded_pat[i] ^ coded_pat[i + 4] {
 480                 if coded_pat[i] {
 481                     bw.write1();
 482                 } else {
 483                     bw.write0();
 484                 }
 485             }
 486         }
 487     }
 488     pub fn write_coeffs(&mut self, bw: &mut BitWriter, coeffs: &[Block; 25]) {
 489         let mut fset = if self.intra { &self.iset[self.set_idx] } else { &self.pset[self.set_idx] };
 490
 491         const CODED_ORDER: [usize; 24] = [0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23];
 492         let cbp_code = &fset.cbp[if self.is16 { 1 } else { 0 }];
 493         let mut coded_blk = [false; 24];
 494         let mut coded_pat = [false; 24];
 495         for (i, ((cpat, cblk), &seq)) in coded_pat.iter_mut().zip(coded_blk.iter_mut())
 496                 .zip(CODED_ORDER.iter()).enumerate() {
 497             *cpat = !coeffs[seq].is_empty();
 498             *cblk = !coeffs[i].is_empty();
 499         }
 500         self.write_cbp(bw, coded_pat, cbp_code);
 501
 502         if self.is16 {
 503             fset.write_block(bw, &coeffs[24], 3, true);
 504         }
 505         let (luma_set, chroma_set) = if self.intra {
 506                 (if self.is16 { 2 } else { 1 }, if !self.is_p16 { 0 } else { 1 })
 507             } else {
 508                 (0, 1)
 509             };
 510         let mut citer = coded_blk.iter();
 511         for blk in coeffs[..16].iter() {
 512             if let Some(true) = citer.next() {
 513                 fset.write_block(bw, blk, luma_set, true);
 514             }
 515         }
 516         if self.is_p16 {
 517             self.set_idx = RV34_SET_IDX_INTER[self.super_idx];
 518             fset = &self.pset[self.set_idx];
 519         }
 520         for blk in coeffs[16..24].iter() {
 521             if let Some(true) = citer.next() {
 522                 fset.write_block(bw, blk, chroma_set, false);
 523             }
 524         }
 525     }
 526 }
 527
 528 fn get_subblock_index(blk: &[i16; 4]) -> u16 {
 529     let mut idx = blk[0].abs().min(3) as u16;
 530     idx = idx * 3 + (blk[1].abs().min(2) as u16);
 531     idx = idx * 3 + (blk[2].abs().min(2) as u16);
 532     idx = idx * 3 + (blk[3].abs().min(2) as u16);
 533     idx
 534 }