From 4ffdb358ed675f6b96fed03c1110739126b60bcc Mon Sep 17 00:00:00 2001
From: Yusuke Nakamura <muken.the.vfrmaniac@gmail.com>
Date: Mon, 3 Feb 2014 22:31:58 +0900
Subject: [PATCH] nal: Add the writing of an additional zero_byte before the
 start code.

Enables the output of spec-compliant byte streams, as the specification
notes that an additional zero_byte has to be added under certain
circuimstances.
---
 src/encoder.c | 12 ++++++------
 src/nal.c     |  9 ++++++++-
 src/nal.h     |  3 ++-
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/encoder.c b/src/encoder.c
index 28b1799a..0d701956 100644
--- a/src/encoder.c
+++ b/src/encoder.c
@@ -304,7 +304,7 @@ void encode_one_frame(encoder_control* encoder)
     bitstream_align(encoder->stream);
     bitstream_flush(encoder->stream);
     nal_write(encoder->output, encoder->stream->buffer,
-              encoder->stream->buffer_pos, 0, NAL_VPS_NUT, 0);
+              encoder->stream->buffer_pos, 0, NAL_VPS_NUT, 0, 1);
     bitstream_clear_buffer(encoder->stream);
 
     // Sequence Parameter Set (SPS)
@@ -312,7 +312,7 @@ void encode_one_frame(encoder_control* encoder)
     bitstream_align(encoder->stream);
     bitstream_flush(encoder->stream);
     nal_write(encoder->output, encoder->stream->buffer,
-              encoder->stream->buffer_pos, 0, NAL_SPS_NUT, 0);
+              encoder->stream->buffer_pos, 0, NAL_SPS_NUT, 0, 1);
     bitstream_clear_buffer(encoder->stream);
         
     // Picture Parameter Set (PPS)
@@ -320,7 +320,7 @@ void encode_one_frame(encoder_control* encoder)
     bitstream_align(encoder->stream);
     bitstream_flush(encoder->stream);
     nal_write(encoder->output, encoder->stream->buffer,
-              encoder->stream->buffer_pos, 0, NAL_PPS_NUT, 0);
+              encoder->stream->buffer_pos, 0, NAL_PPS_NUT, 0, 1);
     bitstream_clear_buffer(encoder->stream);
 
     // First slice is IDR
@@ -337,7 +337,7 @@ void encode_one_frame(encoder_control* encoder)
     bitstream_align(encoder->stream);
     bitstream_flush(encoder->stream);
     nal_write(encoder->output, encoder->stream->buffer,
-              encoder->stream->buffer_pos, 0, NAL_IDR_W_RADL, 0);
+              encoder->stream->buffer_pos, 0, NAL_IDR_W_RADL, 0, 0);
     bitstream_clear_buffer(encoder->stream);
   } else {
     cabac_start(&cabac);
@@ -354,7 +354,7 @@ void encode_one_frame(encoder_control* encoder)
     bitstream_align(encoder->stream);
     bitstream_flush(encoder->stream);
     nal_write(encoder->output, encoder->stream->buffer,
-              encoder->stream->buffer_pos, 0, NAL_TRAIL_R, 0);
+              encoder->stream->buffer_pos, 0, NAL_TRAIL_R, 0, 1);
     bitstream_clear_buffer(encoder->stream);
   }  
   
@@ -462,7 +462,7 @@ static void add_checksum(encoder_control* encoder)
   bitstream_align(encoder->stream);
   bitstream_flush(encoder->stream);
   nal_write(encoder->output, encoder->stream->buffer,
-            encoder->stream->buffer_pos, 0, NAL_SUFFIT_SEI_NUT, 0);
+            encoder->stream->buffer_pos, 0, NAL_SUFFIT_SEI_NUT, 0, 0);
   bitstream_clear_buffer(encoder->stream);
 }
 
diff --git a/src/nal.c b/src/nal.c
index 7f536b8d..8f036070 100644
--- a/src/nal.c
+++ b/src/nal.c
@@ -35,7 +35,9 @@
 /**
  * \brief Write a Network Abstraction Layer (NAL) packet to the output.
  */
-void nal_write(FILE *output, uint8_t *buffer, uint32_t buffer_len, uint8_t nal_ref, uint8_t nal_type, uint8_t temporal_id)
+void nal_write(FILE *output, uint8_t *buffer, uint32_t buffer_len,
+               uint8_t nal_ref, uint8_t nal_type, uint8_t temporal_id,
+               int long_start_code)
 {
   uint8_t byte;
   uint32_t i;
@@ -46,6 +48,11 @@ void nal_write(FILE *output, uint8_t *buffer, uint32_t buffer_len, uint8_t nal_r
   const uint8_t start_code_prefix_one_3bytes = 0x01;
   const uint8_t zero = 0x00;
 
+  // zero_byte (0x00) shall be present in the byte stream NALU of VPS, SPS
+  // and PPS, or the first NALU of an access unit
+  if(long_start_code)
+    fwrite(&zero, 1, 1, output);
+
   // start_code_prefix_one_3bytes
   fwrite(&zero, 1, 1, output);
   fwrite(&zero, 1, 1, output);
diff --git a/src/nal.h b/src/nal.h
index 8ccbb2a8..1395e4e5 100644
--- a/src/nal.h
+++ b/src/nal.h
@@ -93,7 +93,8 @@ enum {
 //////////////////////////////////////////////////////////////////////////
 // FUNCTIONS
 void nal_write(FILE *output, uint8_t *buffer, uint32_t buffer_len,
-               uint8_t nal_ref, uint8_t nal_type, uint8_t temporal_id);
+               uint8_t nal_ref, uint8_t nal_type, uint8_t temporal_id,
+               int long_start_code);
 void picture_checksum(const picture *pic,
                       unsigned char checksum_out[][SEI_HASH_MAX_LENGTH]);