cbs: Add a table of all supported codec IDs
[ffmpeg.git] / libavcodec / cbs.c
1 /*
2  * This file is part of FFmpeg.
3  *
4  * FFmpeg is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU Lesser General Public
6  * License as published by the Free Software Foundation; either
7  * version 2.1 of the License, or (at your option) any later version.
8  *
9  * FFmpeg is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
12  * Lesser General Public License for more details.
13  *
14  * You should have received a copy of the GNU Lesser General Public
15  * License along with FFmpeg; if not, write to the Free Software
16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
17  */
18
19 #include <string.h>
20
21 #include "config.h"
22
23 #include "libavutil/avassert.h"
24 #include "libavutil/buffer.h"
25 #include "libavutil/common.h"
26
27 #include "cbs.h"
28 #include "cbs_internal.h"
29
30
31 static const CodedBitstreamType *cbs_type_table[] = {
32 #if CONFIG_CBS_H264
33     &ff_cbs_type_h264,
34 #endif
35 #if CONFIG_CBS_H265
36     &ff_cbs_type_h265,
37 #endif
38 #if CONFIG_CBS_MPEG2
39     &ff_cbs_type_mpeg2,
40 #endif
41 };
42
43 const enum AVCodecID ff_cbs_all_codec_ids[] = {
44 #if CONFIG_CBS_H264
45     AV_CODEC_ID_H264,
46 #endif
47 #if CONFIG_CBS_H265
48     AV_CODEC_ID_H265,
49 #endif
50 #if CONFIG_CBS_MPEG2
51     AV_CODEC_ID_MPEG2VIDEO,
52 #endif
53     AV_CODEC_ID_NONE
54 };
55
56 int ff_cbs_init(CodedBitstreamContext **ctx_ptr,
57                 enum AVCodecID codec_id, void *log_ctx)
58 {
59     CodedBitstreamContext *ctx;
60     const CodedBitstreamType *type;
61     int i;
62
63     type = NULL;
64     for (i = 0; i < FF_ARRAY_ELEMS(cbs_type_table); i++) {
65         if (cbs_type_table[i]->codec_id == codec_id) {
66             type = cbs_type_table[i];
67             break;
68         }
69     }
70     if (!type)
71         return AVERROR(EINVAL);
72
73     ctx = av_mallocz(sizeof(*ctx));
74     if (!ctx)
75         return AVERROR(ENOMEM);
76
77     ctx->log_ctx = log_ctx;
78     ctx->codec   = type;
79
80     ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
81     if (!ctx->priv_data) {
82         av_freep(&ctx);
83         return AVERROR(ENOMEM);
84     }
85
86     ctx->decompose_unit_types = NULL;
87
88     ctx->trace_enable = 0;
89     ctx->trace_level  = AV_LOG_TRACE;
90
91     *ctx_ptr = ctx;
92     return 0;
93 }
94
95 void ff_cbs_close(CodedBitstreamContext **ctx_ptr)
96 {
97     CodedBitstreamContext *ctx = *ctx_ptr;
98
99     if (!ctx)
100         return;
101
102     if (ctx->codec && ctx->codec->close)
103         ctx->codec->close(ctx);
104
105     av_freep(&ctx->priv_data);
106     av_freep(ctx_ptr);
107 }
108
109 static void cbs_unit_uninit(CodedBitstreamContext *ctx,
110                             CodedBitstreamUnit *unit)
111 {
112     av_buffer_unref(&unit->content_ref);
113     unit->content = NULL;
114
115     av_buffer_unref(&unit->data_ref);
116     unit->data             = NULL;
117     unit->data_size        = 0;
118     unit->data_bit_padding = 0;
119 }
120
121 void ff_cbs_fragment_uninit(CodedBitstreamContext *ctx,
122                             CodedBitstreamFragment *frag)
123 {
124     int i;
125
126     for (i = 0; i < frag->nb_units; i++)
127         cbs_unit_uninit(ctx, &frag->units[i]);
128     av_freep(&frag->units);
129     frag->nb_units = 0;
130
131     av_buffer_unref(&frag->data_ref);
132     frag->data             = NULL;
133     frag->data_size        = 0;
134     frag->data_bit_padding = 0;
135 }
136
137 static int cbs_read_fragment_content(CodedBitstreamContext *ctx,
138                                      CodedBitstreamFragment *frag)
139 {
140     int err, i, j;
141
142     for (i = 0; i < frag->nb_units; i++) {
143         if (ctx->decompose_unit_types) {
144             for (j = 0; j < ctx->nb_decompose_unit_types; j++) {
145                 if (ctx->decompose_unit_types[j] == frag->units[i].type)
146                     break;
147             }
148             if (j >= ctx->nb_decompose_unit_types)
149                 continue;
150         }
151
152         av_buffer_unref(&frag->units[i].content_ref);
153         frag->units[i].content = NULL;
154
155         err = ctx->codec->read_unit(ctx, &frag->units[i]);
156         if (err == AVERROR(ENOSYS)) {
157             av_log(ctx->log_ctx, AV_LOG_VERBOSE,
158                    "Decomposition unimplemented for unit %d "
159                    "(type %"PRIu32").\n", i, frag->units[i].type);
160         } else if (err < 0) {
161             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to read unit %d "
162                    "(type %"PRIu32").\n", i, frag->units[i].type);
163             return err;
164         }
165     }
166
167     return 0;
168 }
169
170 int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
171                           CodedBitstreamFragment *frag,
172                           const AVCodecParameters *par)
173 {
174     int err;
175
176     memset(frag, 0, sizeof(*frag));
177
178     frag->data      = par->extradata;
179     frag->data_size = par->extradata_size;
180
181     err = ctx->codec->split_fragment(ctx, frag, 1);
182     if (err < 0)
183         return err;
184
185     frag->data      = NULL;
186     frag->data_size = 0;
187
188     return cbs_read_fragment_content(ctx, frag);
189 }
190
191 static int cbs_fill_fragment_data(CodedBitstreamContext *ctx,
192                                   CodedBitstreamFragment *frag,
193                                   const uint8_t *data, size_t size)
194 {
195     av_assert0(!frag->data && !frag->data_ref);
196
197     frag->data_ref =
198         av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
199     if (!frag->data_ref)
200         return AVERROR(ENOMEM);
201
202     frag->data      = frag->data_ref->data;
203     frag->data_size = size;
204
205     memcpy(frag->data, data, size);
206     memset(frag->data + size, 0,
207            AV_INPUT_BUFFER_PADDING_SIZE);
208
209     return 0;
210 }
211
212 int ff_cbs_read_packet(CodedBitstreamContext *ctx,
213                        CodedBitstreamFragment *frag,
214                        const AVPacket *pkt)
215 {
216     int err;
217
218     memset(frag, 0, sizeof(*frag));
219
220     if (pkt->buf) {
221         frag->data_ref = av_buffer_ref(pkt->buf);
222         if (!frag->data_ref)
223             return AVERROR(ENOMEM);
224
225         frag->data      = pkt->data;
226         frag->data_size = pkt->size;
227
228     } else {
229         err = cbs_fill_fragment_data(ctx, frag, pkt->data, pkt->size);
230         if (err < 0)
231             return err;
232     }
233
234     err = ctx->codec->split_fragment(ctx, frag, 0);
235     if (err < 0)
236         return err;
237
238     return cbs_read_fragment_content(ctx, frag);
239 }
240
241 int ff_cbs_read(CodedBitstreamContext *ctx,
242                 CodedBitstreamFragment *frag,
243                 const uint8_t *data, size_t size)
244 {
245     int err;
246
247     memset(frag, 0, sizeof(*frag));
248
249     err = cbs_fill_fragment_data(ctx, frag, data, size);
250     if (err < 0)
251         return err;
252
253     err = ctx->codec->split_fragment(ctx, frag, 0);
254     if (err < 0)
255         return err;
256
257     return cbs_read_fragment_content(ctx, frag);
258 }
259
260
261 int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
262                                CodedBitstreamFragment *frag)
263 {
264     int err, i;
265
266     for (i = 0; i < frag->nb_units; i++) {
267         CodedBitstreamUnit *unit = &frag->units[i];
268
269         if (!unit->content)
270             continue;
271
272         av_buffer_unref(&unit->data_ref);
273         unit->data = NULL;
274
275         err = ctx->codec->write_unit(ctx, unit);
276         if (err < 0) {
277             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to write unit %d "
278                    "(type %"PRIu32").\n", i, unit->type);
279             return err;
280         }
281     }
282
283     av_buffer_unref(&frag->data_ref);
284     frag->data = NULL;
285
286     err = ctx->codec->assemble_fragment(ctx, frag);
287     if (err < 0) {
288         av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to assemble fragment.\n");
289         return err;
290     }
291
292     return 0;
293 }
294
295 int ff_cbs_write_extradata(CodedBitstreamContext *ctx,
296                            AVCodecParameters *par,
297                            CodedBitstreamFragment *frag)
298 {
299     int err;
300
301     err = ff_cbs_write_fragment_data(ctx, frag);
302     if (err < 0)
303         return err;
304
305     av_freep(&par->extradata);
306
307     par->extradata = av_malloc(frag->data_size +
308                                AV_INPUT_BUFFER_PADDING_SIZE);
309     if (!par->extradata)
310         return AVERROR(ENOMEM);
311
312     memcpy(par->extradata, frag->data, frag->data_size);
313     memset(par->extradata + frag->data_size, 0,
314            AV_INPUT_BUFFER_PADDING_SIZE);
315     par->extradata_size = frag->data_size;
316
317     return 0;
318 }
319
320 int ff_cbs_write_packet(CodedBitstreamContext *ctx,
321                         AVPacket *pkt,
322                         CodedBitstreamFragment *frag)
323 {
324     AVBufferRef *buf;
325     int err;
326
327     err = ff_cbs_write_fragment_data(ctx, frag);
328     if (err < 0)
329         return err;
330
331     av_assert0(frag->data_ref);
332     buf = av_buffer_ref(frag->data_ref);
333     if (!buf)
334         return AVERROR(ENOMEM);
335
336     av_init_packet(pkt);
337     pkt->buf  = buf;
338     pkt->data = frag->data;
339     pkt->size = frag->data_size;
340
341     return 0;
342 }
343
344
345 void ff_cbs_trace_header(CodedBitstreamContext *ctx,
346                          const char *name)
347 {
348     if (!ctx->trace_enable)
349         return;
350
351     av_log(ctx->log_ctx, ctx->trace_level, "%s\n", name);
352 }
353
354 void ff_cbs_trace_syntax_element(CodedBitstreamContext *ctx, int position,
355                                  const char *name, const char *bits,
356                                  int64_t value)
357 {
358     size_t name_len, bits_len;
359     int pad;
360
361     if (!ctx->trace_enable)
362         return;
363
364     av_assert0(value >= INT_MIN && value <= UINT32_MAX);
365
366     name_len = strlen(name);
367     bits_len = strlen(bits);
368
369     if (name_len + bits_len > 60)
370         pad = bits_len + 2;
371     else
372         pad = 61 - name_len;
373
374     av_log(ctx->log_ctx, ctx->trace_level, "%-10d  %s%*s = %"PRId64"\n",
375            position, name, pad, bits, value);
376 }
377
378 int ff_cbs_read_unsigned(CodedBitstreamContext *ctx, GetBitContext *gbc,
379                          int width, const char *name, uint32_t *write_to,
380                          uint32_t range_min, uint32_t range_max)
381 {
382     uint32_t value;
383     int position;
384
385     av_assert0(width > 0 && width <= 32);
386
387     if (get_bits_left(gbc) < width) {
388         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid value at "
389                "%s: bitstream ended.\n", name);
390         return AVERROR_INVALIDDATA;
391     }
392
393     if (ctx->trace_enable)
394         position = get_bits_count(gbc);
395
396     value = get_bits_long(gbc, width);
397
398     if (ctx->trace_enable) {
399         char bits[33];
400         int i;
401         for (i = 0; i < width; i++)
402             bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
403         bits[i] = 0;
404
405         ff_cbs_trace_syntax_element(ctx, position, name, bits, value);
406     }
407
408     if (value < range_min || value > range_max) {
409         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
410                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
411                name, value, range_min, range_max);
412         return AVERROR_INVALIDDATA;
413     }
414
415     *write_to = value;
416     return 0;
417 }
418
419 int ff_cbs_write_unsigned(CodedBitstreamContext *ctx, PutBitContext *pbc,
420                           int width, const char *name, uint32_t value,
421                           uint32_t range_min, uint32_t range_max)
422 {
423     av_assert0(width > 0 && width <= 32);
424
425     if (value < range_min || value > range_max) {
426         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
427                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
428                name, value, range_min, range_max);
429         return AVERROR_INVALIDDATA;
430     }
431
432     if (put_bits_left(pbc) < width)
433         return AVERROR(ENOSPC);
434
435     if (ctx->trace_enable) {
436         char bits[33];
437         int i;
438         for (i = 0; i < width; i++)
439             bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
440         bits[i] = 0;
441
442         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc), name, bits, value);
443     }
444
445     if (width < 32)
446         put_bits(pbc, width, value);
447     else
448         put_bits32(pbc, value);
449
450     return 0;
451 }
452
453
454 int ff_cbs_alloc_unit_content(CodedBitstreamContext *ctx,
455                               CodedBitstreamUnit *unit,
456                               size_t size,
457                               void (*free)(void *opaque, uint8_t *data))
458 {
459     av_assert0(!unit->content && !unit->content_ref);
460
461     unit->content = av_mallocz(size);
462     if (!unit->content)
463         return AVERROR(ENOMEM);
464
465     unit->content_ref = av_buffer_create(unit->content, size,
466                                          free, ctx, 0);
467     if (!unit->content_ref) {
468         av_freep(&unit->content);
469         return AVERROR(ENOMEM);
470     }
471
472     return 0;
473 }
474
475 int ff_cbs_alloc_unit_data(CodedBitstreamContext *ctx,
476                            CodedBitstreamUnit *unit,
477                            size_t size)
478 {
479     av_assert0(!unit->data && !unit->data_ref);
480
481     unit->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
482     if (!unit->data_ref)
483         return AVERROR(ENOMEM);
484
485     unit->data      = unit->data_ref->data;
486     unit->data_size = size;
487
488     memset(unit->data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
489
490     return 0;
491 }
492
493 static int cbs_insert_unit(CodedBitstreamContext *ctx,
494                            CodedBitstreamFragment *frag,
495                            int position)
496 {
497     CodedBitstreamUnit *units;
498
499     units = av_malloc_array(frag->nb_units + 1, sizeof(*units));
500     if (!units)
501         return AVERROR(ENOMEM);
502
503     if (position > 0)
504         memcpy(units, frag->units, position * sizeof(*units));
505     if (position < frag->nb_units)
506         memcpy(units + position + 1, frag->units + position,
507                (frag->nb_units - position) * sizeof(*units));
508
509     memset(units + position, 0, sizeof(*units));
510
511     av_freep(&frag->units);
512     frag->units = units;
513     ++frag->nb_units;
514
515     return 0;
516 }
517
518 int ff_cbs_insert_unit_content(CodedBitstreamContext *ctx,
519                                CodedBitstreamFragment *frag,
520                                int position,
521                                CodedBitstreamUnitType type,
522                                void *content,
523                                AVBufferRef *content_buf)
524 {
525     CodedBitstreamUnit *unit;
526     AVBufferRef *content_ref;
527     int err;
528
529     if (position == -1)
530         position = frag->nb_units;
531     av_assert0(position >= 0 && position <= frag->nb_units);
532
533     if (content_buf) {
534         content_ref = av_buffer_ref(content_buf);
535         if (!content_ref)
536             return AVERROR(ENOMEM);
537     } else {
538         content_ref = NULL;
539     }
540
541     err = cbs_insert_unit(ctx, frag, position);
542     if (err < 0) {
543         av_buffer_unref(&content_ref);
544         return err;
545     }
546
547     unit = &frag->units[position];
548     unit->type        = type;
549     unit->content     = content;
550     unit->content_ref = content_ref;
551
552     return 0;
553 }
554
555 int ff_cbs_insert_unit_data(CodedBitstreamContext *ctx,
556                             CodedBitstreamFragment *frag,
557                             int position,
558                             CodedBitstreamUnitType type,
559                             uint8_t *data, size_t data_size,
560                             AVBufferRef *data_buf)
561 {
562     CodedBitstreamUnit *unit;
563     AVBufferRef *data_ref;
564     int err;
565
566     if (position == -1)
567         position = frag->nb_units;
568     av_assert0(position >= 0 && position <= frag->nb_units);
569
570     if (data_buf)
571         data_ref = av_buffer_ref(data_buf);
572     else
573         data_ref = av_buffer_create(data, data_size, NULL, NULL, 0);
574     if (!data_ref)
575         return AVERROR(ENOMEM);
576
577     err = cbs_insert_unit(ctx, frag, position);
578     if (err < 0) {
579         av_buffer_unref(&data_ref);
580         return err;
581     }
582
583     unit = &frag->units[position];
584     unit->type      = type;
585     unit->data      = data;
586     unit->data_size = data_size;
587     unit->data_ref  = data_ref;
588
589     return 0;
590 }
591
592 int ff_cbs_delete_unit(CodedBitstreamContext *ctx,
593                        CodedBitstreamFragment *frag,
594                        int position)
595 {
596     if (position < 0 || position >= frag->nb_units)
597         return AVERROR(EINVAL);
598
599     cbs_unit_uninit(ctx, &frag->units[position]);
600
601     --frag->nb_units;
602
603     if (frag->nb_units == 0) {
604         av_freep(&frag->units);
605
606     } else {
607         memmove(frag->units + position,
608                 frag->units + position + 1,
609                 (frag->nb_units - position) * sizeof(*frag->units));
610
611         // Don't bother reallocating the unit array.
612     }
613
614     return 0;
615 }