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