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