libhb: add generic code for converting NAL

unit bitstreams (H.264, HEVC) to and from
Annex B format.

Only used by QSV for now, but could
always come in handy down the road.

May also fix an issue introduced
by the recent QSV cleanup patchset.


git-svn-id: svn://svn.handbrake.fr/HandBrake/trunk@6156 b64f7644-9d1e-0410-96f1-a4d463321fa5

3 files changed, 254 insertions, 89 deletions

diff --git a/libhb/enc_qsv.c b/libhb/enc_qsv.c
index e036c8115..bc9bb4305 100644
--- a/libhb/enc_qsv.c
+++ b/libhb/enc_qsv.c
@@ -29,6 +29,7 @@ THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #ifdef USE_QSV
 
 #include "hb.h"
+#include "nal_units.h"
 #include "qsv_common.h"
 #include "qsv_memory.h"
 #include "h264_common.h"
@@ -50,9 +51,6 @@ THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #define FRAME_INFO_SIZE (1 << (FRAME_INFO_MIN2 - FRAME_INFO_MAX2 + 1))
 #define FRAME_INFO_MASK (FRAME_INFO_SIZE - 1)
 
-int  nal_find_start_code(uint8_t**, size_t*);
-void parse_nalus        (uint8_t*,  size_t, hb_buffer_t*);
-
 int  encqsvInit (hb_work_object_t*, hb_job_t*);
 int  encqsvWork (hb_work_object_t*, hb_buffer_t**, hb_buffer_t**);
 void encqsvClose(hb_work_object_t*);
@@ -1461,20 +1459,17 @@ static int qsv_frame_is_key(mfxU16 FrameType)
 
 static void qsv_bitstream_slurp(hb_work_private_t *pv, mfxBitstream *bs)
 {
-    /* allocate additional data for parse_nalus */
-    hb_buffer_t *buf = hb_buffer_init(bs->DataLength * 2);
-    if (buf == NULL)
-    {
-        hb_error("encqsv: hb_buffer_init failed");
-        goto fail;
-    }
-    buf->size = 0;
-
     /*
      * we need to convert the encoder's Annex B output
      * to an MP4-compatible format (ISO/IEC 14496-15).
      */
-    parse_nalus(bs->Data + bs->DataOffset, bs->DataLength, buf);
+    hb_buffer_t *buf = hb_nal_bitstream_annexb_to_mp4(bs->Data + bs->DataOffset,
+                                                      bs->DataLength);
+    if (buf == NULL)
+    {
+        hb_error("encqsv: hb_nal_bitstream_annexb_to_mp4 failed");
+        goto fail;
+    }
     bs->DataLength = bs->DataOffset = 0;
     bs->MaxLength  = pv->job->qsv.ctx->enc_space->p_buf_max_size;
 
@@ -1870,80 +1865,4 @@ fail:
     return HB_WORK_ERROR;
 }
 
-int nal_find_start_code(uint8_t **pb, size_t *size)
-{
-    if ((int)*size < 4)
-    {
-        return 0;
-    }
-
-    // find start code by MSDK, see ff_prefix_code[]
-    while ((4 <= *size) && ((*pb)[0] || (*pb)[1] || (*pb)[2] != 1))
-    {
-        *pb   += 1;
-        *size -= 1;
-    }
-
-    if (4 <= *size)
-    {
-        return (((*pb)[0] << 24) |
-                ((*pb)[1] << 16) |
-                ((*pb)[2] <<  8) |
-                ((*pb)[3]));
-    }
-
-    return 0;
-}
-
-void parse_nalus(uint8_t *nal_inits, size_t length, hb_buffer_t *buf)
-{
-    uint8_t *offset = nal_inits;
-    size_t   size   = length;
-
-    if (!nal_find_start_code(&offset, &size))
-    {
-        size = 0;
-    }
-
-    while (size > 0)
-    {
-        uint8_t *current_nal  = offset + sizeof(ff_prefix_code) - 1;
-        size_t   current_size = size   - sizeof(ff_prefix_code);
-        uint8_t *next_offset  = current_nal + 1;
-        size_t   next_size    = current_size;
-
-        if (!nal_find_start_code(&next_offset, &next_size))
-        {
-            size          = 0;
-            current_size += 1;
-        }
-        else
-        {
-            if (next_offset > 0 && *(next_offset - 1))
-            {
-                current_size += 1;
-            }
-            current_size -= next_size;
-        }
-
-        char size_position[4];
-        size_position[1] = (current_size >> 24) & 0xFF;
-        size_position[1] = (current_size >> 16) & 0xFF;
-        size_position[2] = (current_size >>  8) & 0xFF;
-        size_position[3] = (current_size      ) & 0xFF;
-
-        memcpy(buf->data + buf->size, &size_position, sizeof(size_position));
-        buf->size += sizeof(size_position);
-
-        memcpy(buf->data + buf->size, current_nal, current_size);
-        buf->size += current_size;
-
-        if (size)
-        {
-            size   = next_size;
-            offset = next_offset;
-        }
-    }
-}
-
 #endif // USE_QSV
