From f1fc0de2bfd73147d72a06dd8fbfc5f051a5cae5 Mon Sep 17 00:00:00 2001
From: Ari Koivula <ari@koivu.la>
Date: Tue, 31 Jan 2017 13:26:56 +0200
Subject: [PATCH] Write slice headers to the parent stream

Appending to the child stream doesn't work is the child is a leaf
slice state.

Simplifies flow by removing distinction between tile and slice. Now
that slice headers are written in the parent stream, there is zero
difference between tiles and slices from bitstream point of view.
---
 src/encoder_state-bitstream.c | 43 ++++++++++++++++++-----------------
 src/encoder_state-bitstream.h |  4 +++-
 src/encoderstate.c            |  6 -----
 3 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/src/encoder_state-bitstream.c b/src/encoder_state-bitstream.c
index 2347ff76..35bf77f2 100644
--- a/src/encoder_state-bitstream.c
+++ b/src/encoder_state-bitstream.c
@@ -674,10 +674,11 @@ static void encoder_state_write_bitstream_entry_points_write(bitstream_t * const
   }
 }
 
-void kvz_encoder_state_write_bitstream_slice_header(encoder_state_t * const state)
+void kvz_encoder_state_write_bitstream_slice_header(
+    struct bitstream_t * const stream,
+    struct encoder_state_t * const state)
 {
   const encoder_control_t * const encoder = state->encoder_control;
-  bitstream_t * const stream = &state->stream;
   int j;
   int ref_negative = 0;
   int ref_positive = 0;
@@ -893,12 +894,30 @@ static void add_checksum(encoder_state_t * const state)
   kvz_bitstream_add_rbsp_trailing_bits(stream);
 }
 
+static void encoder_state_write_slice_header(
+  bitstream_t * stream,
+  encoder_state_t * state)
+{
+  uint8_t nal_type = (state->frame->is_idr_frame ? KVZ_NAL_IDR_W_RADL : KVZ_NAL_TRAIL_R);
+
+  kvz_nal_write(stream, nal_type, 0, state->frame->first_nal);
+  state->frame->first_nal = false;
+
+  kvz_encoder_state_write_bitstream_slice_header(stream, state);
+  kvz_bitstream_add_rbsp_trailing_bits(stream);
+}
+
 /**
  * \brief Move child state bitstreams to the parent stream.
  */
 static void encoder_state_write_bitstream_children(encoder_state_t * const state)
 {
+  // Write Slice headers to the parent stream instead of the child stream
+  // in case the child stream is a leaf with something in it already.
   for (int i = 0; state->children[i].encoder_control; ++i) {
+    if (state->children[i].type == ENCODER_STATE_TYPE_SLICE) {
+      encoder_state_write_slice_header(&state->stream, &state->children[i]);
+    }
     kvz_encoder_state_write_bitstream(&state->children[i]);
     kvz_bitstream_move(&state->stream, &state->children[i].stream);
   }
@@ -976,22 +995,6 @@ static void encoder_state_write_bitstream_main(encoder_state_t * const state)
   state->frame->cur_gop_bits_coded += newpos - curpos;
 }
 
-static void encoder_state_write_bitstream_tile(encoder_state_t * const state)
-{
-  encoder_state_write_bitstream_children(state);
-}
-
-static void encoder_state_write_bitstream_slice(encoder_state_t * const state)
-{
-  uint8_t nal_type = (state->frame->is_idr_frame ? KVZ_NAL_IDR_W_RADL : KVZ_NAL_TRAIL_R);
-  kvz_nal_write(state->cabac.stream, nal_type, 0, state->frame->first_nal);
-
-  kvz_encoder_state_write_bitstream_slice_header(state);
-  kvz_bitstream_add_rbsp_trailing_bits(&state->stream);
-
-  encoder_state_write_bitstream_children(state);
-}
-
 void kvz_encoder_state_write_bitstream(encoder_state_t * const state)
 {
   if (!state->is_leaf) {
@@ -1000,10 +1003,8 @@ void kvz_encoder_state_write_bitstream(encoder_state_t * const state)
         encoder_state_write_bitstream_main(state);
         break;
       case ENCODER_STATE_TYPE_TILE:
-        encoder_state_write_bitstream_tile(state);
-        break;
       case ENCODER_STATE_TYPE_SLICE:
-        encoder_state_write_bitstream_slice(state);
+        encoder_state_write_bitstream_children(state);
         break;
       default:
         fprintf(stderr, "Unsupported node type %c!\n", state->type);
diff --git a/src/encoder_state-bitstream.h b/src/encoder_state-bitstream.h
index 947c56c8..0c724349 100644
--- a/src/encoder_state-bitstream.h
+++ b/src/encoder_state-bitstream.h
@@ -35,7 +35,9 @@ struct encoder_state_t;
 
 struct bitstream_t;
 
-void kvz_encoder_state_write_bitstream_slice_header(struct encoder_state_t * const state);
+void kvz_encoder_state_write_bitstream_slice_header(
+    struct bitstream_t * const stream,
+    struct encoder_state_t * const state);
 void kvz_encoder_state_write_bitstream(struct encoder_state_t * const state);
 void kvz_encoder_state_write_bitstream_leaf(struct encoder_state_t * const state);
 void kvz_encoder_state_worker_write_bitstream(void * opaque);
diff --git a/src/encoderstate.c b/src/encoderstate.c
index b232d712..1308037b 100644
--- a/src/encoderstate.c
+++ b/src/encoderstate.c
@@ -316,12 +316,6 @@ static void encoder_state_worker_encode_lcu(void * opaque)
   //Now write data to bitstream (required to have a correct CABAC state)
   const uint64_t existing_bits = kvz_bitstream_tell(&state->stream);
   
-  //First LCU, and we are in a slice. We need a slice header
-  if (state->type == ENCODER_STATE_TYPE_SLICE && lcu->index == 0) {
-    kvz_encoder_state_write_bitstream_slice_header(state);
-    kvz_bitstream_add_rbsp_trailing_bits(&state->stream); 
-  }
-  
   //Encode SAO
   if (encoder->sao_enable) {
     encode_sao(state, lcu->position.x, lcu->position.y, &frame->sao_luma[lcu->position.y * frame->width_in_lcu + lcu->position.x], &frame->sao_chroma[lcu->position.y * frame->width_in_lcu + lcu->position.x]);