Cdf update and stats collection for palette color index

Change-Id: I546c6c3733a3efa2d0b10f003a9c935f8239b70f

av1/common/entropymode.h

diff --git a/av1/common/entropymode.h b/av1/common/entropymode.h
index 7719929..51da597 100644
--- a/av1/common/entropymode.h
+++ b/av1/common/entropymode.h
@@ -185,6 +185,12 @@
   unsigned int palette_uv_mode[PALETTE_UV_MODE_CONTEXTS][2];
   unsigned int palette_y_size[PALATTE_BSIZE_CTXS][PALETTE_SIZES];
   unsigned int palette_uv_size[PALATTE_BSIZE_CTXS][PALETTE_SIZES];
+  unsigned int palette_y_color_index[PALETTE_SIZES]
+                                    [PALETTE_COLOR_INDEX_CONTEXTS]
+                                    [PALETTE_COLORS];
+  unsigned int palette_uv_color_index[PALETTE_SIZES]
+                                     [PALETTE_COLOR_INDEX_CONTEXTS]
+                                     [PALETTE_COLORS];
 #endif  // CONFIG_ENTROPY_STATS
   unsigned int partition[PARTITION_CONTEXTS][EXT_PARTITION_TYPES];
   unsigned int switchable_interp[SWITCHABLE_FILTER_CONTEXTS]

av1/encoder/encodeframe.c

diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 6128db5..b27f4e7 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -4746,7 +4746,8 @@
         if (mbmi->palette_mode_info.palette_size[plane] > 0) {
           if (!dry_run) {
             av1_tokenize_color_map(x, plane, t, bsize, mbmi->tx_size,
-                                   PALETTE_MAP);
+                                   PALETTE_MAP, tile_data->allow_update_cdf,
+                                   td->counts);
           } else if (dry_run == DRY_RUN_COSTCOEFFS) {
             rate +=
                 av1_cost_color_map(x, plane, bsize, mbmi->tx_size, PALETTE_MAP);

av1/encoder/tokenize.c

diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index 8d35df5..3729e30 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -28,7 +28,8 @@
 #include "av1/encoder/tokenize.h"
 
 static int cost_and_tokenize_map(Av1ColorMapParam *param, TOKENEXTRA **t,
-                                 int calc_rate) {
+                                 int plane, int calc_rate, int allow_update_cdf,
+                                 FRAME_COUNTS *counts) {
   const uint8_t *const color_map = param->color_map;
   MapCdf map_cdf = param->map_cdf;
   ColorCost color_cost = param->color_cost;
@@ -36,9 +37,13 @@
   const int rows = param->rows;
   const int cols = param->cols;
   const int n = param->n_colors;
-
+  const int palette_size_idx = n - PALETTE_MIN_SIZE;
   int this_rate = 0;
   uint8_t color_order[PALETTE_MAX_SIZE];
+
+  (void)plane;
+  (void)counts;
+
   for (int k = 1; k < rows + cols - 1; ++k) {
     for (int j = AOMMIN(k, cols - 1); j >= AOMMAX(0, k - rows + 1); --j) {
       int i = k - j;
@@ -47,12 +52,22 @@
           color_map, plane_block_width, i, j, n, color_order, &color_new_idx);
       assert(color_new_idx >= 0 && color_new_idx < n);
       if (calc_rate) {
-        this_rate +=
-            (*color_cost)[n - PALETTE_MIN_SIZE][color_ctx][color_new_idx];
+        this_rate += (*color_cost)[palette_size_idx][color_ctx][color_new_idx];
       } else {
         (*t)->token = color_new_idx;
-        (*t)->color_map_cdf = map_cdf[n - PALETTE_MIN_SIZE][color_ctx];
+        (*t)->color_map_cdf = map_cdf[palette_size_idx][color_ctx];
         ++(*t);
+        if (allow_update_cdf)
+          update_cdf(map_cdf[palette_size_idx][color_ctx], color_new_idx, n);
+#if CONFIG_ENTROPY_STATS
+        if (plane) {
+          ++counts->palette_uv_color_index[palette_size_idx][color_ctx]
+                                          [color_new_idx];
+        } else {
+          ++counts->palette_y_color_index[palette_size_idx][color_ctx]
+                                         [color_new_idx];
+        }
+#endif
       }
     }
   }
@@ -92,12 +107,13 @@
   assert(plane == 0 || plane == 1);
   Av1ColorMapParam color_map_params;
   get_color_map_params(x, plane, bsize, tx_size, type, &color_map_params);
-  return cost_and_tokenize_map(&color_map_params, NULL, 1);
+  return cost_and_tokenize_map(&color_map_params, NULL, plane, 1, 0, NULL);
 }
 
 void av1_tokenize_color_map(const MACROBLOCK *const x, int plane,
                             TOKENEXTRA **t, BLOCK_SIZE bsize, TX_SIZE tx_size,
-                            COLOR_MAP_TYPE type) {
+                            COLOR_MAP_TYPE type, int allow_update_cdf,
+                            FRAME_COUNTS *counts) {
   assert(plane == 0 || plane == 1);
   Av1ColorMapParam color_map_params;
   get_color_map_params(x, plane, bsize, tx_size, type, &color_map_params);
@@ -105,7 +121,8 @@
   (*t)->token = color_map_params.color_map[0];
   (*t)->color_map_cdf = NULL;
   ++(*t);
-  cost_and_tokenize_map(&color_map_params, t, 0);
+  cost_and_tokenize_map(&color_map_params, t, plane, 0, allow_update_cdf,
+                        counts);
 }
 
 void tokenize_vartx(ThreadData *td, TOKENEXTRA **t, RUN_TYPE dry_run,

av1/encoder/tokenize.h

diff --git a/av1/encoder/tokenize.h b/av1/encoder/tokenize.h
index 26be5a8..5d277f6 100644
--- a/av1/encoder/tokenize.h
+++ b/av1/encoder/tokenize.h
@@ -56,7 +56,8 @@
 
 void av1_tokenize_color_map(const MACROBLOCK *const x, int plane,
                             TOKENEXTRA **t, BLOCK_SIZE bsize, TX_SIZE tx_size,
-                            COLOR_MAP_TYPE type);
+                            COLOR_MAP_TYPE type, int allow_update_cdf,
+                            FRAME_COUNTS *counts);
 
 static INLINE int av1_get_tx_eob(const struct segmentation *seg, int segment_id,
                                  TX_SIZE tx_size) {