Add plane to tx_type read/write/update functions

This is for lv_map experiment

Change-Id: Ie000f7850efac32ffb46b9a4679cff2814c6246a
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index e6e88b5..7b2f4c4 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -1424,7 +1424,7 @@
                        const int supertx_enabled,
 #endif
 #if CONFIG_LV_MAP
-                       int block,
+                       int block, int plane,
 #endif
                        aom_writer *w) {
   MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
@@ -1444,7 +1444,9 @@
   TX_TYPE tx_type = mbmi->tx_type;
 #else
   // Only y plane's tx_type is transmitted
-  TX_TYPE tx_type = get_tx_type(PLANE_TYPE_Y, xd, block, tx_size);
+  if (plane > 0) return;
+  PLANE_TYPE plane_type = get_plane_type(plane);
+  TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
 #endif
 
   if (!FIXED_TX_TYPE) {
diff --git a/av1/encoder/bitstream.h b/av1/encoder/bitstream.h
index 59b4ce3..f2ab301 100644
--- a/av1/encoder/bitstream.h
+++ b/av1/encoder/bitstream.h
@@ -42,7 +42,7 @@
                        const int supertx_enabled,
 #endif
 #if CONFIG_LV_MAP
-                       int block,
+                       int block, int plane,
 #endif
                        aom_writer *w);
 
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 5084216..0b53c02 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -5756,7 +5756,7 @@
 
 void av1_update_tx_type_count(const AV1_COMMON *cm, MACROBLOCKD *xd,
 #if CONFIG_LV_MAP
-                              int block,
+                              int block, int plane,
 #endif
                               BLOCK_SIZE bsize, TX_SIZE tx_size,
                               FRAME_COUNTS *counts) {
@@ -5765,6 +5765,8 @@
 #if !CONFIG_LV_MAP
   TX_TYPE tx_type = mbmi->tx_type;
 #else
+  // Only y plane's tx_type is updated
+  if (plane > 0) return;
   TX_TYPE tx_type = get_tx_type(PLANE_TYPE_Y, xd, block, tx_size);
 #endif
 #if CONFIG_EXT_TX
diff --git a/av1/encoder/encodeframe.h b/av1/encoder/encodeframe.h
index 0ea577b..ad3213e 100644
--- a/av1/encoder/encodeframe.h
+++ b/av1/encoder/encodeframe.h
@@ -46,7 +46,7 @@
 
 void av1_update_tx_type_count(const struct AV1Common *cm, MACROBLOCKD *xd,
 #if CONFIG_LV_MAP
-                              int block,
+                              int block, int plane,
 #endif
                               BLOCK_SIZE bsize, TX_SIZE tx_size,
                               FRAME_COUNTS *counts);
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 55753a3..be7eb88 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -90,7 +90,7 @@
   aom_write(w, eob == 0, cm->fc->txb_skip[tx_size][txb_ctx->txb_skip_ctx]);
 
   if (eob == 0) return;
-  av1_write_tx_type(cm, xd, block, w);
+  av1_write_tx_type(cm, xd, block, plane, w);
 
   nz_map = cm->fc->nz_map[tx_size][plane_type];
   eob_flag = cm->fc->eob_flag[tx_size][plane_type];
@@ -456,7 +456,8 @@
     return;
   }
 
-  av1_update_tx_type_count(cm, xd, block, mbmi->sb_type, tx_size, td->counts);
+  av1_update_tx_type_count(cm, xd, block, plane, mbmi->sb_type, tx_size,
+                           td->counts);
 
   for (c = 0; c < eob; ++c) {
     tran_low_t v = qcoeff[scan[c]];