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