nalutils: Introduce NAL writer helper methods

Add helper methods for writing h264 and h265 NAL

gst-libs/gst/codecparsers/nalutils.c

@@ -36,6 +36,7 @@
 #endif
 
 #include "nalutils.h"
+#include <string.h>
 
 /* Compute Ceil(Log2(v)) */
 /* Derived from branchless code for integer log2(v) from:
@@ -304,3 +305,220 @@ scan_for_start_codes (const guint8 * data, guint size)
   return gst_byte_reader_masked_scan_uint32 (&br, 0xffffff00, 0x00000100,
       0, size);
 }
+
+void
+nal_writer_init (NalWriter * nw, guint nal_prefix_size, gboolean packetized)
+{
+  g_return_if_fail (nw != NULL);
+  g_return_if_fail ((packetized && nal_prefix_size > 1 && nal_prefix_size < 5)
+      || (!packetized && (nal_prefix_size == 3 || nal_prefix_size == 4)));
+
+  gst_bit_writer_init (&nw->bw);
+  nw->nal_prefix_size = nal_prefix_size;
+  nw->packetized = packetized;
+}
+
+void
+nal_writer_reset (NalWriter * nw)
+{
+  g_return_if_fail (nw != NULL);
+
+  gst_bit_writer_reset (&nw->bw);
+  memset (nw, 0, sizeof (NalWriter));
+}
+
+gboolean
+nal_writer_do_rbsp_trailing_bits (NalWriter * nw)
+{
+  g_return_val_if_fail (nw != NULL, FALSE);
+
+  if (!gst_bit_writer_put_bits_uint8 (&nw->bw, 1, 1)) {
+    GST_WARNING ("Cannot put trailing bits");
+    return FALSE;
+  }
+
+  if (!gst_bit_writer_align_bytes (&nw->bw, 0)) {
+    GST_WARNING ("Cannot put align bits");
+    return FALSE;
+  }
+
+  return TRUE;
+}
+
+GstMemory *
+nal_writer_reset_and_get_memory (NalWriter * nw)
+{
+  GstBitWriter bw;
+  gint i;
+  guint8 *src, *dst;
+  gsize size;
+  GstMemory *ret = NULL;
+  gpointer data;
+
+  g_return_val_if_fail (nw != NULL, NULL);
+
+  if ((GST_BIT_WRITER_BIT_SIZE (&nw->bw) >> 3) == 0) {
+    GST_WARNING ("No written byte");
+    goto done;
+  }
+
+  if ((GST_BIT_WRITER_BIT_SIZE (&nw->bw) & 0x7) != 0) {
+    GST_WARNING ("Written stream is not byte aligned");
+    if (!nal_writer_do_rbsp_trailing_bits (nw))
+      goto done;
+  }
+
+  /* scan to put emulation_prevention_three_byte */
+  size = GST_BIT_WRITER_BIT_SIZE (&nw->bw) >> 3;
+  src = GST_BIT_WRITER_DATA (&nw->bw);
+
+  gst_bit_writer_init_with_size (&bw, size + nw->nal_prefix_size, FALSE);
+  for (i = 0; i < nw->nal_prefix_size - 1; i++)
+    gst_bit_writer_put_bits_uint8 (&bw, 0, 8);
+  gst_bit_writer_put_bits_uint8 (&bw, 1, 8);
+
+  for (i = 0; i < size; i++) {
+    guint pos = (GST_BIT_WRITER_BIT_SIZE (&bw) >> 3);
+    dst = GST_BIT_WRITER_DATA (&bw);
+    if (pos >= nw->nal_prefix_size + 2 &&
+        dst[pos - 2] == 0 && dst[pos - 1] == 0 && src[i] <= 0x3) {
+      gst_bit_writer_put_bits_uint8 (&bw, 0x3, 8);
+    }
+
+    gst_bit_writer_put_bits_uint8 (&bw, src[i], 8);
+  }
+
+  size = bw.bit_size >> 3;
+  data = gst_bit_writer_reset_and_get_data (&bw);
+  ret = gst_memory_new_wrapped (0, data, size, 0, size, data, g_free);
+
+  if (nw->packetized) {
+    GstMapInfo info;
+
+    gst_memory_map (ret, &info, GST_MAP_WRITE);
+
+    size = info.size - nw->nal_prefix_size;
+
+    switch (nw->nal_prefix_size) {
+      case 1:
+        GST_WRITE_UINT8 (info.data, size);
+        break;
+      case 2:
+        GST_WRITE_UINT16_BE (info.data, size);
+        break;
+      case 3:
+        GST_WRITE_UINT24_BE (info.data, size);
+        break;
+      case 4:
+        GST_WRITE_UINT32_BE (info.data, size);
+        break;
+      default:
+        g_assert_not_reached ();
+        break;
+    }
+
+    gst_memory_unmap (ret, &info);
+  }
+
+done:
+  gst_bit_writer_reset (&nw->bw);
+
+  return ret;
+}
+
+gboolean
+nal_writer_put_bits_uint8 (NalWriter * nw, guint8 value, guint nbits)
+{
+  g_return_val_if_fail (nw != NULL, FALSE);
+
+  if (!gst_bit_writer_put_bits_uint8 (&nw->bw, value, nbits))
+    return FALSE;
+
+  return TRUE;
+}
+
+gboolean
+nal_writer_put_bits_uint16 (NalWriter * nw, guint16 value, guint nbits)
+{
+  g_return_val_if_fail (nw != NULL, FALSE);
+
+  if (!gst_bit_writer_put_bits_uint16 (&nw->bw, value, nbits))
+    return FALSE;
+
+  return TRUE;
+}
+
+gboolean
+nal_writer_put_bits_uint32 (NalWriter * nw, guint32 value, guint nbits)
+{
+  g_return_val_if_fail (nw != NULL, FALSE);
+
+  if (!gst_bit_writer_put_bits_uint32 (&nw->bw, value, nbits))
+    return FALSE;
+
+  return TRUE;
+}
+
+gboolean
+nal_writer_put_bytes (NalWriter * nw, const guint8 * data, guint nbytes)
+{
+  g_return_val_if_fail (nw != NULL, FALSE);
+  g_return_val_if_fail (data != NULL, FALSE);
+  g_return_val_if_fail (nbytes != 0, FALSE);
+
+  if (!gst_bit_writer_put_bytes (&nw->bw, data, nbytes))
+    return FALSE;
+
+  return TRUE;
+}
+
+gboolean
+nal_writer_put_ue (NalWriter * nw, guint32 value)
+{
+  guint leading_zeros;
+  guint rest;
+
+  g_return_val_if_fail (nw != NULL, FALSE);
+
+  count_exp_golomb_bits (value, &leading_zeros, &rest);
+
+  /* write leading zeros */
+  if (leading_zeros) {
+    if (!nal_writer_put_bits_uint32 (nw, 0, leading_zeros))
+      return FALSE;
+  }
+
+  /* write the rest */
+  if (!nal_writer_put_bits_uint32 (nw, value + 1, rest))
+    return FALSE;
+
+  return TRUE;
+}
+
+gboolean
+count_exp_golomb_bits (guint32 value, guint * leading_zeros, guint * rest)
+{
+  guint32 x;
+  guint count = 0;
+
+  /* https://en.wikipedia.org/wiki/Exponential-Golomb_coding */
+  /* count bits of value + 1 */
+  x = value + 1;
+  while (x) {
+    count++;
+    x >>= 1;
+  }
+
+  if (leading_zeros) {
+    if (count > 1)
+      *leading_zeros = count - 1;
+    else
+      *leading_zeros = 0;
+  }
+
+  if (rest) {
+    *rest = count;
+  }
+
+  return TRUE;
+}

