xref: /third_party/ffmpeg/libavcodec/cbs.c (revision cabdff1a)
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
33static 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
54const 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
76av_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
121av_cold void ff_cbs_flush(CodedBitstreamContext *ctx)
122{
123    if (ctx->codec->flush)
124        ctx->codec->flush(ctx);
125}
126
127av_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
146static 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
157void 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
171av_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
179static 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
222static 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
242static 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
271int 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
280int 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
289int 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
297int 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
310int 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 */
323static 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
340static 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
394int 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
430int 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
455int 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
480void 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
489void 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
539int 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
582int 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
618int 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
661int 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
698int 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
718static 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
756int 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
792static 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
829int 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
839void 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
855static 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
868static 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
896int 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
924static 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
982fail:
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
990int 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
1038int 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}
1086