lavc: Add coded bitstream read/write support for H.264
[libav.git] / libavcodec / cbs_h2645.c
1 /*
2 * This file is part of Libav.
3 *
4 * Libav 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 * Libav 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 Libav; if not, write to the Free Software
16 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
17 */
18
19 #include "libavutil/attributes.h"
20 #include "libavutil/avassert.h"
21
22 #include "bytestream.h"
23 #include "cbs.h"
24 #include "cbs_internal.h"
25 #include "cbs_h264.h"
26 #include "golomb.h"
27 #include "h264.h"
28 #include "h264_sei.h"
29 #include "h2645_parse.h"
30
31
32 static int cbs_read_ue_golomb(CodedBitstreamContext *ctx, BitstreamContext *bc,
33 const char *name, uint32_t *write_to,
34 uint32_t range_min, uint32_t range_max)
35 {
36 uint32_t value;
37 int position;
38
39 if (ctx->trace_enable) {
40 char bits[65];
41 unsigned int k;
42 int i, j;
43
44 position = bitstream_tell(bc);
45
46 for (i = 0; i < 32; i++) {
47 k = bitstream_read_bit(bc);
48 bits[i] = k ? '1' : '0';
49 if (k)
50 break;
51 }
52 if (i >= 32) {
53 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid ue-golomb "
54 "code found while reading %s: "
55 "more than 31 zeroes.\n", name);
56 return AVERROR_INVALIDDATA;
57 }
58 value = 1;
59 for (j = 0; j < i; j++) {
60 k = bitstream_read_bit(bc);
61 bits[i + j + 1] = k ? '1' : '0';
62 value = value << 1 | k;
63 }
64 bits[i + j + 1] = 0;
65 --value;
66
67 ff_cbs_trace_syntax_element(ctx, position, name, bits, value);
68 } else {
69 value = get_ue_golomb_long(bc);
70 }
71
72 if (value < range_min || value > range_max) {
73 av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
74 "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
75 name, value, range_min, range_max);
76 return AVERROR_INVALIDDATA;
77 }
78
79 *write_to = value;
80 return 0;
81 }
82
83 static int cbs_read_se_golomb(CodedBitstreamContext *ctx, BitstreamContext *bc,
84 const char *name, int32_t *write_to,
85 int32_t range_min, int32_t range_max)
86 {
87 int32_t value;
88 int position;
89
90 if (ctx->trace_enable) {
91 char bits[65];
92 uint32_t v;
93 unsigned int k;
94 int i, j;
95
96 position = bitstream_tell(bc);
97
98 for (i = 0; i < 32; i++) {
99 k = bitstream_read_bit(bc);
100 bits[i] = k ? '1' : '0';
101 if (k)
102 break;
103 }
104 if (i >= 32) {
105 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid se-golomb "
106 "code found while reading %s: "
107 "more than 31 zeroes.\n", name);
108 return AVERROR_INVALIDDATA;
109 }
110 v = 1;
111 for (j = 0; j < i; j++) {
112 k = bitstream_read_bit(bc);
113 bits[i + j + 1] = k ? '1' : '0';
114 v = v << 1 | k;
115 }
116 bits[i + j + 1] = 0;
117 if (v & 1)
118 value = -(int32_t)(v / 2);
119 else
120 value = v / 2;
121
122 ff_cbs_trace_syntax_element(ctx, position, name, bits, value);
123 } else {
124 value = get_se_golomb_long(bc);
125 }
126
127 if (value < range_min || value > range_max) {
128 av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
129 "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
130 name, value, range_min, range_max);
131 return AVERROR_INVALIDDATA;
132 }
133
134 *write_to = value;
135 return 0;
136 }
137
138 static int cbs_write_ue_golomb(CodedBitstreamContext *ctx, PutBitContext *pbc,
139 const char *name, uint32_t value,
140 uint32_t range_min, uint32_t range_max)
141 {
142 int len;
143
144 if (value < range_min || value > range_max) {
145 av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
146 "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
147 name, value, range_min, range_max);
148 return AVERROR_INVALIDDATA;
149 }
150 av_assert0(value != UINT32_MAX);
151
152 len = av_log2(value + 1);
153 if (put_bits_left(pbc) < 2 * len + 1)
154 return AVERROR(ENOSPC);
155
156 if (ctx->trace_enable) {
157 char bits[65];
158 int i;
159
160 for (i = 0; i < len; i++)
161 bits[i] = '0';
162 bits[len] = '1';
163 for (i = 0; i < len; i++)
164 bits[len + i + 1] = (value + 1) >> (len - i - 1) & 1 ? '1' : '0';
165 bits[len + len + 1] = 0;
166
167 ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc), name, bits, value);
168 }
169
170 put_bits(pbc, len, 0);
171 if (len + 1 < 32)
172 put_bits(pbc, len + 1, value + 1);
173 else
174 put_bits32(pbc, value + 1);
175
176 return 0;
177 }
178
179 static int cbs_write_se_golomb(CodedBitstreamContext *ctx, PutBitContext *pbc,
180 const char *name, int32_t value,
181 int32_t range_min, int32_t range_max)
182 {
183 int len;
184 uint32_t uvalue;
185
186 if (value < range_min || value > range_max) {
187 av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
188 "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
189 name, value, range_min, range_max);
190 return AVERROR_INVALIDDATA;
191 }
192 av_assert0(value != INT32_MIN);
193
194 if (value == 0)
195 uvalue = 0;
196 else if (value > 0)
197 uvalue = 2 * (uint32_t)value - 1;
198 else
199 uvalue = 2 * (uint32_t)-value;
200
201 len = av_log2(uvalue + 1);
202 if (put_bits_left(pbc) < 2 * len + 1)
203 return AVERROR(ENOSPC);
204
205 if (ctx->trace_enable) {
206 char bits[65];
207 int i;
208
209 for (i = 0; i < len; i++)
210 bits[i] = '0';
211 bits[len] = '1';
212 for (i = 0; i < len; i++)
213 bits[len + i + 1] = (uvalue + 1) >> (len - i - 1) & 1 ? '1' : '0';
214 bits[len + len + 1] = 0;
215
216 ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc), name, bits, value);
217 }
218
219 put_bits(pbc, len, 0);
220 if (len + 1 < 32)
221 put_bits(pbc, len + 1, uvalue + 1);
222 else
223 put_bits32(pbc, uvalue + 1);
224
225 return 0;
226 }
227
228 #define HEADER(name) do { \
229 ff_cbs_trace_header(ctx, name); \
230 } while (0)
231
232 #define CHECK(call) do { \
233 err = (call); \
234 if (err < 0) \
235 return err; \
236 } while (0)
237
238 #define FUNC_NAME(rw, codec, name) cbs_ ## codec ## _ ## rw ## _ ## name
239 #define FUNC_H264(rw, name) FUNC_NAME(rw, h264, name)
240
241
242 #define READ
243 #define READWRITE read
244 #define RWContext BitstreamContext
245
246 #define xu(width, name, var, range_min, range_max) do { \
247 uint32_t value = range_min; \
248 CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
249 &value, range_min, range_max)); \
250 var = value; \
251 } while (0)
252 #define xue(name, var, range_min, range_max) do { \
253 uint32_t value = range_min; \
254 CHECK(cbs_read_ue_golomb(ctx, rw, #name, \
255 &value, range_min, range_max)); \
256 var = value; \
257 } while (0)
258 #define xse(name, var, range_min, range_max) do { \
259 int32_t value = range_min; \
260 CHECK(cbs_read_se_golomb(ctx, rw, #name, \
261 &value, range_min, range_max)); \
262 var = value; \
263 } while (0)
264
265
266 #define u(width, name, range_min, range_max) \
267 xu(width, name, current->name, range_min, range_max)
268 #define flag(name) u(1, name, 0, 1)
269 #define ue(name, range_min, range_max) \
270 xue(name, current->name, range_min, range_max)
271 #define se(name, range_min, range_max) \
272 xse(name, current->name, range_min, range_max)
273
274 #define infer(name, value) do { \
275 current->name = value; \
276 } while (0)
277
278 static int cbs_h2645_read_more_rbsp_data(BitstreamContext *bc)
279 {
280 int bits_left = bitstream_bits_left(bc);
281 if (bits_left > 8)
282 return 1;
283 if (bitstream_peek(bc, bits_left) == 1 << (bits_left - 1))
284 return 0;
285 return 1;
286 }
287
288 #define more_rbsp_data(var) ((var) = cbs_h2645_read_more_rbsp_data(rw))
289
290 #define byte_alignment(rw) (bitstream_tell(rw) % 8)
291
292 #define allocate(name, size) do { \
293 name = av_mallocz(size); \
294 if (!name) \
295 return AVERROR(ENOMEM); \
296 } while (0)
297
298 #define FUNC(name) FUNC_H264(READWRITE, name)
299 #include "cbs_h264_syntax_template.c"
300 #undef FUNC
301
302 #undef READ
303 #undef READWRITE
304 #undef RWContext
305 #undef xu
306 #undef xue
307 #undef xse
308 #undef u
309 #undef flag
310 #undef ue
311 #undef se
312 #undef infer
313 #undef more_rbsp_data
314 #undef byte_alignment
315 #undef allocate
316
317
318 #define WRITE
319 #define READWRITE write
320 #define RWContext PutBitContext
321
322 #define xu(width, name, var, range_min, range_max) do { \
323 uint32_t value = var; \
324 CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
325 value, range_min, range_max)); \
326 } while (0)
327 #define xue(name, var, range_min, range_max) do { \
328 uint32_t value = var; \
329 CHECK(cbs_write_ue_golomb(ctx, rw, #name, \
330 value, range_min, range_max)); \
331 } while (0)
332 #define xse(name, var, range_min, range_max) do { \
333 int32_t value = var; \
334 CHECK(cbs_write_se_golomb(ctx, rw, #name, \
335 value, range_min, range_max)); \
336 } while (0)
337
338 #define u(width, name, range_min, range_max) \
339 xu(width, name, current->name, range_min, range_max)
340 #define flag(name) u(1, name, 0, 1)
341 #define ue(name, range_min, range_max) \
342 xue(name, current->name, range_min, range_max)
343 #define se(name, range_min, range_max) \
344 xse(name, current->name, range_min, range_max)
345
346 #define infer(name, value) do { \
347 if (current->name != (value)) { \
348 av_log(ctx->log_ctx, AV_LOG_WARNING, "Warning: " \
349 "%s does not match inferred value: " \
350 "%"PRId64", but should be %"PRId64".\n", \
351 #name, (int64_t)current->name, (int64_t)(value)); \
352 } \
353 } while (0)
354
355 #define more_rbsp_data(var) (var)
356
357 #define byte_alignment(rw) (put_bits_count(rw) % 8)
358
359 #define allocate(name, size) do { \
360 if (!name) { \
361 av_log(ctx->log_ctx, AV_LOG_ERROR, "%s must be set " \
362 "for writing.\n", #name); \
363 return AVERROR_INVALIDDATA; \
364 } \
365 } while (0)
366
367 #define FUNC(name) FUNC_H264(READWRITE, name)
368 #include "cbs_h264_syntax_template.c"
369 #undef FUNC
370
371 #undef WRITE
372 #undef READWRITE
373 #undef RWContext
374 #undef xu
375 #undef xue
376 #undef xse
377 #undef u
378 #undef flag
379 #undef ue
380 #undef se
381 #undef infer
382 #undef more_rbsp_data
383 #undef byte_alignment
384 #undef allocate
385
386
387 static void cbs_h264_free_sei(H264RawSEI *sei)
388 {
389 int i;
390 for (i = 0; i < sei->payload_count; i++) {
391 H264RawSEIPayload *payload = &sei->payload[i];
392
393 switch (payload->payload_type) {
394 case H264_SEI_TYPE_BUFFERING_PERIOD:
395 case H264_SEI_TYPE_PIC_TIMING:
396 case H264_SEI_TYPE_RECOVERY_POINT:
397 case H264_SEI_TYPE_DISPLAY_ORIENTATION:
398 break;
399 case H264_SEI_TYPE_USER_DATA_REGISTERED:
400 av_freep(&payload->payload.user_data_registered.data);
401 break;
402 case H264_SEI_TYPE_USER_DATA_UNREGISTERED:
403 av_freep(&payload->payload.user_data_unregistered.data);
404 break;
405 default:
406 av_freep(&payload->payload.other.data);
407 break;
408 }
409 }
410 }
411
412 static void cbs_h264_free_slice(H264RawSlice *slice)
413 {
414 av_freep(&slice->data);
415 }
416
417 static void cbs_h264_free_nal_unit(CodedBitstreamUnit *unit)
418 {
419 switch (unit->type) {
420 case H264_NAL_SEI:
421 cbs_h264_free_sei(unit->content);
422 break;
423 case H264_NAL_IDR_SLICE:
424 case H264_NAL_SLICE:
425 cbs_h264_free_slice(unit->content);
426 break;
427 }
428 av_freep(&unit->content);
429 }
430
431 static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
432 CodedBitstreamFragment *frag,
433 const H2645Packet *packet)
434 {
435 int err, i;
436
437 for (i = 0; i < packet->nb_nals; i++) {
438 const H2645NAL *nal = &packet->nals[i];
439 uint8_t *data;
440
441 data = av_malloc(nal->size);
442 if (!data)
443 return AVERROR(ENOMEM);
444 memcpy(data, nal->data, nal->size);
445
446 err = ff_cbs_insert_unit_data(ctx, frag, -1, nal->type,
447 data, nal->size);
448 if (err < 0) {
449 av_freep(&data);
450 return err;
451 }
452 }
453
454 return 0;
455 }
456
457 static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
458 CodedBitstreamFragment *frag,
459 int header)
460 {
461 enum AVCodecID codec_id = ctx->codec->codec_id;
462 CodedBitstreamH2645Context *priv = ctx->priv_data;
463 GetByteContext gbc;
464 int err;
465
466 av_assert0(frag->data && frag->nb_units == 0);
467 if (frag->data_size == 0)
468 return 0;
469
470 if (header && frag->data[0] && codec_id == AV_CODEC_ID_H264) {
471 // AVCC header.
472 size_t size, start, end;
473 int i, count, version;
474
475 priv->mp4 = 1;
476
477 bytestream2_init(&gbc, frag->data, frag->data_size);
478
479 if (bytestream2_get_bytes_left(&gbc) < 6)
480 return AVERROR_INVALIDDATA;
481
482 version = bytestream2_get_byte(&gbc);
483 if (version != 1) {
484 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid AVCC header: "
485 "first byte %u.", version);
486 return AVERROR_INVALIDDATA;
487 }
488
489 bytestream2_skip(&gbc, 3);
490 priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
491
492 // SPS array.
493 count = bytestream2_get_byte(&gbc) & 0x1f;
494 start = bytestream2_tell(&gbc);
495 for (i = 0; i < count; i++) {
496 if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
497 return AVERROR_INVALIDDATA;
498 size = bytestream2_get_be16(&gbc);
499 if (bytestream2_get_bytes_left(&gbc) < size)
500 return AVERROR_INVALIDDATA;
501 bytestream2_skip(&gbc, size);
502 }
503 end = bytestream2_tell(&gbc);
504
505 err = ff_h2645_packet_split(&priv->read_packet,
506 frag->data + start, end - start,
507 ctx->log_ctx, 1, 2, AV_CODEC_ID_H264);
508 if (err < 0) {
509 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC SPS array.\n");
510 return err;
511 }
512 err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
513 if (err < 0)
514 return err;
515
516 // PPS array.
517 count = bytestream2_get_byte(&gbc);
518 start = bytestream2_tell(&gbc);
519 for (i = 0; i < count; i++) {
520 if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
521 return AVERROR_INVALIDDATA;
522 size = bytestream2_get_be16(&gbc);
523 if (bytestream2_get_bytes_left(&gbc) < size)
524 return AVERROR_INVALIDDATA;
525 bytestream2_skip(&gbc, size);
526 }
527 end = bytestream2_tell(&gbc);
528
529 err = ff_h2645_packet_split(&priv->read_packet,
530 frag->data + start, end - start,
531 ctx->log_ctx, 1, 2, AV_CODEC_ID_H264);
532 if (err < 0) {
533 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC PPS array.\n");
534 return err;
535 }
536 err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
537 if (err < 0)
538 return err;
539
540 if (bytestream2_get_bytes_left(&gbc) > 0) {
541 av_log(ctx->log_ctx, AV_LOG_WARNING, "%u bytes left at end of AVCC "
542 "header.\n", bytestream2_get_bytes_left(&gbc));
543 }
544
545 } else {
546 // Annex B, or later MP4 with already-known parameters.
547
548 err = ff_h2645_packet_split(&priv->read_packet,
549 frag->data, frag->data_size,
550 ctx->log_ctx,
551 priv->mp4, priv->nal_length_size,
552 codec_id);
553 if (err < 0)
554 return err;
555
556 err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
557 if (err < 0)
558 return err;
559 }
560
561 return 0;
562 }
563
564 #define cbs_h2645_replace_ps(h26n, ps_name, ps_var, id_element) \
565 static int cbs_h26 ## h26n ## _replace_ ## ps_var(CodedBitstreamContext *ctx, \
566 const H26 ## h26n ## Raw ## ps_name *ps_var) \
567 { \
568 CodedBitstreamH26 ## h26n ## Context *priv = ctx->priv_data; \
569 unsigned int id = ps_var->id_element; \
570 if (id > FF_ARRAY_ELEMS(priv->ps_var)) { \
571 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid " #ps_name \
572 " id : %d.\n", id); \
573 return AVERROR_INVALIDDATA; \
574 } \
575 av_freep(&priv->ps_var[id]); \
576 priv->ps_var[id] = av_malloc(sizeof(*ps_var)); \
577 if (!priv->ps_var[id]) \
578 return AVERROR(ENOMEM); \
579 memcpy(priv->ps_var[id], ps_var, sizeof(*ps_var)); \
580 return 0; \
581 }
582
583 cbs_h2645_replace_ps(4, SPS, sps, seq_parameter_set_id)
584 cbs_h2645_replace_ps(4, PPS, pps, pic_parameter_set_id)
585
586 static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
587 CodedBitstreamUnit *unit)
588 {
589 BitstreamContext bc;
590 int err;
591
592 err = bitstream_init(&bc, unit->data, 8 * unit->data_size);
593 if (err < 0)
594 return err;
595
596 switch (unit->type) {
597 case H264_NAL_SPS:
598 {
599 H264RawSPS *sps;
600
601 sps = av_mallocz(sizeof(*sps));
602 if (!sps)
603 return AVERROR(ENOMEM);
604 err = cbs_h264_read_sps(ctx, &bc, sps);
605 if (err >= 0)
606 err = cbs_h264_replace_sps(ctx, sps);
607 if (err < 0) {
608 av_free(sps);
609 return err;
610 }
611
612 unit->content = sps;
613 }
614 break;
615
616 case H264_NAL_SPS_EXT:
617 {
618 H264RawSPSExtension *sps_ext;
619
620 sps_ext = av_mallocz(sizeof(*sps_ext));
621 if (!sps_ext)
622 return AVERROR(ENOMEM);
623 err = cbs_h264_read_sps_extension(ctx, &bc, sps_ext);
624 if (err < 0) {
625 av_free(sps_ext);
626 return err;
627 }
628
629 unit->content = sps_ext;
630 }
631 break;
632
633 case H264_NAL_PPS:
634 {
635 H264RawPPS *pps;
636
637 pps = av_mallocz(sizeof(*pps));
638 if (!pps)
639 return AVERROR(ENOMEM);
640 err = cbs_h264_read_pps(ctx, &bc, pps);
641 if (err >= 0)
642 err = cbs_h264_replace_pps(ctx, pps);
643 if (err < 0) {
644 av_free(pps);
645 return err;
646 }
647
648 unit->content = pps;
649 }
650 break;
651
652 case H264_NAL_SLICE:
653 case H264_NAL_IDR_SLICE:
654 case H264_NAL_AUXILIARY_SLICE:
655 {
656 H264RawSlice *slice;
657 int pos, len;
658
659 slice = av_mallocz(sizeof(*slice));
660 if (!slice)
661 return AVERROR(ENOMEM);
662 err = cbs_h264_read_slice_header(ctx, &bc, &slice->header);
663 if (err < 0) {
664 av_free(slice);
665 return err;
666 }
667
668 pos = bitstream_tell(&bc);
669 len = unit->data_size;
670 if (!unit->data[len - 1]) {
671 int z;
672 for (z = 0; z < len && !unit->data[len - z - 1]; z++);
673 av_log(ctx->log_ctx, AV_LOG_DEBUG, "Deleted %d trailing zeroes "
674 "from slice data.\n", z);
675 len -= z;
676 }
677
678 slice->data_size = len - pos / 8;
679 slice->data = av_malloc(slice->data_size);
680 if (!slice->data) {
681 av_free(slice);
682 return AVERROR(ENOMEM);
683 }
684 memcpy(slice->data,
685 unit->data + pos / 8, slice->data_size);
686 slice->data_bit_start = pos % 8;
687
688 unit->content = slice;
689 }
690 break;
691
692 case H264_NAL_AUD:
693 {
694 H264RawAUD *aud;
695
696 aud = av_mallocz(sizeof(*aud));
697 if (!aud)
698 return AVERROR(ENOMEM);
699 err = cbs_h264_read_aud(ctx, &bc, aud);
700 if (err < 0) {
701 av_free(aud);
702 return err;
703 }
704
705 unit->content = aud;
706 }
707 break;
708
709 case H264_NAL_SEI:
710 {
711 H264RawSEI *sei;
712
713 sei = av_mallocz(sizeof(*sei));
714 if (!sei)
715 return AVERROR(ENOMEM);
716 err = cbs_h264_read_sei(ctx, &bc, sei);
717 if (err < 0) {
718 cbs_h264_free_sei(sei);
719 return err;
720 }
721
722 unit->content = sei;
723 }
724 break;
725
726 default:
727 return AVERROR(ENOSYS);
728 }
729
730 return 0;
731 }
732
733 static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
734 CodedBitstreamUnit *unit,
735 PutBitContext *pbc)
736 {
737 int err;
738
739 switch (unit->type) {
740 case H264_NAL_SPS:
741 {
742 H264RawSPS *sps = unit->content;
743
744 err = cbs_h264_write_sps(ctx, pbc, sps);
745 if (err < 0)
746 return err;
747
748 err = cbs_h264_replace_sps(ctx, sps);
749 if (err < 0)
750 return err;
751 }
752 break;
753
754 case H264_NAL_SPS_EXT:
755 {
756 H264RawSPSExtension *sps_ext;
757
758 err = cbs_h264_write_sps_extension(ctx, pbc, sps_ext);
759 if (err < 0)
760 return err;
761 }
762 break;
763
764 case H264_NAL_PPS:
765 {
766 H264RawPPS *pps = unit->content;
767
768 err = cbs_h264_write_pps(ctx, pbc, pps);
769 if (err < 0)
770 return err;
771
772 err = cbs_h264_replace_pps(ctx, pps);
773 if (err < 0)
774 return err;
775 }
776 break;
777
778 case H264_NAL_SLICE:
779 case H264_NAL_IDR_SLICE:
780 {
781 H264RawSlice *slice = unit->content;
782 BitstreamContext bc;
783 int bits_left, end, zeroes;
784
785 err = cbs_h264_write_slice_header(ctx, pbc, &slice->header);
786 if (err < 0)
787 return err;
788
789 if (slice->data) {
790 if (slice->data_size * 8 + 8 > put_bits_left(pbc))
791 return AVERROR(ENOSPC);
792
793 bitstream_init(&bc, slice->data, slice->data_size * 8);
794 bitstream_skip(&bc, slice->data_bit_start);
795
796 // Copy in two-byte blocks, but stop before copying the
797 // rbsp_stop_one_bit in the final byte.
798 while (bitstream_bits_left(&bc) > 23)
799 put_bits(pbc, 16, bitstream_read(&bc, 16));
800
801 bits_left = bitstream_bits_left(&bc);
802 end = bitstream_read(&bc, bits_left);
803
804 // rbsp_stop_one_bit must be present here.
805 av_assert0(end);
806 zeroes = ff_ctz(end);
807 if (bits_left > zeroes + 1)
808 put_bits(pbc, bits_left - zeroes - 1,
809 end >> (zeroes + 1));
810 put_bits(pbc, 1, 1);
811 while (put_bits_count(pbc) % 8 != 0)
812 put_bits(pbc, 1, 0);
813 } else {
814 // No slice data - that was just the header.
815 // (Bitstream may be unaligned!)
816 }
817 }
818 break;
819
820 case H264_NAL_AUD:
821 {
822 err = cbs_h264_write_aud(ctx, pbc, unit->content);
823 if (err < 0)
824 return err;
825 }
826 break;
827
828 case H264_NAL_SEI:
829 {
830 err = cbs_h264_write_sei(ctx, pbc, unit->content);
831 if (err < 0)
832 return err;
833 }
834 break;
835
836 default:
837 av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
838 "NAL unit type %"PRIu32".\n", unit->type);
839 return AVERROR_PATCHWELCOME;
840 }
841
842 return 0;
843 }
844
845 static int cbs_h2645_write_nal_unit(CodedBitstreamContext *ctx,
846 CodedBitstreamUnit *unit)
847 {
848 CodedBitstreamH2645Context *priv = ctx->priv_data;
849 enum AVCodecID codec_id = ctx->codec->codec_id;
850 PutBitContext pbc;
851 int err;
852
853 if (!priv->write_buffer) {
854 // Initial write buffer size is 1MB.
855 priv->write_buffer_size = 1024 * 1024;
856
857 reallocate_and_try_again:
858 err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
859 if (err < 0) {
860 av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
861 "sufficiently large write buffer (last attempt "
862 "%zu bytes).\n", priv->write_buffer_size);
863 return err;
864 }
865 }
866
867 init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
868
869 err = cbs_h264_write_nal_unit(ctx, unit, &pbc);
870
871 if (err == AVERROR(ENOSPC)) {
872 // Overflow.
873 priv->write_buffer_size *= 2;
874 goto reallocate_and_try_again;
875 }
876 // Overflow but we didn't notice.
877 av_assert0(put_bits_count(&pbc) <= 8 * priv->write_buffer_size);
878
879 if (put_bits_count(&pbc) % 8)
880 unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
881 else
882 unit->data_bit_padding = 0;
883
884 unit->data_size = (put_bits_count(&pbc) + 7) / 8;
885 flush_put_bits(&pbc);
886
887 err = av_reallocp(&unit->data, unit->data_size);
888 if (err < 0)
889 return err;
890
891 memcpy(unit->data, priv->write_buffer, unit->data_size);
892
893 return 0;
894 }
895
896 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
897 CodedBitstreamFragment *frag)
898 {
899 uint8_t *data;
900 size_t max_size, dp, sp;
901 int err, i, zero_run;
902
903 for (i = 0; i < frag->nb_units; i++) {
904 // Data should already all have been written when we get here.
905 av_assert0(frag->units[i].data);
906 }
907
908 max_size = 0;
909 for (i = 0; i < frag->nb_units; i++) {
910 // Start code + content with worst-case emulation prevention.
911 max_size += 3 + frag->units[i].data_size * 3 / 2;
912 }
913
914 data = av_malloc(max_size);
915 if (!data)
916 return AVERROR(ENOMEM);
917
918 dp = 0;
919 for (i = 0; i < frag->nb_units; i++) {
920 CodedBitstreamUnit *unit = &frag->units[i];
921
922 if (unit->data_bit_padding > 0) {
923 if (i < frag->nb_units - 1)
924 av_log(ctx->log_ctx, AV_LOG_WARNING, "Probably invalid "
925 "unaligned padding on non-final NAL unit.\n");
926 else
927 frag->data_bit_padding = unit->data_bit_padding;
928 }
929
930 if (unit->type == H264_NAL_SPS ||
931 unit->type == H264_NAL_PPS ||
932 i == 0 /* (Assume this is the start of an access unit.) */) {
933 // zero_byte
934 data[dp++] = 0;
935 }
936 // start_code_prefix_one_3bytes
937 data[dp++] = 0;
938 data[dp++] = 0;
939 data[dp++] = 1;
940
941 zero_run = 0;
942 for (sp = 0; sp < unit->data_size; sp++) {
943 if (zero_run < 2) {
944 if (unit->data[sp] == 0)
945 ++zero_run;
946 else
947 zero_run = 0;
948 } else {
949 if ((unit->data[sp] & ~3) == 0) {
950 // emulation_prevention_three_byte
951 data[dp++] = 3;
952 }
953 zero_run = unit->data[sp] == 0;
954 }
955 data[dp++] = unit->data[sp];
956 }
957 }
958
959 av_assert0(dp <= max_size);
960 err = av_reallocp(&data, dp);
961 if (err)
962 return err;
963
964 frag->data = data;
965 frag->data_size = dp;
966
967 return 0;
968 }
969
970 static void cbs_h264_close(CodedBitstreamContext *ctx)
971 {
972 CodedBitstreamH264Context *h264 = ctx->priv_data;
973 int i;
974
975 ff_h2645_packet_uninit(&h264->common.read_packet);
976
977 av_freep(&h264->common.write_buffer);
978
979 for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
980 av_freep(&h264->sps[i]);
981 for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
982 av_freep(&h264->pps[i]);
983 }
984
985 const CodedBitstreamType ff_cbs_type_h264 = {
986 .codec_id = AV_CODEC_ID_H264,
987
988 .priv_data_size = sizeof(CodedBitstreamH264Context),
989
990 .split_fragment = &cbs_h2645_split_fragment,
991 .read_unit = &cbs_h264_read_nal_unit,
992 .write_unit = &cbs_h2645_write_nal_unit,
993 .assemble_fragment = &cbs_h2645_assemble_fragment,
994
995 .free_unit = &cbs_h264_free_nal_unit,
996 .close = &cbs_h264_close,
997 };