avcodec: rename prores encoders
[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 Libav.
10  *
11  * Libav 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  * Libav 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 Libav; 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 "avcodec.h"
28 #include "dsputil.h"
29 #include "put_bits.h"
30 #include "bytestream.h"
31 #include "internal.h"
32 #include "proresdsp.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 3 // should be increased to 4 when there's AV_PIX_FMT_YUV444AP10
41
42 enum {
43     PRORES_PROFILE_PROXY = 0,
44     PRORES_PROFILE_LT,
45     PRORES_PROFILE_STANDARD,
46     PRORES_PROFILE_HQ,
47 };
48
49 enum {
50     QUANT_MAT_PROXY = 0,
51     QUANT_MAT_LT,
52     QUANT_MAT_STANDARD,
53     QUANT_MAT_HQ,
54     QUANT_MAT_DEFAULT,
55 };
56
57 static const uint8_t prores_quant_matrices[][64] = {
58     { // proxy
59          4,  7,  9, 11, 13, 14, 15, 63,
60          7,  7, 11, 12, 14, 15, 63, 63,
61          9, 11, 13, 14, 15, 63, 63, 63,
62         11, 11, 13, 14, 63, 63, 63, 63,
63         11, 13, 14, 63, 63, 63, 63, 63,
64         13, 14, 63, 63, 63, 63, 63, 63,
65         13, 63, 63, 63, 63, 63, 63, 63,
66         63, 63, 63, 63, 63, 63, 63, 63,
67     },
68     { // LT
69          4,  5,  6,  7,  9, 11, 13, 15,
70          5,  5,  7,  8, 11, 13, 15, 17,
71          6,  7,  9, 11, 13, 15, 15, 17,
72          7,  7,  9, 11, 13, 15, 17, 19,
73          7,  9, 11, 13, 14, 16, 19, 23,
74          9, 11, 13, 14, 16, 19, 23, 29,
75          9, 11, 13, 15, 17, 21, 28, 35,
76         11, 13, 16, 17, 21, 28, 35, 41,
77     },
78     { // standard
79          4,  4,  5,  5,  6,  7,  7,  9,
80          4,  4,  5,  6,  7,  7,  9,  9,
81          5,  5,  6,  7,  7,  9,  9, 10,
82          5,  5,  6,  7,  7,  9,  9, 10,
83          5,  6,  7,  7,  8,  9, 10, 12,
84          6,  7,  7,  8,  9, 10, 12, 15,
85          6,  7,  7,  9, 10, 11, 14, 17,
86          7,  7,  9, 10, 11, 14, 17, 21,
87     },
88     { // high quality
89          4,  4,  4,  4,  4,  4,  4,  4,
90          4,  4,  4,  4,  4,  4,  4,  4,
91          4,  4,  4,  4,  4,  4,  4,  4,
92          4,  4,  4,  4,  4,  4,  4,  5,
93          4,  4,  4,  4,  4,  4,  5,  5,
94          4,  4,  4,  4,  4,  5,  5,  6,
95          4,  4,  4,  4,  5,  5,  6,  7,
96          4,  4,  4,  4,  5,  6,  7,  7,
97     },
98     { // codec default
99          4,  4,  4,  4,  4,  4,  4,  4,
100          4,  4,  4,  4,  4,  4,  4,  4,
101          4,  4,  4,  4,  4,  4,  4,  4,
102          4,  4,  4,  4,  4,  4,  4,  4,
103          4,  4,  4,  4,  4,  4,  4,  4,
104          4,  4,  4,  4,  4,  4,  4,  4,
105          4,  4,  4,  4,  4,  4,  4,  4,
106          4,  4,  4,  4,  4,  4,  4,  4,
107     },
108 };
109
110 #define NUM_MB_LIMITS 4
111 static const int prores_mb_limits[NUM_MB_LIMITS] = {
112     1620, // up to 720x576
113     2700, // up to 960x720
114     6075, // up to 1440x1080
115     9216, // up to 2048x1152
116 };
117
118 static const struct prores_profile {
119     const char *full_name;
120     uint32_t    tag;
121     int         min_quant;
122     int         max_quant;
123     int         br_tab[NUM_MB_LIMITS];
124     int         quant;
125 } prores_profile_info[4] = {
126     {
127         .full_name = "proxy",
128         .tag       = MKTAG('a', 'p', 'c', 'o'),
129         .min_quant = 4,
130         .max_quant = 8,
131         .br_tab    = { 300, 242, 220, 194 },
132         .quant     = QUANT_MAT_PROXY,
133     },
134     {
135         .full_name = "LT",
136         .tag       = MKTAG('a', 'p', 'c', 's'),
137         .min_quant = 1,
138         .max_quant = 9,
139         .br_tab    = { 720, 560, 490, 440 },
140         .quant     = QUANT_MAT_LT,
141     },
142     {
143         .full_name = "standard",
144         .tag       = MKTAG('a', 'p', 'c', 'n'),
145         .min_quant = 1,
146         .max_quant = 6,
147         .br_tab    = { 1050, 808, 710, 632 },
148         .quant     = QUANT_MAT_STANDARD,
149     },
150     {
151         .full_name = "high quality",
152         .tag       = MKTAG('a', 'p', 'c', 'h'),
153         .min_quant = 1,
154         .max_quant = 6,
155         .br_tab    = { 1566, 1216, 1070, 950 },
156         .quant     = QUANT_MAT_HQ,
157     }
158 // for 4444 profile bitrate numbers are { 2350, 1828, 1600, 1425 }
159 };
160
161 #define TRELLIS_WIDTH 16
162 #define SCORE_LIMIT   INT_MAX / 2
163
164 struct TrellisNode {
165     int prev_node;
166     int quant;
167     int bits;
168     int score;
169 };
170
171 #define MAX_STORED_Q 16
172
173 typedef struct ProresThreadData {
174     DECLARE_ALIGNED(16, int16_t, blocks)[MAX_PLANES][64 * 4 * MAX_MBS_PER_SLICE];
175     DECLARE_ALIGNED(16, uint16_t, emu_buf)[16 * 16];
176     int16_t custom_q[64];
177     struct TrellisNode *nodes;
178 } ProresThreadData;
179
180 typedef struct ProresContext {
181     AVClass *class;
182     DECLARE_ALIGNED(16, int16_t, blocks)[MAX_PLANES][64 * 4 * MAX_MBS_PER_SLICE];
183     DECLARE_ALIGNED(16, uint16_t, emu_buf)[16*16];
184     int16_t quants[MAX_STORED_Q][64];
185     int16_t custom_q[64];
186     const uint8_t *quant_mat;
187
188     ProresDSPContext dsp;
189     ScanTable  scantable;
190
191     int mb_width, mb_height;
192     int mbs_per_slice;
193     int num_chroma_blocks, chroma_factor;
194     int slices_width;
195     int slices_per_picture;
196     int pictures_per_frame; // 1 for progressive, 2 for interlaced
197     int cur_picture_idx;
198     int num_planes;
199     int bits_per_mb;
200     int force_quant;
201
202     char *vendor;
203     int quant_sel;
204
205     int frame_size_upper_bound;
206
207     int profile;
208     const struct prores_profile *profile_info;
209
210     int *slice_q;
211
212     ProresThreadData *tdata;
213 } ProresContext;
214
215 static void get_slice_data(ProresContext *ctx, const uint16_t *src,
216                            int linesize, int x, int y, int w, int h,
217                            int16_t *blocks, uint16_t *emu_buf,
218                            int mbs_per_slice, int blocks_per_mb, int is_chroma)
219 {
220     const uint16_t *esrc;
221     const int mb_width = 4 * blocks_per_mb;
222     int elinesize;
223     int i, j, k;
224
225     for (i = 0; i < mbs_per_slice; i++, src += mb_width) {
226         if (x >= w) {
227             memset(blocks, 0, 64 * (mbs_per_slice - i) * blocks_per_mb
228                               * sizeof(*blocks));
229             return;
230         }
231         if (x + mb_width <= w && y + 16 <= h) {
232             esrc      = src;
233             elinesize = linesize;
234         } else {
235             int bw, bh, pix;
236
237             esrc      = emu_buf;
238             elinesize = 16 * sizeof(*emu_buf);
239
240             bw = FFMIN(w - x, mb_width);
241             bh = FFMIN(h - y, 16);
242
243             for (j = 0; j < bh; j++) {
244                 memcpy(emu_buf + j * 16,
245                        (const uint8_t*)src + j * linesize,
246                        bw * sizeof(*src));
247                 pix = emu_buf[j * 16 + bw - 1];
248                 for (k = bw; k < mb_width; k++)
249                     emu_buf[j * 16 + k] = pix;
250             }
251             for (; j < 16; j++)
252                 memcpy(emu_buf + j * 16,
253                        emu_buf + (bh - 1) * 16,
254                        mb_width * sizeof(*emu_buf));
255         }
256         if (!is_chroma) {
257             ctx->dsp.fdct(esrc, elinesize, blocks);
258             blocks += 64;
259             if (blocks_per_mb > 2) {
260                 ctx->dsp.fdct(esrc + 8, elinesize, blocks);
261                 blocks += 64;
262             }
263             ctx->dsp.fdct(esrc + elinesize * 4, elinesize, blocks);
264             blocks += 64;
265             if (blocks_per_mb > 2) {
266                 ctx->dsp.fdct(esrc + elinesize * 4 + 8, elinesize, blocks);
267                 blocks += 64;
268             }
269         } else {
270             ctx->dsp.fdct(esrc, elinesize, blocks);
271             blocks += 64;
272             ctx->dsp.fdct(esrc + elinesize * 4, elinesize, blocks);
273             blocks += 64;
274             if (blocks_per_mb > 2) {
275                 ctx->dsp.fdct(esrc + 8, elinesize, blocks);
276                 blocks += 64;
277                 ctx->dsp.fdct(esrc + elinesize * 4 + 8, elinesize, blocks);
278                 blocks += 64;
279             }
280         }
281
282         x += mb_width;
283     }
284 }
285
286 /**
287  * Write an unsigned rice/exp golomb codeword.
288  */
289 static inline void encode_vlc_codeword(PutBitContext *pb, unsigned codebook, int val)
290 {
291     unsigned int rice_order, exp_order, switch_bits, switch_val;
292     int exponent;
293
294     /* number of prefix bits to switch between Rice and expGolomb */
295     switch_bits = (codebook & 3) + 1;
296     rice_order  =  codebook >> 5;       /* rice code order */
297     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
298
299     switch_val  = switch_bits << rice_order;
300
301     if (val >= switch_val) {
302         val -= switch_val - (1 << exp_order);
303         exponent = av_log2(val);
304
305         put_bits(pb, exponent - exp_order + switch_bits, 0);
306         put_bits(pb, exponent + 1, val);
307     } else {
308         exponent = val >> rice_order;
309
310         if (exponent)
311             put_bits(pb, exponent, 0);
312         put_bits(pb, 1, 1);
313         if (rice_order)
314             put_sbits(pb, rice_order, val);
315     }
316 }
317
318 #define GET_SIGN(x)  ((x) >> 31)
319 #define MAKE_CODE(x) (((x) << 1) ^ GET_SIGN(x))
320
321 static void encode_dcs(PutBitContext *pb, int16_t *blocks,
322                        int blocks_per_slice, int scale)
323 {
324     int i;
325     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
326
327     prev_dc = (blocks[0] - 0x4000) / scale;
328     encode_vlc_codeword(pb, FIRST_DC_CB, MAKE_CODE(prev_dc));
329     sign     = 0;
330     codebook = 3;
331     blocks  += 64;
332
333     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
334         dc       = (blocks[0] - 0x4000) / scale;
335         delta    = dc - prev_dc;
336         new_sign = GET_SIGN(delta);
337         delta    = (delta ^ sign) - sign;
338         code     = MAKE_CODE(delta);
339         encode_vlc_codeword(pb, ff_prores_dc_codebook[codebook], code);
340         codebook = (code + (code & 1)) >> 1;
341         codebook = FFMIN(codebook, 3);
342         sign     = new_sign;
343         prev_dc  = dc;
344     }
345 }
346
347 static void encode_acs(PutBitContext *pb, int16_t *blocks,
348                        int blocks_per_slice,
349                        int plane_size_factor,
350                        const uint8_t *scan, const int16_t *qmat)
351 {
352     int idx, i;
353     int run, level, run_cb, lev_cb;
354     int max_coeffs, abs_level;
355
356     max_coeffs = blocks_per_slice << 6;
357     run_cb     = ff_prores_run_to_cb_index[4];
358     lev_cb     = ff_prores_lev_to_cb_index[2];
359     run        = 0;
360
361     for (i = 1; i < 64; i++) {
362         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
363             level = blocks[idx] / qmat[scan[i]];
364             if (level) {
365                 abs_level = FFABS(level);
366                 encode_vlc_codeword(pb, ff_prores_ac_codebook[run_cb], run);
367                 encode_vlc_codeword(pb, ff_prores_ac_codebook[lev_cb],
368                                     abs_level - 1);
369                 put_sbits(pb, 1, GET_SIGN(level));
370
371                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
372                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
373                 run    = 0;
374             } else {
375                 run++;
376             }
377         }
378     }
379 }
380
381 static int encode_slice_plane(ProresContext *ctx, PutBitContext *pb,
382                               const uint16_t *src, int linesize,
383                               int mbs_per_slice, int16_t *blocks,
384                               int blocks_per_mb, int plane_size_factor,
385                               const int16_t *qmat)
386 {
387     int blocks_per_slice, saved_pos;
388
389     saved_pos = put_bits_count(pb);
390     blocks_per_slice = mbs_per_slice * blocks_per_mb;
391
392     encode_dcs(pb, blocks, blocks_per_slice, qmat[0]);
393     encode_acs(pb, blocks, blocks_per_slice, plane_size_factor,
394                ctx->scantable.permutated, qmat);
395     flush_put_bits(pb);
396
397     return (put_bits_count(pb) - saved_pos) >> 3;
398 }
399
400 static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
401                         PutBitContext *pb,
402                         int sizes[4], int x, int y, int quant,
403                         int mbs_per_slice)
404 {
405     ProresContext *ctx = avctx->priv_data;
406     int i, xp, yp;
407     int total_size = 0;
408     const uint16_t *src;
409     int slice_width_factor = av_log2(mbs_per_slice);
410     int num_cblocks, pwidth, linesize, line_add;
411     int plane_factor, is_chroma;
412     uint16_t *qmat;
413
414     if (ctx->pictures_per_frame == 1)
415         line_add = 0;
416     else
417         line_add = ctx->cur_picture_idx ^ !pic->top_field_first;
418
419     if (ctx->force_quant) {
420         qmat = ctx->quants[0];
421     } else if (quant < MAX_STORED_Q) {
422         qmat = ctx->quants[quant];
423     } else {
424         qmat = ctx->custom_q;
425         for (i = 0; i < 64; i++)
426             qmat[i] = ctx->quant_mat[i] * quant;
427     }
428
429     for (i = 0; i < ctx->num_planes; i++) {
430         is_chroma    = (i == 1 || i == 2);
431         plane_factor = slice_width_factor + 2;
432         if (is_chroma)
433             plane_factor += ctx->chroma_factor - 3;
434         if (!is_chroma || ctx->chroma_factor == CFACTOR_Y444) {
435             xp          = x << 4;
436             yp          = y << 4;
437             num_cblocks = 4;
438             pwidth      = avctx->width;
439         } else {
440             xp          = x << 3;
441             yp          = y << 4;
442             num_cblocks = 2;
443             pwidth      = avctx->width >> 1;
444         }
445
446         linesize = pic->linesize[i] * ctx->pictures_per_frame;
447         src = (const uint16_t*)(pic->data[i] + yp * linesize +
448                                 line_add * pic->linesize[i]) + xp;
449
450         get_slice_data(ctx, src, linesize, xp, yp,
451                        pwidth, avctx->height / ctx->pictures_per_frame,
452                        ctx->blocks[0], ctx->emu_buf,
453                        mbs_per_slice, num_cblocks, is_chroma);
454         sizes[i] = encode_slice_plane(ctx, pb, src, linesize,
455                                       mbs_per_slice, ctx->blocks[0],
456                                       num_cblocks, plane_factor,
457                                       qmat);
458         total_size += sizes[i];
459     }
460     return total_size;
461 }
462
463 static inline int estimate_vlc(unsigned codebook, int val)
464 {
465     unsigned int rice_order, exp_order, switch_bits, switch_val;
466     int exponent;
467
468     /* number of prefix bits to switch between Rice and expGolomb */
469     switch_bits = (codebook & 3) + 1;
470     rice_order  =  codebook >> 5;       /* rice code order */
471     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
472
473     switch_val  = switch_bits << rice_order;
474
475     if (val >= switch_val) {
476         val -= switch_val - (1 << exp_order);
477         exponent = av_log2(val);
478
479         return exponent * 2 - exp_order + switch_bits + 1;
480     } else {
481         return (val >> rice_order) + rice_order + 1;
482     }
483 }
484
485 static int estimate_dcs(int *error, int16_t *blocks, int blocks_per_slice,
486                         int scale)
487 {
488     int i;
489     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
490     int bits;
491
492     prev_dc  = (blocks[0] - 0x4000) / scale;
493     bits     = estimate_vlc(FIRST_DC_CB, MAKE_CODE(prev_dc));
494     sign     = 0;
495     codebook = 3;
496     blocks  += 64;
497     *error  += FFABS(blocks[0] - 0x4000) % scale;
498
499     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
500         dc       = (blocks[0] - 0x4000) / scale;
501         *error  += FFABS(blocks[0] - 0x4000) % scale;
502         delta    = dc - prev_dc;
503         new_sign = GET_SIGN(delta);
504         delta    = (delta ^ sign) - sign;
505         code     = MAKE_CODE(delta);
506         bits    += estimate_vlc(ff_prores_dc_codebook[codebook], code);
507         codebook = (code + (code & 1)) >> 1;
508         codebook = FFMIN(codebook, 3);
509         sign     = new_sign;
510         prev_dc  = dc;
511     }
512
513     return bits;
514 }
515
516 static int estimate_acs(int *error, int16_t *blocks, int blocks_per_slice,
517                         int plane_size_factor,
518                         const uint8_t *scan, const int16_t *qmat)
519 {
520     int idx, i;
521     int run, level, run_cb, lev_cb;
522     int max_coeffs, abs_level;
523     int bits = 0;
524
525     max_coeffs = blocks_per_slice << 6;
526     run_cb     = ff_prores_run_to_cb_index[4];
527     lev_cb     = ff_prores_lev_to_cb_index[2];
528     run        = 0;
529
530     for (i = 1; i < 64; i++) {
531         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
532             level   = blocks[idx] / qmat[scan[i]];
533             *error += FFABS(blocks[idx]) % qmat[scan[i]];
534             if (level) {
535                 abs_level = FFABS(level);
536                 bits += estimate_vlc(ff_prores_ac_codebook[run_cb], run);
537                 bits += estimate_vlc(ff_prores_ac_codebook[lev_cb],
538                                      abs_level - 1) + 1;
539
540                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
541                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
542                 run    = 0;
543             } else {
544                 run++;
545             }
546         }
547     }
548
549     return bits;
550 }
551
552 static int estimate_slice_plane(ProresContext *ctx, int *error, int plane,
553                                 const uint16_t *src, int linesize,
554                                 int mbs_per_slice,
555                                 int blocks_per_mb, int plane_size_factor,
556                                 const int16_t *qmat, ProresThreadData *td)
557 {
558     int blocks_per_slice;
559     int bits;
560
561     blocks_per_slice = mbs_per_slice * blocks_per_mb;
562
563     bits  = estimate_dcs(error, td->blocks[plane], blocks_per_slice, qmat[0]);
564     bits += estimate_acs(error, td->blocks[plane], blocks_per_slice,
565                          plane_size_factor, ctx->scantable.permutated, qmat);
566
567     return FFALIGN(bits, 8);
568 }
569
570 static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
571                             int trellis_node, int x, int y, int mbs_per_slice,
572                             ProresThreadData *td)
573 {
574     ProresContext *ctx = avctx->priv_data;
575     int i, q, pq, xp, yp;
576     const uint16_t *src;
577     int slice_width_factor = av_log2(mbs_per_slice);
578     int num_cblocks[MAX_PLANES], pwidth;
579     int plane_factor[MAX_PLANES], is_chroma[MAX_PLANES];
580     const int min_quant = ctx->profile_info->min_quant;
581     const int max_quant = ctx->profile_info->max_quant;
582     int error, bits, bits_limit;
583     int mbs, prev, cur, new_score;
584     int slice_bits[TRELLIS_WIDTH], slice_score[TRELLIS_WIDTH];
585     int overquant;
586     uint16_t *qmat;
587     int linesize[4], line_add;
588
589     if (ctx->pictures_per_frame == 1)
590         line_add = 0;
591     else
592         line_add = ctx->cur_picture_idx ^ !pic->top_field_first;
593     mbs = x + mbs_per_slice;
594
595     for (i = 0; i < ctx->num_planes; i++) {
596         is_chroma[i]    = (i == 1 || i == 2);
597         plane_factor[i] = slice_width_factor + 2;
598         if (is_chroma[i])
599             plane_factor[i] += ctx->chroma_factor - 3;
600         if (!is_chroma[i] || ctx->chroma_factor == CFACTOR_Y444) {
601             xp             = x << 4;
602             yp             = y << 4;
603             num_cblocks[i] = 4;
604             pwidth         = avctx->width;
605         } else {
606             xp             = x << 3;
607             yp             = y << 4;
608             num_cblocks[i] = 2;
609             pwidth         = avctx->width >> 1;
610         }
611
612         linesize[i] = pic->linesize[i] * ctx->pictures_per_frame;
613         src = (const uint16_t*)(pic->data[i] + yp * linesize[i] +
614                                 line_add * pic->linesize[i]) + xp;
615
616         get_slice_data(ctx, src, linesize[i], xp, yp,
617                        pwidth, avctx->height / ctx->pictures_per_frame,
618                        td->blocks[i], td->emu_buf,
619                        mbs_per_slice, num_cblocks[i], is_chroma[i]);
620     }
621
622     for (q = min_quant; q < max_quant + 2; q++) {
623         td->nodes[trellis_node + q].prev_node = -1;
624         td->nodes[trellis_node + q].quant     = q;
625     }
626
627     // todo: maybe perform coarser quantising to fit into frame size when needed
628     for (q = min_quant; q <= max_quant; q++) {
629         bits  = 0;
630         error = 0;
631         for (i = 0; i < ctx->num_planes; i++) {
632             bits += estimate_slice_plane(ctx, &error, i,
633                                          src, linesize[i],
634                                          mbs_per_slice,
635                                          num_cblocks[i], plane_factor[i],
636                                          ctx->quants[q], td);
637         }
638         if (bits > 65000 * 8) {
639             error = SCORE_LIMIT;
640             break;
641         }
642         slice_bits[q]  = bits;
643         slice_score[q] = error;
644     }
645     if (slice_bits[max_quant] <= ctx->bits_per_mb * mbs_per_slice) {
646         slice_bits[max_quant + 1]  = slice_bits[max_quant];
647         slice_score[max_quant + 1] = slice_score[max_quant] + 1;
648         overquant = max_quant;
649     } else {
650         for (q = max_quant + 1; q < 128; q++) {
651             bits  = 0;
652             error = 0;
653             if (q < MAX_STORED_Q) {
654                 qmat = ctx->quants[q];
655             } else {
656                 qmat = td->custom_q;
657                 for (i = 0; i < 64; i++)
658                     qmat[i] = ctx->quant_mat[i] * q;
659             }
660             for (i = 0; i < ctx->num_planes; i++) {
661                 bits += estimate_slice_plane(ctx, &error, i,
662                                              src, linesize[i],
663                                              mbs_per_slice,
664                                              num_cblocks[i], plane_factor[i],
665                                              qmat, td);
666             }
667             if (bits <= ctx->bits_per_mb * mbs_per_slice)
668                 break;
669         }
670
671         slice_bits[max_quant + 1]  = bits;
672         slice_score[max_quant + 1] = error;
673         overquant = q;
674     }
675     td->nodes[trellis_node + max_quant + 1].quant = overquant;
676
677     bits_limit = mbs * ctx->bits_per_mb;
678     for (pq = min_quant; pq < max_quant + 2; pq++) {
679         prev = trellis_node - TRELLIS_WIDTH + pq;
680
681         for (q = min_quant; q < max_quant + 2; q++) {
682             cur = trellis_node + q;
683
684             bits  = td->nodes[prev].bits + slice_bits[q];
685             error = slice_score[q];
686             if (bits > bits_limit)
687                 error = SCORE_LIMIT;
688
689             if (td->nodes[prev].score < SCORE_LIMIT && error < SCORE_LIMIT)
690                 new_score = td->nodes[prev].score + error;
691             else
692                 new_score = SCORE_LIMIT;
693             if (td->nodes[cur].prev_node == -1 ||
694                 td->nodes[cur].score >= new_score) {
695
696                 td->nodes[cur].bits      = bits;
697                 td->nodes[cur].score     = new_score;
698                 td->nodes[cur].prev_node = prev;
699             }
700         }
701     }
702
703     error = td->nodes[trellis_node + min_quant].score;
704     pq    = trellis_node + min_quant;
705     for (q = min_quant + 1; q < max_quant + 2; q++) {
706         if (td->nodes[trellis_node + q].score <= error) {
707             error = td->nodes[trellis_node + q].score;
708             pq    = trellis_node + q;
709         }
710     }
711
712     return pq;
713 }
714
715 static int find_quant_thread(AVCodecContext *avctx, void *arg,
716                              int jobnr, int threadnr)
717 {
718     ProresContext *ctx = avctx->priv_data;
719     ProresThreadData *td = ctx->tdata + threadnr;
720     int mbs_per_slice = ctx->mbs_per_slice;
721     int x, y = jobnr, mb, q = 0;
722
723     for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
724         while (ctx->mb_width - x < mbs_per_slice)
725             mbs_per_slice >>= 1;
726         q = find_slice_quant(avctx, avctx->coded_frame,
727                              (mb + 1) * TRELLIS_WIDTH, x, y,
728                              mbs_per_slice, td);
729     }
730
731     for (x = ctx->slices_width - 1; x >= 0; x--) {
732         ctx->slice_q[x + y * ctx->slices_width] = td->nodes[q].quant;
733         q = td->nodes[q].prev_node;
734     }
735
736     return 0;
737 }
738
739 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
740                         const AVFrame *pic, int *got_packet)
741 {
742     ProresContext *ctx = avctx->priv_data;
743     uint8_t *orig_buf, *buf, *slice_hdr, *slice_sizes, *tmp;
744     uint8_t *picture_size_pos;
745     PutBitContext pb;
746     int x, y, i, mb, q = 0;
747     int sizes[4] = { 0 };
748     int slice_hdr_size = 2 + 2 * (ctx->num_planes - 1);
749     int frame_size, picture_size, slice_size;
750     int pkt_size, ret;
751     uint8_t frame_flags;
752
753     *avctx->coded_frame           = *pic;
754     avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
755     avctx->coded_frame->key_frame = 1;
756
757     pkt_size = ctx->frame_size_upper_bound + FF_MIN_BUFFER_SIZE;
758
759     if ((ret = ff_alloc_packet2(avctx, pkt, pkt_size)) < 0)
760         return ret;
761
762     orig_buf = pkt->data;
763
764     // frame atom
765     orig_buf += 4;                              // frame size
766     bytestream_put_be32  (&orig_buf, FRAME_ID); // frame container ID
767     buf = orig_buf;
768
769     // frame header
770     tmp = buf;
771     buf += 2;                                   // frame header size will be stored here
772     bytestream_put_be16  (&buf, 0);             // version 1
773     bytestream_put_buffer(&buf, ctx->vendor, 4);
774     bytestream_put_be16  (&buf, avctx->width);
775     bytestream_put_be16  (&buf, avctx->height);
776
777     frame_flags = ctx->chroma_factor << 6;
778     if (avctx->flags & CODEC_FLAG_INTERLACED_DCT)
779         frame_flags |= pic->top_field_first ? 0x04 : 0x08;
780     bytestream_put_byte  (&buf, frame_flags);
781
782     bytestream_put_byte  (&buf, 0);             // reserved
783     bytestream_put_byte  (&buf, avctx->color_primaries);
784     bytestream_put_byte  (&buf, avctx->color_trc);
785     bytestream_put_byte  (&buf, avctx->colorspace);
786     bytestream_put_byte  (&buf, 0x40);          // source format and alpha information
787     bytestream_put_byte  (&buf, 0);             // reserved
788     if (ctx->quant_sel != QUANT_MAT_DEFAULT) {
789         bytestream_put_byte  (&buf, 0x03);      // matrix flags - both matrices are present
790         // luma quantisation matrix
791         for (i = 0; i < 64; i++)
792             bytestream_put_byte(&buf, ctx->quant_mat[i]);
793         // chroma quantisation matrix
794         for (i = 0; i < 64; i++)
795             bytestream_put_byte(&buf, ctx->quant_mat[i]);
796     } else {
797         bytestream_put_byte  (&buf, 0x00);      // matrix flags - default matrices are used
798     }
799     bytestream_put_be16  (&tmp, buf - orig_buf); // write back frame header size
800
801     for (ctx->cur_picture_idx = 0;
802          ctx->cur_picture_idx < ctx->pictures_per_frame;
803          ctx->cur_picture_idx++) {
804         // picture header
805         picture_size_pos = buf + 1;
806         bytestream_put_byte  (&buf, 0x40);          // picture header size (in bits)
807         buf += 4;                                   // picture data size will be stored here
808         bytestream_put_be16  (&buf, ctx->slices_per_picture);
809         bytestream_put_byte  (&buf, av_log2(ctx->mbs_per_slice) << 4); // slice width and height in MBs
810
811         // seek table - will be filled during slice encoding
812         slice_sizes = buf;
813         buf += ctx->slices_per_picture * 2;
814
815         // slices
816         if (!ctx->force_quant) {
817             ret = avctx->execute2(avctx, find_quant_thread, NULL, NULL,
818                                   ctx->mb_height);
819             if (ret)
820                 return ret;
821         }
822
823         for (y = 0; y < ctx->mb_height; y++) {
824             int mbs_per_slice = ctx->mbs_per_slice;
825             for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
826                 q = ctx->force_quant ? ctx->force_quant
827                                      : ctx->slice_q[mb + y * ctx->slices_width];
828
829                 while (ctx->mb_width - x < mbs_per_slice)
830                     mbs_per_slice >>= 1;
831
832                 bytestream_put_byte(&buf, slice_hdr_size << 3);
833                 slice_hdr = buf;
834                 buf += slice_hdr_size - 1;
835                 init_put_bits(&pb, buf, (pkt_size - (buf - orig_buf)) * 8);
836                 encode_slice(avctx, pic, &pb, sizes, x, y, q, mbs_per_slice);
837
838                 bytestream_put_byte(&slice_hdr, q);
839                 slice_size = slice_hdr_size + sizes[ctx->num_planes - 1];
840                 for (i = 0; i < ctx->num_planes - 1; i++) {
841                     bytestream_put_be16(&slice_hdr, sizes[i]);
842                     slice_size += sizes[i];
843                 }
844                 bytestream_put_be16(&slice_sizes, slice_size);
845                 buf += slice_size - slice_hdr_size;
846             }
847         }
848
849         picture_size = buf - (picture_size_pos - 1);
850         bytestream_put_be32(&picture_size_pos, picture_size);
851     }
852
853     orig_buf -= 8;
854     frame_size = buf - orig_buf;
855     bytestream_put_be32(&orig_buf, frame_size);
856
857     pkt->size   = frame_size;
858     pkt->flags |= AV_PKT_FLAG_KEY;
859     *got_packet = 1;
860
861     return 0;
862 }
863
864 static av_cold int encode_close(AVCodecContext *avctx)
865 {
866     ProresContext *ctx = avctx->priv_data;
867     int i;
868
869     av_freep(&avctx->coded_frame);
870
871     if (ctx->tdata) {
872         for (i = 0; i < avctx->thread_count; i++)
873             av_free(ctx->tdata[i].nodes);
874     }
875     av_freep(&ctx->tdata);
876     av_freep(&ctx->slice_q);
877
878     return 0;
879 }
880
881 static av_cold int encode_init(AVCodecContext *avctx)
882 {
883     ProresContext *ctx = avctx->priv_data;
884     int mps;
885     int i, j;
886     int min_quant, max_quant;
887     int interlaced = !!(avctx->flags & CODEC_FLAG_INTERLACED_DCT);
888
889     avctx->bits_per_raw_sample = 10;
890     avctx->coded_frame = avcodec_alloc_frame();
891     if (!avctx->coded_frame)
892         return AVERROR(ENOMEM);
893
894     ff_proresdsp_init(&ctx->dsp, avctx);
895     ff_init_scantable(ctx->dsp.dct_permutation, &ctx->scantable,
896                       interlaced ? ff_prores_interlaced_scan
897                                  : ff_prores_progressive_scan);
898
899     mps = ctx->mbs_per_slice;
900     if (mps & (mps - 1)) {
901         av_log(avctx, AV_LOG_ERROR,
902                "there should be an integer power of two MBs per slice\n");
903         return AVERROR(EINVAL);
904     }
905
906     ctx->chroma_factor = avctx->pix_fmt == AV_PIX_FMT_YUV422P10
907                          ? CFACTOR_Y422
908                          : CFACTOR_Y444;
909     ctx->profile_info  = prores_profile_info + ctx->profile;
910     ctx->num_planes    = 3;
911
912     ctx->mb_width      = FFALIGN(avctx->width,  16) >> 4;
913
914     if (interlaced)
915         ctx->mb_height = FFALIGN(avctx->height, 32) >> 5;
916     else
917         ctx->mb_height = FFALIGN(avctx->height, 16) >> 4;
918
919     ctx->slices_width  = ctx->mb_width / mps;
920     ctx->slices_width += av_popcount(ctx->mb_width - ctx->slices_width * mps);
921     ctx->slices_per_picture = ctx->mb_height * ctx->slices_width;
922     ctx->pictures_per_frame = 1 + interlaced;
923
924     if (ctx->quant_sel == -1)
925         ctx->quant_mat = prores_quant_matrices[ctx->profile_info->quant];
926     else
927         ctx->quant_mat = prores_quant_matrices[ctx->quant_sel];
928
929     if (strlen(ctx->vendor) != 4) {
930         av_log(avctx, AV_LOG_ERROR, "vendor ID should be 4 bytes\n");
931         return AVERROR_INVALIDDATA;
932     }
933
934     ctx->force_quant = avctx->global_quality / FF_QP2LAMBDA;
935     if (!ctx->force_quant) {
936         if (!ctx->bits_per_mb) {
937             for (i = 0; i < NUM_MB_LIMITS - 1; i++)
938                 if (prores_mb_limits[i] >= ctx->mb_width * ctx->mb_height *
939                                            ctx->pictures_per_frame)
940                     break;
941             ctx->bits_per_mb   = ctx->profile_info->br_tab[i];
942         } else if (ctx->bits_per_mb < 128) {
943             av_log(avctx, AV_LOG_ERROR, "too few bits per MB, please set at least 128\n");
944             return AVERROR_INVALIDDATA;
945         }
946
947         min_quant = ctx->profile_info->min_quant;
948         max_quant = ctx->profile_info->max_quant;
949         for (i = min_quant; i < MAX_STORED_Q; i++) {
950             for (j = 0; j < 64; j++)
951                 ctx->quants[i][j] = ctx->quant_mat[j] * i;
952         }
953
954         ctx->slice_q = av_malloc(ctx->slices_per_picture * sizeof(*ctx->slice_q));
955         if (!ctx->slice_q) {
956             encode_close(avctx);
957             return AVERROR(ENOMEM);
958         }
959
960         ctx->tdata = av_mallocz(avctx->thread_count * sizeof(*ctx->tdata));
961         if (!ctx->tdata) {
962             encode_close(avctx);
963             return AVERROR(ENOMEM);
964         }
965
966         for (j = 0; j < avctx->thread_count; j++) {
967             ctx->tdata[j].nodes = av_malloc((ctx->slices_width + 1)
968                                             * TRELLIS_WIDTH
969                                             * sizeof(*ctx->tdata->nodes));
970             if (!ctx->tdata[j].nodes) {
971                 encode_close(avctx);
972                 return AVERROR(ENOMEM);
973             }
974             for (i = min_quant; i < max_quant + 2; i++) {
975                 ctx->tdata[j].nodes[i].prev_node = -1;
976                 ctx->tdata[j].nodes[i].bits      = 0;
977                 ctx->tdata[j].nodes[i].score     = 0;
978             }
979         }
980     } else {
981         int ls = 0;
982
983         if (ctx->force_quant > 64) {
984             av_log(avctx, AV_LOG_ERROR, "too large quantiser, maximum is 64\n");
985             return AVERROR_INVALIDDATA;
986         }
987
988         for (j = 0; j < 64; j++) {
989             ctx->quants[0][j] = ctx->quant_mat[j] * ctx->force_quant;
990             ls += av_log2((1 << 11)  / ctx->quants[0][j]) * 2 + 1;
991         }
992
993         ctx->bits_per_mb = ls * 8;
994         if (ctx->chroma_factor == CFACTOR_Y444)
995             ctx->bits_per_mb += ls * 4;
996         if (ctx->num_planes == 4)
997             ctx->bits_per_mb += ls * 4;
998     }
999
1000     ctx->frame_size_upper_bound = ctx->pictures_per_frame *
1001                                   ctx->slices_per_picture *
1002                                   (2 + 2 * ctx->num_planes +
1003                                    (mps * ctx->bits_per_mb) / 8)
1004                                   + 200;
1005
1006     avctx->codec_tag   = ctx->profile_info->tag;
1007
1008     av_log(avctx, AV_LOG_DEBUG,
1009            "profile %d, %d slices, interlacing: %s, %d bits per MB\n",
1010            ctx->profile, ctx->slices_per_picture * ctx->pictures_per_frame,
1011            interlaced ? "yes" : "no", ctx->bits_per_mb);
1012     av_log(avctx, AV_LOG_DEBUG, "frame size upper bound: %d\n",
1013            ctx->frame_size_upper_bound);
1014
1015     return 0;
1016 }
1017
1018 #define OFFSET(x) offsetof(ProresContext, x)
1019 #define VE     AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
1020
1021 static const AVOption options[] = {
1022     { "mbs_per_slice", "macroblocks per slice", OFFSET(mbs_per_slice),
1023         AV_OPT_TYPE_INT, { .i64 = 8 }, 1, MAX_MBS_PER_SLICE, VE },
1024     { "profile",       NULL, OFFSET(profile), AV_OPT_TYPE_INT,
1025         { .i64 = PRORES_PROFILE_STANDARD },
1026         PRORES_PROFILE_PROXY, PRORES_PROFILE_HQ, VE, "profile" },
1027     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_PROXY },
1028         0, 0, VE, "profile" },
1029     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_LT },
1030         0, 0, VE, "profile" },
1031     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_STANDARD },
1032         0, 0, VE, "profile" },
1033     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PRORES_PROFILE_HQ },
1034         0, 0, VE, "profile" },
1035     { "vendor", "vendor ID", OFFSET(vendor),
1036         AV_OPT_TYPE_STRING, { .str = "Lavc" }, CHAR_MIN, CHAR_MAX, VE },
1037     { "bits_per_mb", "desired bits per macroblock", OFFSET(bits_per_mb),
1038         AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 8192, VE },
1039     { "quant_mat", "quantiser matrix", OFFSET(quant_sel), AV_OPT_TYPE_INT,
1040         { .i64 = -1 }, -1, QUANT_MAT_DEFAULT, VE, "quant_mat" },
1041     { "auto",          NULL, 0, AV_OPT_TYPE_CONST, { .i64 = -1 },
1042         0, 0, VE, "quant_mat" },
1043     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_PROXY },
1044         0, 0, VE, "quant_mat" },
1045     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_LT },
1046         0, 0, VE, "quant_mat" },
1047     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_STANDARD },
1048         0, 0, VE, "quant_mat" },
1049     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_HQ },
1050         0, 0, VE, "quant_mat" },
1051     { "default",       NULL, 0, AV_OPT_TYPE_CONST, { .i64 = QUANT_MAT_DEFAULT },
1052         0, 0, VE, "quant_mat" },
1053     { NULL }
1054 };
1055
1056 static const AVClass proresenc_class = {
1057     .class_name = "ProRes encoder",
1058     .item_name  = av_default_item_name,
1059     .option     = options,
1060     .version    = LIBAVUTIL_VERSION_INT,
1061 };
1062
1063 AVCodec ff_prores_ks_encoder = {
1064     .name           = "prores_ks",
1065     .type           = AVMEDIA_TYPE_VIDEO,
1066     .id             = AV_CODEC_ID_PRORES,
1067     .priv_data_size = sizeof(ProresContext),
1068     .init           = encode_init,
1069     .close          = encode_close,
1070     .encode2        = encode_frame,
1071     .capabilities   = CODEC_CAP_SLICE_THREADS,
1072     .long_name      = NULL_IF_CONFIG_SMALL("Apple ProRes (iCodec Pro)"),
1073     .pix_fmts       = (const enum AVPixelFormat[]) {
1074                           AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10, AV_PIX_FMT_NONE
1075                       },
1076     .priv_class     = &proresenc_class,
1077 };