Add read_coeffs_reverse_2d()

Change-Id: I9f6d50b21454b28b7ff76959ee8f3367ffd50d17

av1/common/txb_common.h

@@ -357,6 +357,23 @@ static INLINE int get_br_ctx_from_count_mag(const int row, const int col,
   return ctx;
 }
 
+static INLINE int get_br_ctx_2d(const uint8_t *const levels,
+                                const int c,  // raster order
+                                const int bwl) {
+  const int row = c >> bwl;
+  const int col = c - (row << bwl);
+  const int stride = (1 << bwl) + TX_PAD_HOR;
+  const int pos = row * stride + col;
+  int mag =
+      AOMMIN(levels[pos + 1], COEFF_BASE_RANGE + NUM_BASE_LEVELS + 1) +
+      AOMMIN(levels[pos + stride], COEFF_BASE_RANGE + NUM_BASE_LEVELS + 1) +
+      AOMMIN(levels[pos + 1 + stride], COEFF_BASE_RANGE + NUM_BASE_LEVELS + 1);
+  mag = AOMMIN((mag + 1) >> 1, 6);
+  if (c == 0) return mag;
+  if ((row < 2) && (col < 2)) return mag + 7;
+  return mag + 14;
+}
+
 static INLINE int get_br_ctx(const uint8_t *const levels,
                              const int c,  // raster order
                              const int bwl, const TX_TYPE tx_type) {
@@ -373,13 +390,13 @@ static INLINE int get_br_ctx(const uint8_t *const levels,
         AOMMIN(nb_mag[2], COEFF_BASE_RANGE + NUM_BASE_LEVELS + 1);
   mag = AOMMIN((mag + 1) >> 1, 6);
   if (c == 0) return mag;
-  if (tx_class == 0) {
+  if (tx_class == TX_CLASS_2D) {
     if ((row < 2) && (col < 2)) return mag + 7;
   } else {
-    if (tx_class == 1) {
+    if (tx_class == TX_CLASS_HORIZ) {
       if (col == 0) return mag + 7;
     } else {
-      if (tx_class == 2) {
+      if (tx_class == TX_CLASS_VERT) {
         if (row == 0) return mag + 7;
       }
     }
@@ -518,6 +535,26 @@ static INLINE int get_lower_levels_ctx_eob(int bwl, int height, int scan_idx) {
   return 3;
 }
 
+static INLINE int get_lower_levels_ctx_2d(const uint8_t *levels, int coeff_idx,
+                                          int bwl, TX_SIZE tx_size) {
+  int mag;
+  // Note: AOMMIN(level, 3) is useless for decoder since level < 3.
+  levels = levels + get_padded_idx(coeff_idx, bwl);
+  mag = AOMMIN(levels[1], 3);                                     // { 0, 1 }
+  mag += AOMMIN(levels[(1 << bwl) + TX_PAD_HOR], 3);              // { 1, 0 }
+  mag += AOMMIN(levels[(1 << bwl) + TX_PAD_HOR + 1], 3);          // { 1, 1 }
+  mag += AOMMIN(levels[2], 3);                                    // { 0, 2 }
+  mag += AOMMIN(levels[(2 << bwl) + (2 << TX_PAD_HOR_LOG2)], 3);  // { 2, 0 }
+
+  const int ctx = AOMMIN((mag + 1) >> 1, 4);
+  if (!coeff_idx) {
+    return 0;
+  } else {
+    const int row = coeff_idx >> bwl;
+    const int col = coeff_idx - (row << bwl);
+    return ctx + av1_nz_map_ctx_offset[tx_size][AOMMIN(row, 4)][AOMMIN(col, 4)];
+  }
+}
 static INLINE int get_lower_levels_ctx(const uint8_t *levels, int coeff_idx,
                                        int bwl, TX_SIZE tx_size,
                                        TX_TYPE tx_type) {

av1/decoder/decodetxb.c

@@ -63,6 +63,32 @@ static INLINE int get_dqv(const int16_t *dequant, int coeff_idx,
   return dqv;
 }
 
+static INLINE void read_coeffs_reverse_2d(
+    aom_reader *r, TX_SIZE tx_size, int start_si, int end_si,
+    const int16_t *scan, int bwl, uint8_t *levels, base_cdf_arr base_cdf,
+    br_cdf_arr br_cdf, int *num_updates, uint16_t *update_pos) {
+  for (int c = end_si; c >= start_si; --c) {
+    const int pos = scan[c];
+    const int coeff_ctx = get_lower_levels_ctx_2d(levels, pos, bwl, tx_size);
+    const int nsymbs = 4;
+    int level = aom_read_symbol(r, base_cdf[coeff_ctx], nsymbs, ACCT_STR);
+    if (level > NUM_BASE_LEVELS) {
+      const int br_ctx = get_br_ctx_2d(levels, pos, bwl);
+      aom_cdf_prob *cdf = br_cdf[br_ctx];
+      for (int idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
+        const int k = aom_read_symbol(r, cdf, BR_CDF_SIZE, ACCT_STR);
+        level += k;
+        if (k < BR_CDF_SIZE - 1) break;
+      }
+      if (level > NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+        update_pos[*num_updates] = pos;
+        ++*num_updates;
+      }
+    }
+    levels[get_padded_idx(pos, bwl)] = level;
+  }
+}
+
 static INLINE void read_coeffs_reverse(aom_reader *r, TX_SIZE tx_size,
                                        TX_TYPE tx_type, int start_si,
                                        int end_si, const int16_t *scan, int bwl,
@@ -253,8 +279,14 @@ uint8_t av1_read_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *const xd,
   base_cdf_arr base_cdf = ec_ctx->coeff_base_cdf[txs_ctx][plane_type];
   br_cdf_arr br_cdf =
       ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_32X32)][plane_type];
-  read_coeffs_reverse(r, tx_size, tx_type, 0, *eob - 1 - 1, scan, bwl, levels,
-                      base_cdf, br_cdf, &num_updates, update_pos);
+  const TX_CLASS tx_class = tx_type_to_class[tx_type];
+  if (tx_class == TX_CLASS_2D) {
+    read_coeffs_reverse_2d(r, tx_size, 0, *eob - 1 - 1, scan, bwl, levels,
+                           base_cdf, br_cdf, &num_updates, update_pos);
+  } else {
+    read_coeffs_reverse(r, tx_size, tx_type, 0, *eob - 1 - 1, scan, bwl, levels,
+                        base_cdf, br_cdf, &num_updates, update_pos);
+  }
 
   for (int i = 0; i < num_updates; ++i) {
     const int pos = update_pos[i];