diff --git a/libhb/nal_units.c b/libhb/nal_units.c
new file mode 100644
index 000000000..fc87b7798
--- /dev/null
+++ b/libhb/nal_units.c
@@ -0,0 +1,195 @@
+/* nal_units.c
+ *
+ * Copyright (c) 2003-2014 HandBrake Team
+ * This file is part of the HandBrake source code.
+ * Homepage: <http://handbrake.fr/>.
+ * It may be used under the terms of the GNU General Public License v2.
+ * For full terms see the file COPYING file or visit http://www.gnu.org/licenses/gpl-2.0.html
+ */
+
+#include <stdint.h>
+#include <string.h>
+
+#include "common.h"
+#include "nal_units.h"
+
+static const uint8_t hb_annexb_startcode[] = { 0x00, 0x00, 0x00, 0x01, };
+
+size_t hb_nal_unit_write_annexb(uint8_t *buf,
+                                const uint8_t *nal_unit,
+                                const size_t nal_unit_size)
+{
+    if (buf != NULL)
+    {
+        memcpy(buf, hb_annexb_startcode, sizeof(hb_annexb_startcode));
+        memcpy(buf + sizeof(hb_annexb_startcode), nal_unit, nal_unit_size);
+    }
+
+    return sizeof(hb_annexb_startcode) + nal_unit_size;
+}
+
+size_t hb_nal_unit_write_isomp4(uint8_t *buf,
+                                const uint8_t *nal_unit,
+                                const size_t nal_unit_size)
+{
+    int i;
+    uint8_t length[4]; // 4-byte length replaces Annex B start code prefix
+
+    if (buf != NULL)
+    {
+        for (i = 0; i < sizeof(length); i++)
+        {
+            length[i] = (nal_unit_size >> (8 * (sizeof(length) - 1 - i))) & 0xff;
+        }
+
+        memcpy(buf, &length[0], sizeof(length));
+        memcpy(buf + sizeof(length), nal_unit, nal_unit_size);
+    }
+
+    return sizeof(length) + nal_unit_size;
+}
+
+uint8_t* hb_annexb_find_next_nalu(const uint8_t *start, size_t *size)
+{
+    uint8_t *nal = NULL;
+    uint8_t *buf = (uint8_t*)start;
+    uint8_t *end = (uint8_t*)start + *size;
+
+    /* Look for an Annex B start code prefix (3-byte sequence == 1) */
+    while (end - buf > 3)
+    {
+        if (!buf[0] && !buf[1] && buf[2] == 1)
+        {
+            nal = (buf += 3); // NAL unit begins after start code
+            break;
+        }
+        buf++;
+    }
+
+    if (nal == NULL)
+    {
+        *size = 0;
+        return NULL;
+    }
+
+    /*
+     * Start code prefix found, look for the next one to determine the size
+     *
+     * A 4-byte sequence == 1 is also a start code, so check for a 3-byte
+     * sequence == 0 too (start code emulation prevention will prevent such a
+     * sequence from occurring outside of a start code prefix)
+     */
+    while (end - buf > 3)
+    {
+        if (!buf[0] && !buf[1] && (!buf[2] || buf[2] == 1))
+        {
+            end = buf;
+            break;
+        }
+        buf++;
+    }
+
+    *size = end - nal;
+    return  nal;
+}
+
+hb_buffer_t* hb_nal_bitstream_annexb_to_mp4(const uint8_t *data,
+                                            const size_t size)
+{
+    hb_buffer_t *out;
+    uint8_t *buf, *end;
+    size_t out_size, buf_size;
+
+    out_size = 0;
+    buf_size = size;
+    buf      = (uint8_t*)data;
+    end      = (uint8_t*)data + size;
+
+    while ((buf = hb_annexb_find_next_nalu(buf, &buf_size)) != NULL)
+    {
+        out_size += hb_nal_unit_write_isomp4(NULL, buf, buf_size);
+        buf_size  = end - buf;
+    }
+
+    out = hb_buffer_init(out_size);
+    if (out == NULL)
+    {
+        hb_error("hb_nal_bitstream_annexb_to_mp4: hb_buffer_init failed");
+        return NULL;
+    }
+
+    out_size = 0;
+    buf_size = size;
+    buf      = (uint8_t*)data;
+    end      = (uint8_t*)data + size;
+
+    while ((buf = hb_annexb_find_next_nalu(buf, &buf_size)) != NULL)
+    {
+        out_size += hb_nal_unit_write_isomp4(out->data + out_size, buf, buf_size);
+        buf_size  = end - buf;
+    }
+
+    return out;
+}
+
+static size_t mp4_nal_unit_length(const uint8_t *data,
+                                  const size_t nal_length_size,
+                                  size_t *nal_unit_length)
+{
+    uint8_t i;
+
+    /* In MP4, NAL units are preceded by a 2-4 byte length field */
+    for (i = 0, *nal_unit_length = 0; i < nal_length_size; i++)
+    {
+        *nal_unit_length |= data[i] << (8 * (nal_length_size - 1 - i));
+    }
+
+    return nal_length_size;
+}
+
+hb_buffer_t* hb_nal_bitstream_mp4_to_annexb(const uint8_t *data,
+                                            const size_t size,
+                                            const uint8_t nal_length_size)
+{
+    hb_buffer_t *out;
+    uint8_t *buf, *end;
+    size_t out_size, nal_size;
+
+    out_size = 0;
+    buf      = (uint8_t*)data;
+    end      = (uint8_t*)data + size;
+
+    while (end - buf > nal_length_size)
+    {
+        buf += mp4_nal_unit_length(buf, nal_length_size, &nal_size);
+        if (end - buf < nal_size)
+        {
+            hb_log("hb_nal_bitstream_mp4_to_annexb: truncated bitstream"
+                   " (remaining: %lu, expected: %lu)", end - buf, nal_size);
+            return NULL;
+        }
+
+        out_size += hb_nal_unit_write_annexb(NULL, buf, nal_size);
+        buf      += nal_size;
+    }
+
+    out = hb_buffer_init(out_size);
+    if (out == NULL)
+    {
+        hb_error("hb_nal_bitstream_mp4_to_annexb: hb_buffer_init failed");
+        return NULL;
+    }
+
+    out_size = 0;
+    buf      = (uint8_t*)data;
+    end      = (uint8_t*)data + size;
+
+    while (end - buf > nal_length_size)
+    {
+        buf      += mp4_nal_unit_length(buf, nal_length_size, &nal_size);
+        out_size += hb_nal_unit_write_annexb(out->data + out_size, buf, nal_size);
+        buf      += nal_size;
+    }
+
+    return out;
+}
diff --git a/libhb/nal_units.h b/libhb/nal_units.h
new file mode 100644
index 000000000..734362de9
--- /dev/null
+++ b/libhb/nal_units.h
@@ -0,0 +1,51 @@
+/* nal_units.h
+ *
+ * Copyright (c) 2003-2014 HandBrake Team
+ * This file is part of the HandBrake source code.
+ * Homepage: <http://handbrake.fr/>.
+ * It may be used under the terms of the GNU General Public License v2.
+ * For full terms see the file COPYING file or visit http://www.gnu.org/licenses/gpl-2.0.html
+ */
+
+#ifndef HB_NAL_UNITS_H
+#define HB_NAL_UNITS_H
+
+#include <stdint.h>
+
+#include "common.h"
+
+/*
+ * Write a NAL unit of the specified size to the provided
+ * output buffer, using the requested output format.
+ * Returns the amount (in bytes) of data written to the buffer.
+ *
+ * The provided NAL unit must start with the NAL unit header.
+ *
+ * Note: the buffer is assumed to be large enough to hold the NAL unit
+ * as well as any additonal data the function may prepend/append to it.
+ *
+ * The caller may check the minimum required buffer size by passing a
+ * NULL buffer to the fucntion and checking the returned size value.
+ */
+size_t hb_nal_unit_write_annexb(uint8_t *buf, const uint8_t *nal_unit, const size_t nal_unit_size);
+size_t hb_nal_unit_write_isomp4(uint8_t *buf, const uint8_t *nal_unit, const size_t nal_unit_size);
+
+/*
+ * Search the provided data buffer for NAL units in Annex B format.
+ *
+ * Returns a pointer to the start (start code prefix excluded) of the
+ * first NAL unit found, or NULL if no NAL units were found in the buffer.
+ *
+ * On input,  size holds the length of the provided data buffer.
+ * On output, size holds the length of the returned NAL unit.
+ */
+uint8_t* hb_annexb_find_next_nalu(const uint8_t *start, size_t *size);
+
+/*
+ * Returns a newly-allocated buffer holding a copy of the provided
+ * NAL unit bitstream data, converted to the requested format.
+ */
+hb_buffer_t* hb_nal_bitstream_annexb_to_mp4(const uint8_t *data, const size_t size);
+hb_buffer_t* hb_nal_bitstream_mp4_to_annexb(const uint8_t *data, const size_t size, const uint8_t nal_length_size);
+
+#endif // HB_NAL_UNITS_H