gst-libs/gst/codecparsers/nalutils.h

@@ -37,6 +37,7 @@
 
 #include <gst/base/gstbytereader.h>
 #include <gst/base/gstbitreader.h>
+#include <gst/base/gstbitwriter.h>
 #include <string.h>
 
 guint ceil_log2 (guint32 v);
@@ -54,6 +55,14 @@ typedef struct
   guint64 cache;                /* cached bytes */
 } NalReader;
 
+typedef struct
+{
+  GstBitWriter bw;
+
+  guint nal_prefix_size;
+  gboolean packetized;
+} NalWriter;
+
 G_GNUC_INTERNAL
 void nal_reader_init (NalReader * nr, const guint8 * data, guint size);
 
@@ -182,3 +191,68 @@ gboolean nal_reader_get_se (NalReader * nr, gint32 * val);
 
 G_GNUC_INTERNAL
 gint scan_for_start_codes (const guint8 * data, guint size);
+
+G_GNUC_INTERNAL
+void nal_writer_init (NalWriter * nw, guint nal_prefix_size, gboolean packetized);
+
+G_GNUC_INTERNAL
+void nal_writer_reset (NalWriter * nw);
+
+G_GNUC_INTERNAL
+gboolean nal_writer_do_rbsp_trailing_bits (NalWriter * nw);
+
+G_GNUC_INTERNAL
+GstMemory * nal_writer_reset_and_get_memory (NalWriter * nw);
+
+G_GNUC_INTERNAL
+gboolean nal_writer_put_bits_uint8 (NalWriter * nw, guint8 value, guint nbits);
+
+G_GNUC_INTERNAL
+gboolean nal_writer_put_bits_uint16 (NalWriter * nw, guint16 value, guint nbits);
+
+G_GNUC_INTERNAL
+gboolean nal_writer_put_bits_uint32 (NalWriter * nw, guint32 value, guint nbits);
+
+G_GNUC_INTERNAL
+gboolean nal_writer_put_bytes (NalWriter * nw, const guint8 * data, guint nbytes);
+
+G_GNUC_INTERNAL
+gboolean nal_writer_put_ue (NalWriter * nw, guint32 value);
+
+G_GNUC_INTERNAL
+gboolean count_exp_golomb_bits (guint32 value, guint * leading_zeros, guint * rest);
+
+#define WRITE_UINT8(nw, val, nbits) { \
+  if (!nal_writer_put_bits_uint8 (nw, val, nbits)) { \
+    GST_WARNING ("failed to write uint8, nbits: %d", nbits); \
+    goto error; \
+  } \
+}
+
+#define WRITE_UINT16(nw, val, nbits) { \
+  if (!nal_writer_put_bits_uint16 (nw, val, nbits)) { \
+    GST_WARNING ("failed to write uint16, nbits: %d", nbits); \
+    goto error; \
+  } \
+}
+
+#define WRITE_UINT32(nw, val, nbits) { \
+  if (!nal_writer_put_bits_uint32 (nw, val, nbits)) { \
+    GST_WARNING ("failed to write uint32, nbits: %d", nbits); \
+    goto error; \
+  } \
+}
+
+#define WRITE_BYTES(nw, data, nbytes) { \
+  if (!nal_writer_put_bytes (nw, data, nbytes)) { \
+    GST_WARNING ("failed to write bytes, nbits: %d", nbytes); \
+    goto error; \
+  } \
+}
+
+#define WRITE_UE(nw, val) { \
+  if (!nal_writer_put_ue (nw, val)) { \
+    GST_WARNING ("failed to write ue"); \
+    goto error; \
+  } \
+}

