encodetxb.c 67.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12
/*
 * Copyright (c) 2017, Alliance for Open Media. All rights reserved
 *
 * 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.
 */

#include "av1/common/scan.h"
Angie Chiang's avatar
Angie Chiang committed
13
#include "av1/common/blockd.h"
14
#include "av1/common/idct.h"
Angie Chiang's avatar
Angie Chiang committed
15
#include "av1/common/pred_common.h"
16 17
#include "av1/encoder/bitstream.h"
#include "av1/encoder/encodeframe.h"
18
#include "av1/encoder/cost.h"
19
#include "av1/encoder/encodetxb.h"
Angie Chiang's avatar
Angie Chiang committed
20
#include "av1/encoder/rdopt.h"
21
#include "av1/encoder/subexp.h"
Angie Chiang's avatar
Angie Chiang committed
22
#include "av1/encoder/tokenize.h"
23

Angie Chiang's avatar
Angie Chiang committed
24 25
#define TEST_OPTIMIZE_TXB 0

26
void av1_alloc_txb_buf(AV1_COMP *cpi) {
27
#if 0
28 29 30 31 32 33 34 35 36 37 38 39
  AV1_COMMON *cm = &cpi->common;
  int mi_block_size = 1 << MI_SIZE_LOG2;
  // TODO(angiebird): Make sure cm->subsampling_x/y is set correctly, and then
  // use precise buffer size according to cm->subsampling_x/y
  int pixel_stride = mi_block_size * cm->mi_cols;
  int pixel_height = mi_block_size * cm->mi_rows;
  int i;
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    CHECK_MEM_ERROR(
        cm, cpi->tcoeff_buf[i],
        aom_malloc(sizeof(*cpi->tcoeff_buf[i]) * pixel_stride * pixel_height));
  }
40 41 42
#else
  (void)cpi;
#endif
43 44 45
}

void av1_free_txb_buf(AV1_COMP *cpi) {
46
#if 0
47 48 49 50
  int i;
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    aom_free(cpi->tcoeff_buf[i]);
  }
51 52 53
#else
  (void)cpi;
#endif
54 55
}

56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
static void write_golomb(aom_writer *w, int level) {
  int x = level + 1;
  int i = x;
  int length = 0;

  while (i) {
    i >>= 1;
    ++length;
  }
  assert(length > 0);

  for (i = 0; i < length - 1; ++i) aom_write_bit(w, 0);

  for (i = length - 1; i >= 0; --i) aom_write_bit(w, (x >> i) & 0x01);
}

void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
Jingning Han's avatar
Jingning Han committed
73 74 75
                          aom_writer *w, int blk_row, int blk_col, int block,
                          int plane, TX_SIZE tx_size, const tran_low_t *tcoeff,
                          uint16_t eob, TXB_CTX *txb_ctx) {
76 77 78
  aom_prob *nz_map;
  aom_prob *eob_flag;
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
79
  const PLANE_TYPE plane_type = get_plane_type(plane);
80
  const TX_SIZE txs_ctx = get_txsize_context(tx_size);
81 82
  const TX_TYPE tx_type =
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, block, tx_size);
Angie Chiang's avatar
Angie Chiang committed
83
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
84
  const int16_t *scan = scan_order->scan;
85
  const int16_t *iscan = scan_order->iscan;
86 87 88
  int c;
  int is_nz;
  const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
89
  const int height = tx_size_high[tx_size];
90
  const int seg_eob = tx_size_2d[tx_size];
91 92
  uint16_t update_eob = 0;

Jingning Han's avatar
Jingning Han committed
93 94 95
  (void)blk_row;
  (void)blk_col;

96
  aom_write(w, eob == 0, cm->fc->txb_skip[txs_ctx][txb_ctx->txb_skip_ctx]);
97 98

  if (eob == 0) return;
Angie Chiang's avatar
Angie Chiang committed
99
#if CONFIG_TXK_SEL
100 101
  av1_write_tx_type(cm, xd, blk_row, blk_col, block, plane,
                    get_min_tx_size(tx_size), w);
Angie Chiang's avatar
Angie Chiang committed
102
#endif
103

104 105
  nz_map = cm->fc->nz_map[txs_ctx][plane_type];
  eob_flag = cm->fc->eob_flag[txs_ctx][plane_type];
106 107

  for (c = 0; c < eob; ++c) {
108
    int coeff_ctx = get_nz_map_ctx(tcoeff, scan[c], bwl, height, iscan);
109
    int eob_ctx = get_eob_ctx(tcoeff, scan[c], txs_ctx);
110 111 112 113 114 115 116 117 118 119 120 121 122 123 124

    tran_low_t v = tcoeff[scan[c]];
    is_nz = (v != 0);

    if (c == seg_eob - 1) break;

    aom_write(w, is_nz, nz_map[coeff_ctx]);

    if (is_nz) {
      aom_write(w, c == (eob - 1), eob_flag[eob_ctx]);
    }
  }

  int i;
  for (i = 0; i < NUM_BASE_LEVELS; ++i) {
125
    aom_prob *coeff_base = cm->fc->coeff_base[txs_ctx][plane_type][i];
126 127 128 129 130 131 132 133 134 135

    update_eob = 0;
    for (c = eob - 1; c >= 0; --c) {
      tran_low_t v = tcoeff[scan[c]];
      tran_low_t level = abs(v);
      int sign = (v < 0) ? 1 : 0;
      int ctx;

      if (level <= i) continue;

136
      ctx = get_base_ctx(tcoeff, scan[c], bwl, height, i + 1);
137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167

      if (level == i + 1) {
        aom_write(w, 1, coeff_base[ctx]);
        if (c == 0) {
          aom_write(w, sign, cm->fc->dc_sign[plane_type][txb_ctx->dc_sign_ctx]);
        } else {
          aom_write_bit(w, sign);
        }
        continue;
      }
      aom_write(w, 0, coeff_base[ctx]);
      update_eob = AOMMAX(update_eob, c);
    }
  }

  for (c = update_eob; c >= 0; --c) {
    tran_low_t v = tcoeff[scan[c]];
    tran_low_t level = abs(v);
    int sign = (v < 0) ? 1 : 0;
    int idx;
    int ctx;

    if (level <= NUM_BASE_LEVELS) continue;

    if (c == 0) {
      aom_write(w, sign, cm->fc->dc_sign[plane_type][txb_ctx->dc_sign_ctx]);
    } else {
      aom_write_bit(w, sign);
    }

    // level is above 1.
168
    ctx = get_br_ctx(tcoeff, scan[c], bwl, height);
169 170
    for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
      if (level == (idx + 1 + NUM_BASE_LEVELS)) {
171
        aom_write(w, 1, cm->fc->coeff_lps[txs_ctx][plane_type][ctx]);
172 173
        break;
      }
174
      aom_write(w, 0, cm->fc->coeff_lps[txs_ctx][plane_type][ctx]);
175 176 177 178 179 180 181
    }
    if (idx < COEFF_BASE_RANGE) continue;

    // use 0-th order Golomb code to handle the residual level.
    write_golomb(w, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
  }
}
182

