Enable switchable restoration for chroma

Change-Id: I78a8a1749cd4449c61a106f413c697e4a2df0475
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index c420a7f..c71a571 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -3084,30 +3084,14 @@
 static void encode_restoration_mode(AV1_COMMON *cm,
                                     struct aom_write_bit_buffer *wb) {
   int p;
-  RestorationInfo *rsi = &cm->rst_info[0];
-  switch (rsi->frame_restoration_type) {
-    case RESTORE_NONE:
-      aom_wb_write_bit(wb, 0);
-      aom_wb_write_bit(wb, 0);
-      break;
-    case RESTORE_WIENER:
-      aom_wb_write_bit(wb, 1);
-      aom_wb_write_bit(wb, 0);
-      break;
-    case RESTORE_SGRPROJ:
-      aom_wb_write_bit(wb, 1);
-      aom_wb_write_bit(wb, 1);
-      break;
-    case RESTORE_SWITCHABLE:
-      aom_wb_write_bit(wb, 0);
-      aom_wb_write_bit(wb, 1);
-      break;
-    default: assert(0);
-  }
-  for (p = 1; p < MAX_MB_PLANE; ++p) {
+  RestorationInfo *rsi;
+  for (p = 0; p < MAX_MB_PLANE; ++p) {
     rsi = &cm->rst_info[p];
     switch (rsi->frame_restoration_type) {
-      case RESTORE_NONE: aom_wb_write_bit(wb, 0); break;
+      case RESTORE_NONE:
+        aom_wb_write_bit(wb, 0);
+        aom_wb_write_bit(wb, 0);
+        break;
       case RESTORE_WIENER:
         aom_wb_write_bit(wb, 1);
         aom_wb_write_bit(wb, 0);
@@ -3116,15 +3100,19 @@
         aom_wb_write_bit(wb, 1);
         aom_wb_write_bit(wb, 1);
         break;
+      case RESTORE_SWITCHABLE:
+        aom_wb_write_bit(wb, 0);
+        aom_wb_write_bit(wb, 1);
+        break;
       default: assert(0);
     }
   }
   if (cm->rst_info[0].frame_restoration_type != RESTORE_NONE ||
       cm->rst_info[1].frame_restoration_type != RESTORE_NONE ||
       cm->rst_info[2].frame_restoration_type != RESTORE_NONE) {
-    rsi = &cm->rst_info[0];
-    aom_wb_write_bit(wb, rsi->restoration_tilesize != RESTORATION_TILESIZE_MAX);
-    if (rsi->restoration_tilesize != RESTORATION_TILESIZE_MAX) {
+    aom_wb_write_bit(
+        wb, rsi->restoration_tilesize != (RESTORATION_TILESIZE_MAX >> 2));
+    if (rsi->restoration_tilesize != (RESTORATION_TILESIZE_MAX >> 2)) {
       aom_wb_write_bit(
           wb, rsi->restoration_tilesize != (RESTORATION_TILESIZE_MAX >> 1));
     }
@@ -3219,7 +3207,6 @@
   SgrprojInfo *sgrproj_info = xd->sgrproj_info + plane;
 
   if (rsi->frame_restoration_type == RESTORE_SWITCHABLE) {
-    assert(plane == 0);
     aom_write_symbol(w, rsi->restoration_type[rtile_idx],
                      xd->tile_ctx->switchable_restore_cdf,
                      RESTORE_SWITCHABLE_TYPES);
diff --git a/av1/encoder/pickrst.c b/av1/encoder/pickrst.c
index 5030836..f90f3f1 100644
--- a/av1/encoder/pickrst.c
+++ b/av1/encoder/pickrst.c
@@ -1435,6 +1435,7 @@
   }
 
   for (int plane = AOM_PLANE_Y; plane <= AOM_PLANE_V; ++plane) {
+    const int ntiles = (plane == AOM_PLANE_Y ? ntiles_y : ntiles_uv);
     for (r = 0; r < RESTORE_SWITCHABLE_TYPES; ++r) {
       cost_restore[r] = DBL_MAX;
       if (force_restore_type != RESTORE_TYPES)
@@ -1443,7 +1444,7 @@
           src, cpi, plane, &cm->rst_info[plane], restore_types[r], tile_cost[r],
           &cpi->trial_frame_rst);
     }
-    if (plane == AOM_PLANE_Y)
+    if (ntiles > 1)
       cost_restore[RESTORE_SWITCHABLE] = search_switchable_restoration(
           src, cpi, plane, restore_types, tile_cost, &cm->rst_info[plane]);
     else
@@ -1463,9 +1464,8 @@
       assert(best_restore == force_restore_type ||
              best_restore == RESTORE_NONE);
     if (best_restore != RESTORE_SWITCHABLE) {
-      const int nt = (plane == AOM_PLANE_Y ? ntiles_y : ntiles_uv);
       memcpy(cm->rst_info[plane].restoration_type, restore_types[best_restore],
-             nt * sizeof(restore_types[best_restore][0]));
+             ntiles * sizeof(restore_types[best_restore][0]));
     }
   }
   /*