git.sesse.net Git - ffmpeg/blob - libavcodec/cbs.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 <string.h>
  20
  21 #include "config.h"
  22
  23 #include "libavutil/avassert.h"
  24 #include "libavutil/buffer.h"
  25 #include "libavutil/common.h"
  26
  27 #include "cbs.h"
  28 #include "cbs_internal.h"
  29
  30
  31 static const CodedBitstreamType *cbs_type_table[] = {
  32 #if CONFIG_CBS_AV1
  33     &ff_cbs_type_av1,
  34 #endif
  35 #if CONFIG_CBS_H264
  36     &ff_cbs_type_h264,
  37 #endif
  38 #if CONFIG_CBS_H265
  39     &ff_cbs_type_h265,
  40 #endif
  41 #if CONFIG_CBS_JPEG
  42     &ff_cbs_type_jpeg,
  43 #endif
  44 #if CONFIG_CBS_MPEG2
  45     &ff_cbs_type_mpeg2,
  46 #endif
  47 #if CONFIG_CBS_VP9
  48     &ff_cbs_type_vp9,
  49 #endif
  50 };
  51
  52 const enum AVCodecID ff_cbs_all_codec_ids[] = {
  53 #if CONFIG_CBS_AV1
  54     AV_CODEC_ID_AV1,
  55 #endif
  56 #if CONFIG_CBS_H264
  57     AV_CODEC_ID_H264,
  58 #endif
  59 #if CONFIG_CBS_H265
  60     AV_CODEC_ID_H265,
  61 #endif
  62 #if CONFIG_CBS_JPEG
  63     AV_CODEC_ID_MJPEG,
  64 #endif
  65 #if CONFIG_CBS_MPEG2
  66     AV_CODEC_ID_MPEG2VIDEO,
  67 #endif
  68 #if CONFIG_CBS_VP9
  69     AV_CODEC_ID_VP9,
  70 #endif
  71     AV_CODEC_ID_NONE
  72 };
  73
  74 int ff_cbs_init(CodedBitstreamContext **ctx_ptr,
  75                 enum AVCodecID codec_id, void *log_ctx)
  76 {
  77     CodedBitstreamContext *ctx;
  78     const CodedBitstreamType *type;
  79     int i;
  80
  81     type = NULL;
  82     for (i = 0; i < FF_ARRAY_ELEMS(cbs_type_table); i++) {
  83         if (cbs_type_table[i]->codec_id == codec_id) {
  84             type = cbs_type_table[i];
  85             break;
  86         }
  87     }
  88     if (!type)
  89         return AVERROR(EINVAL);
  90
  91     ctx = av_mallocz(sizeof(*ctx));
  92     if (!ctx)
  93         return AVERROR(ENOMEM);
  94
  95     ctx->log_ctx = log_ctx;
  96     ctx->codec   = type;
  97
  98     if (type->priv_data_size) {
  99         ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
 100         if (!ctx->priv_data) {
 101             av_freep(&ctx);
 102             return AVERROR(ENOMEM);
 103         }
 104     }
 105
 106     ctx->decompose_unit_types = NULL;
 107
 108     ctx->trace_enable = 0;
 109     ctx->trace_level  = AV_LOG_TRACE;
 110
 111     *ctx_ptr = ctx;
 112     return 0;
 113 }
 114
 115 void ff_cbs_flush(CodedBitstreamContext *ctx)
 116 {
 117     if (ctx->codec && ctx->codec->flush)
 118         ctx->codec->flush(ctx);
 119 }
 120
 121 void ff_cbs_close(CodedBitstreamContext **ctx_ptr)
 122 {
 123     CodedBitstreamContext *ctx = *ctx_ptr;
 124
 125     if (!ctx)
 126         return;
 127
 128     if (ctx->codec && ctx->codec->close)
 129         ctx->codec->close(ctx);
 130
 131     av_freep(&ctx->write_buffer);
 132     av_freep(&ctx->priv_data);
 133     av_freep(ctx_ptr);
 134 }
 135
 136 static void cbs_unit_uninit(CodedBitstreamUnit *unit)
 137 {
 138     av_buffer_unref(&unit->content_ref);
 139     unit->content = NULL;
 140
 141     av_buffer_unref(&unit->data_ref);
 142     unit->data             = NULL;
 143     unit->data_size        = 0;
 144     unit->data_bit_padding = 0;
 145 }
 146
 147 void ff_cbs_fragment_reset(CodedBitstreamFragment *frag)
 148 {
 149     int i;
 150
 151     for (i = 0; i < frag->nb_units; i++)
 152         cbs_unit_uninit(&frag->units[i]);
 153     frag->nb_units = 0;
 154
 155     av_buffer_unref(&frag->data_ref);
 156     frag->data             = NULL;
 157     frag->data_size        = 0;
 158     frag->data_bit_padding = 0;
 159 }
 160
 161 void ff_cbs_fragment_free(CodedBitstreamFragment *frag)
 162 {
 163     ff_cbs_fragment_reset(frag);
 164
 165     av_freep(&frag->units);
 166     frag->nb_units_allocated = 0;
 167 }
 168
 169 static int cbs_read_fragment_content(CodedBitstreamContext *ctx,
 170                                      CodedBitstreamFragment *frag)
 171 {
 172     int err, i, j;
 173
 174     for (i = 0; i < frag->nb_units; i++) {
 175         CodedBitstreamUnit *unit = &frag->units[i];
 176
 177         if (ctx->decompose_unit_types) {
 178             for (j = 0; j < ctx->nb_decompose_unit_types; j++) {
 179                 if (ctx->decompose_unit_types[j] == unit->type)
 180                     break;
 181             }
 182             if (j >= ctx->nb_decompose_unit_types)
 183                 continue;
 184         }
 185
 186         av_buffer_unref(&unit->content_ref);
 187         unit->content = NULL;
 188
 189         av_assert0(unit->data && unit->data_ref);
 190
 191         err = ctx->codec->read_unit(ctx, unit);
 192         if (err == AVERROR(ENOSYS)) {
 193             av_log(ctx->log_ctx, AV_LOG_VERBOSE,
 194                    "Decomposition unimplemented for unit %d "
 195                    "(type %"PRIu32").\n", i, unit->type);
 196         } else if (err < 0) {
 197             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to read unit %d "
 198                    "(type %"PRIu32").\n", i, unit->type);
 199             return err;
 200         }
 201     }
 202
 203     return 0;
 204 }
 205
 206 static int cbs_fill_fragment_data(CodedBitstreamFragment *frag,
 207                                   const uint8_t *data, size_t size)
 208 {
 209     av_assert0(!frag->data && !frag->data_ref);
 210
 211     frag->data_ref =
 212         av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
 213     if (!frag->data_ref)
 214         return AVERROR(ENOMEM);
 215
 216     frag->data      = frag->data_ref->data;
 217     frag->data_size = size;
 218
 219     memcpy(frag->data, data, size);
 220     memset(frag->data + size, 0,
 221            AV_INPUT_BUFFER_PADDING_SIZE);
 222
 223     return 0;
 224 }
 225
 226 static int cbs_read_data(CodedBitstreamContext *ctx,
 227                          CodedBitstreamFragment *frag,
 228                          AVBufferRef *buf,
 229                          const uint8_t *data, size_t size,
 230                          int header)
 231 {
 232     int err;
 233
 234     if (buf) {
 235         frag->data_ref = av_buffer_ref(buf);
 236         if (!frag->data_ref)
 237             return AVERROR(ENOMEM);
 238
 239         frag->data      = (uint8_t *)data;
 240         frag->data_size = size;
 241
 242     } else {
 243         err = cbs_fill_fragment_data(frag, data, size);
 244         if (err < 0)
 245             return err;
 246     }
 247
 248     err = ctx->codec->split_fragment(ctx, frag, header);
 249     if (err < 0)
 250         return err;
 251
 252     return cbs_read_fragment_content(ctx, frag);
 253 }
 254
 255 int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
 256                           CodedBitstreamFragment *frag,
 257                           const AVCodecParameters *par)
 258 {
 259     return cbs_read_data(ctx, frag, NULL,
 260                          par->extradata,
 261                          par->extradata_size, 1);
 262 }
 263
 264 int ff_cbs_read_extradata_from_codec(CodedBitstreamContext *ctx,
 265                                      CodedBitstreamFragment *frag,
 266                                      const AVCodecContext *avctx)
 267 {
 268     return cbs_read_data(ctx, frag, NULL,
 269                          avctx->extradata,
 270                          avctx->extradata_size, 1);
 271 }
 272
 273 int ff_cbs_read_packet(CodedBitstreamContext *ctx,
 274                        CodedBitstreamFragment *frag,
 275                        const AVPacket *pkt)
 276 {
 277     return cbs_read_data(ctx, frag, pkt->buf,
 278                          pkt->data, pkt->size, 0);
 279 }
 280
 281 int ff_cbs_read(CodedBitstreamContext *ctx,
 282                 CodedBitstreamFragment *frag,
 283                 const uint8_t *data, size_t size)
 284 {
 285     return cbs_read_data(ctx, frag, NULL,
 286                          data, size, 0);
 287 }
 288
 289 static int cbs_write_unit_data(CodedBitstreamContext *ctx,
 290                                CodedBitstreamUnit *unit)
 291 {
 292     PutBitContext pbc;
 293     int ret;
 294
 295     if (!ctx->write_buffer) {
 296         // Initial write buffer size is 1MB.
 297         ctx->write_buffer_size = 1024 * 1024;
 298
 299     reallocate_and_try_again:
 300         ret = av_reallocp(&ctx->write_buffer, ctx->write_buffer_size);
 301         if (ret < 0) {
 302             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
 303                    "sufficiently large write buffer (last attempt "
 304                    "%"SIZE_SPECIFIER" bytes).\n", ctx->write_buffer_size);
 305             return ret;
 306         }
 307     }
 308
 309     init_put_bits(&pbc, ctx->write_buffer, ctx->write_buffer_size);
 310
 311     ret = ctx->codec->write_unit(ctx, unit, &pbc);
 312     if (ret < 0) {
 313         if (ret == AVERROR(ENOSPC)) {
 314             // Overflow.
 315             if (ctx->write_buffer_size == INT_MAX / 8)
 316                 return AVERROR(ENOMEM);
 317             ctx->write_buffer_size = FFMIN(2 * ctx->write_buffer_size, INT_MAX / 8);
 318             goto reallocate_and_try_again;
 319         }
 320         // Write failed for some other reason.
 321         return ret;
 322     }
 323
 324     // Overflow but we didn't notice.
 325     av_assert0(put_bits_count(&pbc) <= 8 * ctx->write_buffer_size);
 326
 327     if (put_bits_count(&pbc) % 8)
 328         unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
 329     else
 330         unit->data_bit_padding = 0;
 331
 332     flush_put_bits(&pbc);
 333
 334     ret = ff_cbs_alloc_unit_data(unit, put_bits_count(&pbc) / 8);
 335     if (ret < 0)
 336         return ret;
 337
 338     memcpy(unit->data, ctx->write_buffer, unit->data_size);
 339
 340     return 0;
 341 }
 342
 343 int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
 344                                CodedBitstreamFragment *frag)
 345 {
 346     int err, i;
 347
 348     for (i = 0; i < frag->nb_units; i++) {
 349         CodedBitstreamUnit *unit = &frag->units[i];
 350
 351         if (!unit->content)
 352             continue;
 353
 354         av_buffer_unref(&unit->data_ref);
 355         unit->data = NULL;
 356
 357         err = cbs_write_unit_data(ctx, unit);
 358         if (err < 0) {
 359             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to write unit %d "
 360                    "(type %"PRIu32").\n", i, unit->type);
 361             return err;
 362         }
 363         av_assert0(unit->data && unit->data_ref);
 364     }
 365
 366     av_buffer_unref(&frag->data_ref);
 367     frag->data = NULL;
 368
 369     err = ctx->codec->assemble_fragment(ctx, frag);
 370     if (err < 0) {
 371         av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to assemble fragment.\n");
 372         return err;
 373     }
 374     av_assert0(frag->data && frag->data_ref);
 375
 376     return 0;
 377 }
 378
 379 int ff_cbs_write_extradata(CodedBitstreamContext *ctx,
 380                            AVCodecParameters *par,
 381                            CodedBitstreamFragment *frag)
 382 {
 383     int err;
 384
 385     err = ff_cbs_write_fragment_data(ctx, frag);
 386     if (err < 0)
 387         return err;
 388
 389     av_freep(&par->extradata);
 390
 391     par->extradata = av_malloc(frag->data_size +
 392                                AV_INPUT_BUFFER_PADDING_SIZE);
 393     if (!par->extradata)
 394         return AVERROR(ENOMEM);
 395
 396     memcpy(par->extradata, frag->data, frag->data_size);
 397     memset(par->extradata + frag->data_size, 0,
 398            AV_INPUT_BUFFER_PADDING_SIZE);
 399     par->extradata_size = frag->data_size;
 400
 401     return 0;
 402 }
 403
 404 int ff_cbs_write_packet(CodedBitstreamContext *ctx,
 405                         AVPacket *pkt,
 406                         CodedBitstreamFragment *frag)
 407 {
 408     AVBufferRef *buf;
 409     int err;
 410
 411     err = ff_cbs_write_fragment_data(ctx, frag);
 412     if (err < 0)
 413         return err;
 414
 415     buf = av_buffer_ref(frag->data_ref);
 416     if (!buf)
 417         return AVERROR(ENOMEM);
 418
 419     av_buffer_unref(&pkt->buf);
 420
 421     pkt->buf  = buf;
 422     pkt->data = frag->data;
 423     pkt->size = frag->data_size;
 424
 425     return 0;
 426 }
 427
 428
 429 void ff_cbs_trace_header(CodedBitstreamContext *ctx,
 430                          const char *name)
 431 {
 432     if (!ctx->trace_enable)
 433         return;
 434
 435     av_log(ctx->log_ctx, ctx->trace_level, "%s\n", name);
 436 }
 437
 438 void ff_cbs_trace_syntax_element(CodedBitstreamContext *ctx, int position,
 439                                  const char *str, const int *subscripts,
 440                                  const char *bits, int64_t value)
 441 {
 442     char name[256];
 443     size_t name_len, bits_len;
 444     int pad, subs, i, j, k, n;
 445
 446     if (!ctx->trace_enable)
 447         return;
 448
 449     av_assert0(value >= INT_MIN && value <= UINT32_MAX);
 450
 451     subs = subscripts ? subscripts[0] : 0;
 452     n = 0;
 453     for (i = j = 0; str[i];) {
 454         if (str[i] == '[') {
 455             if (n < subs) {
 456                 ++n;
 457                 k = snprintf(name + j, sizeof(name) - j, "[%d", subscripts[n]);
 458                 av_assert0(k > 0 && j + k < sizeof(name));
 459                 j += k;
 460                 for (++i; str[i] && str[i] != ']'; i++);
 461                 av_assert0(str[i] == ']');
 462             } else {
 463                 while (str[i] && str[i] != ']')
 464                     name[j++] = str[i++];
 465                 av_assert0(str[i] == ']');
 466             }
 467         } else {
 468             av_assert0(j + 1 < sizeof(name));
 469             name[j++] = str[i++];
 470         }
 471     }
 472     av_assert0(j + 1 < sizeof(name));
 473     name[j] = 0;
 474     av_assert0(n == subs);
 475
 476     name_len = strlen(name);
 477     bits_len = strlen(bits);
 478
 479     if (name_len + bits_len > 60)
 480         pad = bits_len + 2;
 481     else
 482         pad = 61 - name_len;
 483
 484     av_log(ctx->log_ctx, ctx->trace_level, "%-10d  %s%*s = %"PRId64"\n",
 485            position, name, pad, bits, value);
 486 }
 487
 488 int ff_cbs_read_unsigned(CodedBitstreamContext *ctx, GetBitContext *gbc,
 489                          int width, const char *name,
 490                          const int *subscripts, uint32_t *write_to,
 491                          uint32_t range_min, uint32_t range_max)
 492 {
 493     uint32_t value;
 494     int position;
 495
 496     av_assert0(width > 0 && width <= 32);
 497
 498     if (get_bits_left(gbc) < width) {
 499         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid value at "
 500                "%s: bitstream ended.\n", name);
 501         return AVERROR_INVALIDDATA;
 502     }
 503
 504     if (ctx->trace_enable)
 505         position = get_bits_count(gbc);
 506
 507     value = get_bits_long(gbc, width);
 508
 509     if (ctx->trace_enable) {
 510         char bits[33];
 511         int i;
 512         for (i = 0; i < width; i++)
 513             bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
 514         bits[i] = 0;
 515
 516         ff_cbs_trace_syntax_element(ctx, position, name, subscripts,
 517                                     bits, value);
 518     }
 519
 520     if (value < range_min || value > range_max) {
 521         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 522                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
 523                name, value, range_min, range_max);
 524         return AVERROR_INVALIDDATA;
 525     }
 526
 527     *write_to = value;
 528     return 0;
 529 }
 530
 531 int ff_cbs_write_unsigned(CodedBitstreamContext *ctx, PutBitContext *pbc,
 532                           int width, const char *name,
 533                           const int *subscripts, uint32_t value,
 534                           uint32_t range_min, uint32_t range_max)
 535 {
 536     av_assert0(width > 0 && width <= 32);
 537
 538     if (value < range_min || value > range_max) {
 539         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 540                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
 541                name, value, range_min, range_max);
 542         return AVERROR_INVALIDDATA;
 543     }
 544
 545     if (put_bits_left(pbc) < width)
 546         return AVERROR(ENOSPC);
 547
 548     if (ctx->trace_enable) {
 549         char bits[33];
 550         int i;
 551         for (i = 0; i < width; i++)
 552             bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
 553         bits[i] = 0;
 554
 555         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc),
 556                                     name, subscripts, bits, value);
 557     }
 558
 559     if (width < 32)
 560         put_bits(pbc, width, value);
 561     else
 562         put_bits32(pbc, value);
 563
 564     return 0;
 565 }
 566
 567 int ff_cbs_read_signed(CodedBitstreamContext *ctx, GetBitContext *gbc,
 568                        int width, const char *name,
 569                        const int *subscripts, int32_t *write_to,
 570                        int32_t range_min, int32_t range_max)
 571 {
 572     int32_t value;
 573     int position;
 574
 575     av_assert0(width > 0 && width <= 32);
 576
 577     if (get_bits_left(gbc) < width) {
 578         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid value at "
 579                "%s: bitstream ended.\n", name);
 580         return AVERROR_INVALIDDATA;
 581     }
 582
 583     if (ctx->trace_enable)
 584         position = get_bits_count(gbc);
 585
 586     value = get_sbits_long(gbc, width);
 587
 588     if (ctx->trace_enable) {
 589         char bits[33];
 590         int i;
 591         for (i = 0; i < width; i++)
 592             bits[i] = value & (1U << (width - i - 1)) ? '1' : '0';
 593         bits[i] = 0;
 594
 595         ff_cbs_trace_syntax_element(ctx, position, name, subscripts,
 596                                     bits, value);
 597     }
 598
 599     if (value < range_min || value > range_max) {
 600         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 601                "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
 602                name, value, range_min, range_max);
 603         return AVERROR_INVALIDDATA;
 604     }
 605
 606     *write_to = value;
 607     return 0;
 608 }
 609
 610 int ff_cbs_write_signed(CodedBitstreamContext *ctx, PutBitContext *pbc,
 611                         int width, const char *name,
 612                         const int *subscripts, int32_t value,
 613                         int32_t range_min, int32_t range_max)
 614 {
 615     av_assert0(width > 0 && width <= 32);
 616
 617     if (value < range_min || value > range_max) {
 618         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 619                "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
 620                name, value, range_min, range_max);
 621         return AVERROR_INVALIDDATA;
 622     }
 623
 624     if (put_bits_left(pbc) < width)
 625         return AVERROR(ENOSPC);
 626
 627     if (ctx->trace_enable) {
 628         char bits[33];
 629         int i;
 630         for (i = 0; i < width; i++)
 631             bits[i] = value & (1U << (width - i - 1)) ? '1' : '0';
 632         bits[i] = 0;
 633
 634         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc),
 635                                     name, subscripts, bits, value);
 636     }
 637
 638     if (width < 32)
 639         put_sbits(pbc, width, value);
 640     else
 641         put_bits32(pbc, value);
 642
 643     return 0;
 644 }
 645
 646
 647 int ff_cbs_alloc_unit_content(CodedBitstreamUnit *unit,
 648                               size_t size,
 649                               void (*free)(void *opaque, uint8_t *data))
 650 {
 651     av_assert0(!unit->content && !unit->content_ref);
 652
 653     unit->content = av_mallocz(size);
 654     if (!unit->content)
 655         return AVERROR(ENOMEM);
 656
 657     unit->content_ref = av_buffer_create(unit->content, size,
 658                                          free, NULL, 0);
 659     if (!unit->content_ref) {
 660         av_freep(&unit->content);
 661         return AVERROR(ENOMEM);
 662     }
 663
 664     return 0;
 665 }
 666
 667 int ff_cbs_alloc_unit_data(CodedBitstreamUnit *unit,
 668                            size_t size)
 669 {
 670     av_assert0(!unit->data && !unit->data_ref);
 671
 672     unit->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
 673     if (!unit->data_ref)
 674         return AVERROR(ENOMEM);
 675
 676     unit->data      = unit->data_ref->data;
 677     unit->data_size = size;
 678
 679     memset(unit->data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 680
 681     return 0;
 682 }
 683
 684 static int cbs_insert_unit(CodedBitstreamFragment *frag,
 685                            int position)
 686 {
 687     CodedBitstreamUnit *units;
 688
 689     if (frag->nb_units < frag->nb_units_allocated) {
 690         units = frag->units;
 691
 692         if (position < frag->nb_units)
 693             memmove(units + position + 1, units + position,
 694                     (frag->nb_units - position) * sizeof(*units));
 695     } else {
 696         units = av_malloc_array(frag->nb_units*2 + 1, sizeof(*units));
 697         if (!units)
 698             return AVERROR(ENOMEM);
 699
 700         frag->nb_units_allocated = 2*frag->nb_units_allocated + 1;
 701
 702         if (position > 0)
 703             memcpy(units, frag->units, position * sizeof(*units));
 704
 705         if (position < frag->nb_units)
 706             memcpy(units + position + 1, frag->units + position,
 707                    (frag->nb_units - position) * sizeof(*units));
 708     }
 709
 710     memset(units + position, 0, sizeof(*units));
 711
 712     if (units != frag->units) {
 713         av_free(frag->units);
 714         frag->units = units;
 715     }
 716
 717     ++frag->nb_units;
 718
 719     return 0;
 720 }
 721
 722 int ff_cbs_insert_unit_content(CodedBitstreamFragment *frag,
 723                                int position,
 724                                CodedBitstreamUnitType type,
 725                                void *content,
 726                                AVBufferRef *content_buf)
 727 {
 728     CodedBitstreamUnit *unit;
 729     AVBufferRef *content_ref;
 730     int err;
 731
 732     if (position == -1)
 733         position = frag->nb_units;
 734     av_assert0(position >= 0 && position <= frag->nb_units);
 735
 736     if (content_buf) {
 737         content_ref = av_buffer_ref(content_buf);
 738         if (!content_ref)
 739             return AVERROR(ENOMEM);
 740     } else {
 741         content_ref = NULL;
 742     }
 743
 744     err = cbs_insert_unit(frag, position);
 745     if (err < 0) {
 746         av_buffer_unref(&content_ref);
 747         return err;
 748     }
 749
 750     unit = &frag->units[position];
 751     unit->type        = type;
 752     unit->content     = content;
 753     unit->content_ref = content_ref;
 754
 755     return 0;
 756 }
 757
 758 int ff_cbs_insert_unit_data(CodedBitstreamFragment *frag,
 759                             int position,
 760                             CodedBitstreamUnitType type,
 761                             uint8_t *data, size_t data_size,
 762                             AVBufferRef *data_buf)
 763 {
 764     CodedBitstreamUnit *unit;
 765     AVBufferRef *data_ref;
 766     int err;
 767
 768     if (position == -1)
 769         position = frag->nb_units;
 770     av_assert0(position >= 0 && position <= frag->nb_units);
 771
 772     if (data_buf)
 773         data_ref = av_buffer_ref(data_buf);
 774     else
 775         data_ref = av_buffer_create(data, data_size, NULL, NULL, 0);
 776     if (!data_ref) {
 777         if (!data_buf)
 778             av_free(data);
 779         return AVERROR(ENOMEM);
 780     }
 781
 782     err = cbs_insert_unit(frag, position);
 783     if (err < 0) {
 784         av_buffer_unref(&data_ref);
 785         return err;
 786     }
 787
 788     unit = &frag->units[position];
 789     unit->type      = type;
 790     unit->data      = data;
 791     unit->data_size = data_size;
 792     unit->data_ref  = data_ref;
 793
 794     return 0;
 795 }
 796
 797 void ff_cbs_delete_unit(CodedBitstreamFragment *frag,
 798                         int position)
 799 {
 800     av_assert0(0 <= position && position < frag->nb_units
 801                              && "Unit to be deleted not in fragment.");
 802
 803     cbs_unit_uninit(&frag->units[position]);
 804
 805     --frag->nb_units;
 806
 807     if (frag->nb_units > 0)
 808         memmove(frag->units + position,
 809                 frag->units + position + 1,
 810                 (frag->nb_units - position) * sizeof(*frag->units));
 811 }
 812
 813 static void cbs_default_free_unit_content(void *opaque, uint8_t *data)
 814 {
 815     const CodedBitstreamUnitTypeDescriptor *desc = opaque;
 816     if (desc->content_type == CBS_CONTENT_TYPE_INTERNAL_REFS) {
 817         int i;
 818         for (i = 0; i < desc->nb_ref_offsets; i++) {
 819             void **ptr = (void**)(data + desc->ref_offsets[i]);
 820             av_buffer_unref((AVBufferRef**)(ptr + 1));
 821         }
 822     }
 823     av_free(data);
 824 }
 825
 826 static const CodedBitstreamUnitTypeDescriptor
 827     *cbs_find_unit_type_desc(CodedBitstreamContext *ctx,
 828                              CodedBitstreamUnit *unit)
 829 {
 830     const CodedBitstreamUnitTypeDescriptor *desc;
 831     int i, j;
 832
 833     if (!ctx->codec->unit_types)
 834         return NULL;
 835
 836     for (i = 0;; i++) {
 837         desc = &ctx->codec->unit_types[i];
 838         if (desc->nb_unit_types == 0)
 839             break;
 840         if (desc->nb_unit_types == CBS_UNIT_TYPE_RANGE) {
 841             if (unit->type >= desc->unit_type_range_start &&
 842                 unit->type <= desc->unit_type_range_end)
 843                 return desc;
 844         } else {
 845             for (j = 0; j < desc->nb_unit_types; j++) {
 846                 if (desc->unit_types[j] == unit->type)
 847                     return desc;
 848             }
 849         }
 850     }
 851     return NULL;
 852 }
 853
 854 int ff_cbs_alloc_unit_content2(CodedBitstreamContext *ctx,
 855                                CodedBitstreamUnit *unit)
 856 {
 857     const CodedBitstreamUnitTypeDescriptor *desc;
 858
 859     av_assert0(!unit->content && !unit->content_ref);
 860
 861     desc = cbs_find_unit_type_desc(ctx, unit);
 862     if (!desc)
 863         return AVERROR(ENOSYS);
 864
 865     unit->content = av_mallocz(desc->content_size);
 866     if (!unit->content)
 867         return AVERROR(ENOMEM);
 868
 869     unit->content_ref =
 870         av_buffer_create(unit->content, desc->content_size,
 871                          desc->content_free ? desc->content_free
 872                                             : cbs_default_free_unit_content,
 873                          (void*)desc, 0);
 874     if (!unit->content_ref) {
 875         av_freep(&unit->content);
 876         return AVERROR(ENOMEM);
 877     }
 878
 879     return 0;
 880 }
 881
 882 static int cbs_clone_unit_content(AVBufferRef **clone_ref,
 883                                   CodedBitstreamUnit *unit,
 884                                   const CodedBitstreamUnitTypeDescriptor *desc)
 885 {
 886     uint8_t *src, *copy;
 887     uint8_t **src_ptr, **copy_ptr;
 888     AVBufferRef **src_buf, **copy_buf;
 889     int err, i;
 890
 891     av_assert0(unit->content);
 892     src = unit->content;
 893
 894     copy = av_memdup(src, desc->content_size);
 895     if (!copy)
 896         return AVERROR(ENOMEM);
 897
 898     for (i = 0; i < desc->nb_ref_offsets; i++) {
 899         src_ptr  = (uint8_t**)(src + desc->ref_offsets[i]);
 900         src_buf  = (AVBufferRef**)(src_ptr + 1);
 901         copy_ptr = (uint8_t**)(copy + desc->ref_offsets[i]);
 902         copy_buf = (AVBufferRef**)(copy_ptr + 1);
 903
 904         if (!*src_ptr) {
 905             av_assert0(!*src_buf);
 906             continue;
 907         }
 908         if (!*src_buf) {
 909             // We can't handle a non-refcounted pointer here - we don't
 910             // have enough information to handle whatever structure lies
 911             // at the other end of it.
 912             err = AVERROR(EINVAL);
 913             goto fail;
 914         }
 915
 916         // src_ptr is required to point somewhere inside src_buf.  If it
 917         // doesn't, there is a bug somewhere.
 918         av_assert0(*src_ptr >= (*src_buf)->data &&
 919                    *src_ptr <  (*src_buf)->data + (*src_buf)->size);
 920
 921         *copy_buf = av_buffer_ref(*src_buf);
 922         if (!*copy_buf) {
 923             err = AVERROR(ENOMEM);
 924             goto fail;
 925         }
 926         *copy_ptr = (*copy_buf)->data + (*src_ptr - (*src_buf)->data);
 927     }
 928
 929     *clone_ref = av_buffer_create(copy, desc->content_size,
 930                                   desc->content_free ? desc->content_free :
 931                                   cbs_default_free_unit_content,
 932                                   (void*)desc, 0);
 933     if (!*clone_ref) {
 934         err = AVERROR(ENOMEM);
 935         goto fail;
 936     }
 937
 938     return 0;
 939
 940 fail:
 941     for (--i; i >= 0; i--)
 942         av_buffer_unref((AVBufferRef**)(copy + desc->ref_offsets[i]));
 943     av_freep(&copy);
 944     *clone_ref = NULL;
 945     return err;
 946 }
 947
 948 int ff_cbs_make_unit_refcounted(CodedBitstreamContext *ctx,
 949                                 CodedBitstreamUnit *unit)
 950 {
 951     const CodedBitstreamUnitTypeDescriptor *desc;
 952     AVBufferRef *ref;
 953     int err;
 954
 955     av_assert0(unit->content);
 956     if (unit->content_ref) {
 957         // Already refcounted, nothing to do.
 958         return 0;
 959     }
 960
 961     desc = cbs_find_unit_type_desc(ctx, unit);
 962     if (!desc)
 963         return AVERROR(ENOSYS);
 964
 965     switch (desc->content_type) {
 966     case CBS_CONTENT_TYPE_POD:
 967         ref = av_buffer_alloc(desc->content_size);
 968         if (!ref)
 969             return AVERROR(ENOMEM);
 970         memcpy(ref->data, unit->content, desc->content_size);
 971         err = 0;
 972         break;
 973
 974     case CBS_CONTENT_TYPE_INTERNAL_REFS:
 975         err = cbs_clone_unit_content(&ref, unit, desc);
 976         break;
 977
 978     case CBS_CONTENT_TYPE_COMPLEX:
 979         if (!desc->content_clone)
 980             return AVERROR_PATCHWELCOME;
 981         err = desc->content_clone(&ref, unit);
 982         break;
 983
 984     default:
 985         av_assert0(0 && "Invalid content type.");
 986     }
 987
 988     if (err < 0)
 989         return err;
 990
 991     unit->content_ref = ref;
 992     unit->content     = ref->data;
 993     return 0;
 994 }
 995
 996 int ff_cbs_make_unit_writable(CodedBitstreamContext *ctx,
 997                               CodedBitstreamUnit *unit)
 998 {
 999     const CodedBitstreamUnitTypeDescriptor *desc;
1000     AVBufferRef *ref;
1001     int err;
1002
1003     // This can only be applied to refcounted units.
1004     err = ff_cbs_make_unit_refcounted(ctx, unit);
1005     if (err < 0)
1006         return err;
1007     av_assert0(unit->content && unit->content_ref);
1008
1009     if (av_buffer_is_writable(unit->content_ref))
1010         return 0;
1011
1012     desc = cbs_find_unit_type_desc(ctx, unit);
1013     if (!desc)
1014         return AVERROR(ENOSYS);
1015
1016     switch (desc->content_type) {
1017     case CBS_CONTENT_TYPE_POD:
1018         err = av_buffer_make_writable(&unit->content_ref);
1019         break;
1020
1021     case CBS_CONTENT_TYPE_INTERNAL_REFS:
1022         err = cbs_clone_unit_content(&ref, unit, desc);
1023         break;
1024
1025     case CBS_CONTENT_TYPE_COMPLEX:
1026         if (!desc->content_clone)
1027             return AVERROR_PATCHWELCOME;
1028         err = desc->content_clone(&ref, unit);
1029         break;
1030
1031     default:
1032         av_assert0(0 && "Invalid content type.");
1033     }
1034     if (err < 0)
1035         return err;
1036
1037     if (desc->content_type != CBS_CONTENT_TYPE_POD) {
1038         av_buffer_unref(&unit->content_ref);
1039         unit->content_ref = ref;
1040     }
1041     unit->content = unit->content_ref->data;
1042     return 0;
1043 }