RealVideo 4 encoder
[nihav.git] / nihav-realmedia / src / codecs / rv40enc / types.rs
diff --git a/nihav-realmedia/src/codecs/rv40enc/types.rs b/nihav-realmedia/src/codecs/rv40enc/types.rs
new file mode 100644 (file)
index 0000000..0d3f11f
--- /dev/null
@@ -0,0 +1,586 @@
+use nihav_codec_support::codecs::{MV, ZERO_MV};
+
+pub trait RV34MVOps {
+    fn scale(&self, trd: u32, trb: u32) -> (MV, MV);
+    fn diff_gt_3(self, other: Self) -> bool;
+}
+
+impl RV34MVOps for MV {
+    fn scale(&self, trd: u32, trb: u32) -> (MV, MV) {
+        const TR_SHIFT: u8 = 14;
+        const TR_BIAS: i32 = 1 << (TR_SHIFT - 1);
+
+        let ratio = ((trb as i32) << TR_SHIFT) / (trd as i32);
+        let mv_f = MV {
+                x: (((self.x as i32) * ratio + TR_BIAS) >> TR_SHIFT) as i16,
+                y: (((self.y as i32) * ratio + TR_BIAS) >> TR_SHIFT) as i16
+            };
+        let mv_b = mv_f - *self;
+        (mv_f, mv_b)
+    }
+    fn diff_gt_3(self, other: Self) -> bool {
+        let diff = self - other;
+        diff.x.abs() > 3 || diff.y.abs() > 3
+    }
+}
+
+#[derive(Debug,Clone,Copy)]
+pub enum PredType4x4 {
+    Ver,
+    Hor,
+    DC,
+    DiagDownLeft,
+    DiagDownRight,
+    VerRight,
+    HorDown,
+    VerLeft,
+    HorUp,
+    LeftDC,
+    TopDC,
+    DC128,
+    DiagDownLeftNoDown,
+    HorUpNoDown,
+    VerLeftNoDown
+}
+
+#[derive(Debug,Clone,Copy)]
+pub enum PredType8x8 {
+    DC,
+    Hor,
+    Ver,
+    Plane,
+    LeftDC,
+    TopDC,
+    DC128
+}
+
+pub trait ToIndex {
+    fn to_index(self) -> i8;
+}
+
+impl ToIndex for PredType8x8 {
+    fn to_index(self) -> i8 {
+        match self {
+            PredType8x8::Ver => 1,
+            PredType8x8::Hor => 2,
+            PredType8x8::Plane => 3,
+            _ => 0,
+        }
+    }
+}
+
+impl ToIndex for PredType4x4 {
+    fn to_index(self) -> i8 {
+        match self {
+            PredType4x4::Ver => 1,
+            PredType4x4::Hor => 2,
+            PredType4x4::DiagDownRight => 3,
+            PredType4x4::DiagDownLeft | PredType4x4::DiagDownLeftNoDown => 4,
+            PredType4x4::VerRight => 5,
+            PredType4x4::VerLeft | PredType4x4::VerLeftNoDown => 6,
+            PredType4x4::HorUp |PredType4x4::HorUpNoDown => 7,
+            PredType4x4::HorDown => 8,
+            _ => 0, // DC predictions
+        }
+    }
+}
+
+#[derive(Clone,Copy,Default)]
+pub struct Block {
+    pub coeffs: [i16; 16],
+}
+
+impl Block {
+    pub fn new() -> Self { Self::default() }
+    pub fn is_empty(&self) -> bool {
+        for &el in self.coeffs.iter() {
+            if el != 0 {
+                return false;
+            }
+        }
+        true
+    }
+    pub fn count_nz(&self) -> usize {
+        self.coeffs.iter().filter(|&&x| x != 0).count()
+    }
+}
+impl std::fmt::Display for Block {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        let mut out = String::new();
+        for row in self.coeffs.chunks(4) {
+            out += format!(" {:3} {:3} {:3} {:3}\n", row[0], row[1], row[2], row[3]).as_str();
+        }
+        write!(f, "{}", out)
+    }
+}
+
+#[derive(Clone,Copy,Default)]
+pub struct DeblockInfo {
+    pub is_strong:  bool,
+    pub q:          u8,
+    pub cbp_y:      u16,
+    pub cbp_c:      u8,
+    pub deblock_y:  u16,
+}
+
+#[derive(Debug,Clone,Copy,PartialEq)]
+pub enum MBType {
+    Intra,
+    Intra16,
+    Skip,
+    P16x16,
+    P16x16Mix,
+    P16x8,
+    P8x16,
+    P8x8,
+    Direct,
+    Bidir,
+    Forward,
+    Backward,
+    Invalid,
+}
+
+impl MBType {
+    pub fn is_intra(self) -> bool { matches!(self, MBType::Intra | MBType::Intra16) }
+    fn get_weight(self) -> u8 {
+        match self {
+            MBType::Intra     => 0,
+            MBType::Intra16   => 1,
+            MBType::Skip      => unreachable!(),
+            MBType::P16x16    => 2,
+            MBType::P16x16Mix => 10,
+            MBType::P16x8     => 7,
+            MBType::P8x16     => 8,
+            MBType::P8x8      => 3,
+            MBType::Direct    => 6,
+            MBType::Bidir     => 9,
+            MBType::Forward   => 4,
+            MBType::Backward  => 5,
+            MBType::Invalid     => unreachable!(),
+        }
+    }
+    pub fn to_code(self) -> usize {
+        match self {
+            MBType::Intra                     => 0,
+            MBType::Intra16                   => 1,
+            MBType::P16x16 | MBType::Forward  => 2,
+            MBType::P8x8   | MBType::Backward => 3,
+            MBType::P16x8  | MBType::Bidir    => 4,
+            MBType::P8x16  | MBType::Direct   => 5,
+            MBType::P16x16Mix                 => 6,
+            _ => unreachable!(),
+        }
+    }
+    pub fn has_dir_mv(self, fwd: bool) -> bool {
+        match self {
+            MBType::Bidir => true,
+            MBType::Forward if fwd => true,
+            MBType::Backward if !fwd => true,
+            _ => false,
+        }
+    }
+}
+
+#[derive(Default)]
+pub struct SliceState {
+    pub has_t:      bool,
+    pub has_l:      bool,
+    pub has_tl:     bool,
+    pub has_tr:     bool,
+    pub mb_x:       usize,
+    pub mb_y:       usize,
+}
+
+impl SliceState {
+    pub fn new() -> Self { Self::default() }
+}
+
+#[derive(Default)]
+pub struct MBState {
+    pub mb_type:        Vec<MBType>,
+    pub ipred:          Vec<i8>,
+    pub fwd_mv:         Vec<MV>,
+    pub bwd_mv:         Vec<MV>,
+    pub ref_mv:         Vec<MV>,
+    pub mb_stride:      usize,
+    pub blk8_stride:    usize,
+    pub blk4_stride:    usize,
+}
+
+impl MBState {
+    pub fn new() -> Self { Self::default() }
+    pub fn resize(&mut self, mb_w: usize, mb_h: usize) {
+        self.mb_stride   = mb_w + 2;
+        self.blk8_stride = mb_w * 2 + 2;
+        self.blk4_stride = mb_w * 4 + 2;
+
+        self.mb_type.resize(self.mb_stride * (mb_h + 1), MBType::Invalid);
+        self.ipred.resize(self.blk4_stride * (mb_h * 4 + 1), -1);
+        self.fwd_mv.resize(self.blk8_stride * (mb_w * 2 + 1), ZERO_MV);
+        self.bwd_mv.resize(self.blk8_stride * (mb_w * 2 + 1), ZERO_MV);
+        self.ref_mv.resize(self.blk8_stride * (mb_w * 2 + 1), ZERO_MV);
+    }
+    pub fn reset(&mut self) {
+        for el in self.mb_type.iter_mut() {
+            *el = MBType::Invalid;
+        }
+        for el in self.ipred.iter_mut() {
+            *el = -1;
+        }
+    }
+    fn set_mv(&mut self, blk8_idx: usize, fwd: bool, mv: MV) {
+        if fwd {
+            self.fwd_mv[blk8_idx] = mv;
+            self.fwd_mv[blk8_idx + 1] = mv;
+            self.fwd_mv[blk8_idx + self.blk8_stride] = mv;
+            self.fwd_mv[blk8_idx + self.blk8_stride + 1] = mv;
+        } else {
+            self.bwd_mv[blk8_idx] = mv;
+            self.bwd_mv[blk8_idx + 1] = mv;
+            self.bwd_mv[blk8_idx + self.blk8_stride] = mv;
+            self.bwd_mv[blk8_idx + self.blk8_stride + 1] = mv;
+        }
+    }
+    pub fn get_mb_idx(&self, mb_x: usize, mb_y: usize) -> usize {
+        mb_x + 1 + (mb_y + 1) * self.mb_stride
+    }
+    pub fn get_blk8_idx(&self, mb_x: usize, mb_y: usize) -> usize {
+        mb_x * 2 + 1 + (mb_y * 2 + 1) * self.blk8_stride
+    }
+    pub fn get_blk4_idx(&self, mb_x: usize, mb_y: usize) -> usize {
+        mb_x * 4 + 1 + (mb_y * 4 + 1) * self.blk4_stride
+    }
+    pub fn update(&mut self, mb_type: &MacroblockType, mb_x: usize, mb_y: usize) {
+        let mb_idx = self.get_mb_idx(mb_x, mb_y);
+        let blk8_idx = self.get_blk8_idx(mb_x, mb_y);
+        let blk4_idx = self.get_blk4_idx(mb_x, mb_y);
+
+        for row in self.ipred[blk4_idx..].chunks_mut(self.blk4_stride).take(4) {
+            for el in row[..4].iter_mut() {
+                *el = 0;
+            }
+        }
+
+        match *mb_type {
+            MacroblockType::Intra16x16(ptype) => {
+                self.mb_type[mb_idx] = MBType::Intra16;
+                let pred_id = ptype.to_index();
+                for row in self.ipred[blk4_idx..].chunks_mut(self.blk4_stride).take(4) {
+                    for el in row[..4].iter_mut() {
+                        *el = pred_id;
+                    }
+                }
+                self.set_mv(blk8_idx, true,  ZERO_MV);
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::Intra4x4(ptypes) => {
+                self.mb_type[mb_idx] = MBType::Intra;
+                for (dst, src) in self.ipred[blk4_idx..].chunks_mut(self.blk4_stride).zip(ptypes.chunks(4)) {
+                    for (dst, &ptype) in dst.iter_mut().zip(src.iter()) {
+                        *dst = ptype.to_index();
+                    }
+                }
+                self.set_mv(blk8_idx, true,  ZERO_MV);
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::PSkip => {
+                self.mb_type[mb_idx] = MBType::Skip;
+                self.set_mv(blk8_idx, true,  ZERO_MV);
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::Inter16x16(mv) => {
+                self.mb_type[mb_idx] = MBType::P16x16;
+                self.set_mv(blk8_idx, true,  mv);
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::InterMix(mv) => {
+                self.mb_type[mb_idx] = MBType::P16x16Mix;
+                self.set_mv(blk8_idx, true,  mv);
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::Inter16x8(mvs) => {
+                self.mb_type[mb_idx] = MBType::P16x8;
+                self.fwd_mv[blk8_idx] = mvs[0];
+                self.fwd_mv[blk8_idx + 1] = mvs[0];
+                self.fwd_mv[blk8_idx + self.blk8_stride] = mvs[1];
+                self.fwd_mv[blk8_idx + self.blk8_stride + 1] = mvs[1];
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::Inter8x16(mvs) => {
+                self.mb_type[mb_idx] = MBType::P8x16;
+                self.fwd_mv[blk8_idx] = mvs[0];
+                self.fwd_mv[blk8_idx + 1] = mvs[1];
+                self.fwd_mv[blk8_idx + self.blk8_stride] = mvs[0];
+                self.fwd_mv[blk8_idx + self.blk8_stride + 1] = mvs[1];
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::Inter8x8(mvs) => {
+                self.mb_type[mb_idx] = MBType::P8x8;
+                self.fwd_mv[blk8_idx] = mvs[0];
+                self.fwd_mv[blk8_idx + 1] = mvs[1];
+                self.fwd_mv[blk8_idx + self.blk8_stride] = mvs[2];
+                self.fwd_mv[blk8_idx + self.blk8_stride + 1] = mvs[3];
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::BSkip(fmvs, bmvs) => {
+                self.mb_type[mb_idx] = MBType::Skip;
+                self.fwd_mv[blk8_idx] = fmvs[0];
+                self.fwd_mv[blk8_idx + 1] = fmvs[1];
+                self.fwd_mv[blk8_idx + self.blk8_stride] = fmvs[0];
+                self.fwd_mv[blk8_idx + self.blk8_stride + 1] = fmvs[1];
+                self.bwd_mv[blk8_idx] = bmvs[0];
+                self.bwd_mv[blk8_idx + 1] = bmvs[1];
+                self.bwd_mv[blk8_idx + self.blk8_stride] = bmvs[0];
+                self.bwd_mv[blk8_idx + self.blk8_stride + 1] = bmvs[1];
+            },
+            /*MacroblockType::Direct(fmv, bmv) => {
+                self.mb_type[mb_idx] = MBType::Direct;
+                self.set_mv(blk8_idx, true,  fmv);
+                self.set_mv(blk8_idx, false, bmv);
+            },*/
+            MacroblockType::Bidir(fmv, bmv) => {
+                self.mb_type[mb_idx] = MBType::Bidir;
+                self.set_mv(blk8_idx, true,  fmv);
+                self.set_mv(blk8_idx, false, bmv);
+            },
+            MacroblockType::Forward(mv) => {
+                self.mb_type[mb_idx] = MBType::Forward;
+                self.set_mv(blk8_idx, true,  mv);
+                self.set_mv(blk8_idx, false, ZERO_MV);
+            },
+            MacroblockType::Backward(mv) => {
+                self.mb_type[mb_idx] = MBType::Backward;
+                self.set_mv(blk8_idx, true,  ZERO_MV);
+                self.set_mv(blk8_idx, false, mv);
+            },
+        };
+    }
+    pub fn get_pred_mbtype(&self, sstate: &SliceState, is_b: bool) -> MBType {
+        let mut cand = [MBType::Invalid; 4];
+        let mut ccount = 0;
+
+        let mb_idx = self.get_mb_idx(sstate.mb_x, sstate.mb_y);
+        if sstate.has_t {
+            cand[ccount] = self.mb_type[mb_idx - self.mb_stride];
+            ccount += 1;
+            if sstate.has_tr {
+                cand[ccount] = self.mb_type[mb_idx - self.mb_stride + 1];
+                ccount += 1;
+            }
+        }
+        if sstate.has_l {
+            cand[ccount] = self.mb_type[mb_idx - 1];
+            ccount += 1;
+        }
+        if sstate.has_tl {
+            cand[ccount] = self.mb_type[mb_idx - self.mb_stride - 1];
+            ccount += 1;
+        }
+        if !is_b {
+            for el in cand[..ccount].iter_mut() {
+                if *el == MBType::Skip {
+                    *el = MBType::P16x16;
+                }
+            }
+        } else {
+            for el in cand[..ccount].iter_mut() {
+                if *el == MBType::Skip {
+                    *el = MBType::Direct;
+                }
+            }
+        }
+        match ccount {
+            0 => MBType::Intra,
+            1 => cand[0],
+            2 => if cand[0].get_weight() <= cand[1].get_weight() { cand[0] } else { cand[1] },
+            _ => {
+                const MBTYPE_FROM_WEIGHT: [MBType; 11] = [
+                    MBType::Intra,   MBType::Intra16,  MBType::P16x16,   MBType::P8x8,
+                    MBType::Forward, MBType::Backward, MBType::Direct,   MBType::P16x8,
+                    MBType::P8x16,   MBType::Bidir,    MBType::P16x16Mix
+                ];
+
+                let mut counts = [0; 12];
+                for el in cand[..ccount].iter() {
+                    counts[usize::from(el.get_weight())] += 1;
+                }
+                let mut best_idx = 0;
+                let mut best_wgt = 0;
+                for (idx, &weight) in counts.iter().enumerate() {
+                    if weight > best_wgt {
+                        best_idx = idx;
+                        best_wgt = weight;
+                    }
+                }
+                MBTYPE_FROM_WEIGHT[best_idx]
+            },
+        }
+    }
+    pub fn get_ipred4x4_ctx(&self, mb_x: usize, mb_y: usize, x: usize, y: usize) -> (i8, i8, i8) {
+        let blk4_idx = self.get_blk4_idx(mb_x, mb_y) + x + y * self.blk4_stride;
+        (self.ipred[blk4_idx - 1],
+         self.ipred[blk4_idx - self.blk4_stride],
+         self.ipred[blk4_idx - self.blk4_stride + 1])
+    }
+    pub fn set_ipred4x4(&mut self, mb_x: usize, mb_y: usize, modes: &[PredType4x4; 16]) {
+        let blk4_idx = self.get_blk4_idx(mb_x, mb_y);
+        for (dst, src) in self.ipred[blk4_idx..].chunks_mut(self.blk4_stride).zip(modes.chunks(4)) {
+            for (dst, src) in dst.iter_mut().zip(src.iter()) {
+                *dst = src.to_index();
+            }
+        }
+    }
+    fn get_mv(&self, idx: usize, fwd: bool) -> MV {
+        if fwd {
+            self.fwd_mv[idx]
+        } else {
+            self.bwd_mv[idx]
+        }
+    }
+    pub fn get_diff_mv(&self, sstate: &SliceState, w16: bool, xoff: usize, yoff: usize) -> MV {
+        let blk8_idx = self.get_blk8_idx(sstate.mb_x, sstate.mb_y) + xoff + yoff * self.blk8_stride;
+
+        let cur_mv = self.get_mv(blk8_idx, true);
+
+        if (yoff == 0 && !sstate.has_t) && (xoff == 0 && !sstate.has_l) {
+            return cur_mv;
+        }
+
+        let left_mv = if sstate.has_l || (xoff != 0) { self.get_mv(blk8_idx - 1, true) } else { ZERO_MV };
+        let top_mv = if sstate.has_t || (yoff != 0) { self.get_mv(blk8_idx - self.blk8_stride, true) } else { left_mv };
+        let has_tr = match xoff + yoff * 2 {
+                0 if w16 => sstate.has_tr,
+                0 => sstate.has_t,
+                1 => sstate.has_tr,
+                2 if w16 => false,
+                2 => true,
+                _ => false,
+            };
+        let has_tl = match xoff + yoff * 2 {
+                0 => sstate.has_tl,
+                1 => sstate.has_t,
+                2 => sstate.has_l,
+                _ => true,
+            };
+        let mv_c = if has_tr {
+                self.get_mv(blk8_idx - self.blk8_stride + if w16 { 2 } else { 1 }, true)
+            } else if has_tl {
+                self.get_mv(blk8_idx - self.blk8_stride - 1, true)
+            } else {
+                return cur_mv - left_mv;
+            };
+
+        cur_mv - MV::pred(left_mv, top_mv, mv_c)
+    }
+    pub fn get_diff_mv_b(&self, sstate: &SliceState, fwd: bool) -> MV {
+        let mb_idx = self.get_mb_idx(sstate.mb_x, sstate.mb_y);
+        let blk8_idx = self.get_blk8_idx(sstate.mb_x, sstate.mb_y);
+
+        let mut pred_mv = [ZERO_MV; 3];
+        let mut pcount = 0;
+
+        let cur_mv = self.get_mv(blk8_idx, fwd);
+
+        if sstate.has_l && self.mb_type[mb_idx - 1].has_dir_mv(fwd) {
+            pred_mv[pcount] = self.get_mv(blk8_idx - 1, fwd);
+            pcount += 1;
+        }
+        if !sstate.has_t {
+            return cur_mv - pred_mv[0];
+        }
+        if self.mb_type[mb_idx - self.mb_stride].has_dir_mv(fwd) {
+            pred_mv[pcount] = self.get_mv(blk8_idx - self.blk8_stride, fwd);
+            pcount += 1;
+        }
+        if sstate.has_tr {
+            if self.mb_type[mb_idx - self.mb_stride + 1].has_dir_mv(fwd) {
+                pred_mv[pcount] = self.get_mv(blk8_idx - self.blk8_stride + 2, fwd);
+                pcount += 1;
+            }
+        } else if sstate.has_tl && self.mb_type[mb_idx - self.mb_stride - 1].has_dir_mv(fwd) {
+            pred_mv[pcount] = self.get_mv(blk8_idx - self.blk8_stride - 1, fwd);
+            pcount += 1;
+        }
+        let pred_mv = match pcount {
+                3 => MV::pred(pred_mv[0], pred_mv[1], pred_mv[2]),
+                2 => MV{ x: (pred_mv[0].x + pred_mv[1].x) / 2, y: (pred_mv[0].y + pred_mv[1].y) / 2 },
+                1 => pred_mv[0],
+                _ => ZERO_MV,
+            };
+        cur_mv - pred_mv
+    }
+    pub fn swap_mvs(&mut self) {
+        std::mem::swap(&mut self.fwd_mv, &mut self.ref_mv);
+    }
+    pub fn fill_deblock(&self, dblk: &mut DeblockInfo, sstate: &SliceState) {
+        if dblk.is_strong {
+            dblk.deblock_y = 0xFFFF;
+            return;
+        }
+        let mut hmvmask = 0;
+        let mut vmvmask = 0;
+
+        let mut blk8_idx = self.get_blk8_idx(sstate.mb_x, sstate.mb_y);
+        for y in 0..2 {
+            for x in 0..2 {
+                let shift = x * 2 + y * 8;
+                let cur_mv = self.get_mv(blk8_idx + x, true);
+                if (x > 0) || (sstate.mb_x > 0) {
+                    let left_mv = self.get_mv(blk8_idx + x - 1, true);
+                    if cur_mv.diff_gt_3(left_mv) {
+                        vmvmask |= 0x11 << shift;
+                    }
+                }
+                if (y > 0) || (sstate.mb_y > 0) {
+                    let top_mv = self.get_mv(blk8_idx + x - self.blk8_stride, true);
+                    if cur_mv.diff_gt_3(top_mv) {
+                        hmvmask |= 0x03 << shift;
+                    }
+                }
+            }
+            blk8_idx += self.blk8_stride;
+        }
+        if sstate.mb_y == 0 { hmvmask &= !0x000F; }
+        if sstate.mb_x == 0 { vmvmask &= !0x1111; }
+
+        dblk.deblock_y = dblk.cbp_y | hmvmask | vmvmask;
+    }
+}
+
+#[derive(Clone)]
+pub enum MacroblockType {
+    Intra16x16(PredType8x8),
+    Intra4x4([PredType4x4; 16]),
+    PSkip,
+    Inter16x16(MV),
+    InterMix(MV),
+    Inter16x8([MV; 2]),
+    Inter8x16([MV; 2]),
+    Inter8x8([MV; 4]),
+    BSkip([MV; 4], [MV; 4]),
+    //Direct(MV, MV),
+    Bidir(MV, MV),
+    Forward(MV),
+    Backward(MV),
+}
+
+impl Default for MacroblockType {
+    fn default() -> Self { Self::Intra16x16(PredType8x8::DC) }
+}
+
+impl MacroblockType {
+    pub fn is_intra(&self) -> bool {
+        matches!(*self, MacroblockType::Intra16x16(_) | MacroblockType::Intra4x4(_))
+    }
+    pub fn is_16(&self) -> bool {
+        matches!(*self, MacroblockType::Intra16x16(_) | MacroblockType::InterMix(_))
+    }
+    pub fn is_skip(&self) -> bool {
+        matches!(*self, MacroblockType::PSkip | MacroblockType::BSkip(_, _))
+    }
+}
+
+pub struct Macroblock {
+    pub mb_type:    MacroblockType,
+    pub coeffs:     [Block; 25],
+}