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