avcodec/agm: Fix overflow of signed shift
[ffmpeg.git] / libavcodec / agm.c
1 /*
2  * Amuse Graphics Movie decoder
3  *
4  * Copyright (c) 2018 Paul B Mahol
5  *
6  * This file is part of FFmpeg.
7  *
8  * FFmpeg 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  * FFmpeg 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 FFmpeg; if not, write to the Free Software
20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
21  */
22
23 #include <stdio.h>
24 #include <stdlib.h>
25 #include <string.h>
26
27 #define BITSTREAM_READER_LE
28
29 #include "avcodec.h"
30 #include "bytestream.h"
31 #include "copy_block.h"
32 #include "get_bits.h"
33 #include "idctdsp.h"
34 #include "internal.h"
35
36 static const uint8_t unscaled_luma[64] = {
37     16, 11, 10, 16, 24, 40, 51, 61, 12, 12, 14, 19,
38     26, 58, 60, 55, 14, 13, 16, 24, 40, 57, 69, 56,
39     14, 17, 22, 29, 51, 87, 80, 62, 18, 22, 37, 56,
40     68,109,103, 77, 24, 35, 55, 64, 81,104,113, 92,
41     49, 64, 78, 87,103,121,120,101, 72, 92, 95, 98,
42     112,100,103,99
43 };
44
45 static const uint8_t unscaled_chroma[64] = {
46     17, 18, 24, 47, 99, 99, 99, 99, 18, 21, 26, 66,
47     99, 99, 99, 99, 24, 26, 56, 99, 99, 99, 99, 99,
48     47, 66, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
49     99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
50     99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
51     99, 99, 99, 99
52 };
53
54 typedef struct MotionVector {
55     int16_t x, y;
56 } MotionVector;
57
58 typedef struct AGMContext {
59     const AVClass  *class;
60     AVCodecContext *avctx;
61     GetBitContext   gb;
62     GetByteContext  gbyte;
63
64     int key_frame;
65     int bitstream_size;
66     int compression;
67     int blocks_w;
68     int blocks_h;
69     int size[3];
70     int plus;
71     int dct;
72     int rgb;
73     unsigned flags;
74     unsigned fflags;
75
76     uint8_t *output;
77     unsigned padded_output_size;
78     unsigned output_size;
79
80     MotionVector *mvectors;
81     unsigned      mvectors_size;
82
83     VLC vlc;
84
85     AVFrame *prev_frame;
86
87     int luma_quant_matrix[64];
88     int chroma_quant_matrix[64];
89
90     ScanTable scantable;
91     DECLARE_ALIGNED(32, int16_t, block)[64];
92
93     int16_t *wblocks;
94     unsigned wblocks_size;
95
96     int      *map;
97     unsigned  map_size;
98
99     IDCTDSPContext idsp;
100 } AGMContext;
101
102 static int read_code(GetBitContext *gb, int *oskip, int *level, int *map, int mode)
103 {
104     int len = 0, skip = 0, max;
105
106     if (get_bits_left(gb) < 2)
107         return AVERROR_INVALIDDATA;
108
109     if (show_bits(gb, 2)) {
110         switch (show_bits(gb, 4)) {
111         case 1:
112         case 9:
113             len = 1;
114             skip = 3;
115             break;
116         case 2:
117             len = 3;
118             skip = 4;
119             break;
120         case 3:
121             len = 7;
122             skip = 4;
123             break;
124         case 5:
125         case 13:
126             len = 2;
127             skip = 3;
128             break;
129         case 6:
130             len = 4;
131             skip = 4;
132             break;
133         case 7:
134             len = 8;
135             skip = 4;
136             break;
137         case 10:
138             len = 5;
139             skip = 4;
140             break;
141         case 11:
142             len = 9;
143             skip = 4;
144             break;
145         case 14:
146             len = 6;
147             skip = 4;
148             break;
149         case 15:
150             len = ((show_bits(gb, 5) & 0x10) | 0xA0) >> 4;
151             skip = 5;
152             break;
153         default:
154             return AVERROR_INVALIDDATA;
155         }
156
157         skip_bits(gb, skip);
158         *level = get_bits(gb, len);
159         *map = 1;
160         *oskip = 0;
161         max = 1 << (len - 1);
162         if (*level < max)
163             *level = -(max + *level);
164     } else if (show_bits(gb, 3) & 4) {
165         skip_bits(gb, 3);
166         if (mode == 1) {
167             if (show_bits(gb, 4)) {
168                 if (show_bits(gb, 4) == 1) {
169                     skip_bits(gb, 4);
170                     *oskip = get_bits(gb, 16);
171                 } else {
172                     *oskip = get_bits(gb, 4);
173                 }
174             } else {
175                 skip_bits(gb, 4);
176                 *oskip = get_bits(gb, 10);
177             }
178         } else if (mode == 0) {
179             *oskip = get_bits(gb, 10);
180         }
181         *level = 0;
182     } else {
183         skip_bits(gb, 3);
184         if (mode == 0)
185             *oskip = get_bits(gb, 4);
186         else if (mode == 1)
187             *oskip = 0;
188         *level = 0;
189     }
190
191     return 0;
192 }
193
194 static int decode_intra_blocks(AGMContext *s, GetBitContext *gb,
195                                const int *quant_matrix, int *skip, int *dc_level)
196 {
197     const uint8_t *scantable = s->scantable.permutated;
198     int level, ret, map = 0;
199
200     memset(s->wblocks, 0, s->wblocks_size);
201
202     for (int i = 0; i < 64; i++) {
203         int16_t *block = s->wblocks + scantable[i];
204
205         for (int j = 0; j < s->blocks_w;) {
206             if (*skip > 0) {
207                 int rskip;
208
209                 rskip = FFMIN(*skip, s->blocks_w - j);
210                 j += rskip;
211                 if (i == 0) {
212                     for (int k = 0; k < rskip; k++)
213                         block[64 * k] = *dc_level * quant_matrix[0];
214                 }
215                 block += rskip * 64;
216                 *skip -= rskip;
217             } else {
218                 ret = read_code(gb, skip, &level, &map, s->flags & 1);
219                 if (ret < 0)
220                     return ret;
221
222                 if (i == 0)
223                     *dc_level += level;
224
225                 block[0] = (i == 0 ? *dc_level : level) * quant_matrix[i];
226                 block += 64;
227                 j++;
228             }
229         }
230     }
231
232     return 0;
233 }
234
235 static int decode_inter_blocks(AGMContext *s, GetBitContext *gb,
236                                const int *quant_matrix, int *skip,
237                                int *map)
238 {
239     const uint8_t *scantable = s->scantable.permutated;
240     int level, ret;
241
242     memset(s->wblocks, 0, s->wblocks_size);
243     memset(s->map, 0, s->map_size);
244
245     for (int i = 0; i < 64; i++) {
246         int16_t *block = s->wblocks + scantable[i];
247
248         for (int j = 0; j < s->blocks_w;) {
249             if (*skip > 0) {
250                 int rskip;
251
252                 rskip = FFMIN(*skip, s->blocks_w - j);
253                 j += rskip;
254                 block += rskip * 64;
255                 *skip -= rskip;
256             } else {
257                 ret = read_code(gb, skip, &level, &map[j], s->flags & 1);
258                 if (ret < 0)
259                     return ret;
260
261                 block[0] = level * quant_matrix[i];
262                 block += 64;
263                 j++;
264             }
265         }
266     }
267
268     return 0;
269 }
270
271 static int decode_intra_block(AGMContext *s, GetBitContext *gb,
272                               const int *quant_matrix, int *skip, int *dc_level)
273 {
274     const uint8_t *scantable = s->scantable.permutated;
275     const int offset = s->plus ? 0 : 1024;
276     int16_t *block = s->block;
277     int level, ret, map = 0;
278
279     memset(block, 0, sizeof(s->block));
280
281     if (*skip > 0) {
282         (*skip)--;
283     } else {
284         ret = read_code(gb, skip, &level, &map, s->flags & 1);
285         if (ret < 0)
286             return ret;
287         *dc_level += level;
288     }
289     block[scantable[0]] = offset + *dc_level * quant_matrix[0];
290
291     for (int i = 1; i < 64;) {
292         if (*skip > 0) {
293             int rskip;
294
295             rskip = FFMIN(*skip, 64 - i);
296             i += rskip;
297             *skip -= rskip;
298         } else {
299             ret = read_code(gb, skip, &level, &map, s->flags & 1);
300             if (ret < 0)
301                 return ret;
302
303             block[scantable[i]] = level * quant_matrix[i];
304             i++;
305         }
306     }
307
308     return 0;
309 }
310
311 static int decode_intra_plane(AGMContext *s, GetBitContext *gb, int size,
312                               const int *quant_matrix, AVFrame *frame,
313                               int plane)
314 {
315     int ret, skip = 0, dc_level = 0;
316     const int offset = s->plus ? 0 : 1024;
317
318     if ((ret = init_get_bits8(gb, s->gbyte.buffer, size)) < 0)
319         return ret;
320
321     if (s->flags & 1) {
322         av_fast_padded_malloc(&s->wblocks, &s->wblocks_size,
323                               64 * s->blocks_w * sizeof(*s->wblocks));
324         if (!s->wblocks)
325             return AVERROR(ENOMEM);
326
327         for (int y = 0; y < s->blocks_h; y++) {
328             ret = decode_intra_blocks(s, gb, quant_matrix, &skip, &dc_level);
329             if (ret < 0)
330                 return ret;
331
332             for (int x = 0; x < s->blocks_w; x++) {
333                 s->wblocks[64 * x] += offset;
334                 s->idsp.idct_put(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
335                                  frame->linesize[plane], s->wblocks + 64 * x);
336             }
337         }
338     } else {
339         for (int y = 0; y < s->blocks_h; y++) {
340             for (int x = 0; x < s->blocks_w; x++) {
341                 ret = decode_intra_block(s, gb, quant_matrix, &skip, &dc_level);
342                 if (ret < 0)
343                     return ret;
344
345                 s->idsp.idct_put(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
346                                  frame->linesize[plane], s->block);
347             }
348         }
349     }
350
351     align_get_bits(gb);
352     if (get_bits_left(gb) < 0)
353         av_log(s->avctx, AV_LOG_WARNING, "overread\n");
354     if (get_bits_left(gb) > 0)
355         av_log(s->avctx, AV_LOG_WARNING, "underread: %d\n", get_bits_left(gb));
356
357     return 0;
358 }
359
360 static int decode_inter_block(AGMContext *s, GetBitContext *gb,
361                               const int *quant_matrix, int *skip,
362                               int *map)
363 {
364     const uint8_t *scantable = s->scantable.permutated;
365     int16_t *block = s->block;
366     int level, ret;
367
368     memset(block, 0, sizeof(s->block));
369
370     for (int i = 0; i < 64;) {
371         if (*skip > 0) {
372             int rskip;
373
374             rskip = FFMIN(*skip, 64 - i);
375             i += rskip;
376             *skip -= rskip;
377         } else {
378             ret = read_code(gb, skip, &level, map, s->flags & 1);
379             if (ret < 0)
380                 return ret;
381
382             block[scantable[i]] = level * quant_matrix[i];
383             i++;
384         }
385     }
386
387     return 0;
388 }
389
390 static int decode_inter_plane(AGMContext *s, GetBitContext *gb, int size,
391                               const int *quant_matrix, AVFrame *frame,
392                               AVFrame *prev, int plane)
393 {
394     int ret, skip = 0;
395
396     if ((ret = init_get_bits8(gb, s->gbyte.buffer, size)) < 0)
397         return ret;
398
399     if (s->flags == 3) {
400         av_fast_padded_malloc(&s->wblocks, &s->wblocks_size,
401                               64 * s->blocks_w * sizeof(*s->wblocks));
402         if (!s->wblocks)
403             return AVERROR(ENOMEM);
404
405         av_fast_padded_malloc(&s->map, &s->map_size,
406                               s->blocks_w * sizeof(*s->map));
407         if (!s->map)
408             return AVERROR(ENOMEM);
409
410         for (int y = 0; y < s->blocks_h; y++) {
411             ret = decode_inter_blocks(s, gb, quant_matrix, &skip, s->map);
412             if (ret < 0)
413                 return ret;
414
415             for (int x = 0; x < s->blocks_w; x++) {
416                 int shift = plane == 0;
417                 int mvpos = (y >> shift) * (s->blocks_w >> shift) + (x >> shift);
418                 int orig_mv_x = s->mvectors[mvpos].x;
419                 int mv_x = s->mvectors[mvpos].x / (1 + !shift);
420                 int mv_y = s->mvectors[mvpos].y / (1 + !shift);
421                 int h = s->avctx->coded_height >> !shift;
422                 int w = s->avctx->coded_width  >> !shift;
423                 int map = s->map[x];
424
425                 if (orig_mv_x >= -32) {
426                     if (y * 8 + mv_y < 0 || y * 8 + mv_y >= h ||
427                         x * 8 + mv_x < 0 || x * 8 + mv_x >= w)
428                         return AVERROR_INVALIDDATA;
429
430                     copy_block8(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
431                                 prev->data[plane] + ((s->blocks_h - 1 - y) * 8 - mv_y) * prev->linesize[plane] + (x * 8 + mv_x),
432                                 frame->linesize[plane], prev->linesize[plane], 8);
433                     if (map) {
434                         s->idsp.idct(s->wblocks + x * 64);
435                         for (int i = 0; i < 64; i++)
436                             s->wblocks[i + x * 64] = (s->wblocks[i + x * 64] + 1) & 0xFFFC;
437                         s->idsp.add_pixels_clamped(&s->wblocks[x*64], frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
438                                                    frame->linesize[plane]);
439                     }
440                 } else if (map) {
441                     s->idsp.idct_put(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
442                                      frame->linesize[plane], s->wblocks + x * 64);
443                 }
444             }
445         }
446     } else if (s->flags & 2) {
447         for (int y = 0; y < s->blocks_h; y++) {
448             for (int x = 0; x < s->blocks_w; x++) {
449                 int shift = plane == 0;
450                 int mvpos = (y >> shift) * (s->blocks_w >> shift) + (x >> shift);
451                 int orig_mv_x = s->mvectors[mvpos].x;
452                 int mv_x = s->mvectors[mvpos].x / (1 + !shift);
453                 int mv_y = s->mvectors[mvpos].y / (1 + !shift);
454                 int h = s->avctx->coded_height >> !shift;
455                 int w = s->avctx->coded_width  >> !shift;
456                 int map = 0;
457
458                 ret = decode_inter_block(s, gb, quant_matrix, &skip, &map);
459                 if (ret < 0)
460                     return ret;
461
462                 if (orig_mv_x >= -32) {
463                     if (y * 8 + mv_y < 0 || y * 8 + mv_y >= h ||
464                         x * 8 + mv_x < 0 || x * 8 + mv_x >= w)
465                         return AVERROR_INVALIDDATA;
466
467                     copy_block8(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
468                                 prev->data[plane] + ((s->blocks_h - 1 - y) * 8 - mv_y) * prev->linesize[plane] + (x * 8 + mv_x),
469                                 frame->linesize[plane], prev->linesize[plane], 8);
470                     if (map) {
471                         s->idsp.idct(s->block);
472                         for (int i = 0; i < 64; i++)
473                             s->block[i] = (s->block[i] + 1) & 0xFFFC;
474                         s->idsp.add_pixels_clamped(s->block, frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
475                                                    frame->linesize[plane]);
476                     }
477                 } else if (map) {
478                     s->idsp.idct_put(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
479                                      frame->linesize[plane], s->block);
480                 }
481             }
482         }
483     } else if (s->flags & 1) {
484         av_fast_padded_malloc(&s->wblocks, &s->wblocks_size,
485                               64 * s->blocks_w * sizeof(*s->wblocks));
486         if (!s->wblocks)
487             return AVERROR(ENOMEM);
488
489         av_fast_padded_malloc(&s->map, &s->map_size,
490                               s->blocks_w * sizeof(*s->map));
491         if (!s->map)
492             return AVERROR(ENOMEM);
493
494         for (int y = 0; y < s->blocks_h; y++) {
495             ret = decode_inter_blocks(s, gb, quant_matrix, &skip, s->map);
496             if (ret < 0)
497                 return ret;
498
499             for (int x = 0; x < s->blocks_w; x++) {
500                 if (!s->map[x])
501                     continue;
502                 s->idsp.idct_add(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
503                                  frame->linesize[plane], s->wblocks + 64 * x);
504             }
505         }
506     } else {
507         for (int y = 0; y < s->blocks_h; y++) {
508             for (int x = 0; x < s->blocks_w; x++) {
509                 int map = 0;
510
511                 ret = decode_inter_block(s, gb, quant_matrix, &skip, &map);
512                 if (ret < 0)
513                     return ret;
514
515                 if (!map)
516                     continue;
517                 s->idsp.idct_add(frame->data[plane] + (s->blocks_h - 1 - y) * 8 * frame->linesize[plane] + x * 8,
518                                  frame->linesize[plane], s->block);
519             }
520         }
521     }
522
523     align_get_bits(gb);
524     if (get_bits_left(gb) < 0)
525         av_log(s->avctx, AV_LOG_WARNING, "overread\n");
526     if (get_bits_left(gb) > 0)
527         av_log(s->avctx, AV_LOG_WARNING, "underread: %d\n", get_bits_left(gb));
528
529     return 0;
530 }
531
532 static void compute_quant_matrix(AGMContext *s, double qscale)
533 {
534     int luma[64], chroma[64];
535     double f = 1.0 - fabs(qscale);
536
537     if (!s->key_frame && (s->flags & 2)) {
538         if (qscale >= 0.0) {
539             for (int i = 0; i < 64; i++) {
540                 luma[i]   = FFMAX(1, 16 * f);
541                 chroma[i] = FFMAX(1, 16 * f);
542             }
543         } else {
544             for (int i = 0; i < 64; i++) {
545                 luma[i]   = FFMAX(1, 16 - qscale * 32);
546                 chroma[i] = FFMAX(1, 16 - qscale * 32);
547             }
548         }
549     } else {
550         if (qscale >= 0.0) {
551             for (int i = 0; i < 64; i++) {
552                 luma[i]   = FFMAX(1, unscaled_luma  [(i & 7) * 8 + (i >> 3)] * f);
553                 chroma[i] = FFMAX(1, unscaled_chroma[(i & 7) * 8 + (i >> 3)] * f);
554             }
555         } else {
556             for (int i = 0; i < 64; i++) {
557                 luma[i]   = FFMAX(1, 255.0 - (255 - unscaled_luma  [(i & 7) * 8 + (i >> 3)]) * f);
558                 chroma[i] = FFMAX(1, 255.0 - (255 - unscaled_chroma[(i & 7) * 8 + (i >> 3)]) * f);
559             }
560         }
561     }
562
563     for (int i = 0; i < 64; i++) {
564         int pos = ff_zigzag_direct[i];
565
566         s->luma_quant_matrix[i]   = luma[pos]   * ((pos / 8) & 1 ? -1 : 1);
567         s->chroma_quant_matrix[i] = chroma[pos] * ((pos / 8) & 1 ? -1 : 1);
568     }
569 }
570
571 static int decode_raw_intra_rgb(AVCodecContext *avctx, GetByteContext *gbyte, AVFrame *frame)
572 {
573     uint8_t *dst = frame->data[0] + (avctx->height - 1) * frame->linesize[0];
574     uint8_t r = 0, g = 0, b = 0;
575
576     for (int y = 0; y < avctx->height; y++) {
577         for (int x = 0; x < avctx->width; x++) {
578             dst[x*3+0] = bytestream2_get_byte(gbyte) + r;
579             r = dst[x*3+0];
580             dst[x*3+1] = bytestream2_get_byte(gbyte) + g;
581             g = dst[x*3+1];
582             dst[x*3+2] = bytestream2_get_byte(gbyte) + b;
583             b = dst[x*3+2];
584         }
585         dst -= frame->linesize[0];
586     }
587
588     return 0;
589 }
590
591 static int fill_pixels(uint8_t **y0, uint8_t **y1,
592                        uint8_t **u, uint8_t **v,
593                        int ylinesize, int ulinesize, int vlinesize,
594                        uint8_t *fill,
595                        int *nx, int *ny, int *np, int w, int h)
596 {
597     uint8_t *y0dst = *y0;
598     uint8_t *y1dst = *y1;
599     uint8_t *udst = *u;
600     uint8_t *vdst = *v;
601     int x = *nx, y = *ny, pos = *np;
602
603     if (pos == 0) {
604         y0dst[2*x+0] += fill[0];
605         y0dst[2*x+1] += fill[1];
606         y1dst[2*x+0] += fill[2];
607         y1dst[2*x+1] += fill[3];
608         pos++;
609     } else if (pos == 1) {
610         udst[x] += fill[0];
611         vdst[x] += fill[1];
612         x++;
613         if (x >= w) {
614             x = 0;
615             y++;
616             if (y >= h)
617                 return 1;
618             y0dst -= 2*ylinesize;
619             y1dst -= 2*ylinesize;
620             udst  -=   ulinesize;
621             vdst  -=   vlinesize;
622         }
623         y0dst[2*x+0] += fill[2];
624         y0dst[2*x+1] += fill[3];
625         pos++;
626     } else if (pos == 2) {
627         y1dst[2*x+0] += fill[0];
628         y1dst[2*x+1] += fill[1];
629         udst[x]      += fill[2];
630         vdst[x]      += fill[3];
631         x++;
632         if (x >= w) {
633             x = 0;
634             y++;
635             if (y >= h)
636                 return 1;
637             y0dst -= 2*ylinesize;
638             y1dst -= 2*ylinesize;
639             udst  -=   ulinesize;
640             vdst  -=   vlinesize;
641         }
642         pos = 0;
643     }
644
645     *y0 = y0dst;
646     *y1 = y1dst;
647     *u = udst;
648     *v = vdst;
649     *np = pos;
650     *nx = x;
651     *ny = y;
652
653     return 0;
654 }
655
656 static int decode_runlen_rgb(AVCodecContext *avctx, GetByteContext *gbyte, AVFrame *frame)
657 {
658     uint8_t *dst = frame->data[0] + (avctx->height - 1) * frame->linesize[0];
659     int runlen, y = 0, x = 0;
660     uint8_t fill[4];
661     unsigned code;
662
663     while (bytestream2_get_bytes_left(gbyte) > 0) {
664         code = bytestream2_peek_le32(gbyte);
665         runlen = code & 0xFFFFFF;
666
667         if (code >> 24 == 0x77) {
668             bytestream2_skip(gbyte, 4);
669
670             for (int i = 0; i < 4; i++)
671                 fill[i] = bytestream2_get_byte(gbyte);
672
673             while (runlen > 0) {
674                 runlen--;
675
676                 for (int i = 0; i < 4; i++) {
677                     dst[x] += fill[i];
678                     x++;
679                     if (x >= frame->width * 3) {
680                         x = 0;
681                         y++;
682                         dst -= frame->linesize[0];
683                         if (y >= frame->height)
684                             return 0;
685                     }
686                 }
687             }
688         } else {
689             for (int i = 0; i < 4; i++)
690                 fill[i] = bytestream2_get_byte(gbyte);
691
692             for (int i = 0; i < 4; i++) {
693                 dst[x] += fill[i];
694                 x++;
695                 if (x >= frame->width * 3) {
696                     x = 0;
697                     y++;
698                     dst -= frame->linesize[0];
699                     if (y >= frame->height)
700                         return 0;
701                 }
702             }
703         }
704     }
705
706     return 0;
707 }
708
709 static int decode_runlen(AVCodecContext *avctx, GetByteContext *gbyte, AVFrame *frame)
710 {
711     uint8_t *y0dst = frame->data[0] + (avctx->height - 1) * frame->linesize[0];
712     uint8_t *y1dst = y0dst - frame->linesize[0];
713     uint8_t *udst = frame->data[1] + ((avctx->height >> 1) - 1) * frame->linesize[1];
714     uint8_t *vdst = frame->data[2] + ((avctx->height >> 1) - 1) * frame->linesize[2];
715     int runlen, y = 0, x = 0, pos = 0;
716     uint8_t fill[4];
717     unsigned code;
718
719     while (bytestream2_get_bytes_left(gbyte) > 0) {
720         code = bytestream2_peek_le32(gbyte);
721         runlen = code & 0xFFFFFF;
722
723         if (code >> 24 == 0x77) {
724             bytestream2_skip(gbyte, 4);
725
726             for (int i = 0; i < 4; i++)
727                 fill[i] = bytestream2_get_byte(gbyte);
728
729             while (runlen > 0) {
730                 runlen--;
731
732                 if (fill_pixels(&y0dst, &y1dst, &udst, &vdst,
733                                 frame->linesize[0],
734                                 frame->linesize[1],
735                                 frame->linesize[2],
736                                 fill, &x, &y, &pos,
737                                 avctx->width / 2,
738                                 avctx->height / 2))
739                     return 0;
740             }
741         } else {
742             for (int i = 0; i < 4; i++)
743                 fill[i] = bytestream2_get_byte(gbyte);
744
745             if (fill_pixels(&y0dst, &y1dst, &udst, &vdst,
746                             frame->linesize[0],
747                             frame->linesize[1],
748                             frame->linesize[2],
749                             fill, &x, &y, &pos,
750                             avctx->width / 2,
751                             avctx->height / 2))
752                 return 0;
753         }
754     }
755
756     return 0;
757 }
758
759 static int decode_raw_intra(AVCodecContext *avctx, GetByteContext *gbyte, AVFrame *frame)
760 {
761     uint8_t *y0dst = frame->data[0] + (avctx->height - 1) * frame->linesize[0];
762     uint8_t *y1dst = y0dst - frame->linesize[0];
763     uint8_t *udst = frame->data[1] + ((avctx->height >> 1) - 1) * frame->linesize[1];
764     uint8_t *vdst = frame->data[2] + ((avctx->height >> 1) - 1) * frame->linesize[2];
765     uint8_t ly0 = 0, ly1 = 0, ly2 = 0, ly3 = 0, lu = 0, lv = 0;
766
767     for (int y = 0; y < avctx->height / 2; y++) {
768         for (int x = 0; x < avctx->width / 2; x++) {
769             y0dst[x*2+0] = bytestream2_get_byte(gbyte) + ly0;
770             ly0 = y0dst[x*2+0];
771             y0dst[x*2+1] = bytestream2_get_byte(gbyte) + ly1;
772             ly1 = y0dst[x*2+1];
773             y1dst[x*2+0] = bytestream2_get_byte(gbyte) + ly2;
774             ly2 = y1dst[x*2+0];
775             y1dst[x*2+1] = bytestream2_get_byte(gbyte) + ly3;
776             ly3 = y1dst[x*2+1];
777             udst[x] = bytestream2_get_byte(gbyte) + lu;
778             lu = udst[x];
779             vdst[x] = bytestream2_get_byte(gbyte) + lv;
780             lv = vdst[x];
781         }
782
783         y0dst -= 2*frame->linesize[0];
784         y1dst -= 2*frame->linesize[0];
785         udst  -= frame->linesize[1];
786         vdst  -= frame->linesize[2];
787     }
788
789     return 0;
790 }
791
792 static int decode_intra(AVCodecContext *avctx, GetBitContext *gb, AVFrame *frame)
793 {
794     AGMContext *s = avctx->priv_data;
795     int ret;
796
797     compute_quant_matrix(s, (2 * s->compression - 100) / 100.0);
798
799     s->blocks_w = avctx->coded_width  >> 3;
800     s->blocks_h = avctx->coded_height >> 3;
801
802     ret = decode_intra_plane(s, gb, s->size[0], s->luma_quant_matrix, frame, 0);
803     if (ret < 0)
804         return ret;
805
806     bytestream2_skip(&s->gbyte, s->size[0]);
807
808     s->blocks_w = avctx->coded_width  >> 4;
809     s->blocks_h = avctx->coded_height >> 4;
810
811     ret = decode_intra_plane(s, gb, s->size[1], s->chroma_quant_matrix, frame, 2);
812     if (ret < 0)
813         return ret;
814
815     bytestream2_skip(&s->gbyte, s->size[1]);
816
817     s->blocks_w = avctx->coded_width  >> 4;
818     s->blocks_h = avctx->coded_height >> 4;
819
820     ret = decode_intra_plane(s, gb, s->size[2], s->chroma_quant_matrix, frame, 1);
821     if (ret < 0)
822         return ret;
823
824     return 0;
825 }
826
827 static int decode_motion_vectors(AVCodecContext *avctx, GetBitContext *gb)
828 {
829     AGMContext *s = avctx->priv_data;
830     int nb_mvs = ((avctx->height + 15) >> 4) * ((avctx->width + 15) >> 4);
831     int ret, skip = 0, value, map;
832
833     av_fast_padded_malloc(&s->mvectors, &s->mvectors_size,
834                           nb_mvs * sizeof(*s->mvectors));
835     if (!s->mvectors)
836         return AVERROR(ENOMEM);
837
838     if ((ret = init_get_bits8(gb, s->gbyte.buffer, bytestream2_get_bytes_left(&s->gbyte) -
839                                                    (s->size[0] + s->size[1] + s->size[2]))) < 0)
840         return ret;
841
842     memset(s->mvectors, 0, sizeof(*s->mvectors) * nb_mvs);
843
844     for (int i = 0; i < nb_mvs; i++) {
845         ret = read_code(gb, &skip, &value, &map, 1);
846         if (ret < 0)
847             return ret;
848         s->mvectors[i].x = value;
849         i += skip;
850     }
851
852     for (int i = 0; i < nb_mvs; i++) {
853         ret = read_code(gb, &skip, &value, &map, 1);
854         if (ret < 0)
855             return ret;
856         s->mvectors[i].y = value;
857         i += skip;
858     }
859
860     if (get_bits_left(gb) <= 0)
861         return AVERROR_INVALIDDATA;
862     skip = (get_bits_count(gb) >> 3) + 1;
863     bytestream2_skip(&s->gbyte, skip);
864
865     return 0;
866 }
867
868 static int decode_inter(AVCodecContext *avctx, GetBitContext *gb,
869                         AVFrame *frame, AVFrame *prev)
870 {
871     AGMContext *s = avctx->priv_data;
872     int ret;
873
874     compute_quant_matrix(s, (2 * s->compression - 100) / 100.0);
875
876     if (s->flags & 2) {
877         ret = decode_motion_vectors(avctx, gb);
878         if (ret < 0)
879             return ret;
880     }
881
882     s->blocks_w = avctx->coded_width  >> 3;
883     s->blocks_h = avctx->coded_height >> 3;
884
885     ret = decode_inter_plane(s, gb, s->size[0], s->luma_quant_matrix, frame, prev, 0);
886     if (ret < 0)
887         return ret;
888
889     bytestream2_skip(&s->gbyte, s->size[0]);
890
891     s->blocks_w = avctx->coded_width  >> 4;
892     s->blocks_h = avctx->coded_height >> 4;
893
894     ret = decode_inter_plane(s, gb, s->size[1], s->chroma_quant_matrix, frame, prev, 2);
895     if (ret < 0)
896         return ret;
897
898     bytestream2_skip(&s->gbyte, s->size[1]);
899
900     s->blocks_w = avctx->coded_width  >> 4;
901     s->blocks_h = avctx->coded_height >> 4;
902
903     ret = decode_inter_plane(s, gb, s->size[2], s->chroma_quant_matrix, frame, prev, 1);
904     if (ret < 0)
905         return ret;
906
907     return 0;
908 }
909
910 typedef struct Node {
911     int parent;
912     int child[2];
913 } Node;
914
915 static void get_tree_codes(uint32_t *codes, Node *nodes, int idx, uint32_t pfx, int bitpos)
916 {
917     if (idx < 256 && idx >= 0) {
918         codes[idx] = pfx;
919     } else if (idx >= 0) {
920         get_tree_codes(codes, nodes, nodes[idx].child[0], pfx + (0 << bitpos), bitpos + 1);
921         get_tree_codes(codes, nodes, nodes[idx].child[1], pfx + (1U << bitpos), bitpos + 1);
922     }
923 }
924
925 static int make_new_tree(const uint8_t *bitlens, uint32_t *codes)
926 {
927     int zlcount = 0, curlen, idx, nindex, last, llast;
928     int blcounts[32] = { 0 };
929     int syms[8192];
930     Node nodes[512];
931     int node_idx[1024];
932     int old_idx[512];
933
934     for (int i = 0; i < 256; i++) {
935         int bitlen = bitlens[i];
936         int blcount = blcounts[bitlen];
937
938         zlcount += bitlen < 1;
939         syms[(bitlen << 8) + blcount] = i;
940         blcounts[bitlen]++;
941     }
942
943     for (int i = 0; i < 512; i++) {
944         nodes[i].child[0] = -1;
945         nodes[i].child[1] = -1;
946     }
947
948     for (int i = 0; i < 256; i++) {
949         node_idx[i] = 257 + i;
950     }
951
952     curlen = 1;
953     node_idx[512] = 256;
954     last = 255;
955     nindex = 1;
956
957     for (curlen = 1; curlen < 32; curlen++) {
958         if (blcounts[curlen] > 0) {
959             int max_zlcount = zlcount + blcounts[curlen];
960
961             for (int i = 0; zlcount < 256 && zlcount < max_zlcount; zlcount++, i++) {
962                 int p = node_idx[nindex - 1 + 512];
963                 int ch = syms[256 * curlen + i];
964
965                 if (nindex <= 0)
966                     return AVERROR_INVALIDDATA;
967
968                 if (nodes[p].child[0] == -1) {
969                     nodes[p].child[0] = ch;
970                 } else {
971                     nodes[p].child[1] = ch;
972                     nindex--;
973                 }
974                 nodes[ch].parent = p;
975             }
976         }
977         llast = last - 1;
978         idx = 0;
979         while (nindex > 0) {
980             int p, ch;
981
982             last = llast - idx;
983             p = node_idx[nindex - 1 + 512];
984             ch = node_idx[last];
985             if (nodes[p].child[0] == -1) {
986                 nodes[p].child[0] = ch;
987             } else {
988                 nodes[p].child[1] = ch;
989                 nindex--;
990             }
991             old_idx[idx] = ch;
992             nodes[ch].parent = p;
993             if (idx == llast)
994                 goto next;
995             idx++;
996             if (nindex <= 0) {
997                 for (int i = 0; i < idx; i++)
998                     node_idx[512 + i] = old_idx[i];
999             }
1000         }
1001         nindex = idx;
1002     }
1003
1004 next:
1005
1006     get_tree_codes(codes, nodes, 256, 0, 0);
1007     return 0;
1008 }
1009
1010 static int build_huff(const uint8_t *bitlen, VLC *vlc)
1011 {
1012     uint32_t new_codes[256];
1013     uint8_t bits[256];
1014     uint8_t symbols[256];
1015     uint32_t codes[256];
1016     int nb_codes = 0;
1017
1018     int ret = make_new_tree(bitlen, new_codes);
1019     if (ret < 0)
1020         return ret;
1021
1022     for (int i = 0; i < 256; i++) {
1023         if (bitlen[i]) {
1024             bits[nb_codes] = bitlen[i];
1025             codes[nb_codes] = new_codes[i];
1026             symbols[nb_codes] = i;
1027             nb_codes++;
1028         }
1029     }
1030
1031     ff_free_vlc(vlc);
1032     return ff_init_vlc_sparse(vlc, 13, nb_codes,
1033                               bits, 1, 1,
1034                               codes, 4, 4,
1035                               symbols, 1, 1,
1036                               INIT_VLC_LE);
1037 }
1038
1039 static int decode_huffman2(AVCodecContext *avctx, int header, int size)
1040 {
1041     AGMContext *s = avctx->priv_data;
1042     GetBitContext *gb = &s->gb;
1043     uint8_t lens[256];
1044     int ret, x, len;
1045
1046     if ((ret = init_get_bits8(gb, s->gbyte.buffer,
1047                               bytestream2_get_bytes_left(&s->gbyte))) < 0)
1048         return ret;
1049
1050     s->output_size = get_bits_long(gb, 32);
1051
1052     if (s->output_size > avctx->width * avctx->height * 9LL + 10000)
1053         return AVERROR_INVALIDDATA;
1054
1055     av_fast_padded_malloc(&s->output, &s->padded_output_size, s->output_size);
1056     if (!s->output)
1057         return AVERROR(ENOMEM);
1058
1059     x = get_bits(gb, 1);
1060     len = 4 + get_bits(gb, 1);
1061     if (x) {
1062         int cb[8] = { 0 };
1063         int count = get_bits(gb, 3) + 1;
1064
1065         for (int i = 0; i < count; i++)
1066             cb[i] = get_bits(gb, len);
1067
1068         for (int i = 0; i < 256; i++) {
1069             int idx = get_bits(gb, 3);
1070             lens[i] = cb[idx];
1071         }
1072     } else {
1073         for (int i = 0; i < 256; i++)
1074             lens[i] = get_bits(gb, len);
1075     }
1076
1077     if ((ret = build_huff(lens, &s->vlc)) < 0)
1078         return ret;
1079
1080     x = 0;
1081     while (get_bits_left(gb) > 0 && x < s->output_size) {
1082         int val = get_vlc2(gb, s->vlc.table, s->vlc.bits, 3);
1083         if (val < 0)
1084             return AVERROR_INVALIDDATA;
1085         s->output[x++] = val;
1086     }
1087
1088     return 0;
1089 }
1090
1091 static int decode_frame(AVCodecContext *avctx, void *data,
1092                         int *got_frame, AVPacket *avpkt)
1093 {
1094     AGMContext *s = avctx->priv_data;
1095     GetBitContext *gb = &s->gb;
1096     GetByteContext *gbyte = &s->gbyte;
1097     AVFrame *frame = data;
1098     int w, h, width, height, header;
1099     unsigned compressed_size;
1100     long skip;
1101     int ret;
1102
1103     if (!avpkt->size)
1104         return 0;
1105
1106     bytestream2_init(gbyte, avpkt->data, avpkt->size);
1107
1108     header = bytestream2_get_le32(gbyte);
1109     s->fflags = bytestream2_get_le32(gbyte);
1110     s->bitstream_size = s->fflags & 0x1FFFFFFF;
1111     s->fflags >>= 29;
1112     av_log(avctx, AV_LOG_DEBUG, "fflags: %X\n", s->fflags);
1113     if (avpkt->size < s->bitstream_size + 8)
1114         return AVERROR_INVALIDDATA;
1115
1116     s->key_frame = (avpkt->flags & AV_PKT_FLAG_KEY);
1117     frame->key_frame = s->key_frame;
1118     frame->pict_type = s->key_frame ? AV_PICTURE_TYPE_I : AV_PICTURE_TYPE_P;
1119
1120     if (header) {
1121         if (avctx->codec_tag == MKTAG('A', 'G', 'M', '0') ||
1122             avctx->codec_tag == MKTAG('A', 'G', 'M', '1'))
1123             return AVERROR_PATCHWELCOME;
1124         else
1125             ret = decode_huffman2(avctx, header, (avpkt->size - s->bitstream_size) - 8);
1126         if (ret < 0)
1127             return ret;
1128         bytestream2_init(gbyte, s->output, s->output_size);
1129     } else if (!s->dct) {
1130         bytestream2_skip(gbyte, 4);
1131     }
1132
1133     if (s->dct) {
1134         s->flags = 0;
1135         w = bytestream2_get_le32(gbyte);
1136         h = bytestream2_get_le32(gbyte);
1137         if (w == INT32_MIN || h == INT32_MIN)
1138             return AVERROR_INVALIDDATA;
1139         if (w < 0) {
1140             w = -w;
1141             s->flags |= 2;
1142         }
1143         if (h < 0) {
1144             h = -h;
1145             s->flags |= 1;
1146         }
1147
1148         width  = avctx->width;
1149         height = avctx->height;
1150         if (w < width || h < height || w & 7 || h & 7)
1151             return AVERROR_INVALIDDATA;
1152
1153         ret = ff_set_dimensions(avctx, w, h);
1154         if (ret < 0)
1155             return ret;
1156         avctx->width = width;
1157         avctx->height = height;
1158
1159         s->compression = bytestream2_get_le32(gbyte);
1160         if (s->compression < 0 || s->compression > 100)
1161             return AVERROR_INVALIDDATA;
1162
1163         for (int i = 0; i < 3; i++)
1164             s->size[i] = bytestream2_get_le32(gbyte);
1165         if (header) {
1166             compressed_size = s->output_size;
1167             skip = 8LL;
1168         } else {
1169             compressed_size = avpkt->size;
1170             skip = 32LL;
1171         }
1172         if (s->size[0] < 0 || s->size[1] < 0 || s->size[2] < 0 ||
1173             skip + s->size[0] + s->size[1] + s->size[2] > compressed_size) {
1174             return AVERROR_INVALIDDATA;
1175         }
1176     }
1177
1178     if ((ret = ff_get_buffer(avctx, frame, AV_GET_BUFFER_FLAG_REF)) < 0)
1179         return ret;
1180
1181     if (frame->key_frame) {
1182         if (!s->dct && !s->rgb)
1183             ret = decode_raw_intra(avctx, gbyte, frame);
1184         else if (!s->dct && s->rgb)
1185             ret = decode_raw_intra_rgb(avctx, gbyte, frame);
1186         else
1187             ret = decode_intra(avctx, gb, frame);
1188     } else {
1189         if (!s->prev_frame->data[0]) {
1190             av_log(avctx, AV_LOG_ERROR, "Missing reference frame.\n");
1191             return AVERROR_INVALIDDATA;
1192         }
1193         if (s->prev_frame-> width != frame->width ||
1194             s->prev_frame->height != frame->height)
1195             return AVERROR_INVALIDDATA;
1196
1197         if (!(s->flags & 2)) {
1198             ret = av_frame_copy(frame, s->prev_frame);
1199             if (ret < 0)
1200                 return ret;
1201         }
1202
1203         if (s->dct) {
1204             ret = decode_inter(avctx, gb, frame, s->prev_frame);
1205         } else if (!s->dct && !s->rgb) {
1206             ret = decode_runlen(avctx, gbyte, frame);
1207         } else {
1208             ret = decode_runlen_rgb(avctx, gbyte, frame);
1209         }
1210     }
1211     if (ret < 0)
1212         return ret;
1213
1214     av_frame_unref(s->prev_frame);
1215     if ((ret = av_frame_ref(s->prev_frame, frame)) < 0)
1216         return ret;
1217
1218     frame->crop_top  = avctx->coded_height - avctx->height;
1219     frame->crop_left = avctx->coded_width  - avctx->width;
1220
1221     *got_frame = 1;
1222
1223     return avpkt->size;
1224 }
1225
1226 static av_cold int decode_init(AVCodecContext *avctx)
1227 {
1228     AGMContext *s = avctx->priv_data;
1229
1230     s->rgb = avctx->codec_tag == MKTAG('A', 'G', 'M', '4');
1231     avctx->pix_fmt = s->rgb ? AV_PIX_FMT_BGR24 : AV_PIX_FMT_YUV420P;
1232     s->avctx = avctx;
1233     s->plus = avctx->codec_tag == MKTAG('A', 'G', 'M', '3') ||
1234               avctx->codec_tag == MKTAG('A', 'G', 'M', '7');
1235
1236     s->dct = avctx->codec_tag != MKTAG('A', 'G', 'M', '4') &&
1237              avctx->codec_tag != MKTAG('A', 'G', 'M', '5');
1238
1239     avctx->idct_algo = FF_IDCT_SIMPLE;
1240     ff_idctdsp_init(&s->idsp, avctx);
1241     ff_init_scantable(s->idsp.idct_permutation, &s->scantable, ff_zigzag_direct);
1242
1243     s->prev_frame = av_frame_alloc();
1244     if (!s->prev_frame)
1245         return AVERROR(ENOMEM);
1246
1247     return 0;
1248 }
1249
1250 static void decode_flush(AVCodecContext *avctx)
1251 {
1252     AGMContext *s = avctx->priv_data;
1253
1254     av_frame_unref(s->prev_frame);
1255 }
1256
1257 static av_cold int decode_close(AVCodecContext *avctx)
1258 {
1259     AGMContext *s = avctx->priv_data;
1260
1261     ff_free_vlc(&s->vlc);
1262     av_frame_free(&s->prev_frame);
1263     av_freep(&s->mvectors);
1264     s->mvectors_size = 0;
1265     av_freep(&s->wblocks);
1266     s->wblocks_size = 0;
1267     av_freep(&s->output);
1268     s->padded_output_size = 0;
1269     av_freep(&s->map);
1270     s->map_size = 0;
1271
1272     return 0;
1273 }
1274
1275 AVCodec ff_agm_decoder = {
1276     .name             = "agm",
1277     .long_name        = NULL_IF_CONFIG_SMALL("Amuse Graphics Movie"),
1278     .type             = AVMEDIA_TYPE_VIDEO,
1279     .id               = AV_CODEC_ID_AGM,
1280     .priv_data_size   = sizeof(AGMContext),
1281     .init             = decode_init,
1282     .close            = decode_close,
1283     .decode           = decode_frame,
1284     .flush            = decode_flush,
1285     .capabilities     = AV_CODEC_CAP_DR1,
1286     .caps_internal    = FF_CODEC_CAP_INIT_THREADSAFE |
1287                         FF_CODEC_CAP_INIT_CLEANUP |
1288                         FF_CODEC_CAP_EXPORTS_CROPPING,
1289 };