Angie Chiang's avatar
Angie Chiang committed
183 184 185 186 187 188 189
void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
                         aom_writer *w, int plane) {
  MACROBLOCKD *xd = &x->e_mbd;
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
  BLOCK_SIZE bsize = mbmi->sb_type;
  struct macroblockd_plane *pd = &xd->plane[plane];

190
#if CONFIG_CHROMA_SUB8X8
191 192
  const BLOCK_SIZE plane_bsize =
      AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
193 194
#elif CONFIG_CB4X4
  const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
Angie Chiang's avatar
Angie Chiang committed
195 196 197 198
#else
  const BLOCK_SIZE plane_bsize =
      get_plane_block_size(AOMMAX(bsize, BLOCK_8X8), pd);
#endif
199 200
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
  const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
201
  const TX_SIZE tx_size = av1_get_tx_size(plane, xd);
Angie Chiang's avatar
Angie Chiang committed
202 203 204 205 206
  const int bkw = tx_size_wide_unit[tx_size];
  const int bkh = tx_size_high_unit[tx_size];
  const int step = tx_size_wide_unit[tx_size] * tx_size_high_unit[tx_size];
  int row, col;
  int block = 0;
207 208
  for (row = 0; row < max_blocks_high; row += bkh) {
    for (col = 0; col < max_blocks_wide; col += bkw) {
Angie Chiang's avatar
Angie Chiang committed
209 210 211 212
      tran_low_t *tcoeff = BLOCK_OFFSET(x->mbmi_ext->tcoeff[plane], block);
      uint16_t eob = x->mbmi_ext->eobs[plane][block];
      TXB_CTX txb_ctx = { x->mbmi_ext->txb_skip_ctx[plane][block],
                          x->mbmi_ext->dc_sign_ctx[plane][block] };
Jingning Han's avatar
Jingning Han committed
213 214
      av1_write_coeffs_txb(cm, xd, w, row, col, block, plane, tx_size, tcoeff,
                           eob, &txb_ctx);
Angie Chiang's avatar
Angie Chiang committed
215 216 217 218 219
      block += step;
    }
  }
}

220 221
static INLINE void get_base_ctx_set(const tran_low_t *tcoeffs,
                                    int c,  // raster order
222
                                    const int bwl, const int height,
223 224 225 226 227 228 229 230 231 232 233 234 235 236
                                    int ctx_set[NUM_BASE_LEVELS]) {
  const int row = c >> bwl;
  const int col = c - (row << bwl);
  const int stride = 1 << bwl;
  int mag[NUM_BASE_LEVELS] = { 0 };
  int idx;
  tran_low_t abs_coeff;
  int i;

  for (idx = 0; idx < BASE_CONTEXT_POSITION_NUM; ++idx) {
    int ref_row = row + base_ref_offset[idx][0];
    int ref_col = col + base_ref_offset[idx][1];
    int pos = (ref_row << bwl) + ref_col;

237
    if (ref_row < 0 || ref_col < 0 || ref_row >= height || ref_col >= stride)
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263
      continue;

    abs_coeff = abs(tcoeffs[pos]);

    for (i = 0; i < NUM_BASE_LEVELS; ++i) {
      ctx_set[i] += abs_coeff > i;
      if (base_ref_offset[idx][0] >= 0 && base_ref_offset[idx][1] >= 0)
        mag[i] |= abs_coeff > (i + 1);
    }
  }

  for (i = 0; i < NUM_BASE_LEVELS; ++i) {
    ctx_set[i] = (ctx_set[i] + 1) >> 1;

    if (row == 0 && col == 0)
      ctx_set[i] = (ctx_set[i] << 1) + mag[i];
    else if (row == 0)
      ctx_set[i] = 8 + (ctx_set[i] << 1) + mag[i];
    else if (col == 0)
      ctx_set[i] = 18 + (ctx_set[i] << 1) + mag[i];
    else
      ctx_set[i] = 28 + (ctx_set[i] << 1) + mag[i];
  }
  return;
}

264
static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
265
                              const int coeff_lps[2]) {
266 267
  const tran_low_t min_level = 1 + NUM_BASE_LEVELS;
  const tran_low_t max_level = 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE;
268
  (void)ctx;
269
  if (abs_qc >= min_level) {
270 271
    const int cost0 = coeff_lps[0];
    const int cost1 = coeff_lps[1];
272 273 274 275 276 277 278 279 280 281
    if (abs_qc >= max_level)
      return COEFF_BASE_RANGE * cost0;
    else
      return (abs_qc - min_level) * cost0 + cost1;
  } else {
    return 0;
  }
}

static INLINE int get_base_cost(tran_low_t abs_qc, int ctx,
282
                                const int coeff_base[2], int base_idx) {
283
  const int level = base_idx + 1;
284
  (void)ctx;
285 286 287
  if (abs_qc < level)
    return 0;
  else
288
    return coeff_base[abs_qc == level];
289 290
}

