bitstream.c 69.4 KB
Newer Older
Jingning Han's avatar
Jingning Han committed
1
/*
2
 * Copyright (c) 2016, Alliance for Open Media. All rights reserved
Jingning Han's avatar
Jingning Han committed
3
 *
4
5
6
7
8
9
 * This source code is subject to the terms of the BSD 2 Clause License and
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
 * was not distributed with this source code in the LICENSE file, you can
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
 * Media Patent License 1.0 was not distributed with this source code in the
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
Jingning Han's avatar
Jingning Han committed
10
11
12
13
14
15
 */

#include <assert.h>
#include <stdio.h>
#include <limits.h>

Adrian Grange's avatar
Adrian Grange committed
16
#include "aom/aom_encoder.h"
Yaowu Xu's avatar
Yaowu Xu committed
17
#include "aom_dsp/bitwriter_buffer.h"
Adrian Grange's avatar
Adrian Grange committed
18
19
#include "aom_dsp/aom_dsp_common.h"
#include "aom_mem/aom_mem.h"
Yaowu Xu's avatar
Yaowu Xu committed
20
21
#include "aom_ports/mem_ops.h"
#include "aom_ports/system_state.h"
Jingning Han's avatar
Jingning Han committed
22

Yaowu Xu's avatar
Yaowu Xu committed
23
#if CONFIG_CLPF
Yaowu Xu's avatar
Yaowu Xu committed
24
#include "av1/common/clpf.h"
Yaowu Xu's avatar
Yaowu Xu committed
25
#endif
26
#if CONFIG_DERING
Yaowu Xu's avatar
Yaowu Xu committed
27
#include "av1/common/dering.h"
28
#endif  // CONFIG_DERING
Yaowu Xu's avatar
Yaowu Xu committed
29
30
31
32
33
#include "av1/common/entropy.h"
#include "av1/common/entropymode.h"
#include "av1/common/entropymv.h"
#include "av1/common/mvref_common.h"
#include "av1/common/pred_common.h"
34
#include "av1/common/reconinter.h"
Yaowu Xu's avatar
Yaowu Xu committed
35
36
#include "av1/common/seg_common.h"
#include "av1/common/tile_common.h"
Jingning Han's avatar
Jingning Han committed
37

Yaowu Xu's avatar
Yaowu Xu committed
38
39
40
41
42
43
44
#include "av1/encoder/cost.h"
#include "av1/encoder/bitstream.h"
#include "av1/encoder/encodemv.h"
#include "av1/encoder/mcomp.h"
#include "av1/encoder/segmentation.h"
#include "av1/encoder/subexp.h"
#include "av1/encoder/tokenize.h"
Jingning Han's avatar
Jingning Han committed
45

46
47
48
static struct av1_token intra_mode_encodings[INTRA_MODES];
static struct av1_token switchable_interp_encodings[SWITCHABLE_FILTERS];
static struct av1_token partition_encodings[PARTITION_TYPES];
49
#if !CONFIG_REF_MV
50
static struct av1_token inter_mode_encodings[INTER_MODES];
51
#endif
52
53
54
#if CONFIG_MOTION_VAR
static struct av1_token motion_mode_encodings[MOTION_MODES];
#endif  // CONFIG_MOTION_VAR
55
static struct av1_token ext_tx_encodings[TX_TYPES];
56

57
void av1_encode_token_init() {
58
59
60
61
62
63
  av1_tokens_from_tree(intra_mode_encodings, av1_intra_mode_tree);
  av1_tokens_from_tree(switchable_interp_encodings, av1_switchable_interp_tree);
  av1_tokens_from_tree(partition_encodings, av1_partition_tree);
#if !CONFIG_REF_MV
  av1_tokens_from_tree(inter_mode_encodings, av1_inter_mode_tree);
#endif
64
65
66
#if CONFIG_MOTION_VAR
  av1_tokens_from_tree(motion_mode_encodings, av1_motion_mode_tree);
#endif  // CONFIG_MOTION_VAR
67
  av1_tokens_from_tree(ext_tx_encodings, av1_ext_tx_tree);
68
69
70
71
72
73
#if CONFIG_DAALA_EC
  /* This hack is necessary when CONFIG_EXT_INTERP is enabled because the five
      SWITCHABLE_FILTERS are not consecutive, e.g., 0, 1, 2, 3, 4, when doing
      an in-order traversal of the av1_switchable_interp_tree structure. */
  av1_indices_from_tree(av1_switchable_interp_ind, av1_switchable_interp_inv,
                        SWITCHABLE_FILTERS, av1_switchable_interp_tree);
74
75
76
77
78
  /* This hack is necessary because the four TX_TYPES are not consecutive,
      e.g., 0, 1, 2, 3, when doing an in-order traversal of the av1_ext_tx_tree
      structure. */
  av1_indices_from_tree(av1_ext_tx_ind, av1_ext_tx_inv, TX_TYPES,
                        av1_ext_tx_tree);
79
#endif
80
81
}

Adrian Grange's avatar
Adrian Grange committed
82
83
static void write_intra_mode(aom_writer *w, PREDICTION_MODE mode,
                             const aom_prob *probs) {
84
  av1_write_token(w, av1_intra_mode_tree, probs, &intra_mode_encodings[mode]);
Jingning Han's avatar
Jingning Han committed
85
86
}

clang-format's avatar
clang-format committed
87
88
static void write_inter_mode(AV1_COMMON *cm, aom_writer *w,
                             PREDICTION_MODE mode, const int16_t mode_ctx) {
89
#if CONFIG_REF_MV
90
  const int16_t newmv_ctx = mode_ctx & NEWMV_CTX_MASK;
91
92
93
94
  const aom_prob newmv_prob = cm->fc->newmv_prob[newmv_ctx];
  aom_write(w, mode != NEWMV, newmv_prob);

  if (mode != NEWMV) {
95
    const int16_t zeromv_ctx = (mode_ctx >> ZEROMV_OFFSET) & ZEROMV_CTX_MASK;
96
    const aom_prob zeromv_prob = cm->fc->zeromv_prob[zeromv_ctx];
97
98
99
100
101

    if (mode_ctx & (1 << ALL_ZERO_FLAG_OFFSET)) {
      assert(mode == ZEROMV);
      return;
    }
102
103
104
    aom_write(w, mode != ZEROMV, zeromv_prob);

    if (mode != ZEROMV) {
105
106
107
      int16_t refmv_ctx = (mode_ctx >> REFMV_OFFSET) & REFMV_CTX_MASK;
      aom_prob refmv_prob;

clang-format's avatar
clang-format committed
108
109
110
      if (mode_ctx & (1 << SKIP_NEARESTMV_OFFSET)) refmv_ctx = 6;
      if (mode_ctx & (1 << SKIP_NEARMV_OFFSET)) refmv_ctx = 7;
      if (mode_ctx & (1 << SKIP_NEARESTMV_SUB8X8_OFFSET)) refmv_ctx = 8;
111
112

      refmv_prob = cm->fc->refmv_prob[refmv_ctx];
113
114
115
116
117
      aom_write(w, mode != NEARESTMV, refmv_prob);
    }
  }
#else
  const aom_prob *const inter_probs = cm->fc->inter_mode_probs[mode_ctx];
Jingning Han's avatar
Jingning Han committed
118
  assert(is_inter_mode(mode));
119
  av1_write_token(w, av1_inter_mode_tree, inter_probs,
clang-format's avatar
clang-format committed
120
                  &inter_mode_encodings[INTER_OFFSET(mode)]);
121
#endif
Jingning Han's avatar
Jingning Han committed
122
123
}

