libavcodec/h264enc.c

   1 /*
   2  * H.264 encoder
   3  *
   4  * This file is part of FFmpeg.
   5  *
   6  * FFmpeg is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU Lesser General Public
   8  * License as published by the Free Software Foundation; either
   9  * version 2.1 of the License, or (at your option) any later version.
  10  *
  11  * FFmpeg is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with FFmpeg; if not, write to the Free Software
  18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  19  */
  20
  21
  22 #include "libavutil/common.h"
  23 #include "bitstream.h"
  24 #include "mpegvideo.h"
  25 #include "h264data.h"
  26
  27 /**
  28  * Write out the provided data into a NAL unit.
  29  * @param nal_ref_idc NAL reference IDC
  30  * @param nal_unit_type NAL unit payload type
  31  * @param dest the target buffer, dst+1 == src is allowed as a special case
  32  * @param destsize the length of the dst array
  33  * @param b2 the data which should be escaped
  34  * @returns pointer to current position in the output buffer or NULL if an error occurred
  35  */
  36 static uint8_t *h264_write_nal_unit(int nal_ref_idc, int nal_unit_type, uint8_t *dest, int *destsize,
  37                           PutBitContext *b2)
  38 {
  39     PutBitContext b;
  40     int i, destpos, rbsplen, escape_count;
  41     uint8_t *rbsp;
  42
  43     if (nal_unit_type != NAL_END_STREAM)
  44         put_bits(b2,1,1); // rbsp_stop_bit
  45
  46     // Align b2 on a byte boundary
  47     align_put_bits(b2);
  48     rbsplen = put_bits_count(b2)/8;
  49     flush_put_bits(b2);
  50     rbsp = b2->buf;
  51
  52     init_put_bits(&b,dest,*destsize);
  53
  54     put_bits(&b,16,0);
  55     put_bits(&b,16,0x01);
  56
  57     put_bits(&b,1,0); // forbidden zero bit
  58     put_bits(&b,2,nal_ref_idc); // nal_ref_idc
  59     put_bits(&b,5,nal_unit_type); // nal_unit_type
  60
  61     flush_put_bits(&b);
  62
  63     destpos = 5;
  64     escape_count= 0;
  65
  66     for (i=0; i<rbsplen; i+=2)
  67     {
  68         if (rbsp[i]) continue;
  69         if (i>0 && rbsp[i-1]==0)
  70             i--;
  71         if (i+2<rbsplen && rbsp[i+1]==0 && rbsp[i+2]<=3)
  72         {
  73             escape_count++;
  74             i+=2;
  75         }
  76     }
  77
  78     if(escape_count==0)
  79     {
  80         if(dest+destpos != rbsp)
  81         {
  82             memcpy(dest+destpos, rbsp, rbsplen);
  83             *destsize -= (rbsplen+destpos);
  84         }
  85         return dest+rbsplen+destpos;
  86     }
  87
  88     if(rbsplen + escape_count + 1> *destsize)
  89     {
  90         av_log(NULL, AV_LOG_ERROR, "Destination buffer too small!\n");
  91         return NULL;
  92     }
  93
  94     // this should be damn rare (hopefully)
  95     for (i = 0 ; i < rbsplen ; i++)
  96     {
  97         if (i + 2 < rbsplen && (rbsp[i] == 0 && rbsp[i+1] == 0 && rbsp[i+2] < 4))
  98         {
  99             dest[destpos++] = rbsp[i++];
 100             dest[destpos++] = rbsp[i];
 101             dest[destpos++] = 0x03; // emulation prevention byte
 102         }
 103         else
 104             dest[destpos++] = rbsp[i];
 105     }
 106     *destsize -= destpos;
 107     return dest+destpos;
 108 }
 109