291
int av1_cost_coeffs_txb(const AV1_COMP *const cpi, MACROBLOCK *x, int plane,
Jingning Han's avatar
Jingning Han committed
292 293
                        int blk_row, int blk_col, int block, TX_SIZE tx_size,
                        TXB_CTX *txb_ctx) {
294
  const AV1_COMMON *const cm = &cpi->common;
295
  MACROBLOCKD *const xd = &x->e_mbd;
296
  TX_SIZE txs_ctx = get_txsize_context(tx_size);
297
  const PLANE_TYPE plane_type = get_plane_type(plane);
298 299
  const TX_TYPE tx_type =
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, block, tx_size);
300 301 302 303 304
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
  const struct macroblock_plane *p = &x->plane[plane];
  const int eob = p->eobs[block];
  const tran_low_t *const qcoeff = BLOCK_OFFSET(p->qcoeff, block);
  int c, cost;
305
  const int seg_eob = AOMMIN(eob, tx_size_2d[tx_size] - 1);
306 307 308
  int txb_skip_ctx = txb_ctx->txb_skip_ctx;

  const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
309 310
  const int height = tx_size_high[tx_size];

Angie Chiang's avatar
Angie Chiang committed
311
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
312
  const int16_t *scan = scan_order->scan;
313
  const int16_t *iscan = scan_order->iscan;
314

315 316
  LV_MAP_COEFF_COST *coeff_costs = &x->coeff_costs[txs_ctx][plane_type];

317 318 319
  cost = 0;

  if (eob == 0) {
320
    cost = coeff_costs->txb_skip_cost[txb_skip_ctx][1];
321 322
    return cost;
  }
323
  cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];
324

Angie Chiang's avatar
Angie Chiang committed
325
#if CONFIG_TXK_SEL
326
  cost += av1_tx_type_cost(cm, x, xd, mbmi->sb_type, plane, tx_size, tx_type);
Angie Chiang's avatar
Angie Chiang committed
327
#endif
328

329 330 331 332 333 334
  for (c = 0; c < eob; ++c) {
    tran_low_t v = qcoeff[scan[c]];
    int is_nz = (v != 0);
    int level = abs(v);

    if (c < seg_eob) {
335
      int coeff_ctx = get_nz_map_ctx(qcoeff, scan[c], bwl, height, iscan);
336
      cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
337 338 339 340 341 342 343 344 345
    }

    if (is_nz) {
      int ctx_ls[NUM_BASE_LEVELS] = { 0 };
      int sign = (v < 0) ? 1 : 0;

      // sign bit cost
      if (c == 0) {
        int dc_sign_ctx = txb_ctx->dc_sign_ctx;
346
        cost += coeff_costs->dc_sign_cost[dc_sign_ctx][sign];
347 348 349 350
      } else {
        cost += av1_cost_bit(128, sign);
      }

351
      get_base_ctx_set(qcoeff, scan[c], bwl, height, ctx_ls);
352 353 354 355 356 357

      int i;
      for (i = 0; i < NUM_BASE_LEVELS; ++i) {
        if (level <= i) continue;

        if (level == i + 1) {
358
          cost += coeff_costs->base_cost[i][ctx_ls[i]][1];
359 360
          continue;
        }
361
        cost += coeff_costs->base_cost[i][ctx_ls[i]][0];
362 363 364 365 366 367
      }

      if (level > NUM_BASE_LEVELS) {
        int idx;
        int ctx;

368
        ctx = get_br_ctx(qcoeff, scan[c], bwl, height);
369 370 371

        for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
          if (level == (idx + 1 + NUM_BASE_LEVELS)) {
372
            cost += coeff_costs->lps_cost[ctx][1];
373 374
            break;
          }
375
          cost += coeff_costs->lps_cost[ctx][0];
376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396
        }

        if (idx >= COEFF_BASE_RANGE) {
          // residual cost
          int r = level - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
          int ri = r;
          int length = 0;

          while (ri) {
            ri >>= 1;
            ++length;
          }

          for (ri = 0; ri < length - 1; ++ri) cost += av1_cost_bit(128, 0);

          for (ri = length - 1; ri >= 0; --ri)
            cost += av1_cost_bit(128, (r >> ri) & 0x01);
        }
      }

      if (c < seg_eob) {
397
        int eob_ctx = get_eob_ctx(qcoeff, scan[c], txs_ctx);
398
        cost += coeff_costs->eob_cost[eob_ctx][c == (eob - 1)];
399 400 401 402 403 404
      }
    }
  }

  return cost;
}
Angie Chiang's avatar
Angie Chiang committed
405

406 407 408 409 410 411 412 413
static INLINE int has_base(tran_low_t qc, int base_idx) {
  const int level = base_idx + 1;
  return abs(qc) >= level;
}

static void gen_base_count_mag_arr(int (*base_count_arr)[MAX_TX_SQUARE],
                                   int (*base_mag_arr)[2],
                                   const tran_low_t *qcoeff, int stride,
414
                                   int height, int eob, const int16_t *scan) {
415 416 417 418 419 420
  for (int c = 0; c < eob; ++c) {
    const int coeff_idx = scan[c];  // raster order
    if (!has_base(qcoeff[coeff_idx], 0)) continue;
    const int row = coeff_idx / stride;
    const int col = coeff_idx % stride;
    int *mag = base_mag_arr[coeff_idx];
421
    get_mag(mag, qcoeff, stride, height, row, col, base_ref_offset,
422 423 424 425
            BASE_CONTEXT_POSITION_NUM);
    for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
      if (!has_base(qcoeff[coeff_idx], i)) continue;
      int *count = base_count_arr[i] + coeff_idx;
426 427
      *count = get_level_count(qcoeff, stride, height, row, col, i,
                               base_ref_offset, BASE_CONTEXT_POSITION_NUM);
428 429 430 431 432
    }
  }
}

static void gen_nz_count_arr(int(*nz_count_arr), const tran_low_t *qcoeff,
433
                             int stride, int height, int eob,
434 435 436 437 438 439 440
                             const SCAN_ORDER *scan_order) {
  const int16_t *scan = scan_order->scan;
  const int16_t *iscan = scan_order->iscan;
  for (int c = 0; c < eob; ++c) {
    const int coeff_idx = scan[c];  // raster order
    const int row = coeff_idx / stride;
    const int col = coeff_idx % stride;
441 442
    nz_count_arr[coeff_idx] =
        get_nz_count(qcoeff, stride, height, row, col, iscan);
443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483
  }
}