124
#if CONFIG_REF_MV
clang-format's avatar
clang-format committed
125
126
static void write_drl_idx(const AV1_COMMON *cm, const MB_MODE_INFO *mbmi,
                          const MB_MODE_INFO_EXT *mbmi_ext, aom_writer *w) {
127
  uint8_t ref_frame_type = av1_ref_frame_type(mbmi->ref_frame);
128
129
130

  assert(mbmi->ref_mv_idx < 3);

131
132
133
134
135
136
137
138
139
  if (mbmi->mode == NEWMV) {
    int idx;
    for (idx = 0; idx < 2; ++idx) {
      if (mbmi_ext->ref_mv_count[ref_frame_type] > idx + 1) {
        uint8_t drl_ctx =
            av1_drl_ctx(mbmi_ext->ref_mv_stack[ref_frame_type], idx);
        aom_prob drl_prob = cm->fc->drl_prob[drl_ctx];

        aom_write(w, mbmi->ref_mv_idx != idx, drl_prob);
clang-format's avatar
clang-format committed
140
        if (mbmi->ref_mv_idx == idx) return;
141
      }
142
143
144
145
    }
    return;
  }

146
147
148
149
150
151
152
153
154
155
  if (mbmi->mode == NEARMV) {
    int idx;
    // TODO(jingning): Temporary solution to compensate the NEARESTMV offset.
    for (idx = 1; idx < 3; ++idx) {
      if (mbmi_ext->ref_mv_count[ref_frame_type] > idx + 1) {
        uint8_t drl_ctx =
            av1_drl_ctx(mbmi_ext->ref_mv_stack[ref_frame_type], idx);
        aom_prob drl_prob = cm->fc->drl_prob[drl_ctx];

        aom_write(w, mbmi->ref_mv_idx != (idx - 1), drl_prob);
clang-format's avatar
clang-format committed
156
        if (mbmi->ref_mv_idx == (idx - 1)) return;
157
      }
158
    }
159
    return;
160
161
162
163
  }
}
#endif

164
165
166
167
168
169
170
171
172
173
#if CONFIG_MOTION_VAR
static void write_motion_mode(const AV1_COMMON *cm, const MB_MODE_INFO *mbmi,
                              aom_writer *w) {
  if (is_motion_variation_allowed(mbmi))
    av1_write_token(w, av1_motion_mode_tree,
                    cm->fc->motion_mode_prob[mbmi->sb_type],
                    &motion_mode_encodings[mbmi->motion_mode]);
}
#endif  // CONFIG_MOTION_VAR

Adrian Grange's avatar
Adrian Grange committed
174
static void encode_unsigned_max(struct aom_write_bit_buffer *wb, int data,
clang-format's avatar
clang-format committed
175
                                int max) {
Adrian Grange's avatar
Adrian Grange committed
176
  aom_wb_write_literal(wb, data, get_unsigned_bits(max));
Jingning Han's avatar
Jingning Han committed
177
178
}

Adrian Grange's avatar
Adrian Grange committed
179
180
static void prob_diff_update(const aom_tree_index *tree,
                             aom_prob probs[/*n - 1*/],
clang-format's avatar
clang-format committed
181
                             const unsigned int counts[/*n - 1*/], int n,
Adrian Grange's avatar
Adrian Grange committed
182
                             aom_writer *w) {
Jingning Han's avatar
Jingning Han committed
183
184
185
186
187
188
  int i;
  unsigned int branch_ct[32][2];

  // Assuming max number of probabilities <= 32
  assert(n <= 32);

189
  av1_tree_probs_from_distribution(tree, branch_ct, counts);
Jingning Han's avatar
Jingning Han committed
190
  for (i = 0; i < n - 1; ++i)
191
    av1_cond_prob_diff_update(w, &probs[i], branch_ct[i]);
Jingning Han's avatar
Jingning Han committed
192
193
}

Adrian Grange's avatar
Adrian Grange committed
194
195
static int prob_diff_update_savings(const aom_tree_index *tree,
                                    aom_prob probs[/*n - 1*/],
196
197
198
199
200
201
202
203
                                    const unsigned int counts[/*n - 1*/],
                                    int n) {
  int i;
  unsigned int branch_ct[32][2];
  int savings = 0;

  // Assuming max number of probabilities <= 32
  assert(n <= 32);
204
  av1_tree_probs_from_distribution(tree, branch_ct, counts);
205
  for (i = 0; i < n - 1; ++i) {
206
    savings += av1_cond_prob_diff_update_savings(&probs[i], branch_ct[i]);
207
208
209
210
  }
  return savings;
}

211
static void write_selected_tx_size(const AV1_COMMON *cm, const MACROBLOCKD *xd,
Adrian Grange's avatar
Adrian Grange committed
212
                                   aom_writer *w) {
Jingning Han's avatar
Jingning Han committed
213
214
215
  TX_SIZE tx_size = xd->mi[0]->mbmi.tx_size;
  BLOCK_SIZE bsize = xd->mi[0]->mbmi.sb_type;
  const TX_SIZE max_tx_size = max_txsize_lookup[bsize];
Adrian Grange's avatar
Adrian Grange committed
216
  const aom_prob *const tx_probs =
clang-format's avatar
clang-format committed
217
      get_tx_probs2(max_tx_size, xd, &cm->fc->tx_probs);
Adrian Grange's avatar
Adrian Grange committed
218
  aom_write(w, tx_size != TX_4X4, tx_probs[0]);
Jingning Han's avatar
Jingning Han committed
219
  if (tx_size != TX_4X4 && max_tx_size >= TX_16X16) {
Adrian Grange's avatar
Adrian Grange committed
220
    aom_write(w, tx_size != TX_8X8, tx_probs[1]);
Jingning Han's avatar
Jingning Han committed
221
    if (tx_size != TX_8X8 && max_tx_size >= TX_32X32)
Adrian Grange's avatar
Adrian Grange committed
222
      aom_write(w, tx_size != TX_16X16, tx_probs[2]);
Jingning Han's avatar
Jingning Han committed
223
224
225
  }
}

