prores: handle 444 chroma in right order

ProRes codes chroma blocks in 444 mode in different order than luma blocks,
so make both decoder and encoder read/write chroma blocks in right order.

Reported by Phil Barrett

3 files changed, 55 insertions, 27 deletions

diff --git a/libavcodec/proresdec.c b/libavcodec/proresdec.c
index 8ca5f3d632..da90912031 100644
--- a/libavcodec/proresdec.c
+++ b/libavcodec/proresdec.c
@@ -411,7 +411,7 @@ static void decode_slice_plane(ProresContext *ctx, ProresThreadData *td,
                                int data_size, uint16_t *out_ptr,
                                int linesize, int mbs_per_slice,
                                int blocks_per_mb, int plane_size_factor,
-                               const int16_t *qmat)
+                               const int16_t *qmat, int is_chroma)
 {
     GetBitContext gb;
     DCTELEM *block_ptr;
@@ -431,18 +431,33 @@ static void decode_slice_plane(ProresContext *ctx, ProresThreadData *td,
     /* inverse quantization, inverse transform and output */
     block_ptr = td->blocks;
 
-    for (mb_num = 0; mb_num < mbs_per_slice; mb_num++, out_ptr += blocks_per_mb * 4) {
-        ctx->dsp.idct_put(out_ptr,                    linesize, block_ptr, qmat);
-        block_ptr += 64;
-        if (blocks_per_mb > 2) {
-            ctx->dsp.idct_put(out_ptr + 8,            linesize, block_ptr, qmat);
+    if (!is_chroma) {
+        for (mb_num = 0; mb_num < mbs_per_slice; mb_num++, out_ptr += blocks_per_mb * 4) {
+            ctx->dsp.idct_put(out_ptr,                    linesize, block_ptr, qmat);
             block_ptr += 64;
+            if (blocks_per_mb > 2) {
+                ctx->dsp.idct_put(out_ptr + 8,            linesize, block_ptr, qmat);
+                block_ptr += 64;
+            }
+            ctx->dsp.idct_put(out_ptr + linesize * 4,     linesize, block_ptr, qmat);
+            block_ptr += 64;
+            if (blocks_per_mb > 2) {
+                ctx->dsp.idct_put(out_ptr + linesize * 4 + 8, linesize, block_ptr, qmat);
+                block_ptr += 64;
+            }
         }
-        ctx->dsp.idct_put(out_ptr + linesize * 4,     linesize, block_ptr, qmat);
-        block_ptr += 64;
-        if (blocks_per_mb > 2) {
-            ctx->dsp.idct_put(out_ptr + linesize * 4 + 8, linesize, block_ptr, qmat);
+    } else {
+        for (mb_num = 0; mb_num < mbs_per_slice; mb_num++, out_ptr += blocks_per_mb * 4) {
+            ctx->dsp.idct_put(out_ptr,                    linesize, block_ptr, qmat);
+            block_ptr += 64;
+            ctx->dsp.idct_put(out_ptr + linesize * 4,     linesize, block_ptr, qmat);
             block_ptr += 64;
+            if (blocks_per_mb > 2) {
+                ctx->dsp.idct_put(out_ptr + 8,            linesize, block_ptr, qmat);
+                block_ptr += 64;
+                ctx->dsp.idct_put(out_ptr + linesize * 4 + 8, linesize, block_ptr, qmat);
+                block_ptr += 64;
+            }
         }
     }
 }
@@ -523,7 +538,7 @@ static int decode_slice(AVCodecContext *avctx, void *tdata)
                        (uint16_t*) (y_data + (mb_y_pos << 4) * y_linesize +
                                     (mb_x_pos << 5)), y_linesize,
                        mbs_per_slice, 4, slice_width_factor + 2,
-                       td->qmat_luma_scaled);
+                       td->qmat_luma_scaled, 0);
 
     /* decode U chroma plane */
     decode_slice_plane(ctx, td, buf + hdr_size + y_data_size, u_data_size,
@@ -531,7 +546,7 @@ static int decode_slice(AVCodecContext *avctx, void *tdata)
                                     (mb_x_pos << ctx->mb_chroma_factor)),
                        u_linesize, mbs_per_slice, ctx->num_chroma_blocks,
                        slice_width_factor + ctx->chroma_factor - 1,
-                       td->qmat_chroma_scaled);
+                       td->qmat_chroma_scaled, 1);
 
     /* decode V chroma plane */
     decode_slice_plane(ctx, td, buf + hdr_size + y_data_size + u_data_size,
@@ -540,7 +555,7 @@ static int decode_slice(AVCodecContext *avctx, void *tdata)
                                     (mb_x_pos << ctx->mb_chroma_factor)),
                        v_linesize, mbs_per_slice, ctx->num_chroma_blocks,
                        slice_width_factor + ctx->chroma_factor - 1,