static void gen_nz_ctx_arr(int (*nz_ctx_arr)[2], int(*nz_count_arr),
                           const tran_low_t *qcoeff, int bwl, int eob,
                           const SCAN_ORDER *scan_order) {
  const int16_t *scan = scan_order->scan;
  const int16_t *iscan = scan_order->iscan;
  for (int c = 0; c < eob; ++c) {
    const int coeff_idx = scan[c];  // raster order
    const int count = nz_count_arr[coeff_idx];
    nz_ctx_arr[coeff_idx][0] =
        get_nz_map_ctx_from_count(count, qcoeff, coeff_idx, bwl, iscan);
  }
}

static void gen_base_ctx_arr(int (*base_ctx_arr)[MAX_TX_SQUARE][2],
                             int (*base_count_arr)[MAX_TX_SQUARE],
                             int (*base_mag_arr)[2], const tran_low_t *qcoeff,
                             int stride, int eob, const int16_t *scan) {
  (void)qcoeff;
  for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
    for (int c = 0; c < eob; ++c) {
      const int coeff_idx = scan[c];  // raster order
      if (!has_base(qcoeff[coeff_idx], i)) continue;
      const int row = coeff_idx / stride;
      const int col = coeff_idx % stride;
      const int count = base_count_arr[i][coeff_idx];
      const int *mag = base_mag_arr[coeff_idx];
      const int level = i + 1;
      base_ctx_arr[i][coeff_idx][0] =
          get_base_ctx_from_count_mag(row, col, count, mag[0], level);
    }
  }
}

static INLINE int has_br(tran_low_t qc) {
  return abs(qc) >= 1 + NUM_BASE_LEVELS;
}

static void gen_br_count_mag_arr(int *br_count_arr, int (*br_mag_arr)[2],
484 485
                                 const tran_low_t *qcoeff, int stride,
                                 int height, int eob, const int16_t *scan) {
486 487 488 489 490 491 492
  for (int c = 0; c < eob; ++c) {
    const int coeff_idx = scan[c];  // raster order
    if (!has_br(qcoeff[coeff_idx])) continue;
    const int row = coeff_idx / stride;
    const int col = coeff_idx % stride;
    int *count = br_count_arr + coeff_idx;
    int *mag = br_mag_arr[coeff_idx];
493
    *count = get_level_count(qcoeff, stride, height, row, col, NUM_BASE_LEVELS,
494
                             br_ref_offset, BR_CONTEXT_POSITION_NUM);
495
    get_mag(mag, qcoeff, stride, height, row, col, br_ref_offset,
496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515
            BR_CONTEXT_POSITION_NUM);
  }
}

static void gen_br_ctx_arr(int (*br_ctx_arr)[2], const int *br_count_arr,
                           int (*br_mag_arr)[2], const tran_low_t *qcoeff,
                           int stride, int eob, const int16_t *scan) {
  (void)qcoeff;
  for (int c = 0; c < eob; ++c) {
    const int coeff_idx = scan[c];  // raster order
    if (!has_br(qcoeff[coeff_idx])) continue;
    const int row = coeff_idx / stride;
    const int col = coeff_idx % stride;
    const int count = br_count_arr[coeff_idx];
    const int *mag = br_mag_arr[coeff_idx];
    br_ctx_arr[coeff_idx][0] =
        get_br_ctx_from_count_mag(row, col, count, mag[0]);
  }
}

516
static INLINE int get_sign_bit_cost(tran_low_t qc, int coeff_idx,
517
                                    const int (*dc_sign_cost)[2],
518 519 520 521
                                    int dc_sign_ctx) {
  const int sign = (qc < 0) ? 1 : 0;
  // sign bit cost
  if (coeff_idx == 0) {
522
    return dc_sign_cost[dc_sign_ctx][sign];
523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544
  } else {
    return av1_cost_bit(128, sign);
  }
}
static INLINE int get_golomb_cost(int abs_qc) {
  if (abs_qc >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
    // residual cost
    int r = abs_qc - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
    int ri = r;
    int length = 0;

    while (ri) {
      ri >>= 1;
      ++length;
    }

    return av1_cost_literal(2 * length - 1);
  } else {
    return 0;
  }
}

545 546 547 548
// TODO(angiebird): add static once this function is called
void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
  const int16_t *scan = txb_info->scan_order->scan;
  gen_nz_count_arr(txb_cache->nz_count_arr, txb_info->qcoeff, txb_info->stride,
549
                   txb_info->height, txb_info->eob, txb_info->scan_order);
550 551 552 553
  gen_nz_ctx_arr(txb_cache->nz_ctx_arr, txb_cache->nz_count_arr,
                 txb_info->qcoeff, txb_info->bwl, txb_info->eob,
                 txb_info->scan_order);
  gen_base_count_mag_arr(txb_cache->base_count_arr, txb_cache->base_mag_arr,
554 555
                         txb_info->qcoeff, txb_info->stride, txb_info->height,
                         txb_info->eob, scan);
556 557 558 559
  gen_base_ctx_arr(txb_cache->base_ctx_arr, txb_cache->base_count_arr,
                   txb_cache->base_mag_arr, txb_info->qcoeff, txb_info->stride,
                   txb_info->eob, scan);
  gen_br_count_mag_arr(txb_cache->br_count_arr, txb_cache->br_mag_arr,
560 561
                       txb_info->qcoeff, txb_info->stride, txb_info->height,
                       txb_info->eob, scan);
562 563 564 565 566
  gen_br_ctx_arr(txb_cache->br_ctx_arr, txb_cache->br_count_arr,
                 txb_cache->br_mag_arr, txb_info->qcoeff, txb_info->stride,
                 txb_info->eob, scan);
}

