JNT_COMP: 4. add context and entropy read/write

Change-Id: I0e6f7ab981e31f7120105515f6204568b6dc82d3
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index bbe80d5..59fb670 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -1605,6 +1605,14 @@
     int16_t mode_ctx;
     write_ref_frames(cm, xd, w);
 
+#if CONFIG_JNT_COMP
+    if (has_second_ref(mbmi)) {
+      const int comp_index_ctx = get_comp_index_context(cm, xd);
+      aom_write(w, mbmi->compound_idx,
+                ec_ctx->compound_index_probs[comp_index_ctx]);
+    }
+#endif  // CONFIG_JNT_COMP
+
 #if CONFIG_COMPOUND_SINGLEREF
     if (!segfeature_active(seg, segment_id, SEG_LVL_REF_FRAME)) {
       // NOTE: Handle single ref comp mode
@@ -4570,6 +4578,11 @@
 
 #if !CONFIG_NEW_MULTISYMBOL
   update_skip_probs(cm, header_bc, counts);
+#if CONFIG_JNT_COMP
+  for (int k = 0; k < COMP_INDEX_CONTEXTS; ++k)
+    av1_cond_prob_diff_update(header_bc, &cm->fc->compound_index_probs[k],
+                              counts->compound_index[k], probwt);
+#endif  // CONFIG_JNT_COMP
 #endif
 
   if (!frame_is_intra_only(cm)) {
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 5403318..9099737 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -590,6 +590,13 @@
   const int x_mis = AOMMIN(bw, cm->mi_cols - mi_col);
   const int y_mis = AOMMIN(bh, cm->mi_rows - mi_row);
   av1_copy_frame_mvs(cm, mi, mi_row, mi_col, x_mis, y_mis);
+
+#if CONFIG_JNT_COMP
+  if (has_second_ref(mbmi)) {
+    const int comp_index_ctx = get_comp_index_context(cm, xd);
+    ++td->counts->compound_index[comp_index_ctx][mbmi->compound_idx];
+  }
+#endif  // CONFIG_JNT_COMP
 }
 
 #if CONFIG_MOTION_VAR && NC_MODE_INFO
@@ -3443,7 +3450,7 @@
   av1_setup_block_planes(xd, cm->subsampling_x, cm->subsampling_y);
 }
 
-#if !CONFIG_REF_ADAPT
+#if !CONFIG_REF_ADAPT && !CONFIG_JNT_COMP
 static int check_dual_ref_flags(AV1_COMP *cpi) {
   const int ref_flags = cpi->ref_frame_flags;
 
@@ -4162,12 +4169,14 @@
     if (is_alt_ref || !cpi->allow_comp_inter_inter)
 #endif  // CONFIG_BGSPRITE
       cm->reference_mode = SINGLE_REFERENCE;
+#if !CONFIG_JNT_COMP
     else if (mode_thrs[COMPOUND_REFERENCE] > mode_thrs[SINGLE_REFERENCE] &&
              mode_thrs[COMPOUND_REFERENCE] > mode_thrs[REFERENCE_MODE_SELECT] &&
              check_dual_ref_flags(cpi) && cpi->static_mb_pct == 100)
       cm->reference_mode = COMPOUND_REFERENCE;
     else if (mode_thrs[SINGLE_REFERENCE] > mode_thrs[REFERENCE_MODE_SELECT])
       cm->reference_mode = SINGLE_REFERENCE;
+#endif  // CONFIG_JNT_COMP
     else
       cm->reference_mode = REFERENCE_MODE_SELECT;
 #endif  // CONFIG_REF_ADAPT
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index ce4c03d..f4d001c 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -8462,6 +8462,14 @@
   }
 #endif  // CONFIG_COMPOUND_SINGLEREF
 
+#if CONFIG_JNT_COMP
+  if (is_comp_pred) {
+    const int comp_index_ctx = get_comp_index_context(cm, xd);
+    rd_stats->rate += av1_cost_bit(cm->fc->compound_index_probs[comp_index_ctx],
+                                   mbmi->compound_idx);
+  }
+#endif  // CONFIG_JNT_COMP
+
   if (this_mode == NEAREST_NEARESTMV) {
     if (mbmi_ext->ref_mv_count[ref_frame_type] > 0) {
       cur_mv[0] = mbmi_ext->ref_mv_stack[ref_frame_type][0].this_mv;