git.sesse.net Git - ffmpeg/blob - libavcodec/proresenc_kostya.c

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