567 568 569
static INLINE const int *get_level_prob(int level, int coeff_idx,
                                        const TxbCache *txb_cache,
                                        const LV_MAP_COEFF_COST *txb_costs) {
570 571
  if (level == 0) {
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx][0];
572
    return txb_costs->nz_map_cost[ctx];
573 574 575
  } else if (level >= 1 && level < 1 + NUM_BASE_LEVELS) {
    const int idx = level - 1;
    const int ctx = txb_cache->base_ctx_arr[idx][coeff_idx][0];
576
    return txb_costs->base_cost[idx][ctx];
577 578 579
  } else if (level >= 1 + NUM_BASE_LEVELS &&
             level < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
    const int ctx = txb_cache->br_ctx_arr[coeff_idx][0];
580
    return txb_costs->lps_cost[ctx];
581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597
  } else if (level >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
    printf("get_level_prob does not support golomb\n");
    assert(0);
    return 0;
  } else {
    assert(0);
    return 0;
  }
}

static INLINE tran_low_t get_lower_coeff(tran_low_t qc) {
  if (qc == 0) {
    return 0;
  }
  return qc > 0 ? qc - 1 : qc + 1;
}

Angie Chiang's avatar
Angie Chiang committed
598 599 600 601 602 603 604
static INLINE void update_mag_arr(int *mag_arr, int abs_qc) {
  if (mag_arr[0] == abs_qc) {
    mag_arr[1] -= 1;
    assert(mag_arr[1] >= 0);
  }
}

Angie Chiang's avatar
Angie Chiang committed
605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654
static INLINE int get_mag_from_mag_arr(const int *mag_arr) {
  int mag;
  if (mag_arr[1] > 0) {
    mag = mag_arr[0];
  } else if (mag_arr[0] > 0) {
    mag = mag_arr[0] - 1;
  } else {
    // no neighbor
    assert(mag_arr[0] == 0 && mag_arr[1] == 0);
    mag = 0;
  }
  return mag;
}

static int neighbor_level_down_update(int *new_count, int *new_mag, int count,
                                      const int *mag, int coeff_idx,
                                      tran_low_t abs_nb_coeff, int nb_coeff_idx,
                                      int level, const TxbInfo *txb_info) {
  *new_count = count;
  *new_mag = get_mag_from_mag_arr(mag);

  int update = 0;
  // check if br_count changes
  if (abs_nb_coeff == level) {
    update = 1;
    *new_count -= 1;
    assert(*new_count >= 0);
  }
  const int row = coeff_idx >> txb_info->bwl;
  const int col = coeff_idx - (row << txb_info->bwl);
  const int nb_row = nb_coeff_idx >> txb_info->bwl;
  const int nb_col = nb_coeff_idx - (nb_row << txb_info->bwl);

  // check if mag changes
  if (nb_row >= row && nb_col >= col) {
    if (abs_nb_coeff == mag[0]) {
      assert(mag[1] > 0);
      if (mag[1] == 1) {
        // the nb is the only qc with max mag
        *new_mag -= 1;
        assert(*new_mag >= 0);
        update = 1;
      }
    }
  }
  return update;
}

static int try_neighbor_level_down_br(int coeff_idx, int nb_coeff_idx,
                                      const TxbCache *txb_cache,
655
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674
                                      const TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const tran_low_t abs_qc = abs(qc);
  const int level = NUM_BASE_LEVELS + 1;
  if (abs_qc < level) return 0;

  const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
  const tran_low_t abs_nb_coeff = abs(nb_coeff);
  const int count = txb_cache->br_count_arr[coeff_idx];
  const int *mag = txb_cache->br_mag_arr[coeff_idx];
  int new_count;
  int new_mag;
  const int update =
      neighbor_level_down_update(&new_count, &new_mag, count, mag, coeff_idx,
                                 abs_nb_coeff, nb_coeff_idx, level, txb_info);
  if (update) {
    const int row = coeff_idx >> txb_info->bwl;
    const int col = coeff_idx - (row << txb_info->bwl);
    const int ctx = txb_cache->br_ctx_arr[coeff_idx][0];
675
    const int org_cost = get_br_cost(abs_qc, ctx, txb_costs->lps_cost[ctx]);
Angie Chiang's avatar
Angie Chiang committed
676 677

    const int new_ctx = get_br_ctx_from_count_mag(row, col, new_count, new_mag);
678 679
    const int new_cost =
        get_br_cost(abs_qc, new_ctx, txb_costs->lps_cost[new_ctx]);
Angie Chiang's avatar
Angie Chiang committed
680 681 682 683 684 685 686 687 688
    const int cost_diff = -org_cost + new_cost;
    return cost_diff;
  } else {
    return 0;
  }
}

static int try_neighbor_level_down_base(int coeff_idx, int nb_coeff_idx,
                                        const TxbCache *txb_cache,
689
                                        const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712
                                        const TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const tran_low_t abs_qc = abs(qc);

  int cost_diff = 0;
  for (int base_idx = 0; base_idx < NUM_BASE_LEVELS; ++base_idx) {
    const int level = base_idx + 1;
    if (abs_qc < level) continue;

    const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
    const tran_low_t abs_nb_coeff = abs(nb_coeff);

    const int count = txb_cache->base_count_arr[base_idx][coeff_idx];
    const int *mag = txb_cache->base_mag_arr[coeff_idx];
    int new_count;
    int new_mag;
    const int update =
        neighbor_level_down_update(&new_count, &new_mag, count, mag, coeff_idx,
                                   abs_nb_coeff, nb_coeff_idx, level, txb_info);
    if (update) {
      const int row = coeff_idx >> txb_info->bwl;
      const int col = coeff_idx - (row << txb_info->bwl);
      const int ctx = txb_cache->base_ctx_arr[base_idx][coeff_idx][0];
713 714
      const int org_cost = get_base_cost(
          abs_qc, ctx, txb_costs->base_cost[base_idx][ctx], base_idx);
Angie Chiang's avatar
Angie Chiang committed
715 716 717

      const int new_ctx =
          get_base_ctx_from_count_mag(row, col, new_count, new_mag, level);
718 719
      const int new_cost = get_base_cost(
          abs_qc, new_ctx, txb_costs->base_cost[base_idx][new_ctx], base_idx);
Angie Chiang's avatar
Angie Chiang committed
720 721 722 723 724 725 726 727
      cost_diff += -org_cost + new_cost;
    }
  }
  return cost_diff;
}

