Add chroma bsize modifications correctly

Change-Id: I807dc86489bc9219a24801841e66b555a2d7a51f
diff --git a/av1/decoder/decodeframe.c b/av1/decoder/decodeframe.c
index 83d3eb1..b7bd2be 100644
--- a/av1/decoder/decodeframe.c
+++ b/av1/decoder/decodeframe.c
@@ -511,15 +511,16 @@
       int eobtotal = 0;
       for (int plane = 0; plane < av1_num_planes(cm); ++plane) {
         const struct macroblockd_plane *const pd = &xd->plane[plane];
-        const BLOCK_SIZE plane_bsize =
-            AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
-        const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
-        const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
-        int row, col;
-
         if (!is_chroma_reference(mi_row, mi_col, bsize, pd->subsampling_x,
                                  pd->subsampling_y))
           continue;
+        const BLOCK_SIZE bsizec =
+            scale_chroma_bsize(bsize, pd->subsampling_x, pd->subsampling_y);
+        const BLOCK_SIZE plane_bsize =
+            AOMMAX(BLOCK_4X4, get_plane_block_size(bsizec, pd));
+        const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
+        const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
+        int row, col;
 
         const BLOCK_SIZE max_unit_bsize = get_plane_block_size(BLOCK_64X64, pd);
         int mu_blocks_wide =
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 1bacf31..ec16bd5 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -1871,18 +1871,19 @@
     assert(*tok < tok_end);
 #endif
     for (plane = 0; plane < num_planes; ++plane) {
-      if (!is_chroma_reference(mi_row, mi_col, mbmi->sb_type,
-                               xd->plane[plane].subsampling_x,
-                               xd->plane[plane].subsampling_y)) {
+      const struct macroblockd_plane *const pd = &xd->plane[plane];
+      if (!is_chroma_reference(mi_row, mi_col, mbmi->sb_type, pd->subsampling_x,
+                               pd->subsampling_y)) {
 #if !CONFIG_LV_MAP
         (*tok)++;
 #endif  // !CONFIG_LV_MAP
         continue;
       }
-      const struct macroblockd_plane *const pd = &xd->plane[plane];
-      BLOCK_SIZE bsize = mbmi->sb_type;
+      const BLOCK_SIZE bsize = mbmi->sb_type;
+      const BLOCK_SIZE bsizec =
+          scale_chroma_bsize(bsize, pd->subsampling_x, pd->subsampling_y);
       const BLOCK_SIZE plane_bsize =
-          AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
+          AOMMAX(BLOCK_4X4, get_plane_block_size(bsizec, pd));
 
       const int num_4x4_w =
           block_size_wide[plane_bsize] >> tx_size_wide_log2[0];
diff --git a/av1/encoder/encodemb.c b/av1/encoder/encodemb.c
index 49c79f7..e097d00 100644
--- a/av1/encoder/encodemb.c
+++ b/av1/encoder/encodemb.c
@@ -821,11 +821,12 @@
                              subsampling_y))
       continue;
 
-    bsize = scale_chroma_bsize(bsize, subsampling_x, subsampling_y);
+    const BLOCK_SIZE bsizec =
+        scale_chroma_bsize(bsize, subsampling_x, subsampling_y);
 
     // TODO(jingning): Clean this up.
     const struct macroblockd_plane *const pd = &xd->plane[plane];
-    const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
+    const BLOCK_SIZE plane_bsize = get_plane_block_size(bsizec, pd);
     const int mi_width = block_size_wide[plane_bsize] >> tx_size_wide_log2[0];
     const int mi_height = block_size_high[plane_bsize] >> tx_size_wide_log2[0];
     const TX_SIZE max_tx_size = get_vartx_max_txsize(
@@ -836,9 +837,9 @@
     int idx, idy;
     int block = 0;
     int step = tx_size_wide_unit[max_tx_size] * tx_size_high_unit[max_tx_size];
-    av1_get_entropy_contexts(bsize, 0, pd, ctx.ta[plane], ctx.tl[plane]);
+    av1_get_entropy_contexts(bsizec, 0, pd, ctx.ta[plane], ctx.tl[plane]);
 
-    av1_subtract_plane(x, bsize, plane);
+    av1_subtract_plane(x, bsizec, plane);
 
     arg.ta = ctx.ta[plane];
     arg.tl = ctx.tl[plane];
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index 9c0ceab..99b1ac9 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -5008,8 +5008,8 @@
   av1_init_rd_stats(rd_stats);
 
   if (x->skip_chroma_rd) return is_cost_valid;
-  bsize = scale_chroma_bsize(mbmi->sb_type, xd->plane[1].subsampling_x,
-                             xd->plane[1].subsampling_y);
+  const BLOCK_SIZE bsizec = scale_chroma_bsize(
+      bsize, xd->plane[1].subsampling_x, xd->plane[1].subsampling_y);
 
 #if 0   // CONFIG_EXT_TX
   if (is_rect_tx(mbmi->tx_size)) {
@@ -5019,13 +5019,13 @@
 
   if (is_inter_block(mbmi) && is_cost_valid) {
     for (plane = 1; plane < MAX_MB_PLANE; ++plane)
-      av1_subtract_plane(x, bsize, plane);
+      av1_subtract_plane(x, bsizec, plane);
   }
 
   if (is_cost_valid) {
     for (plane = 1; plane < MAX_MB_PLANE; ++plane) {
       const struct macroblockd_plane *const pd = &xd->plane[plane];
-      const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
+      const BLOCK_SIZE plane_bsize = get_plane_block_size(bsizec, pd);
       const int mi_width = block_size_wide[plane_bsize] >> tx_size_wide_log2[0];
       const int mi_height =
           block_size_high[plane_bsize] >> tx_size_high_log2[0];
@@ -5040,7 +5040,7 @@
       ENTROPY_CONTEXT tl[2 * MAX_MIB_SIZE];
       RD_STATS pn_rd_stats;
       av1_init_rd_stats(&pn_rd_stats);
-      av1_get_entropy_contexts(bsize, 0, pd, ta, tl);
+      av1_get_entropy_contexts(bsizec, 0, pd, ta, tl);
 
       for (idy = 0; idy < mi_height; idy += bh) {
         for (idx = 0; idx < mi_width; idx += bw) {
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index 1f7d2e1..c18594f 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -653,8 +653,10 @@
       continue;
     }
     const struct macroblockd_plane *const pd = &xd->plane[plane];
+    const BLOCK_SIZE bsizec =
+        scale_chroma_bsize(bsize, pd->subsampling_x, pd->subsampling_y);
     const BLOCK_SIZE plane_bsize =
-        AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
+        AOMMAX(BLOCK_4X4, get_plane_block_size(bsizec, pd));
     const int mi_width = block_size_wide[plane_bsize] >> tx_size_wide_log2[0];
     const int mi_height = block_size_high[plane_bsize] >> tx_size_wide_log2[0];
     const TX_SIZE max_tx_size = get_vartx_max_txsize(