226
227
228
229
230
#if CONFIG_REF_MV
static void update_inter_mode_probs(AV1_COMMON *cm, aom_writer *w,
                                    FRAME_COUNTS *counts) {
  int i;
  for (i = 0; i < NEWMV_MODE_CONTEXTS; ++i)
clang-format's avatar
clang-format committed
231
    av1_cond_prob_diff_update(w, &cm->fc->newmv_prob[i], counts->newmv_mode[i]);
232
233
234
235
  for (i = 0; i < ZEROMV_MODE_CONTEXTS; ++i)
    av1_cond_prob_diff_update(w, &cm->fc->zeromv_prob[i],
                              counts->zeromv_mode[i]);
  for (i = 0; i < REFMV_MODE_CONTEXTS; ++i)
clang-format's avatar
clang-format committed
236
    av1_cond_prob_diff_update(w, &cm->fc->refmv_prob[i], counts->refmv_mode[i]);
237
  for (i = 0; i < DRL_MODE_CONTEXTS; ++i)
clang-format's avatar
clang-format committed
238
    av1_cond_prob_diff_update(w, &cm->fc->drl_prob[i], counts->drl_mode[i]);
239
240
241
}
#endif

242
static int write_skip(const AV1_COMMON *cm, const MACROBLOCKD *xd,
Adrian Grange's avatar
Adrian Grange committed
243
                      int segment_id, const MODE_INFO *mi, aom_writer *w) {
Jingning Han's avatar
Jingning Han committed
244
245
246
247
  if (segfeature_active(&cm->seg, segment_id, SEG_LVL_SKIP)) {
    return 1;
  } else {
    const int skip = mi->mbmi.skip;
248
    aom_write(w, skip, av1_get_skip_prob(cm, xd));
Jingning Han's avatar
Jingning Han committed
249
250
251
252
    return skip;
  }
}

253
static void update_skip_probs(AV1_COMMON *cm, aom_writer *w,
Jingning Han's avatar
Jingning Han committed
254
255
256
257
                              FRAME_COUNTS *counts) {
  int k;

  for (k = 0; k < SKIP_CONTEXTS; ++k)
258
    av1_cond_prob_diff_update(w, &cm->fc->skip_probs[k], counts->skip[k]);
Jingning Han's avatar
Jingning Han committed
259
260
}

261
static void update_switchable_interp_probs(AV1_COMMON *cm, aom_writer *w,
Jingning Han's avatar
Jingning Han committed
262
263
                                           FRAME_COUNTS *counts) {
  int j;
264
  for (j = 0; j < SWITCHABLE_FILTER_CONTEXTS; ++j) {
265
    prob_diff_update(av1_switchable_interp_tree,
Jingning Han's avatar
Jingning Han committed
266
267
                     cm->fc->switchable_interp_prob[j],
                     counts->switchable_interp[j], SWITCHABLE_FILTERS, w);
268
269
270
271
272
273
#if CONFIG_DAALA_EC
    av1_tree_to_cdf(av1_switchable_interp_tree,
                    cm->fc->switchable_interp_prob[j],
                    cm->fc->switchable_interp_cdf[j]);
#endif
  }
Jingning Han's avatar
Jingning Han committed
274
275
}

276
277
278
static void update_ext_tx_probs(AV1_COMMON *cm, aom_writer *w) {
  const int savings_thresh = av1_cost_one(GROUP_DIFF_UPDATE_PROB) -
                             av1_cost_zero(GROUP_DIFF_UPDATE_PROB);
279
280
281
282
283
284
285
  int i, j;

  int savings = 0;
  int do_update = 0;
  for (i = TX_4X4; i < EXT_TX_SIZES; ++i) {
    for (j = 0; j < TX_TYPES; ++j)
      savings += prob_diff_update_savings(
286
          av1_ext_tx_tree, cm->fc->intra_ext_tx_prob[i][j],
287
288
289
          cm->counts.intra_ext_tx[i][j], TX_TYPES);
  }
  do_update = savings > savings_thresh;
Adrian Grange's avatar
Adrian Grange committed
290
  aom_write(w, do_update, GROUP_DIFF_UPDATE_PROB);
291
292
  if (do_update) {
    for (i = TX_4X4; i < EXT_TX_SIZES; ++i) {
293
      for (j = 0; j < TX_TYPES; ++j) {
294
        prob_diff_update(av1_ext_tx_tree, cm->fc->intra_ext_tx_prob[i][j],
clang-format's avatar
clang-format committed
295
                         cm->counts.intra_ext_tx[i][j], TX_TYPES, w);
296
297
298
299
300
#if CONFIG_DAALA_EC
        av1_tree_to_cdf(av1_ext_tx_tree, cm->fc->intra_ext_tx_prob[i][j],
                        cm->fc->intra_ext_tx_cdf[i][j]);
#endif
      }
301
302
303
304
    }
  }
  savings = 0;
  for (i = TX_4X4; i < EXT_TX_SIZES; ++i) {
clang-format's avatar
clang-format committed
305
    savings +=
306
        prob_diff_update_savings(av1_ext_tx_tree, cm->fc->inter_ext_tx_prob[i],
clang-format's avatar
clang-format committed
307
                                 cm->counts.inter_ext_tx[i], TX_TYPES);
308
309
  }
  do_update = savings > savings_thresh;
Adrian Grange's avatar
Adrian Grange committed
310
  aom_write(w, do_update, GROUP_DIFF_UPDATE_PROB);
311
312
  if (do_update) {
    for (i = TX_4X4; i < EXT_TX_SIZES; ++i) {
313
      prob_diff_update(av1_ext_tx_tree, cm->fc->inter_ext_tx_prob[i],
clang-format's avatar
clang-format committed
314
                       cm->counts.inter_ext_tx[i], TX_TYPES, w);
315
316
317
318
#if CONFIG_DAALA_EC
      av1_tree_to_cdf(av1_ext_tx_tree, cm->fc->inter_ext_tx_prob[i],
                      cm->fc->inter_ext_tx_cdf[i]);
#endif
319
320
321
322
    }
  }
}