static int try_neighbor_level_down_nz(int coeff_idx, int nb_coeff_idx,
                                      const TxbCache *txb_cache,
728
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748
                                      TxbInfo *txb_info) {
  // assume eob doesn't change
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const tran_low_t abs_qc = abs(qc);
  const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
  const tran_low_t abs_nb_coeff = abs(nb_coeff);
  if (abs_nb_coeff != 1) return 0;
  const int16_t *iscan = txb_info->scan_order->iscan;
  const int scan_idx = iscan[coeff_idx];
  if (scan_idx == txb_info->seg_eob) return 0;
  const int nb_scan_idx = iscan[nb_coeff_idx];
  if (nb_scan_idx < scan_idx) {
    const int count = txb_cache->nz_count_arr[coeff_idx];
    assert(count > 0);
    txb_info->qcoeff[nb_coeff_idx] = get_lower_coeff(nb_coeff);
    const int new_ctx = get_nz_map_ctx_from_count(
        count - 1, txb_info->qcoeff, coeff_idx, txb_info->bwl, iscan);
    txb_info->qcoeff[nb_coeff_idx] = nb_coeff;
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx][0];
    const int is_nz = abs_qc > 0;
749 750
    const int org_cost = txb_costs->nz_map_cost[ctx][is_nz];
    const int new_cost = txb_costs->nz_map_cost[new_ctx][is_nz];
Angie Chiang's avatar
Angie Chiang committed
751 752 753 754 755 756 757 758 759
    const int cost_diff = new_cost - org_cost;
    return cost_diff;
  } else {
    return 0;
  }
}

static int try_self_level_down(tran_low_t *low_coeff, int coeff_idx,
                               const TxbCache *txb_cache,
760 761
                               const LV_MAP_COEFF_COST *txb_costs,
                               TxbInfo *txb_info) {
762 763 764 765 766 767 768 769 770 771
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  if (qc == 0) {
    *low_coeff = 0;
    return 0;
  }
  const tran_low_t abs_qc = abs(qc);
  *low_coeff = get_lower_coeff(qc);
  int cost_diff;
  if (*low_coeff == 0) {
    const int scan_idx = txb_info->scan_order->iscan[coeff_idx];
772 773 774 775
    const int *level_cost =
        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_costs);
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);
776 777 778 779 780
    if (scan_idx < txb_info->seg_eob) {
      // When level-0, we code the binary of abs_qc > level
      // but when level-k k > 0 we code the binary of abs_qc == level
      // That's why wee need this special treatment for level-0 map
      // TODO(angiebird): make leve-0 consistent to other levels
781
      cost_diff = -level_cost[1] + low_level_cost[0] - low_level_cost[1];
782
    } else {
783
      cost_diff = -level_cost[1];
784 785 786 787
    }

    if (scan_idx < txb_info->seg_eob) {
      const int eob_ctx =
788
          get_eob_ctx(txb_info->qcoeff, coeff_idx, txb_info->txs_ctx);
789 790
      cost_diff -=
          txb_costs->eob_cost[eob_ctx][scan_idx == (txb_info->eob - 1)];
791 792 793
    }

    const int sign_cost = get_sign_bit_cost(
794
        qc, coeff_idx, txb_costs->dc_sign_cost, txb_info->txb_ctx->dc_sign_ctx);
795 796
    cost_diff -= sign_cost;
  } else if (abs_qc < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
797 798 799 800 801 802
    const int *level_cost =
        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_costs);
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);

    cost_diff = -level_cost[1] + low_level_cost[1] - low_level_cost[0];
803
  } else if (abs_qc == 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
804 805 806 807
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);
    cost_diff =
        -get_golomb_cost(abs_qc) + low_level_cost[1] - low_level_cost[0];
808 809 810 811 812 813 814 815
  } else {
    assert(abs_qc > 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE);
    const tran_low_t abs_low_coeff = abs(*low_coeff);
    cost_diff = -get_golomb_cost(abs_qc) + get_golomb_cost(abs_low_coeff);
  }
  return cost_diff;
}

Angie Chiang's avatar
Angie Chiang committed
816 817 818 819 820 821 822 823 824 825 826 827 828 829 830
#define COST_MAP_SIZE 5
#define COST_MAP_OFFSET 2

static INLINE int check_nz_neighbor(tran_low_t qc) { return abs(qc) == 1; }

static INLINE int check_base_neighbor(tran_low_t qc) {
  return abs(qc) <= 1 + NUM_BASE_LEVELS;
}

static INLINE int check_br_neighbor(tran_low_t qc) {
  return abs(qc) > BR_MAG_OFFSET;
}

// TODO(angiebird): add static to this function once it's called
int try_level_down(int coeff_idx, const TxbCache *txb_cache,
831
                   const LV_MAP_COEFF_COST *txb_costs, TxbInfo *txb_info,
Angie Chiang's avatar
Angie Chiang committed
832 833 834 835 836 837 838 839 840 841 842 843 844 845 846
                   int (*cost_map)[COST_MAP_SIZE]) {
  if (cost_map) {
    for (int i = 0; i < COST_MAP_SIZE; ++i) av1_zero(cost_map[i]);
  }

  tran_low_t qc = txb_info->qcoeff[coeff_idx];
  tran_low_t low_coeff;
  if (qc == 0) return 0;
  int accu_cost_diff = 0;

  const int16_t *iscan = txb_info->scan_order->iscan;
  const int eob = txb_info->eob;
  const int scan_idx = iscan[coeff_idx];
  if (scan_idx < eob) {
    const int cost_diff = try_self_level_down(&low_coeff, coeff_idx, txb_cache,
847
                                              txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
848 849 850 851 852 853 854 855 856 857 858 859
    if (cost_map)
      cost_map[0 + COST_MAP_OFFSET][0 + COST_MAP_OFFSET] = cost_diff;
    accu_cost_diff += cost_diff;
  }

  const int row = coeff_idx >> txb_info->bwl;
  const int col = coeff_idx - (row << txb_info->bwl);
  if (check_nz_neighbor(qc)) {
    for (int i = 0; i < SIG_REF_OFFSET_NUM; ++i) {
      const int nb_row = row - sig_ref_offset[i][0];
      const int nb_col = col - sig_ref_offset[i][1];
      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
860

861 862 863 864
      if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
            nb_col < txb_info->stride))
        continue;

Angie Chiang's avatar
Angie Chiang committed
865
      const int nb_scan_idx = iscan[nb_coeff_idx];
866
      if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
867
        const int cost_diff = try_neighbor_level_down_nz(
868
            nb_coeff_idx, coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
869 870 871 872 873 874 875 876 877 878 879 880 881
        if (cost_map)
          cost_map[nb_row - row + COST_MAP_OFFSET]
                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
        accu_cost_diff += cost_diff;
      }
    }
  }

  if (check_base_neighbor(qc)) {
    for (int i = 0; i < BASE_CONTEXT_POSITION_NUM; ++i) {
      const int nb_row = row - base_ref_offset[i][0];
      const int nb_col = col - base_ref_offset[i][1];
      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
882 883 884 885 886

      if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
            nb_col < txb_info->stride))
        continue;

