]> git.sesse.net Git - ffmpeg/blob - libavcodec/proresenc_kostya.c
Merge commit 'd42191c78befc1983f23b1899b2dda513b72f1ed'
[ffmpeg] / libavcodec / proresenc_kostya.c
1 /*
2  * Apple ProRes encoder
3  *
4  * Copyright (c) 2012 Konstantin Shishkov
5  *
6  * This encoder appears to be based on Anatoliy Wassermans considering
7  * similarities in the bugs.
8  *
9  * This file is part of FFmpeg.
10  *
11  * FFmpeg is free software; you can redistribute it and/or
12  * modify it under the terms of the GNU Lesser General Public
13  * License as published by the Free Software Foundation; either
14  * version 2.1 of the License, or (at your option) any later version.
15  *
16  * FFmpeg is distributed in the hope that it will be useful,
17  * but WITHOUT ANY WARRANTY; without even the implied warranty of
18  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
19  * Lesser General Public License for more details.
20  *
21  * You should have received a copy of the GNU Lesser General Public
22  * License along with FFmpeg; if not, write to the Free Software
23  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
24  */
25
26 #include "libavutil/opt.h"
27 #include "libavutil/pixdesc.h"
28 #include "avcodec.h"
29 #include "fdctdsp.h"
30 #include "put_bits.h"
31 #include "bytestream.h"
32 #include "internal.h"
33 #include "proresdata.h"
34
35 #define CFACTOR_Y422 2
36 #define CFACTOR_Y444 3
37
38 #define MAX_MBS_PER_SLICE 8
39
40 #define MAX_PLANES 4
41
42 enum {
43     PRORES_PROFILE_AUTO  = -1,
44     PRORES_PROFILE_PROXY = 0,
45     PRORES_PROFILE_LT,
46     PRORES_PROFILE_STANDARD,
47     PRORES_PROFILE_HQ,
48     PRORES_PROFILE_4444,
49 };
50
51 enum {
52     QUANT_MAT_PROXY = 0,
53     QUANT_MAT_LT,
54     QUANT_MAT_STANDARD,
55     QUANT_MAT_HQ,
56     QUANT_MAT_DEFAULT,
57 };
58
59 static const uint8_t prores_quant_matrices[][64] = {
60     { // proxy
61          4,  7,  9, 11, 13, 14, 15, 63,
62          7,  7, 11, 12, 14, 15, 63, 63,
63          9, 11, 13, 14, 15, 63, 63, 63,
64         11, 11, 13, 14, 63, 63, 63, 63,
65         11, 13, 14, 63, 63, 63, 63, 63,
66         13, 14, 63, 63, 63, 63, 63, 63,
67         13, 63, 63, 63, 63, 63, 63, 63,
68         63, 63, 63, 63, 63, 63, 63, 63,
69     },
70     { // LT
71          4,  5,  6,  7,  9, 11, 13, 15,
72          5,  5,  7,  8, 11, 13, 15, 17,
73          6,  7,  9, 11, 13, 15, 15, 17,
74          7,  7,  9, 11, 13, 15, 17, 19,
75          7,  9, 11, 13, 14, 16, 19, 23,
76          9, 11, 13, 14, 16, 19, 23, 29,
77          9, 11, 13, 15, 17, 21, 28, 35,
78         11, 13, 16, 17, 21, 28, 35, 41,
79     },
80     { // standard
81          4,  4,  5,  5,  6,  7,  7,  9,
82          4,  4,  5,  6,  7,  7,  9,  9,
83          5,  5,  6,  7,  7,  9,  9, 10,
84          5,  5,  6,  7,  7,  9,  9, 10,
85          5,  6,  7,  7,  8,  9, 10, 12,
86          6,  7,  7,  8,  9, 10, 12, 15,
87          6,  7,  7,  9, 10, 11, 14, 17,
88          7,  7,  9, 10, 11, 14, 17, 21,
89     },
90     { // high quality
91          4,  4,  4,  4,  4,  4,  4,  4,
92          4,  4,  4,  4,  4,  4,  4,  4,
93          4,  4,  4,  4,  4,  4,  4,  4,
94          4,  4,  4,  4,  4,  4,  4,  5,
95          4,  4,  4,  4,  4,  4,  5,  5,
96          4,  4,  4,  4,  4,  5,  5,  6,
97          4,  4,  4,  4,  5,  5,  6,  7,
98          4,  4,  4,  4,  5,  6,  7,  7,
99     },
100     { // codec default
101          4,  4,  4,  4,  4,  4,  4,  4,
102          4,  4,  4,  4,  4,  4,  4,  4,
103          4,  4,  4,  4,  4,  4,  4,  4,
104          4,  4,  4,  4,  4,  4,  4,  4,
105          4,  4,  4,  4,  4,  4,  4,  4,
106          4,  4,  4,  4,  4,  4,  4,  4,
107          4,  4,  4,  4,  4,  4,  4,  4,
108          4,  4,  4,  4,  4,  4,  4,  4,
109     },
110 };
111
112 #define NUM_MB_LIMITS 4
113 static const int prores_mb_limits[NUM_MB_LIMITS] = {
114     1620, // up to 720x576
115     2700, // up to 960x720
116     6075, // up to 1440x1080
117     9216, // up to 2048x1152
118 };
119
120 static const struct prores_profile {
121     const char *full_name;
122     uint32_t    tag;
123     int         min_quant;
124     int         max_quant;
125     int         br_tab[NUM_MB_LIMITS];
126     int         quant;
127 } prores_profile_info[5] = {
128     {
129         .full_name = "proxy",
130         .tag       = MKTAG('a', 'p', 'c', 'o'),
131         .min_quant = 4,
132         .max_quant = 8,
133         .br_tab    = { 300, 242, 220, 194 },
134         .quant     = QUANT_MAT_PROXY,
135     },
136     {
137         .full_name = "LT",
138         .tag       = MKTAG('a', 'p', 'c', 's'),
139         .min_quant = 1,
140         .max_quant = 9,
141         .br_tab    = { 720, 560, 490, 440 },
142         .quant     = QUANT_MAT_LT,
143     },
144     {
145         .full_name = "standard",
146         .tag       = MKTAG('a', 'p', 'c', 'n'),
147         .min_quant = 1,
148         .max_quant = 6,
149         .br_tab    = { 1050, 808, 710, 632 },
150         .quant     = QUANT_MAT_STANDARD,
151     },
152     {
153         .full_name = "high quality",
154         .tag       = MKTAG('a', 'p', 'c', 'h'),
155         .min_quant = 1,
156         .max_quant = 6,
157         .br_tab    = { 1566, 1216, 1070, 950 },
158         .quant     = QUANT_MAT_HQ,
159     },
160     {
161         .full_name = "4444",
162         .tag       = MKTAG('a', 'p', '4', 'h'),
163         .min_quant = 1,
164         .max_quant = 6,
165         .br_tab    = { 2350, 1828, 1600, 1425 },
166         .quant     = QUANT_MAT_HQ,
167     }
168 };
169
170 #define TRELLIS_WIDTH 16
171 #define SCORE_LIMIT   INT_MAX / 2
172
173 struct TrellisNode {
174     int prev_node;
175     int quant;
176     int bits;
177     int score;
178 };
179
180 #define MAX_STORED_Q 16
181
182 typedef struct ProresThreadData {
183     DECLARE_ALIGNED(16, int16_t, blocks)[MAX_PLANES][64 * 4 * MAX_MBS_PER_SLICE];
184     DECLARE_ALIGNED(16, uint16_t, emu_buf)[16 * 16];
185     int16_t custom_q[64];
186     struct TrellisNode *nodes;
187 } ProresThreadData;
188
189 typedef struct ProresContext {
190     AVClass *class;
191     DECLARE_ALIGNED(16, int16_t, blocks)[MAX_PLANES][64 * 4 * MAX_MBS_PER_SLICE];
192     DECLARE_ALIGNED(16, uint16_t, emu_buf)[16*16];
193     int16_t quants[MAX_STORED_Q][64];
194     int16_t custom_q[64];
195     const uint8_t *quant_mat;
196     const uint8_t *scantable;
197
198     void (*fdct)(FDCTDSPContext *fdsp, const uint16_t *src,
199                  int linesize, int16_t *block);
200     FDCTDSPContext fdsp;
201
202     int mb_width, mb_height;
203     int mbs_per_slice;
204     int num_chroma_blocks, chroma_factor;
205     int slices_width;
206     int slices_per_picture;
207     int pictures_per_frame; // 1 for progressive, 2 for interlaced
208     int cur_picture_idx;
209     int num_planes;
210     int bits_per_mb;
211     int force_quant;
212     int alpha_bits;
213     int warn;
214
215     char *vendor;
216     int quant_sel;
217
218     int frame_size_upper_bound;
219
220     int profile;
221     const struct prores_profile *profile_info;
222
223     int *slice_q;
224
225     ProresThreadData *tdata;
226 } ProresContext;
227
228 static void get_slice_data(ProresContext *ctx, const uint16_t *src,
229                            int linesize, int x, int y, int w, int h,
230                            int16_t *blocks, uint16_t *emu_buf,
231                            int mbs_per_slice, int blocks_per_mb, int is_chroma)
232 {
233     const uint16_t *esrc;
234     const int mb_width = 4 * blocks_per_mb;
235     int elinesize;
236     int i, j, k;
237
238     for (i = 0; i < mbs_per_slice; i++, src += mb_width) {
239         if (x >= w) {
240             memset(blocks, 0, 64 * (mbs_per_slice - i) * blocks_per_mb
241                               * sizeof(*blocks));
242             return;
243         }
244         if (x + mb_width <= w && y + 16 <= h) {
245             esrc      = src;
246             elinesize = linesize;
247         } else {
248             int bw, bh, pix;
249
250             esrc      = emu_buf;
251             elinesize = 16 * sizeof(*emu_buf);
252
253             bw = FFMIN(w - x, mb_width);
254             bh = FFMIN(h - y, 16);
255
256             for (j = 0; j < bh; j++) {
257                 memcpy(emu_buf + j * 16,
258                        (const uint8_t*)src + j * linesize,
259                        bw * sizeof(*src));
260                 pix = emu_buf[j * 16 + bw - 1];
261                 for (k = bw; k < mb_width; k++)
262                     emu_buf[j * 16 + k] = pix;
263             }
264             for (; j < 16; j++)
265                 memcpy(emu_buf + j * 16,
266                        emu_buf + (bh - 1) * 16,
267                        mb_width * sizeof(*emu_buf));
268         }
269         if (!is_chroma) {
270             ctx->fdct(&ctx->fdsp, esrc, elinesize, blocks);
271             blocks += 64;
272             if (blocks_per_mb > 2) {
273                 ctx->fdct(&ctx->fdsp, esrc + 8, elinesize, blocks);
274                 blocks += 64;
275             }
276             ctx->fdct(&ctx->fdsp, esrc + elinesize * 4, elinesize, blocks);
277             blocks += 64;
278             if (blocks_per_mb > 2) {
279                 ctx->fdct(&ctx->fdsp, esrc + elinesize * 4 + 8, elinesize, blocks);
280                 blocks += 64;
281             }
282         } else {
283             ctx->fdct(&ctx->fdsp, esrc, elinesize, blocks);
284             blocks += 64;
285             ctx->fdct(&ctx->fdsp, esrc + elinesize * 4, elinesize, blocks);
286             blocks += 64;
287             if (blocks_per_mb > 2) {
288                 ctx->fdct(&ctx->fdsp, esrc + 8, elinesize, blocks);
289                 blocks += 64;
290                 ctx->fdct(&ctx->fdsp, esrc + elinesize * 4 + 8, elinesize, blocks);
291                 blocks += 64;
292             }
293         }
294
295         x += mb_width;
296     }
297 }
298
299 static void get_alpha_data(ProresContext *ctx, const uint16_t *src,
300                            int linesize, int x, int y, int w, int h,
301                            int16_t *blocks, int mbs_per_slice, int abits)
302 {
303     const int slice_width = 16 * mbs_per_slice;
304     int i, j, copy_w, copy_h;
305
306     copy_w = FFMIN(w - x, slice_width);
307     copy_h = FFMIN(h - y, 16);
308     for (i = 0; i < copy_h; i++) {
309         memcpy(blocks, src, copy_w * sizeof(*src));
310         if (abits == 8)
311             for (j = 0; j < copy_w; j++)
312                 blocks[j] >>= 2;
313         else
314             for (j = 0; j < copy_w; j++)
315                 blocks[j] = (blocks[j] << 6) | (blocks[j] >> 4);
316         for (j = copy_w; j < slice_width; j++)
317             blocks[j] = blocks[copy_w - 1];
318         blocks += slice_width;
319         src    += linesize >> 1;
320     }
321     for (; i < 16; i++) {
322         memcpy(blocks, blocks - slice_width, slice_width * sizeof(*blocks));
323         blocks += slice_width;
324     }
325 }
326
327 /**
328  * Write an unsigned rice/exp golomb codeword.
329  */
330 static inline void encode_vlc_codeword(PutBitContext *pb, unsigned codebook, int val)
331 {
332     unsigned int rice_order, exp_order, switch_bits, switch_val;
333     int exponent;
334
335     /* number of prefix bits to switch between Rice and expGolomb */
336     switch_bits = (codebook & 3) + 1;
337     rice_order  =  codebook >> 5;       /* rice code order */
338     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
339
340     switch_val  = switch_bits << rice_order;
341
342     if (val >= switch_val) {
343         val -= switch_val - (1 << exp_order);
344         exponent = av_log2(val);
345
346         put_bits(pb, exponent - exp_order + switch_bits, 0);
347         put_bits(pb, exponent + 1, val);
348     } else {
349         exponent = val >> rice_order;
350
351         if (exponent)
352             put_bits(pb, exponent, 0);
353         put_bits(pb, 1, 1);
354         if (rice_order)
355             put_sbits(pb, rice_order, val);
356     }
357 }
358
359 #define GET_SIGN(x)  ((x) >> 31)
360 #define MAKE_CODE(x) (((x) << 1) ^ GET_SIGN(x))
361
362 static void encode_dcs(PutBitContext *pb, int16_t *blocks,
363                        int blocks_per_slice, int scale)
364 {
365     int i;
366     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
367
368     prev_dc = (blocks[0] - 0x4000) / scale;
369     encode_vlc_codeword(pb, FIRST_DC_CB, MAKE_CODE(prev_dc));
370     sign     = 0;
371     codebook = 3;
372     blocks  += 64;
373
374     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
375         dc       = (blocks[0] - 0x4000) / scale;
376         delta    = dc - prev_dc;
377         new_sign = GET_SIGN(delta);
378         delta    = (delta ^ sign) - sign;
379         code     = MAKE_CODE(delta);
380         encode_vlc_codeword(pb, ff_prores_dc_codebook[codebook], code);
381         codebook = (code + (code & 1)) >> 1;
382         codebook = FFMIN(codebook, 3);
383         sign     = new_sign;
384         prev_dc  = dc;
385     }
386 }
387
388 static void encode_acs(PutBitContext *pb, int16_t *blocks,
389                        int blocks_per_slice,
390                        int plane_size_factor,
391                        const uint8_t *scan, const int16_t *qmat)
392 {
393     int idx, i;
394     int run, level, run_cb, lev_cb;
395     int max_coeffs, abs_level;
396
397     max_coeffs = blocks_per_slice << 6;
398     run_cb     = ff_prores_run_to_cb_index[4];
399     lev_cb     = ff_prores_lev_to_cb_index[2];
400     run        = 0;
401
402     for (i = 1; i < 64; i++) {
403         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
404             level = blocks[idx] / qmat[scan[i]];
405             if (level) {
406                 abs_level = FFABS(level);
407                 encode_vlc_codeword(pb, ff_prores_ac_codebook[run_cb], run);
408                 encode_vlc_codeword(pb, ff_prores_ac_codebook[lev_cb],
409                                     abs_level - 1);
410                 put_sbits(pb, 1, GET_SIGN(level));
411
412                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
413                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
414                 run    = 0;
415             } else {
416                 run++;
417             }
418         }
419     }
420 }
421
422 static int encode_slice_plane(ProresContext *ctx, PutBitContext *pb,
423                               const uint16_t *src, int linesize,
424                               int mbs_per_slice, int16_t *blocks,
425                               int blocks_per_mb, int plane_size_factor,
426                               const int16_t *qmat)
427 {
428     int blocks_per_slice, saved_pos;
429
430     saved_pos = put_bits_count(pb);
431     blocks_per_slice = mbs_per_slice * blocks_per_mb;
432
433     encode_dcs(pb, blocks, blocks_per_slice, qmat[0]);
434     encode_acs(pb, blocks, blocks_per_slice, plane_size_factor,
435                ctx->scantable, qmat);
436     flush_put_bits(pb);
437
438     return (put_bits_count(pb) - saved_pos) >> 3;
439 }
440
441 static void put_alpha_diff(PutBitContext *pb, int cur, int prev, int abits)
442 {
443     const int dbits = (abits == 8) ? 4 : 7;
444     const int dsize = 1 << dbits - 1;
445     int diff = cur - prev;
446
447     diff = av_mod_uintp2(diff, abits);
448     if (diff >= (1 << abits) - dsize)
449         diff -= 1 << abits;
450     if (diff < -dsize || diff > dsize || !diff) {
451         put_bits(pb, 1, 1);
452         put_bits(pb, abits, diff);
453     } else {
454         put_bits(pb, 1, 0);
455         put_bits(pb, dbits - 1, FFABS(diff) - 1);
456         put_bits(pb, 1, diff < 0);
457     }
458 }
459
460 static void put_alpha_run(PutBitContext *pb, int run)
461 {
462     if (run) {
463         put_bits(pb, 1, 0);
464         if (run < 0x10)
465             put_bits(pb, 4, run);
466         else
467             put_bits(pb, 15, run);
468     } else {
469         put_bits(pb, 1, 1);
470     }
471 }
472
473 // todo alpha quantisation for high quants
474 static int encode_alpha_plane(ProresContext *ctx, PutBitContext *pb,
475                               int mbs_per_slice, uint16_t *blocks,
476                               int quant)
477 {
478     const int abits = ctx->alpha_bits;
479     const int mask  = (1 << abits) - 1;
480     const int num_coeffs = mbs_per_slice * 256;
481     int saved_pos = put_bits_count(pb);
482     int prev = mask, cur;
483     int idx = 0;
484     int run = 0;
485
486     cur = blocks[idx++];
487     put_alpha_diff(pb, cur, prev, abits);
488     prev = cur;
489     do {
490         cur = blocks[idx++];
491         if (cur != prev) {
492             put_alpha_run (pb, run);
493             put_alpha_diff(pb, cur, prev, abits);
494             prev = cur;
495             run  = 0;
496         } else {
497             run++;
498         }
499     } while (idx < num_coeffs);
500     if (run)
501         put_alpha_run(pb, run);
502     flush_put_bits(pb);
503     return (put_bits_count(pb) - saved_pos) >> 3;
504 }
505
506 static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
507                         PutBitContext *pb,
508                         int sizes[4], int x, int y, int quant,
509                         int mbs_per_slice)
510 {
511     ProresContext *ctx = avctx->priv_data;
512     int i, xp, yp;
513     int total_size = 0;
514     const uint16_t *src;
515     int slice_width_factor = av_log2(mbs_per_slice);
516     int num_cblocks, pwidth, linesize, line_add;
517     int plane_factor, is_chroma;
518     uint16_t *qmat;
519
520     if (ctx->pictures_per_frame == 1)
521         line_add = 0;
522     else
523         line_add = ctx->cur_picture_idx ^ !pic->top_field_first;
524
525     if (ctx->force_quant) {
526         qmat = ctx->quants[0];
527     } else if (quant < MAX_STORED_Q) {
528         qmat = ctx->quants[quant];
529     } else {
530         qmat = ctx->custom_q;
531         for (i = 0; i < 64; i++)
532             qmat[i] = ctx->quant_mat[i] * quant;
533     }
534
535     for (i = 0; i < ctx->num_planes; i++) {
536         is_chroma    = (i == 1 || i == 2);
537         plane_factor = slice_width_factor + 2;
538         if (is_chroma)
539             plane_factor += ctx->chroma_factor - 3;
540         if (!is_chroma || ctx->chroma_factor == CFACTOR_Y444) {
541             xp          = x << 4;
542             yp          = y << 4;
543             num_cblocks = 4;
544             pwidth      = avctx->width;
545         } else {
546             xp          = x << 3;
547             yp          = y << 4;
548             num_cblocks = 2;
549             pwidth      = avctx->width >> 1;
550         }
551
552         linesize = pic->linesize[i] * ctx->pictures_per_frame;
553         src = (const uint16_t*)(pic->data[i] + yp * linesize +
554                                 line_add * pic->linesize[i]) + xp;
555
556         if (i < 3) {
557             get_slice_data(ctx, src, linesize, xp, yp,
558                            pwidth, avctx->height / ctx->pictures_per_frame,
559                            ctx->blocks[0], ctx->emu_buf,
560                            mbs_per_slice, num_cblocks, is_chroma);
561             sizes[i] = encode_slice_plane(ctx, pb, src, linesize,
562                                           mbs_per_slice, ctx->blocks[0],
563                                           num_cblocks, plane_factor,
564                                           qmat);
565         } else {
566             get_alpha_data(ctx, src, linesize, xp, yp,
567                            pwidth, avctx->height / ctx->pictures_per_frame,
568                            ctx->blocks[0], mbs_per_slice, ctx->alpha_bits);
569             sizes[i] = encode_alpha_plane(ctx, pb, mbs_per_slice,
570                                           ctx->blocks[0], quant);
571         }
572         total_size += sizes[i];
573         if (put_bits_left(pb) < 0) {
574             av_log(avctx, AV_LOG_ERROR,
575                    "Underestimated required buffer size.\n");
576             return AVERROR_BUG;
577         }
578     }
579     return total_size;
580 }
581
582 static inline int estimate_vlc(unsigned codebook, int val)
583 {
584     unsigned int rice_order, exp_order, switch_bits, switch_val;
585     int exponent;
586
587     /* number of prefix bits to switch between Rice and expGolomb */
588     switch_bits = (codebook & 3) + 1;
589     rice_order  =  codebook >> 5;       /* rice code order */
590     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
591
592     switch_val  = switch_bits << rice_order;
593
594     if (val >= switch_val) {
595         val -= switch_val - (1 << exp_order);
596         exponent = av_log2(val);
597
598         return exponent * 2 - exp_order + switch_bits + 1;
599     } else {
600         return (val >> rice_order) + rice_order + 1;
601     }
602 }
603
604 static int estimate_dcs(int *error, int16_t *blocks, int blocks_per_slice,
605                         int scale)
606 {
607     int i;
608     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
609     int bits;
610
611     prev_dc  = (blocks[0] - 0x4000) / scale;
612     bits     = estimate_vlc(FIRST_DC_CB, MAKE_CODE(prev_dc));
613     sign     = 0;
614     codebook = 3;
615     blocks  += 64;
616     *error  += FFABS(blocks[0] - 0x4000) % scale;
617
618     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
619         dc       = (blocks[0] - 0x4000) / scale;
620         *error  += FFABS(blocks[0] - 0x4000) % scale;
621         delta    = dc - prev_dc;
622         new_sign = GET_SIGN(delta);
623         delta    = (delta ^ sign) - sign;
624         code     = MAKE_CODE(delta);
625         bits    += estimate_vlc(ff_prores_dc_codebook[codebook], code);
626         codebook = (code + (code & 1)) >> 1;
627         codebook = FFMIN(codebook, 3);
628         sign     = new_sign;
629         prev_dc  = dc;
630     }
631
632     return bits;
633 }
634
635 static int estimate_acs(int *error, int16_t *blocks, int blocks_per_slice,
636                         int plane_size_factor,
637                         const uint8_t *scan, const int16_t *qmat)
638 {
639     int idx, i;
640     int run, level, run_cb, lev_cb;
641     int max_coeffs, abs_level;
642     int bits = 0;
643
644     max_coeffs = blocks_per_slice << 6;
645     run_cb     = ff_prores_run_to_cb_index[4];
646     lev_cb     = ff_prores_lev_to_cb_index[2];
647     run        = 0;
648
649     for (i = 1; i < 64; i++) {
650         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
651             level   = blocks[idx] / qmat[scan[i]];
652             *error += FFABS(blocks[idx]) % qmat[scan[i]];
653             if (level) {
654                 abs_level = FFABS(level);
655                 bits += estimate_vlc(ff_prores_ac_codebook[run_cb], run);
656                 bits += estimate_vlc(ff_prores_ac_codebook[lev_cb],
657                                      abs_level - 1) + 1;
658
659                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
660                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
661                 run    = 0;
662             } else {
663                 run++;
664             }
665         }
666     }
667
668     return bits;
669 }
670
671 static int estimate_slice_plane(ProresContext *ctx, int *error, int plane,
672                                 const uint16_t *src, int linesize,
673                                 int mbs_per_slice,
674                                 int blocks_per_mb, int plane_size_factor,
675                                 const int16_t *qmat, ProresThreadData *td)
676 {
677     int blocks_per_slice;
678     int bits;
679
680     blocks_per_slice = mbs_per_slice * blocks_per_mb;
681
682     bits  = estimate_dcs(error, td->blocks[plane], blocks_per_slice, qmat[0]);
683     bits += estimate_acs(error, td->blocks[plane], blocks_per_slice,
684                          plane_size_factor, ctx->scantable, qmat);
685
686     return FFALIGN(bits, 8);
687 }
688
689 static int est_alpha_diff(int cur, int prev, int abits)
690 {
691     const int dbits = (abits == 8) ? 4 : 7;
692     const int dsize = 1 << dbits - 1;
693     int diff = cur - prev;
694
695     diff = av_mod_uintp2(diff, abits);
696     if (diff >= (1 << abits) - dsize)
697         diff -= 1 << abits;
698     if (diff < -dsize || diff > dsize || !diff)
699         return abits + 1;
700     else
701         return dbits + 1;
702 }
703
704 static int estimate_alpha_plane(ProresContext *ctx, int *error,
705                                 const uint16_t *src, int linesize,
706                                 int mbs_per_slice, int quant,
707                                 int16_t *blocks)
708 {
709     const int abits = ctx->alpha_bits;
710     const int mask  = (1 << abits) - 1;
711     const int num_coeffs = mbs_per_slice * 256;
712     int prev = mask, cur;
713     int idx = 0;
714     int run = 0;
715     int bits;
716
717     *error = 0;
718     cur = blocks[idx++];
719     bits = est_alpha_diff(cur, prev, abits);
720     prev = cur;
721     do {
722         cur = blocks[idx++];
723         if (cur != prev) {
724             if (!run)
725                 bits++;
726             else if (run < 0x10)
727                 bits += 4;
728             else
729                 bits += 15;
730             bits += est_alpha_diff(cur, prev, abits);
731             prev = cur;
732             run  = 0;
733         } else {
734             run++;
735         }
736     } while (idx < num_coeffs);
737
738     if (run) {
739         if (run < 0x10)
740             bits += 4;
741         else
742             bits += 15;
743     }
744
745     return bits;
746 }
747
748 static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
749                             int trellis_node, int x, int y, int mbs_per_slice,
750                             ProresThreadData *td)
751 {
752     ProresContext *ctx = avctx->priv_data;
753     int i, q, pq, xp, yp;
754     const uint16_t *src;
755     int slice_width_factor = av_log2(mbs_per_slice);
756     int num_cblocks[MAX_PLANES], pwidth;
757     int plane_factor[MAX_PLANES], is_chroma[MAX_PLANES];
758     const int min_quant = ctx->profile_info->min_quant;
759     const int max_quant = ctx->profile_info->max_quant;
760     int error, bits, bits_limit;
761     int mbs, prev, cur, new_score;
762     int slice_bits[TRELLIS_WIDTH], slice_score[TRELLIS_WIDTH];
763     int overquant;
764     uint16_t *qmat;
765     int linesize[4], line_add;
766
767     if (ctx->pictures_per_frame == 1)
768         line_add = 0;
769     else
770         line_add = ctx->cur_picture_idx ^ !pic->top_field_first;
771     mbs = x + mbs_per_slice;
772
773     for (i = 0; i < ctx->num_planes; i++) {
774         is_chroma[i]    = (i == 1 || i == 2);
775         plane_factor[i] = slice_width_factor + 2;
776         if (is_chroma[i])
777             plane_factor[i] += ctx->chroma_factor - 3;
778         if (!is_chroma[i] || ctx->chroma_factor == CFACTOR_Y444) {
779             xp             = x << 4;
780             yp             = y << 4;
781             num_cblocks[i] = 4;
782             pwidth         = avctx->width;
783         } else {
784             xp             = x << 3;
785             yp             = y << 4;
786             num_cblocks[i] = 2;
787             pwidth         = avctx->width >> 1;
788         }
789
790         linesize[i] = pic->linesize[i] * ctx->pictures_per_frame;
791         src = (const uint16_t*)(pic->data[i] + yp * linesize[i] +
792                                 line_add * pic->linesize[i]) + xp;
793
794         if (i < 3) {
795             get_slice_data(ctx, src, linesize[i], xp, yp,
796                            pwidth, avctx->height / ctx->pictures_per_frame,
797                            td->blocks[i], td->emu_buf,
798                            mbs_per_slice, num_cblocks[i], is_chroma[i]);
799         } else {
800             get_alpha_data(ctx, src, linesize[i], xp, yp,
801                            pwidth, avctx->height / ctx->pictures_per_frame,
802                            td->blocks[i], mbs_per_slice, ctx->alpha_bits);
803         }
804     }
805
806     for (q = min_quant; q < max_quant + 2; q++) {
807         td->nodes[trellis_node + q].prev_node = -1;
808         td->nodes[trellis_node + q].quant     = q;
809     }
810
811     // todo: maybe perform coarser quantising to fit into frame size when needed
812     for (q = min_quant; q <= max_quant; q++) {
813         bits  = 0;
814         error = 0;
815         for (i = 0; i < ctx->num_planes - !!ctx->alpha_bits; i++) {
816             bits += estimate_slice_plane(ctx, &error, i,
817                                          src, linesize[i],
818                                          mbs_per_slice,
819                                          num_cblocks[i], plane_factor[i],
820                                          ctx->quants[q], td);
821         }
822         if (ctx->alpha_bits)
823             bits += estimate_alpha_plane(ctx, &error, src, linesize[3],
824                                          mbs_per_slice, q, td->blocks[3]);
825         if (bits > 65000 * 8)
826             error = SCORE_LIMIT;
827
828         slice_bits[q]  = bits;
829         slice_score[q] = error;
830     }
831     if (slice_bits[max_quant] <= ctx->bits_per_mb * mbs_per_slice) {
832         slice_bits[max_quant + 1]  = slice_bits[max_quant];
833         slice_score[max_quant + 1] = slice_score[max_quant] + 1;
834         overquant = max_quant;
835     } else {
836         for (q = max_quant + 1; q < 128; q++) {
837             bits  = 0;
838             error = 0;
839             if (q < MAX_STORED_Q) {
840                 qmat = ctx->quants[q];
841             } else {
842                 qmat = td->custom_q;
843                 for (i = 0; i < 64; i++)
844                     qmat[i] = ctx->quant_mat[i] * q;
845             }
846             for (i = 0; i < ctx->num_planes - !!ctx->alpha_bits; i++) {
847                 bits += estimate_slice_plane(ctx, &error, i,
848                                              src, linesize[i],
849                                              mbs_per_slice,
850                                              num_cblocks[i], plane_factor[i],
851                                              qmat, td);
852             }
853             if (ctx->alpha_bits)
854                 bits += estimate_alpha_plane(ctx, &error, src, linesize[3],
855                                              mbs_per_slice, q, td->blocks[3]);
856             if (bits <= ctx->bits_per_mb * mbs_per_slice)
857                 break;
858         }
859
860         slice_bits[max_quant + 1]  = bits;
861         slice_score[max_quant + 1] = error;
862         overquant = q;
863     }
864     td->nodes[trellis_node + max_quant + 1].quant = overquant;
865
866     bits_limit = mbs * ctx->bits_per_mb;
867     for (pq = min_quant; pq < max_quant + 2; pq++) {
868         prev = trellis_node - TRELLIS_WIDTH + pq;
869
870         for (q = min_quant; q < max_quant + 2; q++) {
871             cur = trellis_node + q;
872
873             bits  = td->nodes[prev].bits + slice_bits[q];
874             error = slice_score[q];
875             if (bits > bits_limit)
876                 error = SCORE_LIMIT;
877
878             if (td->nodes[prev].score < SCORE_LIMIT && error < SCORE_LIMIT)
879                 new_score = td->nodes[prev].score + error;
880             else
881                 new_score = SCORE_LIMIT;
882             if (td->nodes[cur].prev_node == -1 ||
883                 td->nodes[cur].score >= new_score) {
884
885                 td->nodes[cur].bits      = bits;
886                 td->nodes[cur].score     = new_score;
887                 td->nodes[cur].prev_node = prev;
888             }
889         }
890     }
891
892     error = td->nodes[trellis_node + min_quant].score;
893     pq    = trellis_node + min_quant;
894     for (q = min_quant + 1; q < max_quant + 2; q++) {
895         if (td->nodes[trellis_node + q].score <= error) {
896             error = td->nodes[trellis_node + q].score;
897             pq    = trellis_node + q;
898         }
899     }
900
901     return pq;
902 }
903
904 static int find_quant_thread(AVCodecContext *avctx, void *arg,
905                              int jobnr, int threadnr)
906 {
907     ProresContext *ctx = avctx->priv_data;
908     ProresThreadData *td = ctx->tdata + threadnr;
909     int mbs_per_slice = ctx->mbs_per_slice;
910     int x, y = jobnr, mb, q = 0;
911
912     for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
913         while (ctx->mb_width - x < mbs_per_slice)
914             mbs_per_slice >>= 1;
915         q = find_slice_quant(avctx, arg,
916                              (mb + 1) * TRELLIS_WIDTH, x, y,
917                              mbs_per_slice, td);
918     }
919
920     for (x = ctx->slices_width - 1; x >= 0; x--) {
921         ctx->slice_q[x + y * ctx->slices_width] = td->nodes[q].quant;
922         q = td->nodes[q].prev_node;
923     }
924
925     return 0;
926 }
927
928 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
929                         const AVFrame *pic, int *got_packet)
930 {
931     ProresContext *ctx = avctx->priv_data;
932     uint8_t *orig_buf, *buf, *slice_hdr, *slice_sizes, *tmp;
933     uint8_t *picture_size_pos;
934     PutBitContext pb;
935     int x, y, i, mb, q = 0;
936     int sizes[4] = { 0 };
937     int slice_hdr_size = 2 + 2 * (ctx->num_planes - 1);
938     int frame_size, picture_size, slice_size;
939     int pkt_size, ret;
940     int max_slice_size = (ctx->frame_size_upper_bound - 200) / (ctx->pictures_per_frame * ctx->slices_per_picture + 1);
941     uint8_t frame_flags;
942
943     pkt_size = ctx->frame_size_upper_bound;
944
945     if ((ret = ff_alloc_packet2(avctx, pkt, pkt_size + FF_MIN_BUFFER_SIZE)) < 0)
946         return ret;
947
948     orig_buf = pkt->data;
949
950     // frame atom
951     orig_buf += 4;                              // frame size
952     bytestream_put_be32  (&orig_buf, FRAME_ID); // frame container ID
953     buf = orig_buf;
954
955     // frame header
956     tmp = buf;
957     buf += 2;                                   // frame header size will be stored here
958     bytestream_put_be16  (&buf, 0);             // version 1
959     bytestream_put_buffer(&buf, ctx->vendor, 4);
960     bytestream_put_be16  (&buf, avctx->width);
961     bytestream_put_be16  (&buf, avctx->height);
962
963     frame_flags = ctx->chroma_factor << 6;
964     if (avctx->flags & CODEC_FLAG_INTERLACED_DCT)
965         frame_flags |= pic->top_field_first ? 0x04 : 0x08;
966     bytestream_put_byte  (&buf, frame_flags);
967
968     bytestream_put_byte  (&buf, 0);             // reserved
969     bytestream_put_byte  (&buf, avctx->color_primaries);
970     bytestream_put_byte  (&buf, avctx->color_trc);
971     bytestream_put_byte  (&buf, avctx->colorspace);
972     bytestream_put_byte  (&buf, 0x40 | (ctx->alpha_bits >> 3));
973     bytestream_put_byte  (&buf, 0);             // reserved
974     if (ctx->quant_sel != QUANT_MAT_DEFAULT) {
975         bytestream_put_byte  (&buf, 0x03);      // matrix flags - both matrices are present
976         // luma quantisation matrix
977         for (i = 0; i < 64; i++)
978             bytestream_put_byte(&buf, ctx->quant_mat[i]);
979         // chroma quantisation matrix
980         for (i = 0; i < 64; i++)
981             bytestream_put_byte(&buf, ctx->quant_mat[i]);
982     } else {
983         bytestream_put_byte  (&buf, 0x00);      // matrix flags - default matrices are used
984     }
985     bytestream_put_be16  (&tmp, buf - orig_buf); // write back frame header size
986
987     for (ctx->cur_picture_idx = 0;
988          ctx->cur_picture_idx < ctx->pictures_per_frame;
989          ctx->cur_picture_idx++) {
990         // picture header
991         picture_size_pos = buf + 1;
992         bytestream_put_byte  (&buf, 0x40);          // picture header size (in bits)
993         buf += 4;                                   // picture data size will be stored here
994         bytestream_put_be16  (&buf, ctx->slices_per_picture);
995         bytestream_put_byte  (&buf, av_log2(ctx->mbs_per_slice) << 4); // slice width and height in MBs
996
997         // seek table - will be filled during slice encoding
998         slice_sizes = buf;
999         buf += ctx->slices_per_picture * 2;
1000
1001         // slices
1002         if (!ctx->force_quant) {
1003             ret = avctx->execute2(avctx, find_quant_thread, (void*)pic, NULL,
1004                                   ctx->mb_height);
1005             if (ret)
1006                 return ret;
1007         }
1008
1009         for (y = 0; y < ctx->mb_height; y++) {
1010             int mbs_per_slice = ctx->mbs_per_slice;
1011             for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
1012                 q = ctx->force_quant ? ctx->force_quant
1013                                      : ctx->slice_q[mb + y * ctx->slices_width];
1014
1015                 while (ctx->mb_width - x < mbs_per_slice)
1016                     mbs_per_slice >>= 1;
1017
1018                 bytestream_put_byte(&buf, slice_hdr_size << 3);
1019                 slice_hdr = buf;
1020                 buf += slice_hdr_size - 1;
1021                 if (pkt_size <= buf - orig_buf + 2 * max_slice_size) {
1022                     uint8_t *start = pkt->data;
1023                     // Recompute new size according to max_slice_size
1024                     // and deduce delta
1025                     int delta = 200 + (ctx->pictures_per_frame *
1026                                 ctx->slices_per_picture + 1) *
1027                                 max_slice_size - pkt_size;
1028
1029                     delta = FFMAX(delta, 2 * max_slice_size);
1030                     ctx->frame_size_upper_bound += delta;
1031
1032                     if (!ctx->warn) {
1033                         avpriv_request_sample(avctx,
1034                                               "Packet too small: is %i,"
1035                                               " needs %i (slice: %i). "
1036                                               "Correct allocation",
1037                                               pkt_size, delta, max_slice_size);
1038                         ctx->warn = 1;
1039                     }
1040
1041                     ret = av_grow_packet(pkt, delta);
1042                     if (ret < 0)
1043                         return ret;
1044
1045                     pkt_size += delta;
1046                     // restore pointers
1047                     orig_buf         = pkt->data + (orig_buf         - start);
1048                     buf              = pkt->data + (buf              - start);
1049                     picture_size_pos = pkt->data + (picture_size_pos - start);
1050                     slice_sizes      = pkt->data + (slice_sizes      - start);
1051                     slice_hdr        = pkt->data + (slice_hdr        - start);
1052                     tmp              = pkt->data + (tmp              - start);
1053                 }
1054                 init_put_bits(&pb, buf, (pkt_size - (buf - orig_buf)));
1055                 ret = encode_slice(avctx, pic, &pb, sizes, x, y, q,
1056                                    mbs_per_slice);
1057                 if (ret < 0)
1058                     return ret;
1059
1060                 bytestream_put_byte(&slice_hdr, q);
1061                 slice_size = slice_hdr_size + sizes[ctx->num_planes - 1];
1062                 for (i = 0; i < ctx->num_planes - 1; i++) {
1063                     bytestream_put_be16(&slice_hdr, sizes[i]);
1064                     slice_size += sizes[i];
1065                 }
1066                 bytestream_put_be16(&slice_sizes, slice_size);
1067                 buf += slice_size - slice_hdr_size;
1068                 if (max_slice_size < slice_size)
1069                     max_slice_size = slice_size;
1070             }
1071         }
1072
1073         picture_size = buf - (picture_size_pos - 1);
1074         bytestream_put_be32(&picture_size_pos, picture_size);
1075     }
1076
1077     orig_buf -= 8;
1078     frame_size = buf - orig_buf;
1079     bytestream_put_be32(&orig_buf, frame_size);
1080
1081     pkt->size   = frame_size;
1082     pkt->flags |= AV_PKT_FLAG_KEY;
1083     *got_packet = 1;
1084
1085     return 0;
1086 }
1087
1088 static av_cold int encode_close(AVCodecContext *avctx)
1089 {
1090     ProresContext *ctx = avctx->priv_data;
1091     int i;
1092
1093     av_frame_free(&avctx->coded_frame);
1094
1095     if (ctx->tdata) {
1096         for (i = 0; i < avctx->thread_count; i++)
1097             av_freep(&ctx->tdata[i].nodes);
1098     }
1099     av_freep(&ctx->tdata);
1100     av_freep(&ctx->slice_q);
1101
1102     return 0;
1103 }
1104
1105 static void prores_fdct(FDCTDSPContext *fdsp, const uint16_t *src,
1106                         int linesize, int16_t *block)
1107 {
1108     int x, y;
1109     const uint16_t *tsrc = src;
1110
1111     for (y = 0; y < 8; y++) {
1112         for (x = 0; x < 8; x++)
1113             block[y * 8 + x] = tsrc[x];
1114         tsrc += linesize >> 1;
1115     }
1116     fdsp->fdct(block);
1117 }
1118
1119 static av_cold int encode_init(AVCodecContext *avctx)
1120 {
1121     ProresContext *ctx = avctx->priv_data;
1122     int mps;
1123     int i, j;
1124     int min_quant, max_quant;
1125     int interlaced = !!(avctx->flags & CODEC_FLAG_INTERLACED_DCT);
1126
1127     avctx->bits_per_raw_sample = 10;
1128     avctx->coded_frame = av_frame_alloc();
1129     if (!avctx->coded_frame)
1130         return AVERROR(ENOMEM);
1131     avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
1132     avctx->coded_frame->key_frame = 1;
1133
1134     ctx->fdct      = prores_fdct;
1135     ctx->scantable = interlaced ? ff_prores_interlaced_scan
1136                                 : ff_prores_progressive_scan;
1137     ff_fdctdsp_init(&ctx->fdsp, avctx);
1138
1139     mps = ctx->mbs_per_slice;
1140     if (mps & (mps - 1)) {
1141         av_log(avctx, AV_LOG_ERROR,
1142                "there should be an integer power of two MBs per slice\n");
1143         return AVERROR(EINVAL);
1144     }
1145     if (ctx->profile == PRORES_PROFILE_AUTO) {
1146         const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(avctx->pix_fmt);
1147         ctx->profile = (desc->flags & AV_PIX_FMT_FLAG_ALPHA ||
1148                         !(desc->log2_chroma_w + desc->log2_chroma_h))
1149                      ? PRORES_PROFILE_4444 : PRORES_PROFILE_HQ;
1150         av_log(avctx, AV_LOG_INFO, "Autoselected %s. It can be overridden "
1151                "through -profile option.\n", ctx->profile == PRORES_PROFILE_4444
1152                ? "4:4:4:4 profile because of the used input colorspace"
1153                : "HQ profile to keep best quality");
1154     }
1155     if (av_pix_fmt_desc_get(avctx->pix_fmt)->flags & AV_PIX_FMT_FLAG_ALPHA) {
1156         if (ctx->profile != PRORES_PROFILE_4444) {
1157             // force alpha and warn
1158             av_log(avctx, AV_LOG_WARNING, "Profile selected will not "
1159                    "encode alpha. Override with -profile if needed.\n");
1160             ctx->alpha_bits = 0;
1161         }
1162         if (ctx->alpha_bits & 7) {
1163             av_log(avctx, AV_LOG_ERROR, "alpha bits should be 0, 8 or 16\n");
1164             return AVERROR(EINVAL);
1165         }
1166         avctx->bits_per_coded_sample = 32;
1167     } else {
1168         ctx->alpha_bits = 0;
1169     }
1170
1171     ctx->chroma_factor = avctx->pix_fmt == AV_PIX_FMT_YUV422P10
1172                          ? CFACTOR_Y422
1173                          : CFACTOR_Y444;
1174     ctx->profile_info  = prores_profile_info + ctx->profile;
1175     ctx->num_planes    = 3 + !!ctx->alpha_bits;
1176
1177     ctx->mb_width      = FFALIGN(avctx->width,  16) >> 4;
1178
1179     if (interlaced)
1180         ctx->mb_height = FFALIGN(avctx->height, 32) >> 5;
1181     else
1182         ctx->mb_height = FFALIGN(avctx->height, 16) >> 4;
1183
1184     ctx->slices_width  = ctx->mb_width / mps;
1185     ctx->slices_width += av_popcount(ctx->mb_width - ctx->slices_width * mps);
1186     ctx->slices_per_picture = ctx->mb_height * ctx->slices_width;
1187     ctx->pictures_per_frame = 1 + interlaced;
1188
1189     if (ctx->quant_sel == -1)
1190         ctx->quant_mat = prores_quant_matrices[ctx->profile_info->quant];
1191     else
1192         ctx->quant_mat = prores_quant_matrices[ctx->quant_sel];
1193
1194     if (strlen(ctx->vendor) != 4) {
1195         av_log(avctx, AV_LOG_ERROR, "vendor ID should be 4 bytes\n");
1196         return AVERROR_INVALIDDATA;
1197     }
1198
1199     ctx->force_quant = avctx->global_quality / FF_QP2LAMBDA;
1200     if (!ctx->force_quant) {
1201         if (!ctx->bits_per_mb) {
1202             for (i = 0; i < NUM_MB_LIMITS - 1; i++)
1203                 if (prores_mb_limits[i] >= ctx->mb_width * ctx->mb_height *
1204                                            ctx->pictures_per_frame)
1205                     break;
1206             ctx->bits_per_mb   = ctx->profile_info->br_tab[i];
1207         } else if (ctx->bits_per_mb < 128) {
1208             av_log(avctx, AV_LOG_ERROR, "too few bits per MB, please set at least 128\n");
1209             return AVERROR_INVALIDDATA;
1210         }
1211
1212         min_quant = ctx->profile_info->min_quant;
1213         max_quant = ctx->profile_info->max_quant;
1214         for (i = min_quant; i < MAX_STORED_Q; i++) {
1215             for (j = 0; j < 64; j++)
1216                 ctx->quants[i][j] = ctx->quant_mat[j] * i;
1217         }
1218
1219         ctx->slice_q = av_malloc(ctx->slices_per_picture * sizeof(*ctx->slice_q));
1220         if (!ctx->slice_q) {
1221             encode_close(avctx);
1222             return AVERROR(ENOMEM);
1223         }
1224
1225         ctx->tdata = av_mallocz(avctx->thread_count * sizeof(*ctx->tdata));
1226         if (!ctx->tdata) {
1227             encode_close(avctx);
1228             return AVERROR(ENOMEM);
1229         }
1230
1231         for (j = 0; j < avctx->thread_count; j++) {
1232             ctx->tdata[j].nodes = av_malloc((ctx->slices_width + 1)
1233                                             * TRELLIS_WIDTH
1234                                             * sizeof(*ctx->tdata->nodes));
1235             if (!ctx->tdata[j].nodes) {
1236                 encode_close(avctx);
1237                 return AVERROR(ENOMEM);
1238             }
1239             for (i = min_quant; i < max_quant + 2; i++) {
1240                 ctx->tdata[j].nodes[i].prev_node = -1;
1241                 ctx->tdata[j].nodes[i].bits      = 0;
1242                 ctx->tdata[j].nodes[i].score     = 0;
1243             }
1244         }
1245     } else {
1246         int ls = 0;
1247
1248         if (ctx->force_quant > 64) {
1249             av_log(avctx, AV_LOG_ERROR, "too large quantiser, maximum is 64\n");
1250             return AVERROR_INVALIDDATA;
1251         }
1252
1253         for (j = 0; j < 64; j++) {
1254             ctx->quants[0][j] = ctx->quant_mat[j] * ctx->force_quant;
1255             ls += av_log2((1 << 11)  / ctx->quants[0][j]) * 2 + 1;
1256         }
1257
1258         ctx->bits_per_mb = ls * 8;
1259         if (ctx->chroma_factor == CFACTOR_Y444)
1260             ctx->bits_per_mb += ls * 4;
1261     }
1262
1263     ctx->frame_size_upper_bound = (ctx->pictures_per_frame *
1264                                    ctx->slices_per_picture + 1) *
1265                                   (2 + 2 * ctx->num_planes +
1266                                    (mps * ctx->bits_per_mb) / 8)
1267                                   + 200;
1268
1269     if (ctx->alpha_bits) {
1270          // The alpha plane is run-coded and might exceed the bit budget.
1271          ctx->frame_size_upper_bound += (ctx->pictures_per_frame *
1272                                          ctx->slices_per_picture + 1) *
1273          /* num pixels per slice */     (ctx->mbs_per_slice * 256 *
1274          /* bits per pixel */            (1 + ctx->alpha_bits + 1) + 7 >> 3);
1275     }
1276
1277     avctx->codec_tag   = ctx->profile_info->tag;
1278
1279     av_log(avctx, AV_LOG_DEBUG,
1280            "profile %d, %d slices, interlacing: %s, %d bits per MB\n",
1281            ctx->profile, ctx->slices_per_picture * ctx->pictures_per_frame,
1282            interlaced ? "yes" : "no", ctx->bits_per_mb);
1283     av_log(avctx, AV_LOG_DEBUG, "frame size upper bound: %d\n",
1284            ctx->frame_size_upper_bound);
1285
1286     return 0;
1287 }
1288
1289 #define OFFSET(x) offsetof(ProresContext, x)
1290 #define VE     AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
1291
1292 static const AVOption options[] = {
1293     { "mbs_per_slice", "macroblocks per slice", OFFSET(mbs_per_slice),
1294         AV_OPT_TYPE_INT, { .i64 = 8 }, 1, MAX_MBS_PER_SLICE, VE },
1295     { "profile",       NULL, OFFSET(profile), AV_OPT_TYPE_INT,
1296         { .i64 = PRORES_PROFILE_AUTO },
1297         PRORES_PROFILE_AUTO, PRORES_PROFILE_4444, VE, "profile" },
1298     { "auto",         NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_AUTO },
1299         0, 0, VE, "profile" },
1300     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_PROXY },
1301         0, 0, VE, "profile" },
1302     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_LT },
1303         0, 0, VE, "profile" },
1304     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_STANDARD },
1305         0, 0, VE, "profile" },
1306     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_HQ },
1307         0, 0, VE, "profile" },
1308     { "4444",          NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_4444 },
1309         0, 0, VE, "profile" },
1310     { "vendor", "vendor ID", OFFSET(vendor),
1311         AV_OPT_TYPE_STRING, { .str = "Lavc" }, CHAR_MIN, CHAR_MAX, VE },
1312     { "bits_per_mb", "desired bits per macroblock", OFFSET(bits_per_mb),
1313         AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 8192, VE },
1314     { "quant_mat", "quantiser matrix", OFFSET(quant_sel), AV_OPT_TYPE_INT,
1315         { .i64 = -1 }, -1, QUANT_MAT_DEFAULT, VE, "quant_mat" },
1316     { "auto",          NULL, 0, AV_OPT_TYPE_CONST, { .i64 = -1 },
1317         0, 0, VE, "quant_mat" },
1318     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_PROXY },
1319         0, 0, VE, "quant_mat" },
1320     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_LT },
1321         0, 0, VE, "quant_mat" },
1322     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_STANDARD },
1323         0, 0, VE, "quant_mat" },
1324     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_HQ },
1325         0, 0, VE, "quant_mat" },
1326     { "default",       NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_DEFAULT },
1327         0, 0, VE, "quant_mat" },
1328     { "alpha_bits", "bits for alpha plane", OFFSET(alpha_bits), AV_OPT_TYPE_INT,
1329         { .i64 = 16 }, 0, 16, VE },
1330     { NULL }
1331 };
1332
1333 static const AVClass proresenc_class = {
1334     .class_name = "ProRes encoder",
1335     .item_name  = av_default_item_name,
1336     .option     = options,
1337     .version    = LIBAVUTIL_VERSION_INT,
1338 };
1339
1340 AVCodec ff_prores_ks_encoder = {
1341     .name           = "prores_ks",
1342     .long_name      = NULL_IF_CONFIG_SMALL("Apple ProRes (iCodec Pro)"),
1343     .type           = AVMEDIA_TYPE_VIDEO,
1344     .id             = AV_CODEC_ID_PRORES,
1345     .priv_data_size = sizeof(ProresContext),
1346     .init           = encode_init,
1347     .close          = encode_close,
1348     .encode2        = encode_frame,
1349     .capabilities   = CODEC_CAP_SLICE_THREADS,
1350     .pix_fmts       = (const enum AVPixelFormat[]) {
1351                           AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10,
1352                           AV_PIX_FMT_YUVA444P10, AV_PIX_FMT_NONE
1353                       },
1354     .priv_class     = &proresenc_class,
1355 };