Adrian Grange's avatar
Adrian Grange committed
323
static void pack_mb_tokens(aom_writer *w, TOKENEXTRA **tp,
clang-format's avatar
clang-format committed
324
                           const TOKENEXTRA *const stop,
Adrian Grange's avatar
Adrian Grange committed
325
                           aom_bit_depth_t bit_depth, const TX_SIZE tx) {
Jingning Han's avatar
Jingning Han committed
326
  TOKENEXTRA *p = *tp;
Alex Converse's avatar
Alex Converse committed
327
328
329
330
331
332
333
334
335
#if CONFIG_AOM_HIGHBITDEPTH
  const av1_extra_bit *const extra_bits_table =
      (bit_depth == AOM_BITS_12)
          ? av1_extra_bits_high12
          : (bit_depth == AOM_BITS_10) ? av1_extra_bits_high10 : av1_extra_bits;
#else
  const av1_extra_bit *const extra_bits_table = av1_extra_bits;
  (void)bit_depth;
#endif  // CONFIG_AOM_HIGHBITDEPTH
336
#if !CONFIG_MISC_FIXES
clang-format's avatar
clang-format committed
337
  (void)tx;
338
#endif
Jingning Han's avatar
Jingning Han committed
339
340

  while (p < stop && p->token != EOSB_TOKEN) {
341
342
    const uint8_t token = p->token;
    aom_tree_index index = 0;
343
#if !CONFIG_RANS
344
345
346
    const struct av1_token *const coef_encoding = &av1_coef_encodings[token];
    int coef_value = coef_encoding->value;
    int coef_length = coef_encoding->len;
347
#endif  // !CONFIG_RANS
Alex Converse's avatar
Alex Converse committed
348
    const av1_extra_bit *const extra_bits = &extra_bits_table[token];
Jingning Han's avatar
Jingning Han committed
349

350
#if CONFIG_RANS
351
    if (!p->skip_eob_node) aom_write(w, token != EOB_TOKEN, p->context_tree[0]);
352

353
354
355
    if (token != EOB_TOKEN) {
      aom_write(w, token != ZERO_TOKEN, p->context_tree[1]);
      if (token != ZERO_TOKEN) {
356
357
        aom_write_symbol(w, token - ONE_TOKEN, *p->token_cdf,
                         CATEGORY6_TOKEN - ONE_TOKEN + 1);
358
359
360
      }
    }
#else
Jingning Han's avatar
Jingning Han committed
361
362
    /* skip one or two nodes */
    if (p->skip_eob_node) {
363
364
      coef_length -= p->skip_eob_node;
      index = 2 * p->skip_eob_node;
Jingning Han's avatar
Jingning Han committed
365
366
367
368
369
370
371
372
373
374
    }

    // TODO(jbb): expanding this can lead to big gains.  It allows
    // much better branch prediction and would enable us to avoid numerous
    // lookups and compares.

    // If we have a token that's in the constrained set, the coefficient tree
    // is split into two treed writes.  The first treed write takes care of the
    // unconstrained nodes.  The second treed write takes care of the
    // constrained nodes.
375
376
377
378
379
380
381
382
383
    if (token >= TWO_TOKEN && token < EOB_TOKEN) {
      const int unconstrained_len = UNCONSTRAINED_NODES - p->skip_eob_node;
      const int unconstrained_bits =
          coef_value >> (coef_length - unconstrained_len);
      // Unconstrained nodes.
      aom_write_tree_bits(w, av1_coef_tree, p->context_tree, unconstrained_bits,
                          unconstrained_len, index);
      coef_value &= (1 << (coef_length - unconstrained_len)) - 1;
      // Constrained nodes.
384
      aom_write_tree(w, av1_coef_con_tree,
385
386
                     av1_pareto8_full[p->context_tree[PIVOT_NODE] - 1],
                     coef_value, coef_length - unconstrained_len, 0);
Jingning Han's avatar
Jingning Han committed
387
    } else {
388
389
      aom_write_tree_bits(w, av1_coef_tree, p->context_tree, coef_value,
                          coef_length, index);
Jingning Han's avatar
Jingning Han committed
390
    }
391
#endif  // CONFIG_RANS
Jingning Han's avatar
Jingning Han committed
392

Alex Converse's avatar
Alex Converse committed
393
394
395
396
397
    if (extra_bits->base_val) {
      const int bit_string = p->extra;
      const int bit_string_length = extra_bits->len;  // Length of extra bits to
                                                      // be written excluding
                                                      // the sign bit.
398
#if CONFIG_MISC_FIXES
399
      int skip_bits =
Alex Converse's avatar
Alex Converse committed
400
          (extra_bits->base_val == CAT6_MIN_VAL) ? TX_SIZES - 1 - tx : 0;
401
402
403
#else
      int skip_bits = 0;
#endif
Jingning Han's avatar
Jingning Han committed
404

Alex Converse's avatar
Alex Converse committed
405
406
407
408
      if (bit_string_length > 0) {
        const unsigned char *pb = extra_bits->prob;
        const int value = bit_string >> 1;
        const int num_bits = bit_string_length;  // number of bits in value
409
        assert(num_bits > 0);
Jingning Han's avatar
Jingning Han committed
410

Alex Converse's avatar
Alex Converse committed
411
412
413
        for (index = 0; index < num_bits; ++index) {
          const int shift = num_bits - index - 1;
          const int bb = (value >> shift) & 1;
414
          if (skip_bits) {
415
            --skip_bits;
416
417
            assert(!bb);
          } else {
Alex Converse's avatar
Alex Converse committed
418
            aom_write(w, bb, pb[index]);
419
          }
Alex Converse's avatar
Alex Converse committed
420
        }
Jingning Han's avatar
Jingning Han committed
421
422
      }

Alex Converse's avatar
Alex Converse committed
423
      aom_write_bit(w, bit_string & 1);
Jingning Han's avatar
Jingning Han committed
424
425
426
427
    }
    ++p;
  }

428
  *tp = p;
Jingning Han's avatar
Jingning Han committed
429
430
}

Adrian Grange's avatar
Adrian Grange committed
431
static void write_segment_id(aom_writer *w, const struct segmentation *seg,
432
                             const struct segmentation_probs *segp,
Jingning Han's avatar
Jingning Han committed
433
434
                             int segment_id) {
  if (seg->enabled && seg->update_map)
435
    aom_write_tree(w, av1_segment_tree, segp->tree_probs, segment_id, 3, 0);
Jingning Han's avatar
Jingning Han committed
436
437
438
}

