proresenc: initialise 'sign' variable
[ffmpeg.git] / libavcodec / proresenc.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 #define NUM_MB_LIMITS 4
46 static const int prores_mb_limits[NUM_MB_LIMITS] = {
47     1620, // up to 720x576
48     2700, // up to 960x720
49     6075, // up to 1440x1080
50     9216, // up to 2048x1152
51 };
52
53 static const struct prores_profile {
54     const char *full_name;
55     uint32_t    tag;
56     int         min_quant;
57     int         max_quant;
58     int         br_tab[NUM_MB_LIMITS];
59     uint8_t     quant[64];
60 } prores_profile_info[4] = {
61     {
62         .full_name = "proxy",
63         .tag       = MKTAG('a', 'p', 'c', 'o'),
64         .min_quant = 4,
65         .max_quant = 8,
66         .br_tab    = { 300, 242, 220, 194 },
67         .quant     = {
68              4,  7,  9, 11, 13, 14, 15, 63,
69              7,  7, 11, 12, 14, 15, 63, 63,
70              9, 11, 13, 14, 15, 63, 63, 63,
71             11, 11, 13, 14, 63, 63, 63, 63,
72             11, 13, 14, 63, 63, 63, 63, 63,
73             13, 14, 63, 63, 63, 63, 63, 63,
74             13, 63, 63, 63, 63, 63, 63, 63,
75             63, 63, 63, 63, 63, 63, 63, 63,
76         },
77     },
78     {
79         .full_name = "LT",
80         .tag       = MKTAG('a', 'p', 'c', 's'),
81         .min_quant = 1,
82         .max_quant = 9,
83         .br_tab    = { 720, 560, 490, 440 },
84         .quant     = {
85              4,  5,  6,  7,  9, 11, 13, 15,
86              5,  5,  7,  8, 11, 13, 15, 17,
87              6,  7,  9, 11, 13, 15, 15, 17,
88              7,  7,  9, 11, 13, 15, 17, 19,
89              7,  9, 11, 13, 14, 16, 19, 23,
90              9, 11, 13, 14, 16, 19, 23, 29,
91              9, 11, 13, 15, 17, 21, 28, 35,
92             11, 13, 16, 17, 21, 28, 35, 41,
93         },
94     },
95     {
96         .full_name = "standard",
97         .tag       = MKTAG('a', 'p', 'c', 'n'),
98         .min_quant = 1,
99         .max_quant = 6,
100         .br_tab    = { 1050, 808, 710, 632 },
101         .quant     = {
102              4,  4,  5,  5,  6,  7,  7,  9,
103              4,  4,  5,  6,  7,  7,  9,  9,
104              5,  5,  6,  7,  7,  9,  9, 10,
105              5,  5,  6,  7,  7,  9,  9, 10,
106              5,  6,  7,  7,  8,  9, 10, 12,
107              6,  7,  7,  8,  9, 10, 12, 15,
108              6,  7,  7,  9, 10, 11, 14, 17,
109              7,  7,  9, 10, 11, 14, 17, 21,
110         },
111     },
112     {
113         .full_name = "high quality",
114         .tag       = MKTAG('a', 'p', 'c', 'h'),
115         .min_quant = 1,
116         .max_quant = 6,
117         .br_tab    = { 1566, 1216, 1070, 950 },
118         .quant     = {
119              4,  4,  4,  4,  4,  4,  4,  4,
120              4,  4,  4,  4,  4,  4,  4,  4,
121              4,  4,  4,  4,  4,  4,  4,  4,
122              4,  4,  4,  4,  4,  4,  4,  5,
123              4,  4,  4,  4,  4,  4,  5,  5,
124              4,  4,  4,  4,  4,  5,  5,  6,
125              4,  4,  4,  4,  5,  5,  6,  7,
126              4,  4,  4,  4,  5,  6,  7,  7,
127         },
128     }
129 // for 4444 profile bitrate numbers are { 2350, 1828, 1600, 1425 }
130 };
131
132 #define TRELLIS_WIDTH 16
133 #define SCORE_LIMIT   INT_MAX / 2
134
135 struct TrellisNode {
136     int prev_node;
137     int quant;
138     int bits;
139     int score;
140 };
141
142 typedef struct ProresContext {
143     AVClass *class;
144     DECLARE_ALIGNED(16, DCTELEM, blocks)[MAX_PLANES][64 * 4 * MAX_MBS_PER_SLICE];
145     DECLARE_ALIGNED(16, uint16_t, emu_buf)[16*16];
146     int16_t quants[16][64];
147
148     ProresDSPContext dsp;
149     ScanTable  scantable;
150
151     int mb_width, mb_height;
152     int mbs_per_slice;
153     int num_chroma_blocks, chroma_factor;
154     int slices_width;
155     int num_slices;
156     int num_planes;
157     int bits_per_mb;
158
159     int profile;
160     const struct prores_profile *profile_info;
161
162     struct TrellisNode *nodes;
163     int *slice_q;
164 } ProresContext;
165
166 static void get_slice_data(ProresContext *ctx, const uint16_t *src,
167                            int linesize, int x, int y, int w, int h,
168                            DCTELEM *blocks,
169                            int mbs_per_slice, int blocks_per_mb)
170 {
171     const uint16_t *esrc;
172     const int mb_width = 4 * blocks_per_mb;
173     int elinesize;
174     int i, j, k;
175
176     for (i = 0; i < mbs_per_slice; i++, src += mb_width) {
177         if (x >= w) {
178             memset(blocks, 0, 64 * (mbs_per_slice - i) * blocks_per_mb
179                               * sizeof(*blocks));
180             return;
181         }
182         if (x + mb_width <= w && y + 16 <= h) {
183             esrc      = src;
184             elinesize = linesize;
185         } else {
186             int bw, bh, pix;
187             const int estride = 16 / sizeof(*ctx->emu_buf);
188
189             esrc      = ctx->emu_buf;
190             elinesize = 16;
191
192             bw = FFMIN(w - x, mb_width);
193             bh = FFMIN(h - y, 16);
194
195             for (j = 0; j < bh; j++) {
196                 memcpy(ctx->emu_buf + j * estride, src + j * linesize,
197                        bw * sizeof(*src));
198                 pix = ctx->emu_buf[j * estride + bw - 1];
199                 for (k = bw; k < mb_width; k++)
200                     ctx->emu_buf[j * estride + k] = pix;
201             }
202             for (; j < 16; j++)
203                 memcpy(ctx->emu_buf + j * estride,
204                        ctx->emu_buf + (bh - 1) * estride,
205                        mb_width * sizeof(*ctx->emu_buf));
206         }
207         ctx->dsp.fdct(esrc, elinesize, blocks);
208         blocks += 64;
209         if (blocks_per_mb > 2) {
210             ctx->dsp.fdct(src + 8, linesize, blocks);
211             blocks += 64;
212         }
213         ctx->dsp.fdct(src + linesize * 4, linesize, blocks);
214         blocks += 64;
215         if (blocks_per_mb > 2) {
216             ctx->dsp.fdct(src + linesize * 4 + 8, linesize, blocks);
217             blocks += 64;
218         }
219
220         x += mb_width;
221     }
222 }
223
224 /**
225  * Write an unsigned rice/exp golomb codeword.
226  */
227 static inline void encode_vlc_codeword(PutBitContext *pb, uint8_t codebook, int val)
228 {
229     unsigned int rice_order, exp_order, switch_bits, switch_val;
230     int exponent;
231
232     /* number of prefix bits to switch between Rice and expGolomb */
233     switch_bits = (codebook & 3) + 1;
234     rice_order  =  codebook >> 5;       /* rice code order */
235     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
236
237     switch_val  = switch_bits << rice_order;
238
239     if (val >= switch_val) {
240         val -= switch_val - (1 << exp_order);
241         exponent = av_log2(val);
242
243         put_bits(pb, exponent - exp_order + switch_bits, 0);
244         put_bits(pb, 1, 1);
245         put_bits(pb, exponent, val);
246     } else {
247         exponent = val >> rice_order;
248
249         if (exponent)
250             put_bits(pb, exponent, 0);
251         put_bits(pb, 1, 1);
252         if (rice_order)
253             put_sbits(pb, rice_order, val);
254     }
255 }
256
257 #define GET_SIGN(x)  ((x) >> 31)
258 #define MAKE_CODE(x) (((x) << 1) ^ GET_SIGN(x))
259
260 static void encode_dcs(PutBitContext *pb, DCTELEM *blocks,
261                        int blocks_per_slice, int scale)
262 {
263     int i;
264     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
265
266     prev_dc = (blocks[0] - 0x4000) / scale;
267     encode_vlc_codeword(pb, FIRST_DC_CB, MAKE_CODE(prev_dc));
268     sign     = 0;
269     codebook = 3;
270     blocks  += 64;
271
272     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
273         dc       = (blocks[0] - 0x4000) / scale;
274         delta    = dc - prev_dc;
275         new_sign = GET_SIGN(delta);
276         delta    = (delta ^ sign) - sign;
277         code     = MAKE_CODE(delta);
278         encode_vlc_codeword(pb, ff_prores_dc_codebook[codebook], code);
279         codebook = (code + (code & 1)) >> 1;
280         codebook = FFMIN(codebook, 3);
281         sign     = new_sign;
282         prev_dc  = dc;
283     }
284 }
285
286 static void encode_acs(PutBitContext *pb, DCTELEM *blocks,
287                        int blocks_per_slice,
288                        int plane_size_factor,
289                        const uint8_t *scan, const int16_t *qmat)
290 {
291     int idx, i;
292     int run, level, run_cb, lev_cb;
293     int max_coeffs, abs_level;
294
295     max_coeffs = blocks_per_slice << 6;
296     run_cb     = ff_prores_run_to_cb_index[4];
297     lev_cb     = ff_prores_lev_to_cb_index[2];
298     run        = 0;
299
300     for (i = 1; i < 64; i++) {
301         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
302             level = blocks[idx] / qmat[scan[i]];
303             if (level) {
304                 abs_level = FFABS(level);
305                 encode_vlc_codeword(pb, ff_prores_ac_codebook[run_cb], run);
306                 encode_vlc_codeword(pb, ff_prores_ac_codebook[lev_cb],
307                                     abs_level - 1);
308                 put_sbits(pb, 1, GET_SIGN(level));
309
310                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
311                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
312                 run    = 0;
313             } else {
314                 run++;
315             }
316         }
317     }
318 }
319
320 static int encode_slice_plane(ProresContext *ctx, PutBitContext *pb,
321                               const uint16_t *src, int linesize,
322                               int mbs_per_slice, DCTELEM *blocks,
323                               int blocks_per_mb, int plane_size_factor,
324                               const int16_t *qmat)
325 {
326     int blocks_per_slice, saved_pos;
327
328     saved_pos = put_bits_count(pb);
329     blocks_per_slice = mbs_per_slice * blocks_per_mb;
330
331     encode_dcs(pb, blocks, blocks_per_slice, qmat[0]);
332     encode_acs(pb, blocks, blocks_per_slice, plane_size_factor,
333                ctx->scantable.permutated, qmat);
334     flush_put_bits(pb);
335
336     return (put_bits_count(pb) - saved_pos) >> 3;
337 }
338
339 static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
340                         PutBitContext *pb,
341                         int sizes[4], int x, int y, int quant,
342                         int mbs_per_slice)
343 {
344     ProresContext *ctx = avctx->priv_data;
345     int i, xp, yp;
346     int total_size = 0;
347     const uint16_t *src;
348     int slice_width_factor = av_log2(mbs_per_slice);
349     int num_cblocks, pwidth;
350     int plane_factor, is_chroma;
351
352     for (i = 0; i < ctx->num_planes; i++) {
353         is_chroma    = (i == 1 || i == 2);
354         plane_factor = slice_width_factor + 2;
355         if (is_chroma)
356             plane_factor += ctx->chroma_factor - 3;
357         if (!is_chroma || ctx->chroma_factor == CFACTOR_Y444) {
358             xp          = x << 4;
359             yp          = y << 4;
360             num_cblocks = 4;
361             pwidth      = avctx->width;
362         } else {
363             xp          = x << 3;
364             yp          = y << 4;
365             num_cblocks = 2;
366             pwidth      = avctx->width >> 1;
367         }
368         src = (const uint16_t*)(pic->data[i] + yp * pic->linesize[i]) + xp;
369
370         get_slice_data(ctx, src, pic->linesize[i], xp, yp,
371                        pwidth, avctx->height, ctx->blocks[0],
372                        mbs_per_slice, num_cblocks);
373         sizes[i] = encode_slice_plane(ctx, pb, src, pic->linesize[i],
374                                       mbs_per_slice, ctx->blocks[0],
375                                       num_cblocks, plane_factor,
376                                       ctx->quants[quant]);
377         total_size += sizes[i];
378     }
379     return total_size;
380 }
381
382 static inline int estimate_vlc(uint8_t codebook, int val)
383 {
384     unsigned int rice_order, exp_order, switch_bits, switch_val;
385     int exponent;
386
387     /* number of prefix bits to switch between Rice and expGolomb */
388     switch_bits = (codebook & 3) + 1;
389     rice_order  =  codebook >> 5;       /* rice code order */
390     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
391
392     switch_val  = switch_bits << rice_order;
393
394     if (val >= switch_val) {
395         val -= switch_val - (1 << exp_order);
396         exponent = av_log2(val);
397
398         return exponent * 2 - exp_order + switch_bits + 1;
399     } else {
400         return (val >> rice_order) + rice_order + 1;
401     }
402 }
403
404 static int estimate_dcs(int *error, DCTELEM *blocks, int blocks_per_slice,
405                         int scale)
406 {
407     int i;
408     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
409     int bits;
410
411     prev_dc  = (blocks[0] - 0x4000) / scale;
412     bits     = estimate_vlc(FIRST_DC_CB, MAKE_CODE(prev_dc));
413     sign     = 0;
414     codebook = 3;
415     blocks  += 64;
416     *error  += FFABS(blocks[0] - 0x4000) % scale;
417
418     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
419         dc       = (blocks[0] - 0x4000) / scale;
420         *error  += FFABS(blocks[0] - 0x4000) % scale;
421         delta    = dc - prev_dc;
422         new_sign = GET_SIGN(delta);
423         delta    = (delta ^ sign) - sign;
424         code     = MAKE_CODE(delta);
425         bits    += estimate_vlc(ff_prores_dc_codebook[codebook], code);
426         codebook = (code + (code & 1)) >> 1;
427         codebook = FFMIN(codebook, 3);
428         sign     = new_sign;
429         prev_dc  = dc;
430     }
431
432     return bits;
433 }
434
435 static int estimate_acs(int *error, DCTELEM *blocks, int blocks_per_slice,
436                         int plane_size_factor,
437                         const uint8_t *scan, const int16_t *qmat)
438 {
439     int idx, i;
440     int run, level, run_cb, lev_cb;
441     int max_coeffs, abs_level;
442     int bits = 0;
443
444     max_coeffs = blocks_per_slice << 6;
445     run_cb     = ff_prores_run_to_cb_index[4];
446     lev_cb     = ff_prores_lev_to_cb_index[2];
447     run        = 0;
448
449     for (i = 1; i < 64; i++) {
450         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
451             level   = blocks[idx] / qmat[scan[i]];
452             *error += FFABS(blocks[idx]) % qmat[scan[i]];
453             if (level) {
454                 abs_level = FFABS(level);
455                 bits += estimate_vlc(ff_prores_ac_codebook[run_cb], run);
456                 bits += estimate_vlc(ff_prores_ac_codebook[lev_cb],
457                                      abs_level - 1) + 1;
458
459                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
460                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
461                 run    = 0;
462             } else {
463                 run++;
464             }
465         }
466     }
467
468     return bits;
469 }
470
471 static int estimate_slice_plane(ProresContext *ctx, int *error, int plane,
472                                 const uint16_t *src, int linesize,
473                                 int mbs_per_slice,
474                                 int blocks_per_mb, int plane_size_factor,
475                                 const int16_t *qmat)
476 {
477     int blocks_per_slice;
478     int bits;
479
480     blocks_per_slice = mbs_per_slice * blocks_per_mb;
481
482     bits  = estimate_dcs(error, ctx->blocks[plane], blocks_per_slice, qmat[0]);
483     bits += estimate_acs(error, ctx->blocks[plane], blocks_per_slice,
484                          plane_size_factor, ctx->scantable.permutated, qmat);
485
486     return FFALIGN(bits, 8);
487 }
488
489 static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
490                             int trellis_node, int x, int y, int mbs_per_slice)
491 {
492     ProresContext *ctx = avctx->priv_data;
493     int i, q, pq, xp, yp;
494     const uint16_t *src;
495     int slice_width_factor = av_log2(mbs_per_slice);
496     int num_cblocks[MAX_PLANES], pwidth;
497     int plane_factor[MAX_PLANES], is_chroma[MAX_PLANES];
498     const int min_quant = ctx->profile_info->min_quant;
499     const int max_quant = ctx->profile_info->max_quant;
500     int error, bits, bits_limit;
501     int mbs, prev, cur, new_score;
502     int slice_bits[TRELLIS_WIDTH], slice_score[TRELLIS_WIDTH];
503
504     mbs = x + mbs_per_slice;
505
506     for (i = 0; i < ctx->num_planes; i++) {
507         is_chroma[i]    = (i == 1 || i == 2);
508         plane_factor[i] = slice_width_factor + 2;
509         if (is_chroma[i])
510             plane_factor[i] += ctx->chroma_factor - 3;
511         if (!is_chroma[i] || ctx->chroma_factor == CFACTOR_Y444) {
512             xp             = x << 4;
513             yp             = y << 4;
514             num_cblocks[i] = 4;
515             pwidth         = avctx->width;
516         } else {
517             xp             = x << 3;
518             yp             = y << 4;
519             num_cblocks[i] = 2;
520             pwidth         = avctx->width >> 1;
521         }
522         src = (const uint16_t*)(pic->data[i] + yp * pic->linesize[i]) + xp;
523
524         get_slice_data(ctx, src, pic->linesize[i], xp, yp,
525                        pwidth, avctx->height, ctx->blocks[i],
526                        mbs_per_slice, num_cblocks[i]);
527     }
528
529     for (q = min_quant; q <= max_quant; q++) {
530         ctx->nodes[trellis_node + q].prev_node = -1;
531         ctx->nodes[trellis_node + q].quant     = q;
532     }
533
534     // todo: maybe perform coarser quantising to fit into frame size when needed
535     for (q = min_quant; q <= max_quant; q++) {
536         bits  = 0;
537         error = 0;
538         for (i = 0; i < ctx->num_planes; i++) {
539             bits += estimate_slice_plane(ctx, &error, i,
540                                          src, pic->linesize[i],
541                                          mbs_per_slice,
542                                          num_cblocks[i], plane_factor[i],
543                                          ctx->quants[q]);
544         }
545         if (bits > 65000 * 8) {
546             error = SCORE_LIMIT;
547             break;
548         }
549         slice_bits[q]  = bits;
550         slice_score[q] = error;
551     }
552
553     bits_limit = mbs * ctx->bits_per_mb;
554     for (pq = min_quant; pq <= max_quant; pq++) {
555         prev = trellis_node - TRELLIS_WIDTH + pq;
556
557         for (q = min_quant; q <= max_quant; q++) {
558             cur = trellis_node + q;
559
560             bits  = ctx->nodes[prev].bits + slice_bits[q];
561             error = slice_score[q];
562             if (bits > bits_limit)
563                 error = SCORE_LIMIT;
564
565             if (ctx->nodes[prev].score < SCORE_LIMIT && error < SCORE_LIMIT)
566                 new_score = ctx->nodes[prev].score + error;
567             else
568                 new_score = SCORE_LIMIT;
569             if (ctx->nodes[cur].prev_node == -1 ||
570                 ctx->nodes[cur].score >= new_score) {
571
572                 ctx->nodes[cur].bits      = bits;
573                 ctx->nodes[cur].score     = new_score;
574                 ctx->nodes[cur].prev_node = prev;
575             }
576         }
577     }
578
579     error = ctx->nodes[trellis_node + min_quant].score;
580     pq    = trellis_node + min_quant;
581     for (q = min_quant + 1; q <= max_quant; q++) {
582         if (ctx->nodes[trellis_node + q].score <= error) {
583             error = ctx->nodes[trellis_node + q].score;
584             pq    = trellis_node + q;
585         }
586     }
587
588     return pq;
589 }
590
591 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
592                         const AVFrame *pic, int *got_packet)
593 {
594     ProresContext *ctx = avctx->priv_data;
595     uint8_t *orig_buf, *buf, *slice_hdr, *slice_sizes, *tmp;
596     uint8_t *picture_size_pos;
597     PutBitContext pb;
598     int x, y, i, mb, q = 0;
599     int sizes[4] = { 0 };
600     int slice_hdr_size = 2 + 2 * (ctx->num_planes - 1);
601     int frame_size, picture_size, slice_size;
602     int mbs_per_slice = ctx->mbs_per_slice;
603     int pkt_size, ret;
604
605     *avctx->coded_frame           = *pic;
606     avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
607     avctx->coded_frame->key_frame = 1;
608
609     pkt_size = ctx->mb_width * ctx->mb_height * 64 * 3 * 12
610                + ctx->num_slices * 2 + 200 + FF_MIN_BUFFER_SIZE;
611
612     if ((ret = ff_alloc_packet(pkt, pkt_size)) < 0) {
613         av_log(avctx, AV_LOG_ERROR, "Error getting output packet.\n");
614         return ret;
615     }
616
617     orig_buf = pkt->data;
618
619     // frame atom
620     orig_buf += 4;                              // frame size
621     bytestream_put_be32  (&orig_buf, FRAME_ID); // frame container ID
622     buf = orig_buf;
623
624     // frame header
625     tmp = buf;
626     buf += 2;                                   // frame header size will be stored here
627     bytestream_put_be16  (&buf, 0);             // version 1
628     bytestream_put_buffer(&buf, "Lavc", 4);     // creator
629     bytestream_put_be16  (&buf, avctx->width);
630     bytestream_put_be16  (&buf, avctx->height);
631     bytestream_put_byte  (&buf, ctx->chroma_factor << 6); // frame flags
632     bytestream_put_byte  (&buf, 0);             // reserved
633     bytestream_put_byte  (&buf, 0);             // primaries
634     bytestream_put_byte  (&buf, 0);             // transfer function
635     bytestream_put_byte  (&buf, 6);             // colour matrix - ITU-R BT.601-4
636     bytestream_put_byte  (&buf, 0x40);          // source format and alpha information
637     bytestream_put_byte  (&buf, 0);             // reserved
638     bytestream_put_byte  (&buf, 0x03);          // matrix flags - both matrices are present
639     // luma quantisation matrix
640     for (i = 0; i < 64; i++)
641         bytestream_put_byte(&buf, ctx->profile_info->quant[i]);
642     // chroma quantisation matrix
643     for (i = 0; i < 64; i++)
644         bytestream_put_byte(&buf, ctx->profile_info->quant[i]);
645     bytestream_put_be16  (&tmp, buf - orig_buf); // write back frame header size
646
647     // picture header
648     picture_size_pos = buf + 1;
649     bytestream_put_byte  (&buf, 0x40);          // picture header size (in bits)
650     buf += 4;                                   // picture data size will be stored here
651     bytestream_put_be16  (&buf, ctx->num_slices); // total number of slices
652     bytestream_put_byte  (&buf, av_log2(ctx->mbs_per_slice) << 4); // slice width and height in MBs
653
654     // seek table - will be filled during slice encoding
655     slice_sizes = buf;
656     buf += ctx->num_slices * 2;
657
658     // slices
659     for (y = 0; y < ctx->mb_height; y++) {
660         mbs_per_slice = ctx->mbs_per_slice;
661         for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
662             while (ctx->mb_width - x < mbs_per_slice)
663                 mbs_per_slice >>= 1;
664             q = find_slice_quant(avctx, pic, (mb + 1) * TRELLIS_WIDTH, x, y,
665                                  mbs_per_slice);
666         }
667
668         for (x = ctx->slices_width - 1; x >= 0; x--) {
669             ctx->slice_q[x] = ctx->nodes[q].quant;
670             q = ctx->nodes[q].prev_node;
671         }
672
673         mbs_per_slice = ctx->mbs_per_slice;
674         for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
675             q = ctx->slice_q[mb];
676
677             while (ctx->mb_width - x < mbs_per_slice)
678                 mbs_per_slice >>= 1;
679
680             bytestream_put_byte(&buf, slice_hdr_size << 3);
681             slice_hdr = buf;
682             buf += slice_hdr_size - 1;
683             init_put_bits(&pb, buf, (pkt_size - (buf - orig_buf)) * 8);
684             encode_slice(avctx, pic, &pb, sizes, x, y, q, mbs_per_slice);
685
686             bytestream_put_byte(&slice_hdr, q);
687             slice_size = slice_hdr_size + sizes[ctx->num_planes - 1];
688             for (i = 0; i < ctx->num_planes - 1; i++) {
689                 bytestream_put_be16(&slice_hdr, sizes[i]);
690                 slice_size += sizes[i];
691             }
692             bytestream_put_be16(&slice_sizes, slice_size);
693             buf += slice_size - slice_hdr_size;
694         }
695     }
696
697     orig_buf -= 8;
698     frame_size = buf - orig_buf;
699     picture_size = buf - picture_size_pos - 6;
700     bytestream_put_be32(&orig_buf, frame_size);
701     bytestream_put_be32(&picture_size_pos, picture_size);
702
703     pkt->size   = frame_size;
704     pkt->flags |= AV_PKT_FLAG_KEY;
705     *got_packet = 1;
706
707     return 0;
708 }
709
710 static av_cold int encode_close(AVCodecContext *avctx)
711 {
712     ProresContext *ctx = avctx->priv_data;
713
714     if (avctx->coded_frame->data[0])
715         avctx->release_buffer(avctx, avctx->coded_frame);
716
717     av_freep(&avctx->coded_frame);
718
719     av_freep(&ctx->nodes);
720     av_freep(&ctx->slice_q);
721
722     return 0;
723 }
724
725 static av_cold int encode_init(AVCodecContext *avctx)
726 {
727     ProresContext *ctx = avctx->priv_data;
728     int mps;
729     int i, j;
730     int min_quant, max_quant;
731
732     avctx->bits_per_raw_sample = 10;
733     avctx->coded_frame = avcodec_alloc_frame();
734     if (!avctx->coded_frame)
735         return AVERROR(ENOMEM);
736
737     ff_proresdsp_init(&ctx->dsp);
738     ff_init_scantable(ctx->dsp.dct_permutation, &ctx->scantable,
739                       ff_prores_progressive_scan);
740
741     mps = ctx->mbs_per_slice;
742     if (mps & (mps - 1)) {
743         av_log(avctx, AV_LOG_ERROR,
744                "there should be an integer power of two MBs per slice\n");
745         return AVERROR(EINVAL);
746     }
747
748     ctx->chroma_factor = avctx->pix_fmt == PIX_FMT_YUV422P10
749                          ? CFACTOR_Y422
750                          : CFACTOR_Y444;
751     ctx->profile_info  = prores_profile_info + ctx->profile;
752     ctx->num_planes    = 3;
753
754     ctx->mb_width      = FFALIGN(avctx->width,  16) >> 4;
755     ctx->mb_height     = FFALIGN(avctx->height, 16) >> 4;
756     ctx->slices_width  = ctx->mb_width / mps;
757     ctx->slices_width += av_popcount(ctx->mb_width - ctx->slices_width * mps);
758     ctx->num_slices    = ctx->mb_height * ctx->slices_width;
759
760     for (i = 0; i < NUM_MB_LIMITS - 1; i++)
761         if (prores_mb_limits[i] >= ctx->mb_width * ctx->mb_height)
762             break;
763     ctx->bits_per_mb   = ctx->profile_info->br_tab[i];
764
765     min_quant = ctx->profile_info->min_quant;
766     max_quant = ctx->profile_info->max_quant;
767     for (i = min_quant; i <= max_quant; i++) {
768         for (j = 0; j < 64; j++)
769             ctx->quants[i][j] = ctx->profile_info->quant[j] * i;
770     }
771
772     avctx->codec_tag   = ctx->profile_info->tag;
773
774     av_log(avctx, AV_LOG_DEBUG, "profile %d, %d slices, %d bits per MB\n",
775            ctx->profile, ctx->num_slices, ctx->bits_per_mb);
776
777     ctx->nodes = av_malloc((ctx->slices_width + 1) * TRELLIS_WIDTH
778                            * sizeof(*ctx->nodes));
779     if (!ctx->nodes) {
780         encode_close(avctx);
781         return AVERROR(ENOMEM);
782     }
783     for (i = min_quant; i <= max_quant; i++) {
784         ctx->nodes[i].prev_node = -1;
785         ctx->nodes[i].bits      = 0;
786         ctx->nodes[i].score     = 0;
787     }
788
789     ctx->slice_q = av_malloc(ctx->slices_width * sizeof(*ctx->slice_q));
790     if (!ctx->slice_q) {
791         encode_close(avctx);
792         return AVERROR(ENOMEM);
793     }
794
795     return 0;
796 }
797
798 #define OFFSET(x) offsetof(ProresContext, x)
799 #define VE     AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
800
801 static const AVOption options[] = {
802     { "mbs_per_slice", "macroblocks per slice", OFFSET(mbs_per_slice),
803         AV_OPT_TYPE_INT, { 8 }, 1, MAX_MBS_PER_SLICE, VE },
804     { "profile",       NULL, OFFSET(profile), AV_OPT_TYPE_INT,
805         { PRORES_PROFILE_STANDARD },
806         PRORES_PROFILE_PROXY, PRORES_PROFILE_HQ, VE, "profile" },
807     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_PROXY },
808         0, 0, VE, "profile" },
809     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_LT },
810         0, 0, VE, "profile" },
811     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_STANDARD },
812         0, 0, VE, "profile" },
813     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_HQ },
814         0, 0, VE, "profile" },
815     { NULL }
816 };
817
818 static const AVClass proresenc_class = {
819     .class_name = "ProRes encoder",
820     .item_name  = av_default_item_name,
821     .option     = options,
822     .version    = LIBAVUTIL_VERSION_INT,
823 };
824
825 AVCodec ff_prores_encoder = {
826     .name           = "prores",
827     .type           = AVMEDIA_TYPE_VIDEO,
828     .id             = CODEC_ID_PRORES,
829     .priv_data_size = sizeof(ProresContext),
830     .init           = encode_init,
831     .close          = encode_close,
832     .encode2        = encode_frame,
833     .long_name      = NULL_IF_CONFIG_SMALL("Apple ProRes (iCodec Pro)"),
834     .pix_fmts       = (const enum PixelFormat[]) {
835                           PIX_FMT_YUV422P10, PIX_FMT_YUV444P10, PIX_FMT_NONE
836                       },
837     .priv_class     = &proresenc_class,
838 };