From f33755ef1f61e589d7a88859b06a3a918158bd16 Mon Sep 17 00:00:00 2001
From: Geza Lore <gezalore@gmail.com>
Date: Tue, 26 Jan 2016 14:40:32 +0000
Subject: [PATCH] Always recode SUPERTX blocks.

There is still an assertion failure when tokenizing transform
coefficients for a supertx block, due to the eob not being set
consistently with the coefficients, so we always recode supertx blocks
for now. Also added further PICK_MODE_CONTEXT instances to avoid
potential clash between horizontal/vertical/split partition SUPERTX
trials.

Change-Id: I5f3da1fa0d8d20fc21face170487e1a285fd1cc6
---
 vp10/encoder/context_tree.c |  8 ++++++--
 vp10/encoder/context_tree.h |  4 +++-
 vp10/encoder/encodeframe.c  | 23 ++++++++++++++++++-----
 vp10/encoder/tokenize.c     |  1 +
 4 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/vp10/encoder/context_tree.c b/vp10/encoder/context_tree.c
index 64552abfc..3cd23ecf6 100644
--- a/vp10/encoder/context_tree.c
+++ b/vp10/encoder/context_tree.c
@@ -82,7 +82,9 @@ static void alloc_tree_contexts(VP10_COMMON *cm, PC_TREE *tree,
   alloc_mode_context(cm, num_4x4_blk/2, &tree->horizontal[0]);
   alloc_mode_context(cm, num_4x4_blk/2, &tree->vertical[0]);
 #ifdef CONFIG_SUPERTX
-  alloc_mode_context(cm, num_4x4_blk, &tree->super_tx);
+  alloc_mode_context(cm, num_4x4_blk, &tree->horizontal_supertx);
+  alloc_mode_context(cm, num_4x4_blk, &tree->vertical_supertx);
+  alloc_mode_context(cm, num_4x4_blk, &tree->split_supertx);
 #endif
 
   if (num_4x4_blk > 4) {
@@ -101,7 +103,9 @@ static void free_tree_contexts(PC_TREE *tree) {
   free_mode_context(&tree->vertical[0]);
   free_mode_context(&tree->vertical[1]);
 #ifdef CONFIG_SUPERTX
-  free_mode_context(&tree->super_tx);
+  free_mode_context(&tree->horizontal_supertx);
+  free_mode_context(&tree->vertical_supertx);
+  free_mode_context(&tree->split_supertx);
 #endif
 }
 
diff --git a/vp10/encoder/context_tree.h b/vp10/encoder/context_tree.h
index 80e1a4796..4fa58064e 100644
--- a/vp10/encoder/context_tree.h
+++ b/vp10/encoder/context_tree.h
@@ -88,7 +88,9 @@ typedef struct PC_TREE {
     PICK_MODE_CONTEXT *leaf_split[4];
   };
 #ifdef CONFIG_SUPERTX
-  PICK_MODE_CONTEXT super_tx;
+  PICK_MODE_CONTEXT horizontal_supertx;
+  PICK_MODE_CONTEXT vertical_supertx;
+  PICK_MODE_CONTEXT split_supertx;
 #endif
 } PC_TREE;
 
diff --git a/vp10/encoder/encodeframe.c b/vp10/encoder/encodeframe.c
index e56f753b9..1df5f740b 100644
--- a/vp10/encoder/encodeframe.c
+++ b/vp10/encoder/encodeframe.c
@@ -1336,6 +1336,7 @@ static void update_state_sb_supertx(VP10_COMP *cpi, ThreadData *td,
   PARTITION_TYPE partition = pc_tree->partitioning;
   BLOCK_SIZE subsize = get_subsize(bsize, partition);
   int i;
+  PICK_MODE_CONTEXT *pmc = NULL;
 
   if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols)
     return;
@@ -1355,6 +1356,7 @@ static void update_state_sb_supertx(VP10_COMP *cpi, ThreadData *td,
         update_state_supertx(cpi, td, &pc_tree->vertical[1],
                              mi_row, mi_col + hbs, subsize, output_enabled);
       }
+      pmc = &pc_tree->vertical_supertx;
       break;
     case PARTITION_HORZ:
       set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize);
@@ -1365,6 +1367,7 @@ static void update_state_sb_supertx(VP10_COMP *cpi, ThreadData *td,
         update_state_supertx(cpi, td, &pc_tree->horizontal[1], mi_row + hbs,
                              mi_col, subsize, output_enabled);
       }
+      pmc = &pc_tree->horizontal_supertx;
       break;
     case PARTITION_SPLIT:
       if (bsize == BLOCK_8X8) {
@@ -1385,16 +1388,25 @@ static void update_state_sb_supertx(VP10_COMP *cpi, ThreadData *td,
         update_state_sb_supertx(cpi, td, tile, mi_row + hbs, mi_col + hbs,
                                 subsize, output_enabled, pc_tree->split[3]);
       }
+      pmc = &pc_tree->split_supertx;
       break;
     default:
       assert(0);
   }
 
   for (i = 0; i < MAX_MB_PLANE; ++i) {
-    p[i].coeff = (&pc_tree->super_tx)->coeff_pbuf[i][1];
-    p[i].qcoeff = (&pc_tree->super_tx)->qcoeff_pbuf[i][1];
-    pd[i].dqcoeff = (&pc_tree->super_tx)->dqcoeff_pbuf[i][1];
-    p[i].eobs = (&pc_tree->super_tx)->eobs_pbuf[i][1];
+    if (pmc != NULL) {
+      p[i].coeff = pmc->coeff_pbuf[i][1];
+      p[i].qcoeff = pmc->qcoeff_pbuf[i][1];
+      pd[i].dqcoeff = pmc->dqcoeff_pbuf[i][1];
+      p[i].eobs = pmc->eobs_pbuf[i][1];
+    } else {
+      // These should never be used
+      p[i].coeff = NULL;
+      p[i].qcoeff = NULL;
+      pd[i].dqcoeff = NULL;
+      p[i].eobs = NULL;
+    }
   }
 }
 
@@ -2012,7 +2024,8 @@ static void encode_sb(VP10_COMP *cpi, ThreadData *td,
 
       set_offsets(cpi, tile, x, mi_row, mi_col, bsize);
       if (!x->skip) {
-        xd->mi[0]->mbmi.skip = 1;
+        // TODO(geza.lore): Investigate if this can be relaxed
+        x->skip_recode = 0;
         vp10_encode_sb_supertx(x, bsize);
         vp10_tokenize_sb_supertx(cpi, td, tp, !output_enabled, bsize);
       } else {
diff --git a/vp10/encoder/tokenize.c b/vp10/encoder/tokenize.c
index 707e42c3d..66bb990ea 100644
--- a/vp10/encoder/tokenize.c
+++ b/vp10/encoder/tokenize.c
@@ -574,6 +574,7 @@ static void tokenize_b(int plane, int block, int blk_row, int blk_col,
       pt = get_coef_context(nb, token_cache, c);
       v = qcoeff[scan[c]];
     }
+    assert(c < eob);
 
     vp10_get_token_extra(v, &token, &extra);
 
-- 
2.40.0