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<'a> ContextWriter<'a> {
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}
276