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