Separate set context from av1_encode_block_intra

Add av1_set_txb_context to perform context setting.
This will make the code cleaner and facilitate
lv_map experiment implementation

Change-Id: Iea7e1e8ff6820e2c9cf72961adaa58346a7eca75

av1/encoder/encodemb.c

@@ -1010,6 +1010,34 @@ void av1_encode_sb_supertx(AV1_COMMON *cm, MACROBLOCK *x, BLOCK_SIZE bsize) {
 }
 #endif  // CONFIG_SUPERTX
 
+void av1_set_txb_context(MACROBLOCK *x, int plane, int block,
+                         BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+                         ENTROPY_CONTEXT *a, ENTROPY_CONTEXT *l) {
+  (void)plane_bsize;
+  (void)tx_size;
+#if !CONFIG_PVQ
+  struct macroblock_plane *p = &x->plane[plane];
+  *a = *l = p->eobs[block] > 0;
+#else   // !CONFIG_PVQ
+  (void)block;
+  *a = *l = !x->pvq_skip[plane];
+#endif  // !CONFIG_PVQ
+}
+
+static void encode_block_intra_and_set_context(int plane, int block,
+                                               int blk_row, int blk_col,
+                                               BLOCK_SIZE plane_bsize,
+                                               TX_SIZE tx_size, void *arg) {
+  av1_encode_block_intra(plane, block, blk_row, blk_col, plane_bsize, tx_size,
+                         arg);
+
+  struct encode_b_args *const args = arg;
+  MACROBLOCK *x = args->x;
+  ENTROPY_CONTEXT *a = &args->ta[blk_col];
+  ENTROPY_CONTEXT *l = &args->tl[blk_row];
+  av1_set_txb_context(x, plane, block, plane_bsize, tx_size, a, l);
+}
+
 void av1_encode_block_intra(int plane, int block, int blk_row, int blk_col,
                             BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
                             void *arg) {
@@ -1034,7 +1062,6 @@ void av1_encode_block_intra(int plane, int block, int blk_row, int blk_col,
   const int dst_stride = pd->dst.stride;
   const int tx1d_width = tx_size_wide[tx_size];
   const int tx1d_height = tx_size_high[tx_size];
-  ENTROPY_CONTEXT *a = NULL, *l = NULL;
   int ctx = 0;
   INV_TXFM_PARAM inv_txfm_param;
 #if CONFIG_PVQ
@@ -1080,9 +1107,9 @@ void av1_encode_block_intra(int plane, int block, int blk_row, int blk_col,
 #endif  // CONFIG_AOM_HIGHBITDEPTH
   }
 
-  a = &args->ta[blk_col];
-  l = &args->tl[blk_row];
 #if !CONFIG_PVQ
+  const ENTROPY_CONTEXT *a = &args->ta[blk_col];
+  const ENTROPY_CONTEXT *l = &args->tl[blk_row];
   ctx = combine_entropy_contexts(*a, *l);
 
   if (args->enable_optimize_b) {
@@ -1094,9 +1121,7 @@ void av1_encode_block_intra(int plane, int block, int blk_row, int blk_col,
                     ctx, AV1_XFORM_QUANT_FP);
 #endif  // CONFIG_NEW_QUANT
     if (p->eobs[block]) {
-      *a = *l = av1_optimize_b(cm, x, plane, block, tx_size, ctx) > 0;
-    } else {
-      *a = *l = 0;
+      av1_optimize_b(cm, x, plane, block, tx_size, ctx);
     }
   } else {
 #if CONFIG_NEW_QUANT
@@ -1106,7 +1131,6 @@ void av1_encode_block_intra(int plane, int block, int blk_row, int blk_col,
     av1_xform_quant(cm, x, plane, block, blk_row, blk_col, plane_bsize, tx_size,
                     ctx, AV1_XFORM_QUANT_B);
 #endif  // CONFIG_NEW_QUANT
-    *a = *l = p->eobs[block] > 0;
   }
 
   if (*eob) {
@@ -1138,8 +1162,6 @@ void av1_encode_block_intra(int plane, int block, int blk_row, int blk_col,
                   ctx, AV1_XFORM_QUANT_FP);
 #endif  // CONFIG_NEW_QUANT
 
-  *a = *l = !x->pvq_skip[plane];
-
   // *(args->skip) == mbmi->skip
   if (!x->pvq_skip[plane]) *(args->skip) = 0;
 
@@ -1213,8 +1235,8 @@ void av1_encode_intra_block_plane(AV1_COMMON *cm, MACROBLOCK *x,
     const TX_SIZE tx_size = get_tx_size(plane, xd);
     av1_get_entropy_contexts(bsize, tx_size, pd, ta, tl);
   }
-  av1_foreach_transformed_block_in_plane(xd, bsize, plane,
-                                         av1_encode_block_intra, &arg);
+  av1_foreach_transformed_block_in_plane(
+      xd, bsize, plane, encode_block_intra_and_set_context, &arg);
 }
 
 #if CONFIG_PVQ

av1/encoder/encodemb.h

@@ -63,6 +63,10 @@ int av1_optimize_b(const AV1_COMMON *cm, MACROBLOCK *mb, int plane, int block,
 
 void av1_subtract_plane(MACROBLOCK *x, BLOCK_SIZE bsize, int plane);
 
+void av1_set_txb_context(MACROBLOCK *x, int plane, int block,
+                         BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+                         ENTROPY_CONTEXT *a, ENTROPY_CONTEXT *l);
+
 void av1_encode_block_intra(int plane, int block, int blk_row, int blk_col,
                             BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg);
 

av1/encoder/rdopt.c

@@ -1475,6 +1475,8 @@ static void block_rd_txfm(int plane, int block, int blk_row, int blk_col,
     };
     av1_encode_block_intra(plane, block, blk_row, blk_col, plane_bsize, tx_size,
                            &b_args);
+    av1_set_txb_context(x, plane, block, plane_bsize, tx_size,
+                        args->t_above + blk_col, args->t_left + blk_row);
     if (args->cpi->sf.use_transform_domain_distortion && !CONFIG_DAALA_DIST) {
       dist_block(args->cpi, x, plane, block, blk_row, blk_col, tx_size,
                  &this_rd_stats.dist, &this_rd_stats.sse);