// This function encodes the reference frame
439
static void write_ref_frames(const AV1_COMMON *cm, const MACROBLOCKD *xd,
Adrian Grange's avatar
Adrian Grange committed
440
                             aom_writer *w) {
Jingning Han's avatar
Jingning Han committed
441
442
443
444
445
446
447
448
449
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
  const int is_compound = has_second_ref(mbmi);
  const int segment_id = mbmi->segment_id;

  // If segment level coding of this signal is disabled...
  // or the segment allows multiple reference frame options
  if (segfeature_active(&cm->seg, segment_id, SEG_LVL_REF_FRAME)) {
    assert(!is_compound);
    assert(mbmi->ref_frame[0] ==
clang-format's avatar
clang-format committed
450
           get_segdata(&cm->seg, segment_id, SEG_LVL_REF_FRAME));
Jingning Han's avatar
Jingning Han committed
451
452
453
454
  } else {
    // does the feature use compound prediction or not
    // (if not specified at the frame/segment level)
    if (cm->reference_mode == REFERENCE_MODE_SELECT) {
455
      aom_write(w, is_compound, av1_get_reference_mode_prob(cm, xd));
Jingning Han's avatar
Jingning Han committed
456
    } else {
457
      assert((!is_compound) == (cm->reference_mode == SINGLE_REFERENCE));
Jingning Han's avatar
Jingning Han committed
458
459
460
    }

    if (is_compound) {
461
462
463
464
#if CONFIG_EXT_REFS
      const int bit_fwd = (mbmi->ref_frame[0] == GOLDEN_FRAME ||
                           mbmi->ref_frame[0] == LAST3_FRAME);
      const int bit_bwd = mbmi->ref_frame[1] == ALTREF_FRAME;
465
466

      // Write forward references.
467
468
469
470
471
472
473
474
      aom_write(w, bit_fwd, av1_get_pred_prob_comp_fwdref_p(cm, xd));
      if (!bit_fwd) {
        const int bit1_fwd = mbmi->ref_frame[0] == LAST_FRAME;
        aom_write(w, bit1_fwd, av1_get_pred_prob_comp_fwdref_p1(cm, xd));
      } else {
        const int bit2_fwd = mbmi->ref_frame[0] == GOLDEN_FRAME;
        aom_write(w, bit2_fwd, av1_get_pred_prob_comp_fwdref_p2(cm, xd));
      }
475
      // Write forward references.
476
477
      aom_write(w, bit_bwd, av1_get_pred_prob_comp_bwdref_p(cm, xd));
#else
Adrian Grange's avatar
Adrian Grange committed
478
      aom_write(w, mbmi->ref_frame[0] == GOLDEN_FRAME,
479
                av1_get_pred_prob_comp_ref_p(cm, xd));
480
#endif  // CONFIG_EXT_REFS
Jingning Han's avatar
Jingning Han committed
481
    } else {
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
#if CONFIG_EXT_REFS
      const int bit0 = (mbmi->ref_frame[0] == ALTREF_FRAME ||
                        mbmi->ref_frame[0] == BWDREF_FRAME);
      aom_write(w, bit0, av1_get_pred_prob_single_ref_p1(cm, xd));
      if (bit0) {
        const int bit1 = mbmi->ref_frame[0] == ALTREF_FRAME;
        aom_write(w, bit1, av1_get_pred_prob_single_ref_p2(cm, xd));
      } else {
        const int bit2 = (mbmi->ref_frame[0] == LAST3_FRAME ||
                          mbmi->ref_frame[0] == GOLDEN_FRAME);
        aom_write(w, bit2, av1_get_pred_prob_single_ref_p3(cm, xd));
        if (!bit2) {
          const int bit3 = mbmi->ref_frame[0] != LAST_FRAME;
          aom_write(w, bit3, av1_get_pred_prob_single_ref_p4(cm, xd));
        } else {
          const int bit4 = mbmi->ref_frame[0] != LAST3_FRAME;
          aom_write(w, bit4, av1_get_pred_prob_single_ref_p5(cm, xd));
        }
      }
#else
Jingning Han's avatar
Jingning Han committed
502
      const int bit0 = mbmi->ref_frame[0] != LAST_FRAME;
503
      aom_write(w, bit0, av1_get_pred_prob_single_ref_p1(cm, xd));
Jingning Han's avatar
Jingning Han committed
504
505
      if (bit0) {
        const int bit1 = mbmi->ref_frame[0] != GOLDEN_FRAME;
506
        aom_write(w, bit1, av1_get_pred_prob_single_ref_p2(cm, xd));
Jingning Han's avatar
Jingning Han committed
507
      }
508
#endif  // CONFIG_EXT_REFS
Jingning Han's avatar
Jingning Han committed
509
510
511
512
    }
  }
}

hui su's avatar
hui su committed
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
#if CONFIG_EXT_INTRA
static INLINE void write_uniform(aom_writer *w, int n, int v) {
  const int l = get_unsigned_bits(n);
  const int m = (1 << l) - n;

  if (l == 0) return;
  if (v < m) {
    aom_write_literal(w, v, l - 1);
  } else {
    aom_write_literal(w, m + ((v - m) >> 1), l - 1);
    aom_write_literal(w, (v - m) & 1, 1);
  }
}

static void write_intra_angle_info(const MB_MODE_INFO *const mbmi,
                                   aom_writer *w) {
  if (mbmi->sb_type < BLOCK_8X8) return;

  if (is_directional_mode(mbmi->mode)) {
    const TX_SIZE max_tx_size = max_txsize_lookup[mbmi->sb_type];
    const int max_angle_delta = av1_max_angle_delta_y[max_tx_size][mbmi->mode];
    write_uniform(w, 2 * max_angle_delta + 1,
                  max_angle_delta + mbmi->intra_angle_delta[0]);
  }

  if (is_directional_mode(mbmi->uv_mode)) {
    write_uniform(w, 2 * MAX_ANGLE_DELTA_UV + 1,
                  MAX_ANGLE_DELTA_UV + mbmi->intra_angle_delta[1]);
  }
}
#endif  // CONFIG_EXT_INTRA

545
546
547
548
549
550
static void write_switchable_interp_filter(AV1_COMP *const cpi,
                                           const MACROBLOCKD *const xd,
                                           aom_writer *w) {
  const AV1_COMMON *const cm = &cpi->common;
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
  if (cm->interp_filter == SWITCHABLE) {
551
    int ctx;
552
#if CONFIG_EXT_INTERP
553
554
555
556
    if (!is_interp_needed(xd)) {
      assert(mbmi->interp_filter == EIGHTTAP);
      return;
    }
557
#endif
558
    ctx = av1_get_pred_context_switchable_interp(xd);
559
#if CONFIG_DAALA_EC
560
561
    aom_write_symbol(w, av1_switchable_interp_ind[mbmi->interp_filter],
                     cm->fc->switchable_interp_cdf[ctx], SWITCHABLE_FILTERS);
562
#else
563
564
565
    av1_write_token(w, av1_switchable_interp_tree,
                    cm->fc->switchable_interp_prob[ctx],
                    &switchable_interp_encodings[mbmi->interp_filter]);
566
#endif
567
    ++cpi->interp_filter_selected[0][mbmi->interp_filter];
568
569
570
  }
}

