git.sesse.net Git - ffmpeg/blob - libavcodec/cbs_mpeg2.c

   1 /*
   2  * This file is part of FFmpeg.
   3  *
   4  * FFmpeg is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * FFmpeg is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with FFmpeg; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17  */
  18
  19 #include "libavutil/avassert.h"
  20
  21 #include "cbs.h"
  22 #include "cbs_internal.h"
  23 #include "cbs_mpeg2.h"
  24 #include "internal.h"
  25
  26
  27 #define HEADER(name) do { \
  28         ff_cbs_trace_header(ctx, name); \
  29     } while (0)
  30
  31 #define CHECK(call) do { \
  32         err = (call); \
  33         if (err < 0) \
  34             return err; \
  35     } while (0)
  36
  37 #define FUNC_NAME(rw, codec, name) cbs_ ## codec ## _ ## rw ## _ ## name
  38 #define FUNC_MPEG2(rw, name) FUNC_NAME(rw, mpeg2, name)
  39 #define FUNC(name) FUNC_MPEG2(READWRITE, name)
  40
  41
  42 #define READ
  43 #define READWRITE read
  44 #define RWContext GetBitContext
  45
  46 #define xui(width, name, var) do { \
  47         uint32_t value = 0; \
  48         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
  49                                    &value, 0, (1 << width) - 1)); \
  50         var = value; \
  51     } while (0)
  52
  53 #define ui(width, name) \
  54         xui(width, name, current->name)
  55
  56 #define marker_bit() do { \
  57         av_unused int one = 1; \
  58         CHECK(ff_cbs_read_unsigned(ctx, rw, 1, "marker_bit", &one, 1, 1)); \
  59     } while (0)
  60
  61 #define nextbits(width, compare, var) \
  62     (get_bits_left(rw) >= width && \
  63      (var = show_bits(rw, width)) == (compare))
  64
  65 #include "cbs_mpeg2_syntax_template.c"
  66
  67 #undef READ
  68 #undef READWRITE
  69 #undef RWContext
  70 #undef xui
  71 #undef ui
  72 #undef marker_bit
  73 #undef nextbits
  74
  75
  76 #define WRITE
  77 #define READWRITE write
  78 #define RWContext PutBitContext
  79
  80 #define xui(width, name, var) do { \
  81         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
  82                                     var, 0, (1 << width) - 1)); \
  83     } while (0)
  84
  85 #define ui(width, name) \
  86         xui(width, name, current->name)
  87
  88 #define marker_bit() do { \
  89         CHECK(ff_cbs_write_unsigned(ctx, rw, 1, "marker_bit", 1, 1, 1)); \
  90     } while (0)
  91
  92 #define nextbits(width, compare, var) (var)
  93
  94 #include "cbs_mpeg2_syntax_template.c"
  95
  96 #undef READ
  97 #undef READWRITE
  98 #undef RWContext
  99 #undef xui
 100 #undef ui
 101 #undef marker_bit
 102 #undef nextbits
 103
 104
 105 static int cbs_mpeg2_split_fragment(CodedBitstreamContext *ctx,
 106                                     CodedBitstreamFragment *frag,
 107                                     int header)
 108 {
 109     const uint8_t *start, *end;
 110     uint8_t *unit_data;
 111     uint32_t start_code = -1, next_start_code = -1;
 112     size_t unit_size;
 113     int err, i, unit_type;
 114
 115     start = avpriv_find_start_code(frag->data, frag->data + frag->data_size,
 116                                    &start_code);
 117     for (i = 0;; i++) {
 118         end = avpriv_find_start_code(start, frag->data + frag->data_size,
 119                                      &next_start_code);
 120
 121         unit_type = start_code & 0xff;
 122
 123         // The start and end pointers point at to the byte following the
 124         // start_code_identifier in the start code that they found.
 125         if (end == frag->data + frag->data_size) {
 126             // We didn't find a start code, so this is the final unit.
 127             unit_size = end - (start - 1);
 128         } else {
 129             // Unit runs from start to the beginning of the start code
 130             // pointed to by end (including any padding zeroes).
 131             unit_size = (end - 4) - (start - 1);
 132         }
 133
 134         unit_data = av_malloc(unit_size + AV_INPUT_BUFFER_PADDING_SIZE);
 135         if (!unit_data)
 136             return AVERROR(ENOMEM);
 137         memcpy(unit_data, start - 1, unit_size);
 138         memset(unit_data + unit_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 139
 140         err = ff_cbs_insert_unit_data(ctx, frag, i, unit_type,
 141                                       unit_data, unit_size);
 142         if (err < 0) {
 143             av_freep(&unit_data);
 144             return err;
 145         }
 146
 147         if (end == frag->data + frag->data_size)
 148             break;
 149
 150         start_code = next_start_code;
 151         start = end;
 152     }
 153
 154     return 0;
 155 }
 156
 157 static int cbs_mpeg2_read_unit(CodedBitstreamContext *ctx,
 158                                CodedBitstreamUnit *unit)
 159 {
 160     GetBitContext gbc;
 161     int err;
 162
 163     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 164     if (err < 0)
 165         return err;
 166
 167     if (MPEG2_START_IS_SLICE(unit->type)) {
 168         MPEG2RawSlice *slice;
 169         int pos, len;
 170
 171         slice = av_mallocz(sizeof(*slice));
 172         if (!slice)
 173             return AVERROR(ENOMEM);
 174         err = cbs_mpeg2_read_slice_header(ctx, &gbc, &slice->header);
 175         if (err < 0) {
 176             av_free(slice);
 177             return err;
 178         }
 179
 180         pos = get_bits_count(&gbc);
 181         len = unit->data_size;
 182
 183         slice->data_size = len - pos / 8;
 184         slice->data = av_malloc(slice->data_size +
 185                                 AV_INPUT_BUFFER_PADDING_SIZE);
 186         if (!slice->data) {
 187             av_free(slice);
 188             return AVERROR(ENOMEM);
 189         }
 190
 191         memcpy(slice->data,
 192                unit->data + pos / 8, slice->data_size);
 193         memset(slice->data + slice->data_size, 0,
 194                AV_INPUT_BUFFER_PADDING_SIZE);
 195         slice->data_bit_start = pos % 8;
 196
 197         unit->content = slice;
 198
 199     } else {
 200         switch (unit->type) {
 201 #define START(start_code, type, func) \
 202         case start_code: \
 203             { \
 204                 type *header; \
 205                 header = av_mallocz(sizeof(*header)); \
 206                 if (!header) \
 207                     return AVERROR(ENOMEM); \
 208                 err = cbs_mpeg2_read_ ## func(ctx, &gbc, header); \
 209                 if (err < 0) { \
 210                     av_free(header); \
 211                     return err; \
 212                 } \
 213                 unit->content = header; \
 214             } \
 215             break;
 216             START(0x00, MPEG2RawPictureHeader,  picture_header);
 217             START(0xb2, MPEG2RawUserData,       user_data);
 218             START(0xb3, MPEG2RawSequenceHeader, sequence_header);
 219             START(0xb5, MPEG2RawExtensionData,  extension_data);
 220             START(0xb8, MPEG2RawGroupOfPicturesHeader, group_of_pictures_header);
 221 #undef START
 222         default:
 223             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unknown start code %02x.\n",
 224                    unit->type);
 225             return AVERROR_INVALIDDATA;
 226         }
 227     }
 228
 229     return 0;
 230 }
 231
 232 static int cbs_mpeg2_write_header(CodedBitstreamContext *ctx,
 233                                   CodedBitstreamUnit *unit,
 234                                   PutBitContext *pbc)
 235 {
 236     int err;
 237
 238     switch (unit->type) {
 239 #define START(start_code, type, func) \
 240     case start_code: \
 241         err = cbs_mpeg2_write_ ## func(ctx, pbc, unit->content); \
 242         break;
 243         START(0x00, MPEG2RawPictureHeader,  picture_header);
 244         START(0xb2, MPEG2RawUserData,       user_data);
 245         START(0xb3, MPEG2RawSequenceHeader, sequence_header);
 246         START(0xb5, MPEG2RawExtensionData,  extension_data);
 247         START(0xb8, MPEG2RawGroupOfPicturesHeader, group_of_pictures_header);
 248 #undef START
 249     default:
 250         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for start "
 251                "code %02"PRIu32".\n", unit->type);
 252         return AVERROR_PATCHWELCOME;
 253     }
 254
 255     return err;
 256 }
 257
 258 static int cbs_mpeg2_write_slice(CodedBitstreamContext *ctx,
 259                                  CodedBitstreamUnit *unit,
 260                                  PutBitContext *pbc)
 261 {
 262     MPEG2RawSlice *slice = unit->content;
 263     GetBitContext gbc;
 264     size_t bits_left;
 265     int err;
 266
 267     err = cbs_mpeg2_write_slice_header(ctx, pbc, &slice->header);
 268     if (err < 0)
 269         return err;
 270
 271     if (slice->data) {
 272         if (slice->data_size * 8 + 8 > put_bits_left(pbc))
 273             return AVERROR(ENOSPC);
 274
 275         init_get_bits(&gbc, slice->data, slice->data_size * 8);
 276         skip_bits_long(&gbc, slice->data_bit_start);
 277
 278         while (get_bits_left(&gbc) > 15)
 279             put_bits(pbc, 16, get_bits(&gbc, 16));
 280
 281         bits_left = get_bits_left(&gbc);
 282         put_bits(pbc, bits_left, get_bits(&gbc, bits_left));
 283
 284         // Align with zeroes.
 285         while (put_bits_count(pbc) % 8 != 0)
 286             put_bits(pbc, 1, 0);
 287     }
 288
 289     return 0;
 290 }
 291
 292 static int cbs_mpeg2_write_unit(CodedBitstreamContext *ctx,
 293                                 CodedBitstreamUnit *unit)
 294 {
 295     CodedBitstreamMPEG2Context *priv = ctx->priv_data;
 296     PutBitContext pbc;
 297     int err;
 298
 299     if (!priv->write_buffer) {
 300         // Initial write buffer size is 1MB.
 301         priv->write_buffer_size = 1024 * 1024;
 302
 303     reallocate_and_try_again:
 304         err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
 305         if (err < 0) {
 306             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
 307                    "sufficiently large write buffer (last attempt "
 308                    "%zu bytes).\n", priv->write_buffer_size);
 309             return err;
 310         }
 311     }
 312
 313     init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
 314
 315     if (unit->type >= 0x01 && unit->type <= 0xaf)
 316         err = cbs_mpeg2_write_slice(ctx, unit, &pbc);
 317     else
 318         err = cbs_mpeg2_write_header(ctx, unit, &pbc);
 319
 320     if (err == AVERROR(ENOSPC)) {
 321         // Overflow.
 322         priv->write_buffer_size *= 2;
 323         goto reallocate_and_try_again;
 324     }
 325     if (err < 0) {
 326         // Write failed for some other reason.
 327         return err;
 328     }
 329
 330     if (put_bits_count(&pbc) % 8)
 331         unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
 332     else
 333         unit->data_bit_padding = 0;
 334
 335     unit->data_size = (put_bits_count(&pbc) + 7) / 8;
 336     flush_put_bits(&pbc);
 337
 338     err = av_reallocp(&unit->data, unit->data_size);
 339     if (err < 0)
 340         return err;
 341
 342     memcpy(unit->data, priv->write_buffer, unit->data_size);
 343
 344     return 0;
 345 }
 346
 347 static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
 348                                        CodedBitstreamFragment *frag)
 349 {
 350     uint8_t *data;
 351     size_t size, dp, sp;
 352     int i;
 353
 354     size = 0;
 355     for (i = 0; i < frag->nb_units; i++)
 356         size += 3 + frag->units[i].data_size;
 357
 358     data = av_malloc(size);
 359     if (!data)
 360         return AVERROR(ENOMEM);
 361
 362     dp = 0;
 363     for (i = 0; i < frag->nb_units; i++) {
 364         CodedBitstreamUnit *unit = &frag->units[i];
 365
 366         data[dp++] = 0;
 367         data[dp++] = 0;
 368         data[dp++] = 1;
 369
 370         for (sp = 0; sp < unit->data_size; sp++)
 371             data[dp++] = unit->data[sp];
 372     }
 373
 374     av_assert0(dp == size);
 375
 376     frag->data      = data;
 377     frag->data_size = size;
 378
 379     return 0;
 380 }
 381
 382 static void cbs_mpeg2_free_unit(CodedBitstreamUnit *unit)
 383 {
 384     if (MPEG2_START_IS_SLICE(unit->type)) {
 385         MPEG2RawSlice *slice = unit->content;
 386         av_freep(&slice->data);
 387         av_freep(&slice->header.extra_information);
 388     } else if (unit->type == MPEG2_START_USER_DATA) {
 389         MPEG2RawUserData *user = unit->content;
 390         av_freep(&user->user_data);
 391     }
 392     av_freep(&unit->content);
 393 }
 394
 395 static void cbs_mpeg2_close(CodedBitstreamContext *ctx)
 396 {
 397     CodedBitstreamMPEG2Context *priv = ctx->priv_data;
 398
 399     av_freep(&priv->write_buffer);
 400 }
 401
 402 const CodedBitstreamType ff_cbs_type_mpeg2 = {
 403     .codec_id          = AV_CODEC_ID_MPEG2VIDEO,
 404
 405     .priv_data_size    = sizeof(CodedBitstreamMPEG2Context),
 406
 407     .split_fragment    = &cbs_mpeg2_split_fragment,
 408     .read_unit         = &cbs_mpeg2_read_unit,
 409     .write_unit        = &cbs_mpeg2_write_unit,
 410     .assemble_fragment = &cbs_mpeg2_assemble_fragment,
 411
 412     .free_unit         = &cbs_mpeg2_free_unit,
 413     .close             = &cbs_mpeg2_close,
 414 };