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