tests/check/libs/nalutils.c

@@ -0,0 +1,119 @@
+/* GStreamer
+ *
+ * Copyright (C) 2020 Seungha Yang <seungha@centricular.com>
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Library General Public
+ * License as published by the Free Software Foundation; either
+ * version 2 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Library General Public License for more details.
+ *
+ * You should have received a copy of the GNU Library General Public
+ * License along with this library; if not, write to the
+ * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
+ * Boston, MA 02110-1301, USA.
+ */
+
+#include <gst/check/gstcheck.h>
+#include <gst/codecparsers/nalutils.h>
+#include <string.h>
+
+GST_START_TEST (test_nal_writer_init)
+{
+  NalWriter nw;
+
+  /* init with invalid params */
+  ASSERT_CRITICAL (nal_writer_init (&nw, 0, TRUE));
+  ASSERT_CRITICAL (nal_writer_init (&nw, 0, FALSE));
+  ASSERT_CRITICAL (nal_writer_init (&nw, 5, TRUE));
+  ASSERT_CRITICAL (nal_writer_init (&nw, 5, FALSE));
+
+  nal_writer_init (&nw, 4, FALSE);
+  nal_writer_reset (&nw);
+}
+
+GST_END_TEST;
+
+GST_START_TEST (test_nal_writer_emulation_preventation)
+{
+  NalWriter nw;
+  gint i, j;
+  static guint8 patterns[][3] = {
+    {0x00, 0x00, 0x00},
+    {0x00, 0x00, 0x01},
+    {0x00, 0x00, 0x02},
+    {0x00, 0x00, 0x03},
+  };
+  static guint8 expected_rst[][4] = {
+    {0x00, 0x00, 0x03, 0x00},
+    {0x00, 0x00, 0x03, 0x01},
+    {0x00, 0x00, 0x03, 0x02},
+    {0x00, 0x00, 0x03, 0x03},
+  };
+
+  /* Within the NAL unit, the following three-byte sequences shall not occur
+   * at any byte-aligned position:
+   * – 0x000000
+   * – 0x000001
+   * – 0x000002
+   * Within the NAL unit, any four-byte sequence that starts with 0x000003
+   * other than the following sequences shall not occur
+   * at any byte-aligned position:
+   * – 0x00000300
+   * – 0x00000301
+   * – 0x00000302
+   * – 0x00000303
+   */
+
+  for (i = 0; i < G_N_ELEMENTS (patterns); i++) {
+    GstMemory *mem;
+    GstMapInfo info;
+
+    nal_writer_init (&nw, 4, FALSE);
+
+    /* forbidden_zero_bit */
+    fail_unless (nal_writer_put_bits_uint8 (&nw, 0, 1));
+    /* nal_ref_idc, just set zero for test */
+    fail_unless (nal_writer_put_bits_uint8 (&nw, 0, 2));
+    /* nal_unit_type, unknown h264 nal type */
+    fail_unless (nal_writer_put_bits_uint8 (&nw, 0x1f, 5));
+
+    for (j = 0; j < 3; j++)
+      fail_unless (nal_writer_put_bits_uint8 (&nw, patterns[i][j], 8));
+
+    mem = nal_writer_reset_and_get_memory (&nw);
+
+    fail_unless (mem != NULL);
+    fail_unless (gst_memory_map (mem, &info, GST_MAP_READ));
+
+    /* start code prefix 4 + nalu header 1 + written bytes 3 +
+     * emulation prevention byte 1 */
+    assert_equals_int (info.size, (4 + 1 + 3 + 1));
+
+    fail_if (memcmp (info.data + 5, expected_rst[i], 4));
+    gst_memory_unmap (mem, &info);
+    gst_memory_unref (mem);
+  }
+}
+
+GST_END_TEST;
+
+static Suite *
+nalutils_suite (void)
+{
+  Suite *s = suite_create ("H264/H264 Nal Utils");
+
+  TCase *tc_chain = tcase_create ("general");
+
+  suite_add_tcase (s, tc_chain);
+  tcase_add_test (tc_chain, test_nal_writer_init);
+  tcase_add_test (tc_chain, test_nal_writer_emulation_preventation);
+
+  return s;
+}
+
+GST_CHECK_MAIN (nalutils);

tests/check/meson.build

@@ -12,6 +12,9 @@ libparser_dep = declare_dependency(link_with: libparser,
 # FIXME: automagic
 exif_dep = dependency('libexif', version : '>= 0.6.16', required : false)
 
+# Since nalutils API is internal, need to build it again
+nalutils_dep = gstcodecparsers_dep.partial_dependency (compile_args: true, includes: true)
+
 enable_gst_player_tests = get_option('gst_player_tests')
 
 # name, condition when to skip the test and extra dependencies
@@ -53,6 +56,7 @@ base_tests = [
   [['libs/h265parser.c'], false, [gstcodecparsers_dep]],
   [['libs/insertbin.c'], false, [gstinsertbin_dep]],
   [['libs/isoff.c'], false, [gstisoff_dep]],
+  [['libs/nalutils.c', meson.current_source_dir() + '/../../gst-libs/gst/codecparsers/nalutils.c'], false, [nalutils_dep]],
   [['libs/mpegts.c'], false, [gstmpegts_dep]],
   [['libs/mpegvideoparser.c'], false, [gstcodecparsers_dep]],
   [['libs/planaraudioadapter.c'], false, [gstbadaudio_dep]],