571
static void pack_inter_mode_mvs(AV1_COMP *cpi, const MODE_INFO *mi,
Adrian Grange's avatar
Adrian Grange committed
572
                                aom_writer *w) {
573
  AV1_COMMON *const cm = &cpi->common;
574
#if !CONFIG_REF_MV
Jingning Han's avatar
Jingning Han committed
575
  const nmv_context *nmvc = &cm->fc->nmvc;
576
#endif
Jingning Han's avatar
Jingning Han committed
577
578
579
  const MACROBLOCK *const x = &cpi->td.mb;
  const MACROBLOCKD *const xd = &x->e_mbd;
  const struct segmentation *const seg = &cm->seg;
580
581
582
583
584
#if CONFIG_MISC_FIXES
  const struct segmentation_probs *const segp = &cm->fc->seg;
#else
  const struct segmentation_probs *const segp = &cm->segp;
#endif
Jingning Han's avatar
Jingning Han committed
585
586
587
588
589
590
591
592
593
594
595
596
597
  const MB_MODE_INFO *const mbmi = &mi->mbmi;
  const MB_MODE_INFO_EXT *const mbmi_ext = x->mbmi_ext;
  const PREDICTION_MODE mode = mbmi->mode;
  const int segment_id = mbmi->segment_id;
  const BLOCK_SIZE bsize = mbmi->sb_type;
  const int allow_hp = cm->allow_high_precision_mv;
  const int is_inter = is_inter_block(mbmi);
  const int is_compound = has_second_ref(mbmi);
  int skip, ref;

  if (seg->update_map) {
    if (seg->temporal_update) {
      const int pred_flag = mbmi->seg_id_predicted;
598
      aom_prob pred_prob = av1_get_pred_prob_seg_id(segp, xd);
Adrian Grange's avatar
Adrian Grange committed
599
      aom_write(w, pred_flag, pred_prob);
clang-format's avatar
clang-format committed
600
      if (!pred_flag) write_segment_id(w, seg, segp, segment_id);
Jingning Han's avatar
Jingning Han committed
601
    } else {
602
      write_segment_id(w, seg, segp, segment_id);
Jingning Han's avatar
Jingning Han committed
603
604
605
606
607
608
    }
  }

  skip = write_skip(cm, xd, segment_id, mi, w);

  if (!segfeature_active(seg, segment_id, SEG_LVL_REF_FRAME))
609
    aom_write(w, is_inter, av1_get_intra_inter_prob(cm, xd));
Jingning Han's avatar
Jingning Han committed
610
611

  if (bsize >= BLOCK_8X8 && cm->tx_mode == TX_MODE_SELECT &&
612
      !(is_inter && skip) && !xd->lossless[segment_id]) {
Jingning Han's avatar
Jingning Han committed
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
    write_selected_tx_size(cm, xd, w);
  }

  if (!is_inter) {
    if (bsize >= BLOCK_8X8) {
      write_intra_mode(w, mode, cm->fc->y_mode_prob[size_group_lookup[bsize]]);
    } else {
      int idx, idy;
      const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize];
      const int num_4x4_h = num_4x4_blocks_high_lookup[bsize];
      for (idy = 0; idy < 2; idy += num_4x4_h) {
        for (idx = 0; idx < 2; idx += num_4x4_w) {
          const PREDICTION_MODE b_mode = mi->bmi[idy * 2 + idx].as_mode;
          write_intra_mode(w, b_mode, cm->fc->y_mode_prob[0]);
        }
      }
    }
    write_intra_mode(w, mbmi->uv_mode, cm->fc->uv_mode_prob[mode]);
hui su's avatar
hui su committed
631
632
633
#if CONFIG_EXT_INTRA
    write_intra_angle_info(mbmi, w);
#endif  // CONFIG_EXT_INTRA
Jingning Han's avatar
Jingning Han committed
634
  } else {
635
    int16_t mode_ctx = mbmi_ext->mode_context[mbmi->ref_frame[0]];
Jingning Han's avatar
Jingning Han committed
636
637
    write_ref_frames(cm, xd, w);

638
#if CONFIG_REF_MV
639
640
    mode_ctx = av1_mode_context_analyzer(mbmi_ext->mode_context,
                                         mbmi->ref_frame, bsize, -1);
641
642
#endif

Jingning Han's avatar
Jingning Han committed
643
644
645
    // If segment skip is not enabled code the mode.
    if (!segfeature_active(seg, segment_id, SEG_LVL_SKIP)) {
      if (bsize >= BLOCK_8X8) {
646
        write_inter_mode(cm, w, mode, mode_ctx);
647
#if CONFIG_REF_MV
648
        if (mode == NEARMV || mode == NEWMV)
649
          write_drl_idx(cm, mbmi, mbmi_ext, w);
650
#endif
Jingning Han's avatar
Jingning Han committed
651
652
653
      }
    }

654
655
656
#if !CONFIG_EXT_INTERP
    write_switchable_interp_filter(cpi, xd, w);
#endif  // CONFIG_EXT_INTERP
Jingning Han's avatar
Jingning Han committed
657
658
659
660
661
662
663
664
665

    if (bsize < BLOCK_8X8) {
      const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize];
      const int num_4x4_h = num_4x4_blocks_high_lookup[bsize];
      int idx, idy;
      for (idy = 0; idy < 2; idy += num_4x4_h) {
        for (idx = 0; idx < 2; idx += num_4x4_w) {
          const int j = idy * 2 + idx;
          const PREDICTION_MODE b_mode = mi->bmi[j].as_mode;
666
667
668
669
#if CONFIG_REF_MV
          mode_ctx = av1_mode_context_analyzer(mbmi_ext->mode_context,
                                               mbmi->ref_frame, bsize, j);
#endif
670
          write_inter_mode(cm, w, b_mode, mode_ctx);
Jingning Han's avatar
Jingning Han committed
671
          if (b_mode == NEWMV) {
672
673
            for (ref = 0; ref < 1 + is_compound; ++ref) {
#if CONFIG_REF_MV
674
675
676
677
              int8_t rf_type = av1_ref_frame_type(mbmi->ref_frame);
              int nmv_ctx = av1_nmv_ctx(mbmi_ext->ref_mv_count[rf_type],
                                        mbmi_ext->ref_mv_stack[rf_type], ref,
                                        mbmi->ref_mv_idx);
678
679
              const nmv_context *nmvc = &cm->fc->nmvc[nmv_ctx];
#endif
680
              av1_encode_mv(cpi, w, &mi->bmi[j].as_mv[ref].as_mv,
clang-format's avatar
clang-format committed
681
682
                            &mbmi_ext->ref_mvs[mbmi->ref_frame[ref]][0].as_mv,
                            nmvc, allow_hp);
683
            }
Jingning Han's avatar
Jingning Han committed
684
685
686
687
688
          }
        }
      }
    } else {
      if (mode == NEWMV) {
689
        int_mv ref_mv;
690
691
        for (ref = 0; ref < 1 + is_compound; ++ref) {
#if CONFIG_REF_MV
692
693
694
695
          int8_t rf_type = av1_ref_frame_type(mbmi->ref_frame);
          int nmv_ctx = av1_nmv_ctx(mbmi_ext->ref_mv_count[rf_type],
                                    mbmi_ext->ref_mv_stack[rf_type], ref,
                                    mbmi->ref_mv_idx);
696
          const nmv_context *nmvc = &cm->fc->nmvc[nmv_ctx];
697
#endif
698
          ref_mv = mbmi_ext->ref_mvs[mbmi->ref_frame[ref]][0];
clang-format's avatar
clang-format committed
699
700
          av1_encode_mv(cpi, w, &mbmi->mv[ref].as_mv, &ref_mv.as_mv, nmvc,
                        allow_hp);
701
        }
Jingning Han's avatar
Jingning Han committed
702
703
      }
    }
704
705
706
#if CONFIG_MOTION_VAR
    write_motion_mode(cm, mbmi, w);
#endif  // CONFIG_MOTION_VAR
707
708
709
#if CONFIG_EXT_INTERP
    write_switchable_interp_filter(cpi, xd, w);
#endif  // CONFIG_EXT_INTERP
Jingning Han's avatar
Jingning Han committed
710
  }
711

