]> git.sesse.net Git - ffmpeg/blob - libavcodec/av1dec.c
avcodec/av1dec: Check for unset obu instead of crashing
[ffmpeg] / libavcodec / av1dec.c
1 /*
2  * AV1 video decoder
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU Lesser General Public
8  * License as published by the Free Software Foundation; either
9  * version 2.1 of the License, or (at your option) any later version.
10  *
11  * FFmpeg is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14  * Lesser General Public License for more details.
15  *
16  * You should have received a copy of the GNU Lesser General Public
17  * License along with FFmpeg; if not, write to the Free Software
18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19  */
20
21 #include "libavutil/pixdesc.h"
22 #include "avcodec.h"
23 #include "av1dec.h"
24 #include "bytestream.h"
25 #include "hwconfig.h"
26 #include "internal.h"
27 #include "profiles.h"
28
29 static void setup_past_independence(AV1Frame *f)
30 {
31     f->loop_filter_delta_enabled = 1;
32
33     f->loop_filter_ref_deltas[AV1_REF_FRAME_INTRA] = 1;
34     f->loop_filter_ref_deltas[AV1_REF_FRAME_LAST] = 0;
35     f->loop_filter_ref_deltas[AV1_REF_FRAME_LAST2] = 0;
36     f->loop_filter_ref_deltas[AV1_REF_FRAME_LAST3] = 0;
37     f->loop_filter_ref_deltas[AV1_REF_FRAME_GOLDEN] = -1;
38     f->loop_filter_ref_deltas[AV1_REF_FRAME_BWDREF] = 0;
39     f->loop_filter_ref_deltas[AV1_REF_FRAME_ALTREF2] = -1;
40     f->loop_filter_ref_deltas[AV1_REF_FRAME_ALTREF] = -1;
41
42     f->loop_filter_mode_deltas[0] = 0;
43     f->loop_filter_mode_deltas[1] = 0;
44 }
45
46 static void load_previous_and_update(AV1DecContext *s)
47 {
48     uint8_t primary_frame, prev_frame;
49
50     primary_frame = s->raw_frame_header->primary_ref_frame;
51     prev_frame = s->raw_frame_header->ref_frame_idx[primary_frame];
52     memcpy(s->cur_frame.loop_filter_ref_deltas,
53            s->ref[prev_frame].loop_filter_ref_deltas,
54            AV1_NUM_REF_FRAMES * sizeof(int8_t));
55     memcpy(s->cur_frame.loop_filter_mode_deltas,
56            s->ref[prev_frame].loop_filter_mode_deltas,
57            2 * sizeof(int8_t));
58
59     if (s->raw_frame_header->loop_filter_delta_update) {
60         for (int i = 0; i < AV1_NUM_REF_FRAMES; i++) {
61             if (s->raw_frame_header->update_ref_delta[i])
62                 s->cur_frame.loop_filter_ref_deltas[i] =
63                     s->raw_frame_header->loop_filter_ref_deltas[i];
64         }
65
66         for (int i = 0; i < 2; i++) {
67             if (s->raw_frame_header->update_mode_delta[i])
68                 s->cur_frame.loop_filter_mode_deltas[i] =
69                     s->raw_frame_header->loop_filter_mode_deltas[i];
70         }
71     }
72
73     s->cur_frame.loop_filter_delta_enabled =
74         s->raw_frame_header->loop_filter_delta_enabled;
75 }
76
77 static uint32_t inverse_recenter(int r, uint32_t v)
78 {
79     if (v > 2 * r)
80         return v;
81     else if (v & 1)
82         return r - ((v + 1) >> 1);
83     else
84         return r + (v >> 1);
85 }
86
87 static uint32_t decode_unsigned_subexp_with_ref(uint32_t sub_exp,
88                                                 int mx, int r)
89 {
90     if ((r << 1) <= mx) {
91         return inverse_recenter(r, sub_exp);
92     } else {
93         return mx - 1 - inverse_recenter(mx - 1 - r, sub_exp);
94     }
95 }
96
97 static int32_t decode_signed_subexp_with_ref(uint32_t sub_exp, int low,
98                                              int high, int r)
99 {
100     int32_t x = decode_unsigned_subexp_with_ref(sub_exp, high - low, r - low);
101     return x + low;
102 }
103
104 static void read_global_param(AV1DecContext *s, int type, int ref, int idx)
105 {
106     uint8_t primary_frame, prev_frame;
107     uint32_t abs_bits, prec_bits, round, prec_diff, sub, mx;
108     int32_t r;
109
110     primary_frame = s->raw_frame_header->primary_ref_frame;
111     prev_frame = s->raw_frame_header->ref_frame_idx[primary_frame];
112     abs_bits = AV1_GM_ABS_ALPHA_BITS;
113     prec_bits = AV1_GM_ALPHA_PREC_BITS;
114
115     if (idx < 2) {
116         if (type == AV1_WARP_MODEL_TRANSLATION) {
117             abs_bits = AV1_GM_ABS_TRANS_ONLY_BITS -
118                 !s->raw_frame_header->allow_high_precision_mv;
119             prec_bits = AV1_GM_TRANS_ONLY_PREC_BITS -
120                 !s->raw_frame_header->allow_high_precision_mv;
121         } else {
122             abs_bits = AV1_GM_ABS_TRANS_BITS;
123             prec_bits = AV1_GM_TRANS_PREC_BITS;
124         }
125     }
126     round = (idx % 3) == 2 ? (1 << AV1_WARPEDMODEL_PREC_BITS) : 0;
127     prec_diff = AV1_WARPEDMODEL_PREC_BITS - prec_bits;
128     sub = (idx % 3) == 2 ? (1 << prec_bits) : 0;
129     mx = 1 << abs_bits;
130     r = (s->ref[prev_frame].gm_params[ref][idx] >> prec_diff) - sub;
131
132     s->cur_frame.gm_params[ref][idx] =
133         (decode_signed_subexp_with_ref(s->raw_frame_header->gm_params[ref][idx],
134                                        -mx, mx + 1, r) << prec_diff) + round;
135 }
136
137 /**
138 * update gm type/params, since cbs already implemented part of this funcation,
139 * so we don't need to full implement spec.
140 */
141 static void global_motion_params(AV1DecContext *s)
142 {
143     const AV1RawFrameHeader *header = s->raw_frame_header;
144     int type, ref;
145
146     for (ref = AV1_REF_FRAME_LAST; ref <= AV1_REF_FRAME_ALTREF; ref++) {
147         s->cur_frame.gm_type[ref] = AV1_WARP_MODEL_IDENTITY;
148         for (int i = 0; i < 6; i++)
149             s->cur_frame.gm_params[ref][i] = (i % 3 == 2) ?
150                                              1 << AV1_WARPEDMODEL_PREC_BITS : 0;
151     }
152     if (header->frame_type == AV1_FRAME_KEY ||
153         header->frame_type == AV1_FRAME_INTRA_ONLY)
154         return;
155
156     for (ref = AV1_REF_FRAME_LAST; ref <= AV1_REF_FRAME_ALTREF; ref++) {
157         if (header->is_global[ref]) {
158             if (header->is_rot_zoom[ref]) {
159                 type = AV1_WARP_MODEL_ROTZOOM;
160             } else {
161                 type = header->is_translation[ref] ? AV1_WARP_MODEL_TRANSLATION
162                                                    : AV1_WARP_MODEL_AFFINE;
163             }
164         } else {
165             type = AV1_WARP_MODEL_IDENTITY;
166         }
167         s->cur_frame.gm_type[ref] = type;
168
169         if (type >= AV1_WARP_MODEL_ROTZOOM) {
170             read_global_param(s, type, ref, 2);
171             read_global_param(s, type, ref, 3);
172             if (type == AV1_WARP_MODEL_AFFINE) {
173                 read_global_param(s, type, ref, 4);
174                 read_global_param(s, type, ref, 5);
175             } else {
176                 s->cur_frame.gm_params[ref][4] = -s->cur_frame.gm_params[ref][3];
177                 s->cur_frame.gm_params[ref][5] = s->cur_frame.gm_params[ref][2];
178             }
179         }
180         if (type >= AV1_WARP_MODEL_TRANSLATION) {
181             read_global_param(s, type, ref, 0);
182             read_global_param(s, type, ref, 1);
183         }
184     }
185 }
186
187 static int init_tile_data(AV1DecContext *s)
188
189 {
190     int cur_tile_num =
191         s->raw_frame_header->tile_cols * s->raw_frame_header->tile_rows;
192     if (s->tile_num < cur_tile_num) {
193         int ret = av_reallocp_array(&s->tile_group_info, cur_tile_num,
194                                     sizeof(TileGroupInfo));
195         if (ret < 0) {
196             s->tile_num = 0;
197             return ret;
198         }
199     }
200     s->tile_num = cur_tile_num;
201
202     return 0;
203 }
204
205 static int get_tiles_info(AVCodecContext *avctx, const AV1RawTileGroup *tile_group)
206 {
207     AV1DecContext *s = avctx->priv_data;
208     GetByteContext gb;
209     uint16_t tile_num, tile_row, tile_col;
210     uint32_t size = 0, size_bytes = 0;
211
212     bytestream2_init(&gb, tile_group->tile_data.data,
213                      tile_group->tile_data.data_size);
214     s->tg_start = tile_group->tg_start;
215     s->tg_end = tile_group->tg_end;
216
217     for (tile_num = tile_group->tg_start; tile_num <= tile_group->tg_end; tile_num++) {
218         tile_row = tile_num / s->raw_frame_header->tile_cols;
219         tile_col = tile_num % s->raw_frame_header->tile_cols;
220
221         if (tile_num == tile_group->tg_end) {
222             s->tile_group_info[tile_num].tile_size = bytestream2_get_bytes_left(&gb);
223             s->tile_group_info[tile_num].tile_offset = bytestream2_tell(&gb);
224             s->tile_group_info[tile_num].tile_row = tile_row;
225             s->tile_group_info[tile_num].tile_column = tile_col;
226             return 0;
227         }
228         size_bytes = s->raw_frame_header->tile_size_bytes_minus1 + 1;
229         if (bytestream2_get_bytes_left(&gb) < size_bytes)
230             return AVERROR_INVALIDDATA;
231         size = 0;
232         for (int i = 0; i < size_bytes; i++)
233             size |= bytestream2_get_byteu(&gb) << 8 * i;
234         if (bytestream2_get_bytes_left(&gb) <= size)
235             return AVERROR_INVALIDDATA;
236         size++;
237
238         s->tile_group_info[tile_num].tile_size = size;
239         s->tile_group_info[tile_num].tile_offset = bytestream2_tell(&gb);
240         s->tile_group_info[tile_num].tile_row = tile_row;
241         s->tile_group_info[tile_num].tile_column = tile_col;
242
243         bytestream2_skipu(&gb, size);
244     }
245
246     return 0;
247
248 }
249
250 static int get_pixel_format(AVCodecContext *avctx)
251 {
252     AV1DecContext *s = avctx->priv_data;
253     const AV1RawSequenceHeader *seq = s->raw_seq;
254     uint8_t bit_depth;
255     int ret;
256     enum AVPixelFormat pix_fmt = AV_PIX_FMT_NONE;
257 #define HWACCEL_MAX (0)
258     enum AVPixelFormat pix_fmts[HWACCEL_MAX + 2], *fmtp = pix_fmts;
259
260     if (seq->seq_profile == 2 && seq->color_config.high_bitdepth)
261         bit_depth = seq->color_config.twelve_bit ? 12 : 10;
262     else if (seq->seq_profile <= 2)
263         bit_depth = seq->color_config.high_bitdepth ? 10 : 8;
264     else {
265         av_log(avctx, AV_LOG_ERROR,
266                "Unknown AV1 profile %d.\n", seq->seq_profile);
267         return -1;
268     }
269
270     if (!seq->color_config.mono_chrome) {
271         // 4:4:4 x:0 y:0, 4:2:2 x:1 y:0, 4:2:0 x:1 y:1
272         if (seq->color_config.subsampling_x == 0 &&
273             seq->color_config.subsampling_y == 0) {
274             if (bit_depth == 8)
275                 pix_fmt = AV_PIX_FMT_YUV444P;
276             else if (bit_depth == 10)
277                 pix_fmt = AV_PIX_FMT_YUV444P10;
278             else if (bit_depth == 12)
279                 pix_fmt = AV_PIX_FMT_YUV444P12;
280             else
281                 av_log(avctx, AV_LOG_WARNING, "Unknown AV1 pixel format.\n");
282         } else if (seq->color_config.subsampling_x == 1 &&
283                    seq->color_config.subsampling_y == 0) {
284             if (bit_depth == 8)
285                 pix_fmt = AV_PIX_FMT_YUV422P;
286             else if (bit_depth == 10)
287                 pix_fmt = AV_PIX_FMT_YUV422P10;
288             else if (bit_depth == 12)
289                 pix_fmt = AV_PIX_FMT_YUV422P12;
290             else
291                 av_log(avctx, AV_LOG_WARNING, "Unknown AV1 pixel format.\n");
292         } else if (seq->color_config.subsampling_x == 1 &&
293                    seq->color_config.subsampling_y == 1) {
294             if (bit_depth == 8)
295                 pix_fmt = AV_PIX_FMT_YUV420P;
296             else if (bit_depth == 10)
297                 pix_fmt = AV_PIX_FMT_YUV420P10;
298             else if (bit_depth == 12)
299                 pix_fmt = AV_PIX_FMT_YUV420P12;
300             else
301                 av_log(avctx, AV_LOG_WARNING, "Unknown AV1 pixel format.\n");
302         }
303     } else {
304         if (seq->color_config.subsampling_x == 1 &&
305             seq->color_config.subsampling_y == 1)
306             pix_fmt = AV_PIX_FMT_YUV440P;
307         else
308             av_log(avctx, AV_LOG_WARNING, "Unknown AV1 pixel format.\n");
309     }
310
311     av_log(avctx, AV_LOG_DEBUG, "AV1 decode get format: %s.\n",
312            av_get_pix_fmt_name(pix_fmt));
313
314     if (pix_fmt == AV_PIX_FMT_NONE)
315         return -1;
316     s->pix_fmt = pix_fmt;
317
318     *fmtp++ = s->pix_fmt;
319     *fmtp = AV_PIX_FMT_NONE;
320
321     ret = ff_thread_get_format(avctx, pix_fmts);
322     if (ret < 0)
323         return ret;
324
325     /**
326      * check if the HW accel is inited correctly. If not, return un-implemented.
327      * Since now the av1 decoder doesn't support native decode, if it will be
328      * implemented in the future, need remove this check.
329      */
330     if (!avctx->hwaccel) {
331         av_log(avctx, AV_LOG_ERROR, "Your platform doesn't suppport"
332                " hardware accelerated AV1 decoding.\n");
333         return AVERROR(ENOSYS);
334     }
335
336     avctx->pix_fmt = ret;
337
338     return 0;
339 }
340
341 static void av1_frame_unref(AVCodecContext *avctx, AV1Frame *f)
342 {
343     ff_thread_release_buffer(avctx, &f->tf);
344     av_buffer_unref(&f->hwaccel_priv_buf);
345     f->hwaccel_picture_private = NULL;
346     f->spatial_id = f->temporal_id = 0;
347 }
348
349 static int av1_frame_ref(AVCodecContext *avctx, AV1Frame *dst, const AV1Frame *src)
350 {
351     int ret;
352
353     ret = ff_thread_ref_frame(&dst->tf, &src->tf);
354     if (ret < 0)
355         return ret;
356
357     if (src->hwaccel_picture_private) {
358         dst->hwaccel_priv_buf = av_buffer_ref(src->hwaccel_priv_buf);
359         if (!dst->hwaccel_priv_buf)
360             goto fail;
361         dst->hwaccel_picture_private = dst->hwaccel_priv_buf->data;
362     }
363
364     dst->spatial_id = src->spatial_id;
365     dst->temporal_id = src->temporal_id;
366     dst->loop_filter_delta_enabled = src->loop_filter_delta_enabled;
367     memcpy(dst->loop_filter_ref_deltas,
368            src->loop_filter_ref_deltas,
369            AV1_NUM_REF_FRAMES * sizeof(int8_t));
370     memcpy(dst->loop_filter_mode_deltas,
371            src->loop_filter_mode_deltas,
372            2 * sizeof(int8_t));
373     memcpy(dst->gm_type,
374            src->gm_type,
375            AV1_NUM_REF_FRAMES * sizeof(uint8_t));
376     memcpy(dst->gm_params,
377            src->gm_params,
378            AV1_NUM_REF_FRAMES * 6 * sizeof(int32_t));
379
380     return 0;
381
382 fail:
383     av1_frame_unref(avctx, dst);
384     return AVERROR(ENOMEM);
385 }
386
387 static av_cold int av1_decode_free(AVCodecContext *avctx)
388 {
389     AV1DecContext *s = avctx->priv_data;
390
391     for (int i = 0; i < FF_ARRAY_ELEMS(s->ref); i++) {
392         av1_frame_unref(avctx, &s->ref[i]);
393         av_frame_free(&s->ref[i].tf.f);
394     }
395     av1_frame_unref(avctx, &s->cur_frame);
396     av_frame_free(&s->cur_frame.tf.f);
397
398     av_buffer_unref(&s->seq_ref);
399     av_buffer_unref(&s->header_ref);
400     av_freep(&s->tile_group_info);
401
402     ff_cbs_fragment_free(&s->current_obu);
403     ff_cbs_close(&s->cbc);
404
405     return 0;
406 }
407
408 static int set_context_with_sequence(AVCodecContext *avctx,
409                                      const AV1RawSequenceHeader *seq)
410 {
411     int width = seq->max_frame_width_minus_1 + 1;
412     int height = seq->max_frame_height_minus_1 + 1;
413
414     avctx->profile = seq->seq_profile;
415     avctx->level = seq->seq_level_idx[0];
416
417     avctx->color_range =
418         seq->color_config.color_range ? AVCOL_RANGE_JPEG : AVCOL_RANGE_MPEG;
419     avctx->color_primaries = seq->color_config.color_primaries;
420     avctx->colorspace = seq->color_config.color_primaries;
421     avctx->color_trc = seq->color_config.transfer_characteristics;
422
423     switch (seq->color_config.chroma_sample_position) {
424     case AV1_CSP_VERTICAL:
425         avctx->chroma_sample_location = AVCHROMA_LOC_LEFT;
426         break;
427     case AV1_CSP_COLOCATED:
428         avctx->chroma_sample_location = AVCHROMA_LOC_TOPLEFT;
429         break;
430     }
431
432     if (avctx->width != width || avctx->height != height) {
433         int ret = ff_set_dimensions(avctx, width, height);
434         if (ret < 0)
435             return ret;
436     }
437     avctx->sample_aspect_ratio = (AVRational) { 1, 1 };
438
439     if (seq->timing_info.num_units_in_display_tick &&
440         seq->timing_info.time_scale) {
441         av_reduce(&avctx->framerate.den, &avctx->framerate.num,
442                   seq->timing_info.num_units_in_display_tick,
443                   seq->timing_info.time_scale,
444                   INT_MAX);
445         if (seq->timing_info.equal_picture_interval)
446             avctx->ticks_per_frame = seq->timing_info.num_ticks_per_picture_minus_1 + 1;
447     }
448
449     return 0;
450 }
451
452 static int update_context_with_frame_header(AVCodecContext *avctx,
453                                             const AV1RawFrameHeader *header)
454 {
455     AVRational aspect_ratio;
456     int width = header->frame_width_minus_1 + 1;
457     int height = header->frame_height_minus_1 + 1;
458     int r_width = header->render_width_minus_1 + 1;
459     int r_height = header->render_height_minus_1 + 1;
460     int ret;
461
462     if (avctx->width != width || avctx->height != height) {
463         ret = ff_set_dimensions(avctx, width, height);
464         if (ret < 0)
465             return ret;
466     }
467
468     av_reduce(&aspect_ratio.num, &aspect_ratio.den,
469               (int64_t)height * r_width,
470               (int64_t)width * r_height,
471               INT_MAX);
472
473     if (av_cmp_q(avctx->sample_aspect_ratio, aspect_ratio)) {
474         ret = ff_set_sar(avctx, aspect_ratio);
475         if (ret < 0)
476             return ret;
477     }
478
479     return 0;
480 }
481
482 static av_cold int av1_decode_init(AVCodecContext *avctx)
483 {
484     AV1DecContext *s = avctx->priv_data;
485     AV1RawSequenceHeader *seq;
486     int ret;
487
488     s->avctx = avctx;
489     s->pix_fmt = AV_PIX_FMT_NONE;
490
491     for (int i = 0; i < FF_ARRAY_ELEMS(s->ref); i++) {
492         s->ref[i].tf.f = av_frame_alloc();
493         if (!s->ref[i].tf.f) {
494             av_log(avctx, AV_LOG_ERROR,
495                    "Failed to allocate reference frame buffer %d.\n", i);
496             return AVERROR(ENOMEM);
497         }
498     }
499
500     s->cur_frame.tf.f = av_frame_alloc();
501     if (!s->cur_frame.tf.f) {
502         av_log(avctx, AV_LOG_ERROR,
503                "Failed to allocate current frame buffer.\n");
504         return AVERROR(ENOMEM);
505     }
506
507     ret = ff_cbs_init(&s->cbc, AV_CODEC_ID_AV1, avctx);
508     if (ret < 0)
509         return ret;
510
511     if (avctx->extradata && avctx->extradata_size) {
512         ret = ff_cbs_read(s->cbc, &s->current_obu, avctx->extradata,
513                           avctx->extradata_size);
514         if (ret < 0) {
515             av_log(avctx, AV_LOG_WARNING, "Failed to read extradata.\n");
516             return ret;
517         }
518
519         seq = ((CodedBitstreamAV1Context *)(s->cbc->priv_data))->sequence_header;
520         if (!seq) {
521             av_log(avctx, AV_LOG_WARNING, "No sequence header available.\n");
522             goto end;
523         }
524
525         ret = set_context_with_sequence(avctx, seq);
526         if (ret < 0) {
527             av_log(avctx, AV_LOG_WARNING, "Failed to set decoder context.\n");
528             goto end;
529         }
530
531         end:
532         ff_cbs_fragment_reset(&s->current_obu);
533     }
534
535     return ret;
536 }
537
538 static int av1_frame_alloc(AVCodecContext *avctx, AV1Frame *f)
539 {
540     AV1DecContext *s = avctx->priv_data;
541     AV1RawFrameHeader *header= s->raw_frame_header;
542     AVFrame *frame;
543     int ret;
544
545     ret = update_context_with_frame_header(avctx, header);
546     if (ret < 0) {
547         av_log(avctx, AV_LOG_ERROR, "Failed to update context with frame header\n");
548         return ret;
549     }
550
551     if ((ret = ff_thread_get_buffer(avctx, &f->tf, AV_GET_BUFFER_FLAG_REF)) < 0)
552         return ret;
553
554     frame = f->tf.f;
555     frame->key_frame = header->frame_type == AV1_FRAME_KEY;
556
557     switch (header->frame_type) {
558     case AV1_FRAME_KEY:
559     case AV1_FRAME_INTRA_ONLY:
560         frame->pict_type = AV_PICTURE_TYPE_I;
561         break;
562     case AV1_FRAME_INTER:
563         frame->pict_type = AV_PICTURE_TYPE_P;
564         break;
565     case AV1_FRAME_SWITCH:
566         frame->pict_type = AV_PICTURE_TYPE_SP;
567         break;
568     }
569
570     if (avctx->hwaccel) {
571         const AVHWAccel *hwaccel = avctx->hwaccel;
572         if (hwaccel->frame_priv_data_size) {
573             f->hwaccel_priv_buf =
574                 av_buffer_allocz(hwaccel->frame_priv_data_size);
575             if (!f->hwaccel_priv_buf)
576                 goto fail;
577             f->hwaccel_picture_private = f->hwaccel_priv_buf->data;
578         }
579     }
580     return 0;
581
582 fail:
583     av1_frame_unref(avctx, f);
584     return AVERROR(ENOMEM);
585 }
586
587 static int set_output_frame(AVCodecContext *avctx, AVFrame *frame,
588                             const AVPacket *pkt, int *got_frame)
589 {
590     AV1DecContext *s = avctx->priv_data;
591     const AVFrame *srcframe = s->cur_frame.tf.f;
592     int ret;
593
594     ret = av_frame_ref(frame, srcframe);
595     if (ret < 0)
596         return ret;
597
598     frame->pts = pkt->pts;
599     frame->pkt_dts = pkt->dts;
600     frame->pkt_size = pkt->size;
601
602     *got_frame = 1;
603
604     return 0;
605 }
606
607 static int update_reference_list(AVCodecContext *avctx)
608 {
609     AV1DecContext *s = avctx->priv_data;
610     const AV1RawFrameHeader *header = s->raw_frame_header;
611     int ret;
612
613     for (int i = 0; i < AV1_NUM_REF_FRAMES; i++) {
614         if (header->refresh_frame_flags & (1 << i)) {
615             if (s->ref[i].tf.f->buf[0])
616                 av1_frame_unref(avctx, &s->ref[i]);
617             if ((ret = av1_frame_ref(avctx, &s->ref[i], &s->cur_frame)) < 0) {
618                 av_log(avctx, AV_LOG_ERROR,
619                        "Failed to update frame %d in reference list\n", i);
620                 return ret;
621             }
622         }
623     }
624     return 0;
625 }
626
627 static int get_current_frame(AVCodecContext *avctx)
628 {
629     AV1DecContext *s = avctx->priv_data;
630     int ret;
631
632     if (s->cur_frame.tf.f->buf[0])
633         av1_frame_unref(avctx, &s->cur_frame);
634
635     ret = av1_frame_alloc(avctx, &s->cur_frame);
636     if (ret < 0) {
637         av_log(avctx, AV_LOG_ERROR,
638                "Failed to allocate space for current frame.\n");
639         return ret;
640     }
641
642     ret = init_tile_data(s);
643     if (ret < 0) {
644         av_log(avctx, AV_LOG_ERROR, "Failed to init tile data.\n");
645         return ret;
646     }
647
648     if (s->raw_frame_header->primary_ref_frame == AV1_PRIMARY_REF_NONE)
649         setup_past_independence(&s->cur_frame);
650     else
651         load_previous_and_update(s);
652
653     global_motion_params(s);
654
655     return ret;
656 }
657
658 static int av1_decode_frame(AVCodecContext *avctx, void *frame,
659                             int *got_frame, AVPacket *pkt)
660 {
661     AV1DecContext *s = avctx->priv_data;
662     AV1RawTileGroup *raw_tile_group = NULL;
663     int ret;
664
665     ret = ff_cbs_read_packet(s->cbc, &s->current_obu, pkt);
666     if (ret < 0) {
667         av_log(avctx, AV_LOG_ERROR, "Failed to read packet.\n");
668         goto end;
669     }
670     av_log(avctx, AV_LOG_DEBUG, "Total obu for this frame:%d.\n",
671            s->current_obu.nb_units);
672
673     for (int i = 0; i < s->current_obu.nb_units; i++) {
674         CodedBitstreamUnit *unit = &s->current_obu.units[i];
675         AV1RawOBU *obu = unit->content;
676         const AV1RawOBUHeader *header;
677
678         if (!obu)
679             continue;
680
681         header = &obu->header;
682         av_log(avctx, AV_LOG_DEBUG, "Obu idx:%d, obu type:%d.\n", i, unit->type);
683
684         switch (unit->type) {
685         case AV1_OBU_SEQUENCE_HEADER:
686             av_buffer_unref(&s->seq_ref);
687             s->seq_ref = av_buffer_ref(unit->content_ref);
688             if (!s->seq_ref) {
689                 ret = AVERROR(ENOMEM);
690                 goto end;
691             }
692
693             s->raw_seq = &obu->obu.sequence_header;
694
695             ret = set_context_with_sequence(avctx, s->raw_seq);
696             if (ret < 0) {
697                 av_log(avctx, AV_LOG_ERROR, "Failed to set context.\n");
698                 s->raw_seq = NULL;
699                 goto end;
700             }
701
702             if (s->pix_fmt == AV_PIX_FMT_NONE) {
703                 ret = get_pixel_format(avctx);
704                 if (ret < 0) {
705                     av_log(avctx, AV_LOG_ERROR,
706                            "Failed to get pixel format.\n");
707                     s->raw_seq = NULL;
708                     goto end;
709                 }
710             }
711
712             if (avctx->hwaccel && avctx->hwaccel->decode_params) {
713                 ret = avctx->hwaccel->decode_params(avctx, unit->type, unit->data,
714                                                     unit->data_size);
715                 if (ret < 0) {
716                     av_log(avctx, AV_LOG_ERROR, "HW accel decode params fail.\n");
717                     s->raw_seq = NULL;
718                     goto end;
719                 }
720             }
721             break;
722         case AV1_OBU_REDUNDANT_FRAME_HEADER:
723             if (s->raw_frame_header)
724                 break;
725         // fall-through
726         case AV1_OBU_FRAME:
727         case AV1_OBU_FRAME_HEADER:
728             if (!s->raw_seq) {
729                 av_log(avctx, AV_LOG_ERROR, "Missing Sequence Header.\n");
730                 ret = AVERROR_INVALIDDATA;
731                 goto end;
732             }
733
734             av_buffer_unref(&s->header_ref);
735             s->header_ref = av_buffer_ref(unit->content_ref);
736             if (!s->header_ref) {
737                 ret = AVERROR(ENOMEM);
738                 goto end;
739             }
740
741             if (unit->type == AV1_OBU_FRAME)
742                 s->raw_frame_header = &obu->obu.frame.header;
743             else
744                 s->raw_frame_header = &obu->obu.frame_header;
745
746             if (s->raw_frame_header->show_existing_frame) {
747                 if (s->cur_frame.tf.f->buf[0])
748                     av1_frame_unref(avctx, &s->cur_frame);
749
750                 ret = av1_frame_ref(avctx, &s->cur_frame,
751                                     &s->ref[s->raw_frame_header->frame_to_show_map_idx]);
752                 if (ret < 0) {
753                     av_log(avctx, AV_LOG_ERROR, "Failed to get reference frame.\n");
754                     goto end;
755                 }
756
757                 ret = update_reference_list(avctx);
758                 if (ret < 0) {
759                     av_log(avctx, AV_LOG_ERROR, "Failed to update reference list.\n");
760                     goto end;
761                 }
762
763                 ret = set_output_frame(avctx, frame, pkt, got_frame);
764                 if (ret < 0)
765                     av_log(avctx, AV_LOG_ERROR, "Set output frame error.\n");
766
767                 s->raw_frame_header = NULL;
768
769                 goto end;
770             }
771
772             ret = get_current_frame(avctx);
773             if (ret < 0) {
774                 av_log(avctx, AV_LOG_ERROR, "Get current frame error\n");
775                 goto end;
776             }
777
778             s->cur_frame.spatial_id  = header->spatial_id;
779             s->cur_frame.temporal_id = header->temporal_id;
780
781             if (avctx->hwaccel) {
782                 ret = avctx->hwaccel->start_frame(avctx, unit->data,
783                                                   unit->data_size);
784                 if (ret < 0) {
785                     av_log(avctx, AV_LOG_ERROR, "HW accel start frame fail.\n");
786                     goto end;
787                 }
788             }
789             if (unit->type != AV1_OBU_FRAME)
790                 break;
791         // fall-through
792         case AV1_OBU_TILE_GROUP:
793             if (!s->raw_frame_header) {
794                 av_log(avctx, AV_LOG_ERROR, "Missing Frame Header.\n");
795                 ret = AVERROR_INVALIDDATA;
796                 goto end;
797             }
798
799             if (unit->type == AV1_OBU_FRAME)
800                 raw_tile_group = &obu->obu.frame.tile_group;
801             else
802                 raw_tile_group = &obu->obu.tile_group;
803
804             ret = get_tiles_info(avctx, raw_tile_group);
805             if (ret < 0)
806                 goto end;
807
808             if (avctx->hwaccel) {
809                 ret = avctx->hwaccel->decode_slice(avctx,
810                                                    raw_tile_group->tile_data.data,
811                                                    raw_tile_group->tile_data.data_size);
812                 if (ret < 0) {
813                     av_log(avctx, AV_LOG_ERROR,
814                            "HW accel decode slice fail.\n");
815                     goto end;
816                 }
817             }
818             break;
819         case AV1_OBU_TILE_LIST:
820         case AV1_OBU_TEMPORAL_DELIMITER:
821         case AV1_OBU_PADDING:
822         case AV1_OBU_METADATA:
823             break;
824         default:
825             av_log(avctx, AV_LOG_DEBUG,
826                    "Unknown obu type: %d (%"SIZE_SPECIFIER" bits).\n",
827                    unit->type, unit->data_size);
828         }
829
830         if (raw_tile_group && (s->tile_num == raw_tile_group->tg_end + 1)) {
831             if (avctx->hwaccel) {
832                 ret = avctx->hwaccel->end_frame(avctx);
833                 if (ret < 0) {
834                     av_log(avctx, AV_LOG_ERROR, "HW accel end frame fail.\n");
835                     goto end;
836                 }
837             }
838
839             ret = update_reference_list(avctx);
840             if (ret < 0) {
841                 av_log(avctx, AV_LOG_ERROR, "Failed to update reference list.\n");
842                 goto end;
843             }
844
845             if (s->raw_frame_header->show_frame) {
846                 ret = set_output_frame(avctx, frame, pkt, got_frame);
847                 if (ret < 0) {
848                     av_log(avctx, AV_LOG_ERROR, "Set output frame error\n");
849                     goto end;
850                 }
851             }
852             raw_tile_group = NULL;
853             s->raw_frame_header = NULL;
854         }
855     }
856
857 end:
858     ff_cbs_fragment_reset(&s->current_obu);
859     if (ret < 0)
860         s->raw_frame_header = NULL;
861     return ret;
862 }
863
864 static void av1_decode_flush(AVCodecContext *avctx)
865 {
866     AV1DecContext *s = avctx->priv_data;
867
868     for (int i = 0; i < FF_ARRAY_ELEMS(s->ref); i++)
869         av1_frame_unref(avctx, &s->ref[i]);
870
871     av1_frame_unref(avctx, &s->cur_frame);
872     s->raw_frame_header = NULL;
873     s->raw_seq = NULL;
874
875     ff_cbs_flush(s->cbc);
876 }
877
878 AVCodec ff_av1_decoder = {
879     .name                  = "av1",
880     .long_name             = NULL_IF_CONFIG_SMALL("Alliance for Open Media AV1"),
881     .type                  = AVMEDIA_TYPE_VIDEO,
882     .id                    = AV_CODEC_ID_AV1,
883     .priv_data_size        = sizeof(AV1DecContext),
884     .init                  = av1_decode_init,
885     .close                 = av1_decode_free,
886     .decode                = av1_decode_frame,
887     .capabilities          = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_AVOID_PROBING,
888     .caps_internal         = FF_CODEC_CAP_INIT_THREADSAFE |
889                              FF_CODEC_CAP_INIT_CLEANUP |
890                              FF_CODEC_CAP_SETS_PKT_DTS,
891     .flush                 = av1_decode_flush,
892     .profiles              = NULL_IF_CONFIG_SMALL(ff_av1_profiles),
893     .hw_configs            = (const AVCodecHWConfigInternal * []) {
894         NULL
895     },
896 };