Angie Chiang's avatar
Angie Chiang committed
887
      const int nb_scan_idx = iscan[nb_coeff_idx];
888
      if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
889
        const int cost_diff = try_neighbor_level_down_base(
890
            nb_coeff_idx, coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
891 892 893 894 895 896 897 898 899 900 901 902 903
        if (cost_map)
          cost_map[nb_row - row + COST_MAP_OFFSET]
                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
        accu_cost_diff += cost_diff;
      }
    }
  }

  if (check_br_neighbor(qc)) {
    for (int i = 0; i < BR_CONTEXT_POSITION_NUM; ++i) {
      const int nb_row = row - br_ref_offset[i][0];
      const int nb_col = col - br_ref_offset[i][1];
      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
904 905 906 907 908

      if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
            nb_col < txb_info->stride))
        continue;

Angie Chiang's avatar
Angie Chiang committed
909
      const int nb_scan_idx = iscan[nb_coeff_idx];
910
      if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
911
        const int cost_diff = try_neighbor_level_down_br(
912
            nb_coeff_idx, coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
913 914 915 916 917 918 919 920 921 922 923
        if (cost_map)
          cost_map[nb_row - row + COST_MAP_OFFSET]
                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
        accu_cost_diff += cost_diff;
      }
    }
  }

  return accu_cost_diff;
}

Angie Chiang's avatar
Angie Chiang committed
924
static int get_low_coeff_cost(int coeff_idx, const TxbCache *txb_cache,
925
                              const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
926 927 928 929 930 931 932
                              const TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const int abs_qc = abs(qc);
  assert(abs_qc <= 1);
  int cost = 0;
  const int scan_idx = txb_info->scan_order->iscan[coeff_idx];
  if (scan_idx < txb_info->seg_eob) {
933 934
    const int *level_cost = get_level_prob(0, coeff_idx, txb_cache, txb_costs);
    cost += level_cost[qc != 0];
Angie Chiang's avatar
Angie Chiang committed
935 936 937 938 939
  }

  if (qc != 0) {
    const int base_idx = 0;
    const int ctx = txb_cache->base_ctx_arr[base_idx][coeff_idx][0];
940 941
    cost += get_base_cost(abs_qc, ctx, txb_costs->base_cost[base_idx][ctx],
                          base_idx);
Angie Chiang's avatar
Angie Chiang committed
942 943
    if (scan_idx < txb_info->seg_eob) {
      const int eob_ctx =
944
          get_eob_ctx(txb_info->qcoeff, coeff_idx, txb_info->txs_ctx);
945
      cost += txb_costs->eob_cost[eob_ctx][scan_idx == (txb_info->eob - 1)];
Angie Chiang's avatar
Angie Chiang committed
946
    }
947
    cost += get_sign_bit_cost(qc, coeff_idx, txb_costs->dc_sign_cost,
Angie Chiang's avatar
Angie Chiang committed
948 949 950 951 952 953 954 955 956 957 958 959
                              txb_info->txb_ctx->dc_sign_ctx);
  }
  return cost;
}

static INLINE void set_eob(TxbInfo *txb_info, int eob) {
  txb_info->eob = eob;
  txb_info->seg_eob = AOMMIN(eob, tx_size_2d[txb_info->tx_size] - 1);
}