clang-format's avatar
clang-format committed
712
  if (mbmi->tx_size < TX_32X32 && cm->base_qindex > 0 && !mbmi->skip &&
713
714
      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) {
    if (is_inter) {
715
#if CONFIG_DAALA_EC
716
717
      aom_write_symbol(w, av1_ext_tx_ind[mbmi->tx_type],
                       cm->fc->inter_ext_tx_cdf[mbmi->tx_size], TX_TYPES);
718
#else
719
      av1_write_token(w, av1_ext_tx_tree,
clang-format's avatar
clang-format committed
720
721
                      cm->fc->inter_ext_tx_prob[mbmi->tx_size],
                      &ext_tx_encodings[mbmi->tx_type]);
722
#endif
723
    } else {
724
#if CONFIG_DAALA_EC
725
      aom_write_symbol(
726
727
728
729
730
          w, av1_ext_tx_ind[mbmi->tx_type],
          cm->fc->intra_ext_tx_cdf[mbmi->tx_size]
                                  [intra_mode_to_tx_type_context[mbmi->mode]],
          TX_TYPES);
#else
731
732
      av1_write_token(
          w, av1_ext_tx_tree,
733
734
735
          cm->fc->intra_ext_tx_prob[mbmi->tx_size]
                                   [intra_mode_to_tx_type_context[mbmi->mode]],
          &ext_tx_encodings[mbmi->tx_type]);
736
#endif
737
738
    }
  } else {
clang-format's avatar
clang-format committed
739
    if (!mbmi->skip) assert(mbmi->tx_type == DCT_DCT);
740
  }
Jingning Han's avatar
Jingning Han committed
741
742
}

743
static void write_mb_modes_kf(const AV1_COMMON *cm, const MACROBLOCKD *xd,
Adrian Grange's avatar
Adrian Grange committed
744
                              MODE_INFO **mi_8x8, aom_writer *w) {
Jingning Han's avatar
Jingning Han committed
745
  const struct segmentation *const seg = &cm->seg;
746
747
748
749
750
#if CONFIG_MISC_FIXES
  const struct segmentation_probs *const segp = &cm->fc->seg;
#else
  const struct segmentation_probs *const segp = &cm->segp;
#endif
Jingning Han's avatar
Jingning Han committed
751
752
753
754
755
756
  const MODE_INFO *const mi = mi_8x8[0];
  const MODE_INFO *const above_mi = xd->above_mi;
  const MODE_INFO *const left_mi = xd->left_mi;
  const MB_MODE_INFO *const mbmi = &mi->mbmi;
  const BLOCK_SIZE bsize = mbmi->sb_type;

clang-format's avatar
clang-format committed
757
  if (seg->update_map) write_segment_id(w, seg, segp, mbmi->segment_id);
Jingning Han's avatar
Jingning Han committed
758
759
760

  write_skip(cm, xd, mbmi->segment_id, mi, w);

761
762
  if (bsize >= BLOCK_8X8 && cm->tx_mode == TX_MODE_SELECT &&
      !xd->lossless[mbmi->segment_id])
Jingning Han's avatar
Jingning Han committed
763
764
765
    write_selected_tx_size(cm, xd, w);

  if (bsize >= BLOCK_8X8) {
766
767
    write_intra_mode(w, mbmi->mode,
                     get_y_mode_probs(cm, mi, above_mi, left_mi, 0));
Jingning Han's avatar
Jingning Han committed
768
769
770
771
772
773
774
775
776
  } else {
    const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize];
    const int num_4x4_h = num_4x4_blocks_high_lookup[bsize];
    int idx, idy;

    for (idy = 0; idy < 2; idy += num_4x4_h) {
      for (idx = 0; idx < 2; idx += num_4x4_w) {
        const int block = idy * 2 + idx;
        write_intra_mode(w, mi->bmi[block].as_mode,
777
                         get_y_mode_probs(cm, mi, above_mi, left_mi, block));
Jingning Han's avatar
Jingning Han committed
778
779
780
781
      }
    }
  }

782
  write_intra_mode(w, mbmi->uv_mode, cm->fc->uv_mode_prob[mbmi->mode]);
hui su's avatar
hui su committed
783
784
785
#if CONFIG_EXT_INTRA
  write_intra_angle_info(mbmi, w);
#endif  // CONFIG_EXT_INTRA
786

clang-format's avatar
clang-format committed
787
  if (mbmi->tx_size < TX_32X32 && cm->base_qindex > 0 && !mbmi->skip &&
788
      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) {
789
790
    av1_write_token(
        w, av1_ext_tx_tree,
791
792
793
794
        cm->fc->intra_ext_tx_prob[mbmi->tx_size]
                                 [intra_mode_to_tx_type_context[mbmi->mode]],
        &ext_tx_encodings[mbmi->tx_type]);
  }
Jingning Han's avatar
Jingning Han committed
795
796
}

797
static void write_modes_b(AV1_COMP *cpi, const TileInfo *const tile,
Adrian Grange's avatar
Adrian Grange committed
798
                          aom_writer *w, TOKENEXTRA **tok,
clang-format's avatar
clang-format committed
799
800
                          const TOKENEXTRA *const tok_end, int mi_row,
                          int mi_col) {
801
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
802
803
  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;
  MODE_INFO *m;
804
  int plane;
Jingning Han's avatar
Jingning Han committed
805
806
807
808

  xd->mi = cm->mi_grid_visible + (mi_row * cm->mi_stride + mi_col);
  m = xd->mi[0];

Jingning Han's avatar
Jingning Han committed
809
  cpi->td.mb.mbmi_ext = cpi->mbmi_ext_base + (mi_row * cm->mi_cols + mi_col);
Jingning Han's avatar
Jingning Han committed
810

clang-format's avatar
clang-format committed
811
  set_mi_row_col(xd, tile, mi_row, num_8x8_blocks_high_lookup[m->mbmi.sb_type],
Jingning Han's avatar
Jingning Han committed
812
813
814
815
816
817
818
819
                 mi_col, num_8x8_blocks_wide_lookup[m->mbmi.sb_type],
                 cm->mi_rows, cm->mi_cols);
  if (frame_is_intra_only(cm)) {
    write_mb_modes_kf(cm, xd, xd->mi, w);
  } else {
    pack_inter_mode_mvs(cpi, m, w);
  }

820
821
822
  if (!m->mbmi.skip) {
    assert(*tok < tok_end);
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
clang-format's avatar
clang-format committed
823
824
      TX_SIZE tx =
          plane ? get_uv_tx_size(&m->mbmi, &xd->plane[plane]) : m->mbmi.tx_size;
825
826
827
828
829
      pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx);
      assert(*tok < tok_end && (*tok)->token == EOSB_TOKEN);
      (*tok)++;
    }
  }
Jingning Han's avatar
Jingning Han committed
830
831
}

