rav1e/context/
frame_header.rs

1// Copyright (c) 2017-2022, The rav1e contributors. All rights reserved
2//
3// This source code is subject to the terms of the BSD 2 Clause License and
4// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5// was not distributed with this source code in the LICENSE file, you can
6// obtain it at www.aomedia.org/license/software. If the Alliance for Open
7// Media Patent License 1.0 was not distributed with this source code in the
8// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9
10use super::*;
11
12impl CDFContext {
13  // rather than test writing and rolling back the cdf, we just count Q8 bits using the current cdf
14  pub fn count_lrf_switchable<W: Writer>(
15    &self, w: &W, rs: &TileRestorationState, filter: RestorationFilter,
16    pli: usize,
17  ) -> u32 {
18    match filter {
19      RestorationFilter::None => w.symbol_bits(0, &self.lrf_switchable_cdf),
20      RestorationFilter::Wiener { .. } => {
21        unreachable!() // for now, not permanently
22      }
23      RestorationFilter::Sgrproj { set, xqd } => {
24        // Does *not* use 'RESTORE_SGRPROJ' but rather just '2'
25        let rp = &rs.planes[pli];
26        let mut bits = w.symbol_bits(2, &self.lrf_switchable_cdf)
27          + ((SGRPROJ_PARAMS_BITS as u32) << OD_BITRES);
28        for i in 0..2 {
29          let s = SGRPROJ_PARAMS_S[set as usize][i];
30          let min = SGRPROJ_XQD_MIN[i] as i32;
31          let max = SGRPROJ_XQD_MAX[i] as i32;
32          if s > 0 {
33            bits += w.count_signed_subexp_with_ref(
34              xqd[i] as i32,
35              min,
36              max + 1,
37              SGRPROJ_PRJ_SUBEXP_K,
38              rp.sgrproj_ref[i] as i32,
39            );
40          }
41        }
42        bits
43      }
44    }
45  }
46}
47
48impl ContextWriter<'_> {
49  fn get_ref_frame_ctx_b0(&self, bo: TileBlockOffset) -> usize {
50    let ref_counts = self.bc.blocks[bo].neighbors_ref_counts;
51
52    let fwd_cnt = ref_counts[LAST_FRAME.to_index()]
53      + ref_counts[LAST2_FRAME.to_index()]
54      + ref_counts[LAST3_FRAME.to_index()]
55      + ref_counts[GOLDEN_FRAME.to_index()];
56
57    let bwd_cnt = ref_counts[BWDREF_FRAME.to_index()]
58      + ref_counts[ALTREF2_FRAME.to_index()]
59      + ref_counts[ALTREF_FRAME.to_index()];
60
61    ContextWriter::ref_count_ctx(fwd_cnt, bwd_cnt)
62  }
63
64  /// # Panics
65  ///
66  /// - If the `comp_mode` setting does not match the reference mode and size.
67  pub fn write_ref_frames<T: Pixel, W: Writer>(
68    &mut self, w: &mut W, fi: &FrameInvariants<T>, bo: TileBlockOffset,
69  ) {
70    let rf = self.bc.blocks[bo].ref_frames;
71    let sz = self.bc.blocks[bo].n4_w.min(self.bc.blocks[bo].n4_h);
72
73    /* TODO: Handle multiple references */
74    let comp_mode = self.bc.blocks[bo].has_second_ref();
75
76    if fi.reference_mode != ReferenceMode::SINGLE && sz >= 2 {
77      let ctx = self.get_comp_mode_ctx(bo);
78      let cdf = &self.fc.comp_mode_cdf[ctx];
79      symbol_with_update!(self, w, comp_mode as u32, cdf);
80    } else {
81      assert!(!comp_mode);
82    }
83
84    if comp_mode {
85      let comp_ref_type: u32 = 1; // bidir
86      let ctx = self.get_comp_ref_type_ctx(bo);
87      let cdf = &self.fc.comp_ref_type_cdf[ctx];
88      symbol_with_update!(self, w, comp_ref_type, cdf);
89
90      if comp_ref_type == 0 {
91        unimplemented!();
92      } else {
93        let compref = rf[0] == GOLDEN_FRAME || rf[0] == LAST3_FRAME;
94        let ctx = self.get_pred_ctx_ll2_or_l3gld(bo);
95        let cdf = &self.fc.comp_ref_cdf[ctx][0];
96        symbol_with_update!(self, w, compref as u32, cdf);
97        if !compref {
98          let compref_p1 = rf[0] == LAST2_FRAME;
99          let ctx = self.get_pred_ctx_last_or_last2(bo);
100          let cdf = &self.fc.comp_ref_cdf[ctx][1];
101          symbol_with_update!(self, w, compref_p1 as u32, cdf);
102        } else {
103          let compref_p2 = rf[0] == GOLDEN_FRAME;
104          let ctx = self.get_pred_ctx_last3_or_gold(bo);
105          let cdf = &self.fc.comp_ref_cdf[ctx][2];
106          symbol_with_update!(self, w, compref_p2 as u32, cdf);
107        }
108        let comp_bwdref = rf[1] == ALTREF_FRAME;
109        let ctx = self.get_pred_ctx_brfarf2_or_arf(bo);
110        let cdf = &self.fc.comp_bwd_ref_cdf[ctx][0];
111        symbol_with_update!(self, w, comp_bwdref as u32, cdf);
112        if !comp_bwdref {
113          let comp_bwdref_p1 = rf[1] == ALTREF2_FRAME;
114          let ctx = self.get_pred_ctx_brf_or_arf2(bo);
115          let cdf = &self.fc.comp_bwd_ref_cdf[ctx][1];
116          symbol_with_update!(self, w, comp_bwdref_p1 as u32, cdf);
117        }
118      }
119    } else {
120      let b0_ctx = self.get_ref_frame_ctx_b0(bo);
121      let b0 = rf[0] != NONE_FRAME && rf[0].is_bwd_ref();
122
123      let cdf = &self.fc.single_ref_cdfs[b0_ctx][0];
124      symbol_with_update!(self, w, b0 as u32, cdf);
125      if b0 {
126        let b1_ctx = self.get_pred_ctx_brfarf2_or_arf(bo);
127        let b1 = rf[0] == ALTREF_FRAME;
128
129        let cdf = &self.fc.single_ref_cdfs[b1_ctx][1];
130        symbol_with_update!(self, w, b1 as u32, cdf);
131        if !b1 {
132          let b5_ctx = self.get_pred_ctx_brf_or_arf2(bo);
133          let b5 = rf[0] == ALTREF2_FRAME;
134
135          let cdf = &self.fc.single_ref_cdfs[b5_ctx][5];
136          symbol_with_update!(self, w, b5 as u32, cdf);
137        }
138      } else {
139        let b2_ctx = self.get_pred_ctx_ll2_or_l3gld(bo);
140        let b2 = rf[0] == LAST3_FRAME || rf[0] == GOLDEN_FRAME;
141
142        let cdf = &self.fc.single_ref_cdfs[b2_ctx][2];
143        symbol_with_update!(self, w, b2 as u32, cdf);
144        if !b2 {
145          let b3_ctx = self.get_pred_ctx_last_or_last2(bo);
146          let b3 = rf[0] != LAST_FRAME;
147
148          let cdf = &self.fc.single_ref_cdfs[b3_ctx][3];
149          symbol_with_update!(self, w, b3 as u32, cdf);
150        } else {
151          let b4_ctx = self.get_pred_ctx_last3_or_gold(bo);
152          let b4 = rf[0] != LAST3_FRAME;
153
154          let cdf = &self.fc.single_ref_cdfs[b4_ctx][4];
155          symbol_with_update!(self, w, b4 as u32, cdf);
156        }
157      }
158    }
159  }
160
161  pub fn count_lrf_switchable<W: Writer>(
162    &self, w: &W, rs: &TileRestorationState, filter: RestorationFilter,
163    pli: usize,
164  ) -> u32 {
165    self.fc.count_lrf_switchable(w, rs, filter, pli)
166  }
167
168  /// # Panics
169  ///
170  /// - If the LRF has already been written for this superblock
171  pub fn write_lrf<W: Writer>(
172    &mut self, w: &mut W, rs: &mut TileRestorationStateMut,
173    sbo: TileSuperBlockOffset, pli: usize,
174  ) {
175    let rp = &mut rs.planes[pli];
176    if let Some(filter) = rp.restoration_unit(sbo, true).map(|ru| ru.filter) {
177      match filter {
178        RestorationFilter::None => match rp.rp_cfg.lrf_type {
179          RESTORE_WIENER => {
180            let cdf = &self.fc.lrf_wiener_cdf;
181            symbol_with_update!(self, w, 0, cdf);
182          }
183          RESTORE_SGRPROJ => {
184            let cdf = &self.fc.lrf_sgrproj_cdf;
185            symbol_with_update!(self, w, 0, cdf);
186          }
187          RESTORE_SWITCHABLE => {
188            let cdf = &self.fc.lrf_switchable_cdf;
189            symbol_with_update!(self, w, 0, cdf);
190          }
191          RESTORE_NONE => {}
192          _ => unreachable!(),
193        },
194        RestorationFilter::Sgrproj { set, xqd } => {
195          match rp.rp_cfg.lrf_type {
196            RESTORE_SGRPROJ => {
197              let cdf = &self.fc.lrf_sgrproj_cdf;
198              symbol_with_update!(self, w, 1, cdf);
199            }
200            RESTORE_SWITCHABLE => {
201              // Does *not* write 'RESTORE_SGRPROJ'
202              let cdf = &self.fc.lrf_switchable_cdf;
203              symbol_with_update!(self, w, 2, cdf);
204            }
205            _ => unreachable!(),
206          }
207          w.literal(SGRPROJ_PARAMS_BITS, set as u32);
208          for i in 0..2 {
209            let s = SGRPROJ_PARAMS_S[set as usize][i];
210            let min = SGRPROJ_XQD_MIN[i] as i32;
211            let max = SGRPROJ_XQD_MAX[i] as i32;
212            if s > 0 {
213              w.write_signed_subexp_with_ref(
214                xqd[i] as i32,
215                min,
216                max + 1,
217                SGRPROJ_PRJ_SUBEXP_K,
218                rp.sgrproj_ref[i] as i32,
219              );
220              rp.sgrproj_ref[i] = xqd[i];
221            } else {
222              // Nothing written, just update the reference
223              if i == 0 {
224                assert!(xqd[i] == 0);
225                rp.sgrproj_ref[0] = 0;
226              } else {
227                rp.sgrproj_ref[1] = 95; // LOL at spec.  The result is always 95.
228              }
229            }
230          }
231        }
232        RestorationFilter::Wiener { coeffs } => {
233          match rp.rp_cfg.lrf_type {
234            RESTORE_WIENER => {
235              let cdf = &self.fc.lrf_wiener_cdf;
236              symbol_with_update!(self, w, 1, cdf);
237            }
238            RESTORE_SWITCHABLE => {
239              // Does *not* write 'RESTORE_WIENER'
240              let cdf = &self.fc.lrf_switchable_cdf;
241              symbol_with_update!(self, w, 1, cdf);
242            }
243            _ => unreachable!(),
244          }
245          for pass in 0..2 {
246            let first_coeff = if pli == 0 {
247              0
248            } else {
249              assert!(coeffs[pass][0] == 0);
250              1
251            };
252            for i in first_coeff..3 {
253              let min = WIENER_TAPS_MIN[i] as i32;
254              let max = WIENER_TAPS_MAX[i] as i32;
255              w.write_signed_subexp_with_ref(
256                coeffs[pass][i] as i32,
257                min,
258                max + 1,
259                (i + 1) as u8,
260                rp.wiener_ref[pass][i] as i32,
261              );
262              rp.wiener_ref[pass][i] = coeffs[pass][i];
263            }
264          }
265        }
266      }
267    }
268  }
269
270  pub fn write_cdef<W: Writer>(
271    &mut self, w: &mut W, strength_index: u8, bits: u8,
272  ) {
273    w.literal(bits, strength_index as u32);
274  }
275}