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