-                       td->qmat_chroma_scaled);
+                       td->qmat_chroma_scaled, 1);
 
     return 0;
 }
diff --git a/libavcodec/proresenc.c b/libavcodec/proresenc.c
index aa96b35e75..552fa7c14a 100644
--- a/libavcodec/proresenc.c
+++ b/libavcodec/proresenc.c
@@ -171,7 +171,7 @@ typedef struct ProresContext {
 static void get_slice_data(ProresContext *ctx, const uint16_t *src,
                            int linesize, int x, int y, int w, int h,
                            DCTELEM *blocks,
-                           int mbs_per_slice, int blocks_per_mb)
+                           int mbs_per_slice, int blocks_per_mb, int is_chroma)
 {
     const uint16_t *esrc;
     const int mb_width = 4 * blocks_per_mb;
@@ -209,17 +209,30 @@ static void get_slice_data(ProresContext *ctx, const uint16_t *src,
                        ctx->emu_buf + (bh - 1) * estride,
                        mb_width * sizeof(*ctx->emu_buf));
         }
-        ctx->dsp.fdct(esrc, elinesize, blocks);
-        blocks += 64;
-        if (blocks_per_mb > 2) {
-            ctx->dsp.fdct(src + 8, linesize, blocks);
+        if (!is_chroma) {
+            ctx->dsp.fdct(esrc, elinesize, blocks);
             blocks += 64;
-        }
-        ctx->dsp.fdct(src + linesize * 4, linesize, blocks);
-        blocks += 64;
-        if (blocks_per_mb > 2) {
-            ctx->dsp.fdct(src + linesize * 4 + 8, linesize, blocks);
+            if (blocks_per_mb > 2) {
+                ctx->dsp.fdct(src + 8, linesize, blocks);
+                blocks += 64;
+            }
+            ctx->dsp.fdct(src + linesize * 4, linesize, blocks);
             blocks += 64;
+            if (blocks_per_mb > 2) {
+                ctx->dsp.fdct(src + linesize * 4 + 8, linesize, blocks);
+                blocks += 64;
+            }
+        } else {
+            ctx->dsp.fdct(esrc, elinesize, blocks);
+            blocks += 64;
+            ctx->dsp.fdct(src + linesize * 4, linesize, blocks);
+            blocks += 64;
+            if (blocks_per_mb > 2) {
+                ctx->dsp.fdct(src + 8, linesize, blocks);
+                blocks += 64;
+                ctx->dsp.fdct(src + linesize * 4 + 8, linesize, blocks);
+                blocks += 64;
+            }
         }
 
         x += mb_width;
@@ -383,7 +396,7 @@ static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
 
         get_slice_data(ctx, src, pic->linesize[i], xp, yp,
                        pwidth, avctx->height, ctx->blocks[0],
-                       mbs_per_slice, num_cblocks);
+                       mbs_per_slice, num_cblocks, is_chroma);
         sizes[i] = encode_slice_plane(ctx, pb, src, pic->linesize[i],
                                       mbs_per_slice, ctx->blocks[0],
                                       num_cblocks, plane_factor,
@@ -539,7 +552,7 @@ static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
 
         get_slice_data(ctx, src, pic->linesize[i], xp, yp,
                        pwidth, avctx->height, ctx->blocks[i],
-                       mbs_per_slice, num_cblocks[i]);
+                       mbs_per_slice, num_cblocks[i], is_chroma[i]);
     }
 
     for (q = min_quant; q < max_quant + 2; q++) {
diff --git a/tests/ref/fate/prores-alpha b/tests/ref/fate/prores-alpha
index 9aa916013b..80442fe7b5 100644
--- a/tests/ref/fate/prores-alpha
+++ b/tests/ref/fate/prores-alpha
@@ -1,3 +1,3 @@
 #tb 0: 1/2997
-0,          0,          0,        0, 12441600, 0x9d3dc525
-0,        100,        100,        0, 12441600, 0x9d3dc525
+0,          0,          0,        0, 12441600, 0x254d8f95
+0,        100,        100,        0, 12441600, 0x254d8f95