From 8282a18c36c672718b7c5c7ee0d6be9f6949792e Mon Sep 17 00:00:00 2001
From: Marko Viitanen <fador@iki.fi>
Date: Wed, 29 May 2019 16:17:02 +0300
Subject: [PATCH] Fixed headers and NAL writing to match the latest VTM master
 988c22cbb9c58584cac3ef0ec7794cafbea6dfd6

---
 src/bitstream.c               | 12 +-------
 src/cabac.h                   |  2 +-
 src/context.c                 |  2 +-
 src/encoder_state-bitstream.c | 36 ++++++++++++++++-------
 src/encoderstate.c            |  4 +--
 src/kvazaar.h                 | 54 ++++++++++++-----------------------
 src/nal.c                     | 23 +++++++++------
 7 files changed, 64 insertions(+), 69 deletions(-)

diff --git a/src/bitstream.c b/src/bitstream.c
index 78a203f2..680230f3 100644
--- a/src/bitstream.c
+++ b/src/bitstream.c
@@ -266,17 +266,7 @@ void kvz_bitstream_put_se(bitstream_t *stream, int32_t data)
 void kvz_bitstream_add_rbsp_trailing_bits(bitstream_t * const stream)
 {
   kvz_bitstream_put(stream, 1, 1);
-#if VERBOSE
-  printf("%-40s u(%d) : %d\n", "rbsp_stop_one_bit", 1, 1);
-#endif
-  if ((stream->cur_bit & 7) != 0) {
-#if VERBOSE
-    for (int i = 0; i < (8 - stream->cur_bit) & 7; i++) {
-      printf("%-40s u(%d) : %d\n", "rbsp_alignment_zero_bit", 1, 0);
-    }
-#endif
-    kvz_bitstream_put(stream, 0, (8 - stream->cur_bit) & 7);
-  }
+  kvz_bitstream_align_zero(stream);
 }
 
 /**
diff --git a/src/cabac.h b/src/cabac.h
index bc21ae5b..0873f709 100644
--- a/src/cabac.h
+++ b/src/cabac.h
@@ -133,7 +133,7 @@ void kvz_cabac_write_unary_max_symbol_ep(cabac_data_t *data, unsigned int symbol
   (ctx)->state[1]=(state >> 1) & (int)CTX_MASK_1;\
 }
 #define CTX_MPS(ctx) (CTX_STATE(ctx)>>7)
-#define CTX_LPS(ctx,range) ( ((((CTX_STATE(ctx)&0x80) ? (CTX_STATE(ctx)^0xff) : (CTX_STATE(ctx))) >>2)*(range>>5)>>1)+4  )
+#define CTX_LPS(ctx,range) ((uint8_t)( ((((CTX_STATE(ctx)&0x80) ? (CTX_STATE(ctx)^0xff) : (CTX_STATE(ctx))) >>2)*(range>>5)>>1)+4  ))
 #define CTX_UPDATE(ctx,bin) { \
   int rate0 = (ctx)->rate >> 4;\
   int rate1 = (ctx)->rate & 15;\
diff --git a/src/context.c b/src/context.c
index 92c3bf15..96dc4447 100644
--- a/src/context.c
+++ b/src/context.c
@@ -305,7 +305,7 @@ void kvz_ctx_init(cabac_ctx_t *ctx, uint32_t qp, uint32_t init_value, uint8_t ra
 {
   int slope = (init_value >> 4) * 5 - 45;
   int offset = ((init_value & 15) << 3) - 16;
-  int init_state = MIN(MAX(1, ((slope * (int)qp) >> 4) + offset), 126);
+  int init_state = ((slope * (int)qp) >> 4) + offset;
     
   const int p1 = g_inistateToCount[init_state < 0 ? 0 : init_state > 127 ? 127 : init_state];
 
diff --git a/src/encoder_state-bitstream.c b/src/encoder_state-bitstream.c
index ceee1d63..3c7e9a8a 100644
--- a/src/encoder_state-bitstream.c
+++ b/src/encoder_state-bitstream.c
@@ -248,8 +248,7 @@ static void encoder_state_write_bitstream_VUI(bitstream_t *stream,
   } else
     WRITE_U(stream, 0, 1, "aspect_ratio_info_present_flag");
 
-  //IF aspect ratio info
-  //ENDIF
+  WRITE_U(stream, 0, 1, "colour_description_present_flag");
 
   if (encoder->cfg.vui.overscan > 0) {
     WRITE_U(stream, 1, 1, "overscan_info_present_flag");
@@ -257,6 +256,7 @@ static void encoder_state_write_bitstream_VUI(bitstream_t *stream,
   } else
     WRITE_U(stream, 0, 1, "overscan_info_present_flag");
 
+
   //IF overscan info
   //ENDIF
 
@@ -353,6 +353,9 @@ static void encoder_state_write_bitstream_seq_parameter_set(bitstream_t* stream,
 
 
   */
