From e39e60a37ea11900204aa1170b78a61fb4faf576 Mon Sep 17 00:00:00 2001
From: Panagiotis Issaris <takis.issaris@uhasselt.be>
Date: Mon, 11 Dec 2006 01:00:50 +0000
Subject: [PATCH] Adding function which enables writing H.264 NAL units.

Originally committed as revision 7273 to svn://svn.ffmpeg.org/ffmpeg/trunk
---
 libavcodec/Makefile  |   1 +
 libavcodec/h264enc.c | 107 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+)
 create mode 100644 libavcodec/h264enc.c

diff --git a/libavcodec/Makefile b/libavcodec/Makefile
index 1aa29950bc..10278089d5 100644
--- a/libavcodec/Makefile
+++ b/libavcodec/Makefile
@@ -86,6 +86,7 @@ OBJS-$(CONFIG_GIF_ENCODER)             += gif.o
 OBJS-$(CONFIG_H261_DECODER)            += h261.o
 OBJS-$(CONFIG_H261_ENCODER)            += h261.o
 OBJS-$(CONFIG_H264_DECODER)            += h264.o
+OBJS-$(CONFIG_H264_ENCODER)            += h264enc.o
 OBJS-$(CONFIG_HUFFYUV_DECODER)         += huffyuv.o
 OBJS-$(CONFIG_HUFFYUV_ENCODER)         += huffyuv.o
 OBJS-$(CONFIG_IDCIN_DECODER)           += idcinvideo.o
diff --git a/libavcodec/h264enc.c b/libavcodec/h264enc.c
new file mode 100644
index 0000000000..15e1823022
--- /dev/null
+++ b/libavcodec/h264enc.c
@@ -0,0 +1,107 @@
+/*
+ * H.264 encoder
+ *
+ * FFmpeg is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * FFmpeg 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
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with FFmpeg; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+
+#include "common.h"
+#include "bitstream.h"
+#include "mpegvideo.h"
+#include "h264data.h"
+
+/**
+ * Write out the provided data into a NAL unit.
+ * @param nal_ref_idc NAL reference IDC
+ * @param nal_unit_type NAL unit payload type
+ * @param dest the target buffer, dst+1 == src is allowed as a special case
+ * @param destsize the length of the dst array
+ * @param b2 the data which should be escaped
+ * @returns pointer to current position in the output buffer or NULL if an error occured
+ */
+static uint8_t *h264_write_nal_unit(int nal_ref_idc, int nal_unit_type, uint8_t *dest, int *destsize,
+                          PutBitContext *b2)
+{
+    PutBitContext b;
+    int i, destpos, rbsplen, escape_count;
+    uint8_t *rbsp;
+
+    if (nal_unit_type != NAL_END_STREAM)
+        put_bits(b2,1,1); // rbsp_stop_bit
+
+    // Align b2 on a byte boundary
+    align_put_bits(b2);
+    rbsplen = put_bits_count(b2)/8;
+    flush_put_bits(b2);
+    rbsp = b2->buf;
+
+    init_put_bits(&b,dest,*destsize);
+
+    put_bits(&b,16,0);
+    put_bits(&b,16,0x01);
+
+    put_bits(&b,1,0); // forbidden zero bit
+    put_bits(&b,2,nal_ref_idc); // nal_ref_idc
+    put_bits(&b,5,nal_unit_type); // nal_unit_type
+
+    flush_put_bits(&b);
+
+    destpos = 5;
+    escape_count= 0;
+
+    for (i=0; i<rbsplen; i+=2)
+    {
+        if (rbsp[i]) continue;
+        if (i>0 && rbsp[i-1]==0)
+            i--;
+        if (i+2<rbsplen && rbsp[i+1]==0 && rbsp[i+2]<=3)
+        {
+            escape_count++;
+            i+=2;
+        }
+    }
+
+    if(escape_count==0)
+    {
+        if(dest+destpos != rbsp)
+        {
+            memcpy(dest+destpos, rbsp, rbsplen);
+            *destsize -= (rbsplen+destpos);
+        }
+        return dest+rbsplen+destpos;
+    }
+
+    if(rbsplen + escape_count + 1> *destsize)
+    {
+        av_log(NULL, AV_LOG_ERROR, "Destination buffer too small!\n");
+        return NULL;
+    }
+
+    // this should be damn rare (hopefully)
+    for (i = 0 ; i < rbsplen ; i++)
+    {
+        if (i + 2 < rbsplen && (rbsp[i] == 0 && rbsp[i+1] == 0 && rbsp[i+2] < 4))
+        {
+            dest[destpos++] = rbsp[i++];
+            dest[destpos++] = rbsp[i];
+            dest[destpos++] = 0x03; // emulation prevention byte
+        }
+        else
+            dest[destpos++] = rbsp[i];
+    }
+    *destsize -= destpos;
+    return dest+destpos;
+}
+