// TODO(angiebird): add static to this function once it's called
int try_change_eob(int *new_eob, int coeff_idx, const TxbCache *txb_cache,
960
                   const LV_MAP_COEFF_COST *txb_costs, TxbInfo *txb_info) {
Angie Chiang's avatar
Angie Chiang committed
961 962 963 964 965 966 967 968 969 970 971 972
  assert(txb_info->eob > 0);
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const int abs_qc = abs(qc);
  if (abs_qc != 1) {
    *new_eob = -1;
    return 0;
  }
  const int16_t *iscan = txb_info->scan_order->iscan;
  const int16_t *scan = txb_info->scan_order->scan;
  const int scan_idx = iscan[coeff_idx];
  *new_eob = 0;
  int cost_diff = 0;
973
  cost_diff -= get_low_coeff_cost(coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
974
  // int coeff_cost =
975
  //     get_coeff_cost(qc, scan_idx, txb_info, txb_probs);
Angie Chiang's avatar
Angie Chiang committed
976 977 978
  // if (-cost_diff != coeff_cost) {
  //   printf("-cost_diff %d coeff_cost %d\n", -cost_diff, coeff_cost);
  //   get_low_coeff_cost(coeff_idx, txb_cache, txb_probs, txb_info);
979
  //   get_coeff_cost(qc, scan_idx, txb_info, txb_probs);
Angie Chiang's avatar
Angie Chiang committed
980 981 982 983 984 985 986
  // }
  for (int si = scan_idx - 1; si >= 0; --si) {
    const int ci = scan[si];
    if (txb_info->qcoeff[ci] != 0) {
      *new_eob = si + 1;
      break;
    } else {
987
      cost_diff -= get_low_coeff_cost(ci, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
988 989 990 991 992
    }
  }

  const int org_eob = txb_info->eob;
  set_eob(txb_info, *new_eob);
993
  cost_diff += try_level_down(coeff_idx, txb_cache, txb_costs, txb_info, NULL);
Angie Chiang's avatar
Angie Chiang committed
994 995 996 997 998 999
  set_eob(txb_info, org_eob);

  if (*new_eob > 0) {
    // Note that get_eob_ctx does NOT actually account for qcoeff, so we don't
    // need to lower down the qcoeff here
    const int eob_ctx =
1000
        get_eob_ctx(txb_info->qcoeff, scan[*new_eob - 1], txb_info->txs_ctx);
1001 1002
    cost_diff -= txb_costs->eob_cost[eob_ctx][0];
    cost_diff += txb_costs->eob_cost[eob_ctx][1];
Angie Chiang's avatar
Angie Chiang committed
1003 1004
  } else {
    const int txb_skip_ctx = txb_info->txb_ctx->txb_skip_ctx;
1005 1006
    cost_diff -= txb_costs->txb_skip_cost[txb_skip_ctx][0];
    cost_diff += txb_costs->txb_skip_cost[txb_skip_ctx][1];
Angie Chiang's avatar
Angie Chiang committed
1007 1008 1009
  }
  return cost_diff;
}
Angie Chiang's avatar
Angie Chiang committed
1010 1011 1012 1013 1014 1015

static INLINE tran_low_t qcoeff_to_dqcoeff(tran_low_t qc, int dqv, int shift) {
  int sgn = qc < 0 ? -1 : 1;
  return sgn * ((abs(qc) * dqv) >> shift);
}

Angie Chiang's avatar
Angie Chiang committed
1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033
// TODO(angiebird): add static to this function it's called
void update_level_down(int coeff_idx, TxbCache *txb_cache, TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const int abs_qc = abs(qc);
  if (qc == 0) return;
  const tran_low_t low_coeff = get_lower_coeff(qc);
  txb_info->qcoeff[coeff_idx] = low_coeff;
  const int dqv = txb_info->dequant[coeff_idx != 0];
  txb_info->dqcoeff[coeff_idx] =
      qcoeff_to_dqcoeff(low_coeff, dqv, txb_info->shift);

  const int row = coeff_idx >> txb_info->bwl;
  const int col = coeff_idx - (row << txb_info->bwl);
  const int eob = txb_info->eob;
  const int16_t *iscan = txb_info->scan_order->iscan;
  for (int i = 0; i < SIG_REF_OFFSET_NUM; ++i) {
    const int nb_row = row - sig_ref_offset[i][0];
    const int nb_col = col - sig_ref_offset[i][1];
1034 1035 1036 1037 1038

    if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
          nb_col < txb_info->stride))
      continue;

Angie Chiang's avatar
Angie Chiang committed
1039 1040
    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
    const int nb_scan_idx = iscan[nb_coeff_idx];
1041
    if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
1042 1043 1044 1045 1046 1047 1048 1049 1050 1051
      const int scan_idx = iscan[coeff_idx];
      if (scan_idx < nb_scan_idx) {
        const int level = 1;
        if (abs_qc == level) {
          txb_cache->nz_count_arr[nb_coeff_idx] -= 1;
          assert(txb_cache->nz_count_arr[nb_coeff_idx] >= 0);
        }
        const int count = txb_cache->nz_count_arr[nb_coeff_idx];
        txb_cache->nz_ctx_arr[nb_coeff_idx][0] = get_nz_map_ctx_from_count(
            count, txb_info->qcoeff, nb_coeff_idx, txb_info->bwl, iscan);
1052
        // int ref_ctx = get_nz_map_ctx(txb_info->qcoeff, nb_coeff_idx,
Angie Chiang's avatar
Angie Chiang committed
1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064
        // txb_info->bwl, iscan);
        // if (ref_ctx != txb_cache->nz_ctx_arr[nb_coeff_idx][0])
        //   printf("nz ctx %d ref_ctx %d\n",
        //   txb_cache->nz_ctx_arr[nb_coeff_idx][0], ref_ctx);
      }
    }
  }

  for (int i = 0; i < BASE_CONTEXT_POSITION_NUM; ++i) {
    const int nb_row = row - base_ref_offset[i][0];
    const int nb_col = col - base_ref_offset[i][1];
    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
1065 1066 1067 1068 1069

    if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
          nb_col < txb_info->stride))
      continue;

Angie Chiang's avatar
Angie Chiang committed
1070 1071 1072
    const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
    if (!has_base(nb_coeff, 0)) continue;
    const int nb_scan_idx = iscan[nb_coeff_idx];
1073
    if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101
      if (row >= nb_row && col >= nb_col)
        update_mag_arr(txb_cache->base_mag_arr[nb_coeff_idx], abs_qc);
      const int mag =
          get_mag_from_mag_arr(txb_cache->base_mag_arr[nb_coeff_idx]);
      for (int base_idx = 0; base_idx < NUM_BASE_LEVELS; ++base_idx) {
        if (!has_base(nb_coeff, base_idx)) continue;
        const int level = base_idx + 1;
        if (abs_qc == level) {
          txb_cache->base_count_arr[base_idx][nb_coeff_idx] -= 1;
          assert(txb_cache->base_count_arr[base_idx][nb_coeff_idx] >= 0);
        }
        const int count = txb_cache->base_count_arr[base_idx][nb_coeff_idx];
        txb_cache->base_ctx_arr[base_idx][nb_coeff_idx][0] =
            get_base_ctx_from_count_mag(nb_row, nb_col, count, mag, level);
        // int ref_ctx = get_base_ctx(txb_info->qcoeff, nb_coeff_idx,
        // txb_info->bwl, level);
        // if (ref_ctx != txb_cache->base_ctx_arr[base_idx][nb_coeff_idx][0]) {
        //   printf("base ctx %d ref_ctx %d\n",
        //   txb_cache->base_ctx_arr[base_idx][nb_coeff_idx][0], ref_ctx);
        // }
      }
    }
  }

  for (int i = 0; i < BR_CONTEXT_POSITION_NUM; ++i) {
    const int nb_row = row - br_ref_offset[i][0];
    const int nb_col = col - br_ref_offset[i][1];
    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
1102 1103 1104 1105 1106

    if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
          nb_col < txb_info->stride))
      continue;

Angie Chiang's avatar
Angie Chiang committed
1107 1108 1109
    const int nb_scan_idx = iscan[nb_coeff_idx];
    const tran_low_t nb_coeff = txb_info