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