+
+  WRITE_U(stream, 0, 4, "sps_decoding_parameter_set_id");
+
   WRITE_U(stream, 1, 3, "sps_max_sub_layers_minus1");
   WRITE_U(stream, 0, 5, "sps_reserved_zero_5bits");
 
@@ -515,10 +518,17 @@ if (encoder->scaling_list.enable) {
   //ENDIF
 
 
-  //WRITE_U(stream, 0, 1, "sps_strong_intra_smoothing_enable_flag");
-  WRITE_U(stream, 1, 1, "vui_parameters_present_flag");
+  WRITE_U(stream, encoder->vui.timing_info_present_flag, 1, "timing_info_present_flag");
+  if (encoder->vui.timing_info_present_flag) {
+    WRITE_U(stream, encoder->vui.num_units_in_tick, 32, "vui_num_units_in_tick");
+    WRITE_U(stream, encoder->vui.time_scale, 32, "vui_time_scale");
 
-  encoder_state_write_bitstream_VUI(stream, state);
+    WRITE_U(stream, 0, 1, "vui_hrd_parameters_present_flag");
+  }
+
+  WRITE_U(stream, 0, 1, "vui_parameters_present_flag");
+
+  //encoder_state_write_bitstream_VUI(stream, state);
 
   encoder_state_write_bitstream_SPS_extension(stream, state);
 
@@ -564,10 +574,8 @@ static void encoder_state_write_bitstream_pic_parameter_set(bitstream_t* stream,
   WRITE_U(stream, encoder->cfg.lossless, 1, "transquant_bypass_enabled_flag");
 
   
-  WRITE_U(stream, encoder->tiles_enable, 1, "tiles_enabled_flag");
-  //wavefronts
-  WRITE_U(stream, encoder->cfg.wpp, 1, "entropy_coding_sync_enabled_flag");
-
+  WRITE_U(stream, encoder->tiles_enable?0:1, 1, "single_tile_in_pic_flag");
+   
   if (encoder->tiles_enable) {
     WRITE_UE(stream, encoder->cfg.tiles_width_count  - 1, "num_tile_columns_minus1");
     WRITE_UE(stream, encoder->cfg.tiles_height_count - 1, "num_tile_rows_minus1");
@@ -586,8 +594,12 @@ static void encoder_state_write_bitstream_pic_parameter_set(bitstream_t* stream,
     WRITE_U(stream, 0, 1, "loop_filter_across_tiles_enabled_flag");
     
   }
-  
-  WRITE_U(stream, 0, 1, "pps_loop_filter_across_slices_enabled_flag");
+
+  WRITE_U(stream, 0, 1, "signalled_slice_id_flag");
+
+  //wavefronts
+  WRITE_U(stream, encoder->cfg.wpp, 1, "entropy_coding_sync_enabled_flag");
+
   WRITE_U(stream, 1, 1, "deblocking_filter_control_present_flag");
 
   //IF deblocking_filter
@@ -601,6 +613,8 @@ static void encoder_state_write_bitstream_pic_parameter_set(bitstream_t* stream,
        WRITE_SE(stream, encoder->cfg.deblock_tc, "tc_offset_div2");
     }
 
+
+  WRITE_U(stream, 0, 1, "pps_loop_filter_across_virtual_boundaries_disabled_flag");
     //ENDIF
   //ENDIF
   //WRITE_U(stream, 0, 1, "pps_scaling_list_data_present_flag");
diff --git a/src/encoderstate.c b/src/encoderstate.c
index 02ee5c68..4279baf2 100644
--- a/src/encoderstate.c
+++ b/src/encoderstate.c
@@ -1220,9 +1220,9 @@ static void encoder_state_init_new_frame(encoder_state_t * const state, kvz_pict
       state->frame->pictype = KVZ_NAL_CRA_NUT;
     }
   } else if (state->frame->poc < state->frame->irap_poc) {
-    state->frame->pictype = KVZ_NAL_RASL_R;
+    state->frame->pictype = KVZ_NAL_RASL;
   } else {
-    state->frame->pictype = KVZ_NAL_TRAIL_R;
+    state->frame->pictype = KVZ_NAL_TRAIL;
   }
 
   encoder_state_remove_refs(state);
diff --git a/src/kvazaar.h b/src/kvazaar.h
index 4c2fa15a..da6cbbe1 100644
--- a/src/kvazaar.h
+++ b/src/kvazaar.h
@@ -407,48 +407,32 @@ typedef struct kvz_picture {
  */
 enum kvz_nal_unit_type {
 
-  // Trailing pictures
+  // 
 
-  KVZ_NAL_TRAIL_N = 0,
-  KVZ_NAL_TRAIL_R = 0,
+  KVZ_NAL_PPS_NUT = 0,
+  KVZ_NAL_AUD_NUT = 1,
+  KVZ_NAL_PREFIX_SEI_NUT = 2,
+  KVZ_NAL_SUFFIX_SEI_NUT = 3,
 
-  KVZ_NAL_STSA_N = 1,
-  KVZ_NAL_STSA_R = 1,
+  // Coded slices
 
-  // Leading pictures
+  KVZ_NAL_TRAIL = 8,
+  KVZ_NAL_STSA = 9,
+  KVZ_NAL_RADL = 10,
+  KVZ_NAL_RASL = 11,
 
-  KVZ_NAL_RADL_N = 2,
-  KVZ_NAL_RADL_R = 2,
-
-  KVZ_NAL_RASL_N = 3,
-  KVZ_NAL_RASL_R = 3,
+  // non-VCL  
+  KVZ_NAL_SPS_NUT = 17,
 
+  KVZ_NAL_EOS_NUT = 18,
+  KVZ_NAL_EOB_NUT = 19,
+  KVZ_NAL_VPS_NUT = 20,
 
   // Intra random access point pictures
-
-  KVZ_NAL_IDR_W_RADL = 8,
-  KVZ_NAL_IDR_N_LP   = 9,
-
-  KVZ_NAL_CRA_NUT    = 10,
-
-  // Reserved IRAP
-
-  KVZ_NAL_RSV_IRAP_VCL11 = 11,
-  KVZ_NAL_RSV_IRAP_VCL12 = 12,
-
-
-  // non-VCL
-
-  KVZ_NAL_VPS_NUT = 16,
-  KVZ_NAL_SPS_NUT = 17,
-  KVZ_NAL_PPS_NUT = 18,
-
-  KVZ_NAL_AUD_NUT = 20,
-  KVZ_NAL_EOS_NUT = 21,
-  KVZ_NAL_EOB_NUT = 22,
-  KVZ_NAL_PREFIX_SEI_NUT = 23,
-  KVZ_NAL_SUFFIX_SEI_NUT = 24,
-  KVZ_NAL_FD_NUT  = 25,
+  KVZ_NAL_IDR_W_RADL = 24,
+  KVZ_NAL_IDR_N_LP = 25,
+  KVZ_NAL_CRA_NUT = 26,
+  KVZ_NAL_GRA_NUT = 27,
 
 };
 
diff --git a/src/nal.c b/src/nal.c
index 27b6675e..989be93c 100644
--- a/src/nal.c
+++ b/src/nal.c
@@ -52,19 +52,26 @@ void kvz_nal_write(bitstream_t * const bitstream, const uint8_t nal_type,
   kvz_bitstream_writebyte(bitstream, start_code_prefix_one_3bytes);
 
   // Handle header bits with full bytes instead of using bitstream
-  // forbidden_zero_flag(1) + nal_unit_type(6) + 1bit of nuh_layer_id
-  byte = nal_type << 1;
+  // forbidden_zero_flag(1) + nuh_temporal_id_plus1(3) + nal_unit_type(4)
+  uint8_t zero_tid_required_flag = 0;
+  if ((nal_type >= 16) && (nal_type <= 31)) {
+    zero_tid_required_flag = 1;
+  }
+  uint8_t nal_type_lsb = nal_type - (zero_tid_required_flag << 4);
+
+  byte = (zero_tid_required_flag<<7) + ((temporal_id + 1) << 4)  + nal_type_lsb;
   kvz_bitstream_writebyte(bitstream, byte);
 
-  // 5bits of nuh_layer_id + nuh_temporal_id_plus1(3)
-  byte = (temporal_id + 1) & 7;
+  // 7bits of nuh_layer_id_plus1
+  byte = 1<<1;
   kvz_bitstream_writebyte(bitstream, byte);
 
 #if VERBOSE
-  printf("%-40s u(%d) : %d\n", "forbidden_zero_bit", 1, 0);
-  printf("%-40s u(%d) : %d\n", "nal_unit_type", 6, nal_type);
-  printf("%-40s u(%d) : %d\n", "nuh_layer_id", 6, temporal_id);
-  printf("%-40s u(%d) : %d\n", "nuh_temporal_id", 3, temporal_id + 1);
+  printf("%-40s u(%d) : %d\n", "zero_tid_required_flag", 1, zero_tid_required_flag);
+  printf("%-40s u(%d) : %d\n", "nuh_temporal_id_plus1", 3, temporal_id + 1);
+  printf("%-40s u(%d) : %d\n", "nal_unit_type_lsb", 4, nal_type_lsb);
+  printf("%-40s u(%d) : %d\n", "nuh_layer_id_plus1", 6, 1);
+  
 #endif
 }