832
static void write_partition(const AV1_COMMON *const cm,
clang-format's avatar
clang-format committed
833
834
                            const MACROBLOCKD *const xd, int hbs, int mi_row,
                            int mi_col, PARTITION_TYPE p, BLOCK_SIZE bsize,
Adrian Grange's avatar
Adrian Grange committed
835
                            aom_writer *w) {
Jingning Han's avatar
Jingning Han committed
836
  const int ctx = partition_plane_context(xd, mi_row, mi_col, bsize);
Adrian Grange's avatar
Adrian Grange committed
837
  const aom_prob *const probs = cm->fc->partition_prob[ctx];
Jingning Han's avatar
Jingning Han committed
838
839
840
841
  const int has_rows = (mi_row + hbs) < cm->mi_rows;
  const int has_cols = (mi_col + hbs) < cm->mi_cols;

  if (has_rows && has_cols) {
842
#if CONFIG_DAALA_EC
843
    aom_write_symbol(w, p, cm->fc->partition_cdf[ctx], PARTITION_TYPES);
844
#else
845
    av1_write_token(w, av1_partition_tree, probs, &partition_encodings[p]);
846
#endif
Jingning Han's avatar
Jingning Han committed
847
848
  } else if (!has_rows && has_cols) {
    assert(p == PARTITION_SPLIT || p == PARTITION_HORZ);
Adrian Grange's avatar
Adrian Grange committed
849
    aom_write(w, p == PARTITION_SPLIT, probs[1]);
Jingning Han's avatar
Jingning Han committed
850
851
  } else if (has_rows && !has_cols) {
    assert(p == PARTITION_SPLIT || p == PARTITION_VERT);
Adrian Grange's avatar
Adrian Grange committed
852
    aom_write(w, p == PARTITION_SPLIT, probs[2]);
Jingning Han's avatar
Jingning Han committed
853
854
855
856
857
  } else {
    assert(p == PARTITION_SPLIT);
  }
}

858
static void write_modes_sb(AV1_COMP *cpi, const TileInfo *const tile,
Adrian Grange's avatar
Adrian Grange committed
859
                           aom_writer *w, TOKENEXTRA **tok,
clang-format's avatar
clang-format committed
860
861
                           const TOKENEXTRA *const tok_end, int mi_row,
                           int mi_col, BLOCK_SIZE bsize) {
862
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
863
864
865
866
867
868
869
870
  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;

  const int bsl = b_width_log2_lookup[bsize];
  const int bs = (1 << bsl) / 4;
  PARTITION_TYPE partition;
  BLOCK_SIZE subsize;
  const MODE_INFO *m = NULL;

clang-format's avatar
clang-format committed
871
  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return;
Jingning Han's avatar
Jingning Han committed
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903

  m = cm->mi_grid_visible[mi_row * cm->mi_stride + mi_col];

  partition = partition_lookup[bsl][m->mbmi.sb_type];
  write_partition(cm, xd, bs, mi_row, mi_col, partition, bsize, w);
  subsize = get_subsize(bsize, partition);
  if (subsize < BLOCK_8X8) {
    write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
  } else {
    switch (partition) {
      case PARTITION_NONE:
        write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
        break;
      case PARTITION_HORZ:
        write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
        if (mi_row + bs < cm->mi_rows)
          write_modes_b(cpi, tile, w, tok, tok_end, mi_row + bs, mi_col);
        break;
      case PARTITION_VERT:
        write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col);
        if (mi_col + bs < cm->mi_cols)
          write_modes_b(cpi, tile, w, tok, tok_end, mi_row, mi_col + bs);
        break;
      case PARTITION_SPLIT:
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row, mi_col, subsize);
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row, mi_col + bs,
                       subsize);
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row + bs, mi_col,
                       subsize);
        write_modes_sb(cpi, tile, w, tok, tok_end, mi_row + bs, mi_col + bs,
                       subsize);
        break;
clang-format's avatar
clang-format committed
904
      default: assert(0);
Jingning Han's avatar
Jingning Han committed
905
906
907
908
909
910
911
    }
  }

  // update partition context
  if (bsize >= BLOCK_8X8 &&
      (bsize == BLOCK_8X8 || partition != PARTITION_SPLIT))
    update_partition_context(xd, mi_row, mi_col, subsize, bsize);
912
913
914
915

#if DERING_REFINEMENT
  if (bsize == BLOCK_64X64 && cm->dering_level != 0 &&
      !sb_all_skip(cm, mi_row, mi_col)) {
Adrian Grange's avatar
Adrian Grange committed
916
    aom_write_literal(
clang-format's avatar
clang-format committed
917
918
        w,
        cm->mi_grid_visible[mi_row * cm->mi_stride + mi_col]->mbmi.dering_gain,
919
        DERING_REFINEMENT_BITS);
920
921
  }
#endif
Jingning Han's avatar
Jingning Han committed
922
923
}

924
static void write_modes(AV1_COMP *cpi, const TileInfo *const tile,
Adrian Grange's avatar
Adrian Grange committed
925
                        aom_writer *w, TOKENEXTRA **tok,
clang-format's avatar
clang-format committed
926
                        const TOKENEXTRA *const tok_end) {
Jingning Han's avatar
Jingning Han committed
927
928
929
930
  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;
  int mi_row, mi_col;

  for (mi_row = tile->mi_row_start; mi_row < tile->mi_row_end;
931
       mi_row += MAX_MIB_SIZE) {
932
    av1_zero(xd->left_seg_context);
Jingning Han's avatar
Jingning Han committed
933
    for (mi_col = tile->mi_col_start; mi_col < tile->mi_col_end;
934
         mi_col += MAX_MIB_SIZE)
clang-format's avatar
clang-format committed
935
      write_modes_sb(cpi, tile, w, tok, tok_end, mi_row, mi_col, BLOCK_64X64);
Jingning Han's avatar
Jingning Han committed
936
937
938
  }
}

939
940
941
942
static void build_tree_distribution(AV1_COMP *cpi, TX_SIZE tx_size,
                                    av1_coeff_stats *coef_branch_ct,
                                    av1_coeff_probs_model *coef_probs) {
  av1_coeff_count *coef_counts = cpi->td.rd_counts.coef_counts[tx_size];
clang-format's avatar
clang-format committed
943
  unsigned int(*eob_branch_ct)[REF_TYPES][COEF_BANDS][COEFF_CONTEXTS] =
Jingning Han's avatar
Jingning Han committed
944
945
946
947
948
949
950
      cpi->common.counts.eob_branch[tx_size];
  int i, j, k, l, m;

  for (i = 0; i < PLANE_TYPES; ++i) {
    for (j = 0; j < REF_TYPES; ++j) {
      for (k = 0; k < COEF_BANDS; ++k) {
        for (l = 0; l < BAND_COEFF_CONTEXTS(k); ++l) {
951
          av1_tree_probs_from_distribution(av1_coef_tree,
clang-format's avatar
clang-format committed
952
953
                                           coef_branch_ct[i][j][k][l],
                                           coef_counts[i][j][k][l]);
clang-format's avatar
clang-format committed
954
955
          coef_branch_ct[i][j][k][l][0][1] =
              eob_branch_ct[i][j][k][l] - coef_branch_ct[i][j][k][l][0][0];
Jingning Han's avatar
Jingning Han committed
956
          for (m = 0; m < UNCONSTRAINED_NODES; ++m)
clang-format's avatar
clang-format committed
957
958
959
            coef_probs[i][j][k][l][m] =
                get_binary_prob(coef_branch_ct[i][j][k][l][m][0],
                                coef_branch_ct[i][j][k][l][m][1]);
Jingning Han's avatar
Jingning Han committed
960
961
962
963
964
965
        }
      }
    }
  }
}

966
static void update_coef_probs_common(aom_writer *const bc, AV1_COMP *cpi,
Jingning Han's avatar
Jingning Han committed
967
                                     TX_SIZE tx_size,