Refactor av1_set_contexts()

Change-Id: I7b3e747f5cd8a06ab359475388b894699f8bd36f
diff --git a/av1/common/blockd.c b/av1/common/blockd.c
index 3fcaa50..fe089a1 100644
--- a/av1/common/blockd.c
+++ b/av1/common/blockd.c
@@ -87,40 +87,31 @@
 }
 
 void av1_set_contexts(const MACROBLOCKD *xd, struct macroblockd_plane *pd,
-                      int plane, TX_SIZE tx_size, int has_eob, int aoff,
-                      int loff) {
+                      int plane, BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+                      int has_eob, int aoff, int loff) {
   ENTROPY_CONTEXT *const a = pd->above_context + aoff;
   ENTROPY_CONTEXT *const l = pd->left_context + loff;
   const int txs_wide = tx_size_wide_unit[tx_size];
   const int txs_high = tx_size_high_unit[tx_size];
-  const BLOCK_SIZE bsize = xd->mi[0]->sb_type;
-  const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
 
   // above
   if (has_eob && xd->mb_to_right_edge < 0) {
-    int i;
     const int blocks_wide = max_block_wide(xd, plane_bsize, plane);
-    int above_contexts = txs_wide;
-    if (above_contexts + aoff > blocks_wide)
-      above_contexts = blocks_wide - aoff;
-
-    for (i = 0; i < above_contexts; ++i) a[i] = has_eob;
-    for (i = above_contexts; i < txs_wide; ++i) a[i] = 0;
+    const int above_contexts = AOMMIN(txs_wide, blocks_wide - aoff);
+    memset(a, has_eob, sizeof(*a) * above_contexts);
+    memset(a + above_contexts, 0, sizeof(*a) * (txs_wide - above_contexts));
   } else {
-    memset(a, has_eob, sizeof(ENTROPY_CONTEXT) * txs_wide);
+    memset(a, has_eob, sizeof(*a) * txs_wide);
   }
 
   // left
   if (has_eob && xd->mb_to_bottom_edge < 0) {
-    int i;
     const int blocks_high = max_block_high(xd, plane_bsize, plane);
-    int left_contexts = txs_high;
-    if (left_contexts + loff > blocks_high) left_contexts = blocks_high - loff;
-
-    for (i = 0; i < left_contexts; ++i) l[i] = has_eob;
-    for (i = left_contexts; i < txs_high; ++i) l[i] = 0;
+    const int left_contexts = AOMMIN(txs_high, blocks_high - loff);
+    memset(l, has_eob, sizeof(*l) * left_contexts);
+    memset(l + left_contexts, 0, sizeof(*l) * (txs_high - left_contexts));
   } else {
-    memset(l, has_eob, sizeof(ENTROPY_CONTEXT) * txs_high);
+    memset(l, has_eob, sizeof(*l) * txs_high);
   }
 }
 void av1_reset_skip_context(MACROBLOCKD *xd, int mi_row, int mi_col,
diff --git a/av1/common/blockd.h b/av1/common/blockd.h
index b9ebf29..fe8f691 100644
--- a/av1/common/blockd.h
+++ b/av1/common/blockd.h
@@ -943,8 +943,8 @@
                                    void *arg, const int num_planes);
 
 void av1_set_contexts(const MACROBLOCKD *xd, struct macroblockd_plane *pd,
-                      int plane, TX_SIZE tx_size, int has_eob, int aoff,
-                      int loff);
+                      int plane, BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+                      int has_eob, int aoff, int loff);
 
 #define MAX_INTERINTRA_SB_SQUARE 32 * 32
 static INLINE int is_interintra_mode(const MB_MODE_INFO *mbmi) {
diff --git a/av1/decoder/decodetxb.c b/av1/decoder/decodetxb.c
index bd1275f..6c1962f 100644
--- a/av1/decoder/decodetxb.c
+++ b/av1/decoder/decodetxb.c
@@ -319,8 +319,8 @@
   TXB_CTX txb_ctx;
   get_txb_ctx(plane_bsize, tx_size, plane, pd->above_context + col,
               pd->left_context + row, &txb_ctx);
-  uint8_t cul_level = av1_read_coeffs_txb(cm, xd, r, row, col, plane, &txb_ctx,
-                                          tx_size, max_scan_line, eob);
-  av1_set_contexts(xd, pd, plane, tx_size, cul_level, col, row);
+  const uint8_t cul_level = av1_read_coeffs_txb(
+      cm, xd, r, row, col, plane, &txb_ctx, tx_size, max_scan_line, eob);
+  av1_set_contexts(xd, pd, plane, plane_bsize, tx_size, cul_level, col, row);
   return cul_level;
 }
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 197f570..8d3770f 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -1688,10 +1688,9 @@
   const TX_TYPE tx_type = av1_get_tx_type(plane_type, xd, blk_row, blk_col,
                                           tx_size, cm->reduced_tx_set_used);
   const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);
-  (void)plane_bsize;
-
-  int cul_level = av1_get_txb_entropy_context(qcoeff, scan_order, eob);
-  av1_set_contexts(xd, pd, plane, tx_size, cul_level, blk_col, blk_row);
+  const int cul_level = av1_get_txb_entropy_context(qcoeff, scan_order, eob);
+  av1_set_contexts(xd, pd, plane, plane_bsize, tx_size, cul_level, blk_col,
+                   blk_row);
 }
 
 void av1_update_and_record_txb_context(int plane, int block, int blk_row,
@@ -1732,7 +1731,7 @@
   x->mbmi_ext->eobs[plane][block] = eob;
 
   if (eob == 0) {
-    av1_set_contexts(xd, pd, plane, tx_size, 0, blk_col, blk_row);
+    av1_set_contexts(xd, pd, plane, plane_bsize, tx_size, 0, blk_col, blk_row);
     return;
   }
 
@@ -1832,7 +1831,8 @@
   }
 
   const int cul_level = av1_get_txb_entropy_context(tcoeff, scan_order, eob);
-  av1_set_contexts(xd, pd, plane, tx_size, cul_level, blk_col, blk_row);
+  av1_set_contexts(xd, pd, plane, plane_bsize, tx_size, cul_level, blk_col,
+                   blk_row);
 }
 
 void av1_update_txb_context(const AV1_COMP *cpi, ThreadData *td,