avcodec/atrac9dec: Check conditions before apply_band_extension() to avoid out of...
[ffmpeg.git] / libavcodec / atrac9dec.c
1 /*
2  * ATRAC9 decoder
3  * Copyright (c) 2018 Rostislav Pehlivanov <atomnuker@gmail.com>
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * FFmpeg is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with FFmpeg; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21
22 #include "internal.h"
23 #include "get_bits.h"
24 #include "fft.h"
25 #include "atrac9tab.h"
26 #include "libavutil/lfg.h"
27 #include "libavutil/float_dsp.h"
28
29 typedef struct ATRAC9ChannelData {
30     int band_ext;
31     int q_unit_cnt;
32     int band_ext_data[4];
33     int32_t scalefactors[31];
34     int32_t scalefactors_prev[31];
35
36     int precision_coarse[30];
37     int precision_fine[30];
38     int precision_mask[30];
39
40     int codebookset[30];
41
42     int32_t q_coeffs_coarse[256];
43     int32_t q_coeffs_fine[256];
44
45     DECLARE_ALIGNED(32, float, coeffs  )[256];
46     DECLARE_ALIGNED(32, float, prev_win)[128];
47 } ATRAC9ChannelData;
48
49 typedef struct ATRAC9BlockData {
50     ATRAC9ChannelData channel[2];
51
52     /* Base */
53     int band_count;
54     int q_unit_cnt;
55     int q_unit_cnt_prev;
56
57     /* Stereo block only */
58     int stereo_q_unit;
59
60     /* Band extension only */
61     int has_band_ext;
62     int has_band_ext_data;
63     int band_ext_q_unit;
64
65     /* Gradient */
66     int grad_mode;
67     int grad_boundary;
68     int gradient[31];
69
70     /* Stereo */
71     int cpe_base_channel;
72     int is_signs[30];
73
74     int reuseable;
75
76 } ATRAC9BlockData;
77
78 typedef struct ATRAC9Context {
79     AVCodecContext *avctx;
80     AVFloatDSPContext *fdsp;
81     FFTContext imdct;
82     ATRAC9BlockData block[5];
83     AVLFG lfg;
84
85     /* Set on init */
86     int frame_log2;
87     int avg_frame_size;
88     int frame_count;
89     int samplerate_idx;
90     const ATRAC9BlockConfig *block_config;
91
92     /* Generated on init */
93     VLC sf_vlc[2][8];            /* Signed/unsigned, length */
94     VLC coeff_vlc[2][8][4];      /* Cookbook, precision, cookbook index */
95     uint8_t alloc_curve[48][48];
96     DECLARE_ALIGNED(32, float, imdct_win)[256];
97
98     DECLARE_ALIGNED(32, float, temp)[256];
99 } ATRAC9Context;
100
101 static inline int parse_gradient(ATRAC9Context *s, ATRAC9BlockData *b,
102                                  GetBitContext *gb)
103 {
104     int grad_range[2];
105     int grad_value[2];
106     int values, sign, base;
107     uint8_t *curve;
108     float scale;
109
110     b->grad_mode = get_bits(gb, 2);
111     if (b->grad_mode) {
112         grad_range[0] = get_bits(gb, 5);
113         grad_range[1] = 31;
114         grad_value[0] = get_bits(gb, 5);
115         grad_value[1] = 31;
116     } else {
117         grad_range[0] = get_bits(gb, 6);
118         grad_range[1] = get_bits(gb, 6) + 1;
119         grad_value[0] = get_bits(gb, 5);
120         grad_value[1] = get_bits(gb, 5);
121     }
122     b->grad_boundary = get_bits(gb, 4);
123
124     if (grad_range[0] >= grad_range[1] || grad_range[1] > 47)
125         return AVERROR_INVALIDDATA;
126
127     if (grad_value[0] > 31 || grad_value[1] > 31)
128         return AVERROR_INVALIDDATA;
129
130     if (b->grad_boundary > b->q_unit_cnt)
131         return AVERROR_INVALIDDATA;
132
133     values    = grad_value[1] - grad_value[0];
134     sign      = 1 - 2*(values < 0);
135     base      = grad_value[0] + sign;
136     scale     = (FFABS(values) - 1) / 31.0f;
137     curve     = s->alloc_curve[grad_range[1] - grad_range[0] - 1];
138
139     for (int i = 0; i <= b->q_unit_cnt; i++)
140         b->gradient[i] = grad_value[i >= grad_range[0]];
141
142     for (int i = grad_range[0]; i < grad_range[1]; i++)
143         b->gradient[i] = base + sign*((int)(scale*curve[i - grad_range[0]]));
144
145     return 0;
146 }
147
148 static inline void calc_precision(ATRAC9Context *s, ATRAC9BlockData *b,
149                                   ATRAC9ChannelData *c)
150 {
151     memset(c->precision_mask, 0, sizeof(c->precision_mask));
152     for (int i = 1; i < b->q_unit_cnt; i++) {
153         const int delta = FFABS(c->scalefactors[i] - c->scalefactors[i - 1]) - 1;
154         if (delta > 0) {
155             const int neg = c->scalefactors[i - 1] > c->scalefactors[i];
156             c->precision_mask[i - neg] += FFMIN(delta, 5);
157         }
158     }
159
160     if (b->grad_mode) {
161         for (int i = 0; i < b->q_unit_cnt; i++) {
162             c->precision_coarse[i] = c->scalefactors[i];
163             c->precision_coarse[i] += c->precision_mask[i] - b->gradient[i];
164             if (c->precision_coarse[i] < 0)
165                 continue;
166             switch (b->grad_mode) {
167             case 1:
168                 c->precision_coarse[i] >>= 1;
169                 break;
170             case 2:
171                 c->precision_coarse[i] = (3 * c->precision_coarse[i]) >> 3;
172                 break;
173             case 3:
174                 c->precision_coarse[i] >>= 2;
175                 break;
176             }
177         }
178     } else {
179         for (int i = 0; i < b->q_unit_cnt; i++)
180             c->precision_coarse[i] = c->scalefactors[i] - b->gradient[i];
181     }
182
183
184     for (int i = 0; i < b->q_unit_cnt; i++)
185         c->precision_coarse[i] = FFMAX(c->precision_coarse[i], 1);
186
187     for (int i = 0; i < b->grad_boundary; i++)
188         c->precision_coarse[i]++;
189
190     for (int i = 0; i < b->q_unit_cnt; i++) {
191         c->precision_fine[i] = 0;
192         if (c->precision_coarse[i] > 15) {
193             c->precision_fine[i] = c->precision_coarse[i] - 15;
194             c->precision_coarse[i] = 15;
195         }
196     }
197 }
198
199 static inline int parse_band_ext(ATRAC9Context *s, ATRAC9BlockData *b,
200                                  GetBitContext *gb, int stereo)
201 {
202     int ext_band = 0;
203
204     if (b->has_band_ext) {
205         if (b->q_unit_cnt < 13)
206             return AVERROR_INVALIDDATA;
207         ext_band = at9_tab_band_ext_group[b->q_unit_cnt - 13][2];
208         if (stereo) {
209             b->channel[1].band_ext = get_bits(gb, 2);
210             b->channel[1].band_ext = ext_band > 2 ? b->channel[1].band_ext : 4;
211         } else {
212             skip_bits1(gb);
213         }
214     }
215
216     b->has_band_ext_data = get_bits1(gb);
217     if (!b->has_band_ext_data)
218         return 0;
219
220     if (!b->has_band_ext) {
221         skip_bits(gb, 2);
222         skip_bits_long(gb, get_bits(gb, 5));
223         return 0;
224     }
225
226     b->channel[0].band_ext = get_bits(gb, 2);
227     b->channel[0].band_ext = ext_band > 2 ? b->channel[0].band_ext : 4;
228
229     if (!get_bits(gb, 5))
230         return 0;
231
232     for (int i = 0; i <= stereo; i++) {
233         ATRAC9ChannelData *c = &b->channel[i];
234         const int count = at9_tab_band_ext_cnt[c->band_ext][ext_band];
235         for (int j = 0; j < count; j++) {
236             int len = at9_tab_band_ext_lengths[c->band_ext][ext_band][j];
237             c->band_ext_data[j] = get_bits(gb, len);
238         }
239     }
240
241     return 0;
242 }
243
244 static inline int read_scalefactors(ATRAC9Context *s, ATRAC9BlockData *b,
245                                     ATRAC9ChannelData *c, GetBitContext *gb,
246                                     int channel_idx, int first_in_pkt)
247 {
248     static const uint8_t mode_map[2][4] = { { 0, 1, 2, 3 }, { 0, 2, 3, 4 } };
249     const int mode = mode_map[channel_idx][get_bits(gb, 2)];
250
251     memset(c->scalefactors, 0, sizeof(c->scalefactors));
252
253     if (first_in_pkt && (mode == 4 || ((mode == 3) && !channel_idx))) {
254         av_log(s->avctx, AV_LOG_ERROR, "Invalid scalefactor coding mode!\n");
255         return AVERROR_INVALIDDATA;
256     }
257
258     switch (mode) {
259     case 0: { /* VLC delta offset */
260         const uint8_t *sf_weights = at9_tab_sf_weights[get_bits(gb, 3)];
261         const int base = get_bits(gb, 5);
262         const int len = get_bits(gb, 2) + 3;
263         const VLC *tab = &s->sf_vlc[0][len];
264
265         c->scalefactors[0] = get_bits(gb, len);
266
267         for (int i = 1; i < b->band_ext_q_unit; i++) {
268             int val = c->scalefactors[i - 1] + get_vlc2(gb, tab->table, 9, 2);
269             c->scalefactors[i] = val & ((1 << len) - 1);
270         }
271
272         for (int i = 0; i < b->band_ext_q_unit; i++)
273             c->scalefactors[i] += base - sf_weights[i];
274
275         break;
276     }
277     case 1: { /* CLC offset */
278         const int len = get_bits(gb, 2) + 2;
279         const int base = len < 5 ? get_bits(gb, 5) : 0;
280         for (int i = 0; i < b->band_ext_q_unit; i++)
281             c->scalefactors[i] = base + get_bits(gb, len);
282         break;
283     }
284     case 2:
285     case 4: { /* VLC dist to baseline */
286         const int *baseline = mode == 4 ? c->scalefactors_prev :
287                               channel_idx ? b->channel[0].scalefactors :
288                               c->scalefactors_prev;
289         const int baseline_len = mode == 4 ? b->q_unit_cnt_prev :
290                                  channel_idx ? b->band_ext_q_unit :
291                                  b->q_unit_cnt_prev;
292
293         const int len = get_bits(gb, 2) + 2;
294         const int unit_cnt = FFMIN(b->band_ext_q_unit, baseline_len);
295         const VLC *tab = &s->sf_vlc[1][len];
296
297         for (int i = 0; i < unit_cnt; i++) {
298             int dist = get_vlc2(gb, tab->table, 9, 2);
299             c->scalefactors[i] = baseline[i] + dist;
300         }
301
302         for (int i = unit_cnt; i < b->band_ext_q_unit; i++)
303             c->scalefactors[i] = get_bits(gb, 5);
304
305         break;
306     }
307     case 3: { /* VLC offset with baseline */
308         const int *baseline = channel_idx ? b->channel[0].scalefactors :
309                               c->scalefactors_prev;
310         const int baseline_len = channel_idx ? b->band_ext_q_unit :
311                                  b->q_unit_cnt_prev;
312
313         const int base = get_bits(gb, 5) - (1 << (5 - 1));
314         const int len = get_bits(gb, 2) + 1;
315         const int unit_cnt = FFMIN(b->band_ext_q_unit, baseline_len);
316         const VLC *tab = &s->sf_vlc[0][len];
317
318         c->scalefactors[0] = get_bits(gb, len);
319
320         for (int i = 1; i < unit_cnt; i++) {
321             int val = c->scalefactors[i - 1] + get_vlc2(gb, tab->table, 9, 2);
322             c->scalefactors[i] = val & ((1 << len) - 1);
323         }
324
325         for (int i = 0; i < unit_cnt; i++)
326             c->scalefactors[i] += base + baseline[i];
327
328         for (int i = unit_cnt; i < b->band_ext_q_unit; i++)
329             c->scalefactors[i] = get_bits(gb, 5);
330         break;
331     }
332     }
333
334     for (int i = 0; i < b->band_ext_q_unit; i++)
335         if (c->scalefactors[i] < 0 || c->scalefactors[i] > 31)
336             return AVERROR_INVALIDDATA;
337
338     memcpy(c->scalefactors_prev, c->scalefactors, sizeof(c->scalefactors));
339
340     return 0;
341 }
342
343 static inline void calc_codebook_idx(ATRAC9Context *s, ATRAC9BlockData *b,
344                                      ATRAC9ChannelData *c)
345 {
346     int avg = 0;
347     const int last_sf = c->scalefactors[c->q_unit_cnt];
348
349     memset(c->codebookset, 0, sizeof(c->codebookset));
350
351     if (c->q_unit_cnt <= 1)
352         return;
353     if (s->samplerate_idx > 7)
354         return;
355
356     c->scalefactors[c->q_unit_cnt] = c->scalefactors[c->q_unit_cnt - 1];
357
358     if (c->q_unit_cnt > 12) {
359         for (int i = 0; i < 12; i++)
360             avg += c->scalefactors[i];
361         avg = (avg + 6) / 12;
362     }
363
364     for (int i = 8; i < c->q_unit_cnt; i++) {
365         const int prev = c->scalefactors[i - 1];
366         const int cur  = c->scalefactors[i    ];
367         const int next = c->scalefactors[i + 1];
368         const int min  = FFMIN(prev, next);
369         if ((cur - min >= 3 || 2*cur - prev - next >= 3))
370             c->codebookset[i] = 1;
371     }
372
373
374     for (int i = 12; i < c->q_unit_cnt; i++) {
375         const int cur = c->scalefactors[i];
376         const int cnd = at9_q_unit_to_coeff_cnt[i] == 16;
377         const int min = FFMIN(c->scalefactors[i + 1], c->scalefactors[i - 1]);
378         if (c->codebookset[i])
379             continue;
380
381         c->codebookset[i] = (((cur - min) >= 2) && (cur >= (avg - cnd)));
382     }
383
384     c->scalefactors[c->q_unit_cnt] = last_sf;
385 }
386
387 static inline void read_coeffs_coarse(ATRAC9Context *s, ATRAC9BlockData *b,
388                                       ATRAC9ChannelData *c, GetBitContext *gb)
389 {
390     const int max_prec = s->samplerate_idx > 7 ? 1 : 7;
391
392     memset(c->q_coeffs_coarse, 0, sizeof(c->q_coeffs_coarse));
393
394     for (int i = 0; i < c->q_unit_cnt; i++) {
395         int *coeffs = &c->q_coeffs_coarse[at9_q_unit_to_coeff_idx[i]];
396         const int bands = at9_q_unit_to_coeff_cnt[i];
397         const int prec = c->precision_coarse[i] + 1;
398
399         if (prec <= max_prec) {
400             const int cb = c->codebookset[i];
401             const int cbi = at9_q_unit_to_codebookidx[i];
402             const VLC *tab = &s->coeff_vlc[cb][prec][cbi];
403             const HuffmanCodebook *huff = &at9_huffman_coeffs[cb][prec][cbi];
404             const int groups = bands >> huff->value_cnt_pow;
405
406             for (int j = 0; j < groups; j++) {
407                 uint16_t val = get_vlc2(gb, tab->table, 9, huff->max_bit_size);
408
409                 for (int k = 0; k < huff->value_cnt; k++) {
410                     coeffs[k] = sign_extend(val, huff->value_bits);
411                     val >>= huff->value_bits;
412                 }
413
414                 coeffs += huff->value_cnt;
415             }
416         } else {
417             for (int j = 0; j < bands; j++)
418                 coeffs[j] = sign_extend(get_bits(gb, prec), prec);
419         }
420     }
421 }
422
423 static inline void read_coeffs_fine(ATRAC9Context *s, ATRAC9BlockData *b,
424                                     ATRAC9ChannelData *c, GetBitContext *gb)
425 {
426     memset(c->q_coeffs_fine, 0, sizeof(c->q_coeffs_fine));
427
428     for (int i = 0; i < c->q_unit_cnt; i++) {
429         const int start = at9_q_unit_to_coeff_idx[i + 0];
430         const int end   = at9_q_unit_to_coeff_idx[i + 1];
431         const int len   = c->precision_fine[i] + 1;
432
433         if (c->precision_fine[i] <= 0)
434             continue;
435
436         for (int j = start; j < end; j++)
437             c->q_coeffs_fine[j] = sign_extend(get_bits(gb, len), len);
438     }
439 }
440
441 static inline void dequantize(ATRAC9Context *s, ATRAC9BlockData *b,
442                               ATRAC9ChannelData *c)
443 {
444     memset(c->coeffs, 0, sizeof(c->coeffs));
445
446     for (int i = 0; i < c->q_unit_cnt; i++) {
447         const int start = at9_q_unit_to_coeff_idx[i + 0];
448         const int end   = at9_q_unit_to_coeff_idx[i + 1];
449
450         const float coarse_c = at9_quant_step_coarse[c->precision_coarse[i]];
451         const float fine_c   = at9_quant_step_fine[c->precision_fine[i]];
452
453         for (int j = start; j < end; j++) {
454             const float vc = c->q_coeffs_coarse[j] * coarse_c;
455             const float vf = c->q_coeffs_fine[j]   * fine_c;
456             c->coeffs[j] = vc + vf;
457         }
458     }
459 }
460
461 static inline void apply_intensity_stereo(ATRAC9Context *s, ATRAC9BlockData *b,
462                                           const int stereo)
463 {
464     float *src = b->channel[ b->cpe_base_channel].coeffs;
465     float *dst = b->channel[!b->cpe_base_channel].coeffs;
466
467     if (!stereo)
468         return;
469
470     if (b->q_unit_cnt <= b->stereo_q_unit)
471         return;
472
473     for (int i = b->stereo_q_unit; i < b->q_unit_cnt; i++) {
474         const int sign  = b->is_signs[i];
475         const int start = at9_q_unit_to_coeff_idx[i + 0];
476         const int end   = at9_q_unit_to_coeff_idx[i + 1];
477         for (int j = start; j < end; j++)
478             dst[j] = sign*src[j];
479     }
480 }
481
482 static inline void apply_scalefactors(ATRAC9Context *s, ATRAC9BlockData *b,
483                                       const int stereo)
484 {
485     for (int i = 0; i <= stereo; i++) {
486         float *coeffs = b->channel[i].coeffs;
487         for (int j = 0; j < b->q_unit_cnt; j++) {
488             const int start = at9_q_unit_to_coeff_idx[j + 0];
489             const int end   = at9_q_unit_to_coeff_idx[j + 1];
490             const int scalefactor = b->channel[i].scalefactors[j];
491             const float scale = at9_scalefactor_c[scalefactor];
492             for (int k = start; k < end; k++)
493                 coeffs[k] *= scale;
494         }
495     }
496 }
497
498 static inline void fill_with_noise(ATRAC9Context *s, ATRAC9ChannelData *c,
499                                    int start, int count)
500 {
501     float maxval = 0.0f;
502     for (int i = 0; i < count; i += 2) {
503         double tmp[2];
504         av_bmg_get(&s->lfg, tmp);
505         c->coeffs[start + i + 0] = tmp[0];
506         c->coeffs[start + i + 1] = tmp[1];
507         maxval = FFMAX(FFMAX(FFABS(tmp[0]), FFABS(tmp[1])), maxval);
508     }
509     /* Normalize */
510     for (int i = 0; i < count; i++)
511         c->coeffs[start + i] /= maxval;
512 }
513
514 static inline void scale_band_ext_coeffs(ATRAC9ChannelData *c, float sf[6],
515                                          const int s_unit, const int e_unit)
516 {
517     for (int i = s_unit; i < e_unit; i++) {
518         const int start = at9_q_unit_to_coeff_idx[i + 0];
519         const int end   = at9_q_unit_to_coeff_idx[i + 1];
520         for (int j = start; j < end; j++)
521             c->coeffs[j] *= sf[i - s_unit];
522     }
523 }
524
525 static inline void apply_band_extension(ATRAC9Context *s, ATRAC9BlockData *b,
526                                        const int stereo)
527 {
528     const int g_units[4] = { /* A, B, C, total units */
529         b->q_unit_cnt,
530         at9_tab_band_ext_group[b->q_unit_cnt - 13][0],
531         at9_tab_band_ext_group[b->q_unit_cnt - 13][1],
532         FFMAX(g_units[2], 22),
533     };
534
535     const int g_bins[4] = { /* A, B, C, total bins */
536         at9_q_unit_to_coeff_idx[g_units[0]],
537         at9_q_unit_to_coeff_idx[g_units[1]],
538         at9_q_unit_to_coeff_idx[g_units[2]],
539         at9_q_unit_to_coeff_idx[g_units[3]],
540     };
541
542     for (int ch = 0; ch <= stereo; ch++) {
543         ATRAC9ChannelData *c = &b->channel[ch];
544
545         /* Mirror the spectrum */
546         for (int i = 0; i < 3; i++)
547             for (int j = 0; j < (g_bins[i + 1] - g_bins[i + 0]); j++)
548                 c->coeffs[g_bins[i] + j] = c->coeffs[g_bins[i] - j - 1];
549
550         switch (c->band_ext) {
551         case 0: {
552             float sf[6] = { 0.0f };
553             const int l = g_units[3] - g_units[0] - 1;
554             const int n_start = at9_q_unit_to_coeff_idx[g_units[3] - 1];
555             const int n_cnt   = at9_q_unit_to_coeff_cnt[g_units[3] - 1];
556             switch (at9_tab_band_ext_group[b->q_unit_cnt - 13][2]) {
557             case 3:
558                 sf[0] = at9_band_ext_scales_m0[0][0][c->band_ext_data[0]];
559                 sf[1] = at9_band_ext_scales_m0[0][1][c->band_ext_data[0]];
560                 sf[2] = at9_band_ext_scales_m0[0][2][c->band_ext_data[1]];
561                 sf[3] = at9_band_ext_scales_m0[0][3][c->band_ext_data[2]];
562                 sf[4] = at9_band_ext_scales_m0[0][4][c->band_ext_data[3]];
563                 break;
564             case 4:
565                 sf[0] = at9_band_ext_scales_m0[1][0][c->band_ext_data[0]];
566                 sf[1] = at9_band_ext_scales_m0[1][1][c->band_ext_data[0]];
567                 sf[2] = at9_band_ext_scales_m0[1][2][c->band_ext_data[1]];
568                 sf[3] = at9_band_ext_scales_m0[1][3][c->band_ext_data[2]];
569                 sf[4] = at9_band_ext_scales_m0[1][4][c->band_ext_data[3]];
570                 break;
571             case 5:
572                 sf[0] = at9_band_ext_scales_m0[2][0][c->band_ext_data[0]];
573                 sf[1] = at9_band_ext_scales_m0[2][1][c->band_ext_data[1]];
574                 sf[2] = at9_band_ext_scales_m0[2][2][c->band_ext_data[1]];
575                 break;
576             }
577
578             sf[l] = at9_scalefactor_c[c->scalefactors[g_units[0]]];
579
580             fill_with_noise(s, c, n_start, n_cnt);
581             scale_band_ext_coeffs(c, sf, g_units[0], g_units[3]);
582             break;
583         }
584         case 1: {
585             float sf[6];
586             for (int i = g_units[0]; i < g_units[3]; i++)
587                 sf[i - g_units[0]] = at9_scalefactor_c[c->scalefactors[i]];
588
589             fill_with_noise(s, c, g_bins[0], g_bins[3] - g_bins[0]);
590             scale_band_ext_coeffs(c, sf, g_units[0], g_units[3]);
591             break;
592         }
593         case 2: {
594             const float g_sf[2] = {
595                 at9_band_ext_scales_m2[c->band_ext_data[0]],
596                 at9_band_ext_scales_m2[c->band_ext_data[1]],
597             };
598
599             for (int i = 0; i < 2; i++)
600                 for (int j = g_bins[i + 0]; j < g_bins[i + 1]; j++)
601                     c->coeffs[j] *= g_sf[i];
602             break;
603         }
604         case 3: {
605             float scale = at9_band_ext_scales_m3[c->band_ext_data[0]][0];
606             float rate  = at9_band_ext_scales_m3[c->band_ext_data[1]][1];
607             rate = pow(2, rate);
608             for (int i = g_bins[0]; i < g_bins[3]; i++) {
609                 scale *= rate;
610                 c->coeffs[i] *= scale;
611             }
612             break;
613         }
614         case 4: {
615             const float m = at9_band_ext_scales_m4[c->band_ext_data[0]];
616             const float g_sf[3] = { 0.7079468f*m, 0.5011902f*m, 0.3548279f*m };
617
618             for (int i = 0; i < 3; i++)
619                 for (int j = g_bins[i + 0]; j < g_bins[i + 1]; j++)
620                     c->coeffs[j] *= g_sf[i];
621             break;
622         }
623         }
624     }
625 }
626
627 static int atrac9_decode_block(ATRAC9Context *s, GetBitContext *gb,
628                                ATRAC9BlockData *b, AVFrame *frame,
629                                int frame_idx, int block_idx)
630 {
631     const int first_in_pkt = !get_bits1(gb);
632     const int reuse_params =  get_bits1(gb);
633     const int stereo = s->block_config->type[block_idx] == ATRAC9_BLOCK_TYPE_CPE;
634
635     if (s->block_config->type[block_idx] == ATRAC9_BLOCK_TYPE_LFE) {
636         ATRAC9ChannelData *c = &b->channel[0];
637         const int precision = reuse_params ? 8 : 4;
638         c->q_unit_cnt = b->q_unit_cnt = 2;
639
640         memset(c->scalefactors, 0, sizeof(c->scalefactors));
641         memset(c->q_coeffs_fine, 0, sizeof(c->q_coeffs_fine));
642         memset(c->q_coeffs_coarse, 0, sizeof(c->q_coeffs_coarse));
643
644         for (int i = 0; i < b->q_unit_cnt; i++) {
645             c->scalefactors[i] = get_bits(gb, 5);
646             c->precision_coarse[i] = precision;
647             c->precision_fine[i] = 0;
648         }
649
650         for (int i = 0; i < c->q_unit_cnt; i++) {
651             const int start = at9_q_unit_to_coeff_idx[i + 0];
652             const int end   = at9_q_unit_to_coeff_idx[i + 1];
653             for (int j = start; j < end; j++)
654                 c->q_coeffs_coarse[j] = get_bits(gb, c->precision_coarse[i] + 1);
655         }
656
657         dequantize        (s, b, c);
658         apply_scalefactors(s, b, 0);
659
660         goto imdct;
661     }
662
663     if (first_in_pkt && reuse_params) {
664         av_log(s->avctx, AV_LOG_ERROR, "Invalid block flags!\n");
665         return AVERROR_INVALIDDATA;
666     }
667
668     /* Band parameters */
669     if (!reuse_params) {
670         int stereo_band, ext_band;
671         const int min_band_count = s->samplerate_idx > 7 ? 1 : 3;
672         b->reuseable = 0;
673         b->band_count = get_bits(gb, 4) + min_band_count;
674         b->q_unit_cnt = at9_tab_band_q_unit_map[b->band_count];
675
676         b->band_ext_q_unit = b->stereo_q_unit = b->q_unit_cnt;
677
678         if (b->band_count > at9_tab_sri_max_bands[s->samplerate_idx]) {
679             av_log(s->avctx, AV_LOG_ERROR, "Invalid band count %i!\n",
680                    b->band_count);
681             return AVERROR_INVALIDDATA;
682         }
683
684         if (stereo) {
685             stereo_band = get_bits(gb, 4) + min_band_count;
686             if (stereo_band > b->band_count) {
687                 av_log(s->avctx, AV_LOG_ERROR, "Invalid stereo band %i!\n",
688                        stereo_band);
689                 return AVERROR_INVALIDDATA;
690             }
691             b->stereo_q_unit = at9_tab_band_q_unit_map[stereo_band];
692         }
693
694         b->has_band_ext = get_bits1(gb);
695         if (b->has_band_ext) {
696             ext_band = get_bits(gb, 4) + min_band_count;
697             if (ext_band < b->band_count) {
698                 av_log(s->avctx, AV_LOG_ERROR, "Invalid extension band %i!\n",
699                        ext_band);
700                 return AVERROR_INVALIDDATA;
701             }
702             b->band_ext_q_unit = at9_tab_band_q_unit_map[ext_band];
703         }
704         b->reuseable = 1;
705     }
706     if (!b->reuseable) {
707         av_log(s->avctx, AV_LOG_ERROR, "invalid block reused!\n");
708         return AVERROR_INVALIDDATA;
709     }
710
711     /* Calculate bit alloc gradient */
712     if (parse_gradient(s, b, gb))
713         return AVERROR_INVALIDDATA;
714
715     /* IS data */
716     b->cpe_base_channel = 0;
717     if (stereo) {
718         b->cpe_base_channel = get_bits1(gb);
719         if (get_bits1(gb)) {
720             for (int i = b->stereo_q_unit; i < b->q_unit_cnt; i++)
721                 b->is_signs[i] = 1 - 2*get_bits1(gb);
722         } else {
723             for (int i = 0; i < FF_ARRAY_ELEMS(b->is_signs); i++)
724                 b->is_signs[i] = 1;
725         }
726     }
727
728     /* Band extension */
729     if (parse_band_ext(s, b, gb, stereo))
730         return AVERROR_INVALIDDATA;
731
732     /* Scalefactors */
733     for (int i = 0; i <= stereo; i++) {
734         ATRAC9ChannelData *c = &b->channel[i];
735         c->q_unit_cnt = i == b->cpe_base_channel ? b->q_unit_cnt :
736                                                    b->stereo_q_unit;
737         if (read_scalefactors(s, b, c, gb, i, first_in_pkt))
738             return AVERROR_INVALIDDATA;
739
740         calc_precision    (s, b, c);
741         calc_codebook_idx (s, b, c);
742         read_coeffs_coarse(s, b, c, gb);
743         read_coeffs_fine  (s, b, c, gb);
744         dequantize        (s, b, c);
745     }
746
747     b->q_unit_cnt_prev = b->has_band_ext ? b->band_ext_q_unit : b->q_unit_cnt;
748
749     apply_intensity_stereo(s, b, stereo);
750     apply_scalefactors    (s, b, stereo);
751
752     if (b->has_band_ext && b->has_band_ext_data)
753         apply_band_extension  (s, b, stereo);
754
755 imdct:
756     for (int i = 0; i <= stereo; i++) {
757         ATRAC9ChannelData *c = &b->channel[i];
758         const int dst_idx = s->block_config->plane_map[block_idx][i];
759         const int wsize = 1 << s->frame_log2;
760         const ptrdiff_t offset = wsize*frame_idx*sizeof(float);
761         float *dst = (float *)(frame->extended_data[dst_idx] + offset);
762
763         s->imdct.imdct_half(&s->imdct, s->temp, c->coeffs);
764         s->fdsp->vector_fmul_window(dst, c->prev_win, s->temp,
765                                     s->imdct_win, wsize >> 1);
766         memcpy(c->prev_win, s->temp + (wsize >> 1), sizeof(float)*wsize >> 1);
767     }
768
769     return 0;
770 }
771
772 static int atrac9_decode_frame(AVCodecContext *avctx, void *data,
773                                int *got_frame_ptr, AVPacket *avpkt)
774 {
775     int ret;
776     GetBitContext gb;
777     AVFrame *frame = data;
778     ATRAC9Context *s = avctx->priv_data;
779     const int frames = FFMIN(avpkt->size / s->avg_frame_size, s->frame_count);
780
781     frame->nb_samples = (1 << s->frame_log2) * frames;
782     ret = ff_get_buffer(avctx, frame, 0);
783     if (ret < 0)
784         return ret;
785
786     init_get_bits8(&gb, avpkt->data, avpkt->size);
787
788     for (int i = 0; i < frames; i++) {
789         for (int j = 0; j < s->block_config->count; j++) {
790             ret = atrac9_decode_block(s, &gb, &s->block[j], frame, i, j);
791             if (ret)
792                 return ret;
793             align_get_bits(&gb);
794         }
795     }
796
797     *got_frame_ptr = 1;
798
799     return avctx->block_align;
800 }
801
802 static void atrac9_decode_flush(AVCodecContext *avctx)
803 {
804     ATRAC9Context *s = avctx->priv_data;
805
806     for (int j = 0; j < s->block_config->count; j++) {
807         ATRAC9BlockData *b = &s->block[j];
808         const int stereo = s->block_config->type[j] == ATRAC9_BLOCK_TYPE_CPE;
809         for (int i = 0; i <= stereo; i++) {
810             ATRAC9ChannelData *c = &b->channel[i];
811             memset(c->prev_win, 0, sizeof(c->prev_win));
812         }
813     }
814 }
815
816 static av_cold int atrac9_decode_close(AVCodecContext *avctx)
817 {
818     ATRAC9Context *s = avctx->priv_data;
819
820     for (int i = 1; i < 7; i++)
821         ff_free_vlc(&s->sf_vlc[0][i]);
822     for (int i = 2; i < 6; i++)
823         ff_free_vlc(&s->sf_vlc[1][i]);
824     for (int i = 0; i < 2; i++)
825         for (int j = 0; j < 8; j++)
826             for (int k = 0; k < 4; k++)
827                 ff_free_vlc(&s->coeff_vlc[i][j][k]);
828
829     ff_mdct_end(&s->imdct);
830     av_free(s->fdsp);
831
832     return 0;
833 }
834
835 static av_cold int atrac9_decode_init(AVCodecContext *avctx)
836 {
837     GetBitContext gb;
838     ATRAC9Context *s = avctx->priv_data;
839     int version, block_config_idx, superframe_idx, alloc_c_len;
840
841     s->avctx = avctx;
842
843     av_lfg_init(&s->lfg, 0xFBADF00D);
844
845     if (avctx->extradata_size != 12) {
846         av_log(avctx, AV_LOG_ERROR, "Invalid extradata length!\n");
847         return AVERROR_INVALIDDATA;
848     }
849
850     version = AV_RL32(avctx->extradata);
851     if (version > 2) {
852         av_log(avctx, AV_LOG_ERROR, "Unsupported version (%i)!\n", version);
853         return AVERROR_INVALIDDATA;
854     }
855
856     init_get_bits8(&gb, avctx->extradata + 4, avctx->extradata_size);
857
858     if (get_bits(&gb, 8) != 0xFE) {
859         av_log(avctx, AV_LOG_ERROR, "Incorrect magic byte!\n");
860         return AVERROR_INVALIDDATA;
861     }
862
863     s->samplerate_idx = get_bits(&gb, 4);
864     avctx->sample_rate = at9_tab_samplerates[s->samplerate_idx];
865
866     block_config_idx = get_bits(&gb, 3);
867     if (block_config_idx > 5) {
868         av_log(avctx, AV_LOG_ERROR, "Incorrect block config!\n");
869         return AVERROR_INVALIDDATA;
870     }
871     s->block_config = &at9_block_layout[block_config_idx];
872
873     avctx->channel_layout = s->block_config->channel_layout;
874     avctx->sample_fmt     = AV_SAMPLE_FMT_FLTP;
875
876     if (get_bits1(&gb)) {
877         av_log(avctx, AV_LOG_ERROR, "Incorrect verification bit!\n");
878         return AVERROR_INVALIDDATA;
879     }
880
881     /* Average frame size in bytes */
882     s->avg_frame_size = get_bits(&gb, 11) + 1;
883
884     superframe_idx = get_bits(&gb, 2);
885     if (superframe_idx & 1) {
886         av_log(avctx, AV_LOG_ERROR, "Invalid superframe index!\n");
887         return AVERROR_INVALIDDATA;
888     }
889
890     s->frame_count = 1 << superframe_idx;
891     s->frame_log2  = at9_tab_sri_frame_log2[s->samplerate_idx];
892
893     if (ff_mdct_init(&s->imdct, s->frame_log2 + 1, 1, 1.0f / 32768.0f))
894         return AVERROR(ENOMEM);
895
896     s->fdsp = avpriv_float_dsp_alloc(avctx->flags & AV_CODEC_FLAG_BITEXACT);
897     if (!s->fdsp)
898         return AVERROR(ENOMEM);
899
900     /* iMDCT window */
901     for (int i = 0; i < (1 << s->frame_log2); i++) {
902         const int   len  = 1 << s->frame_log2;
903         const float sidx = (      i + 0.5f) / len;
904         const float eidx = (len - i - 0.5f) / len;
905         const float s_c  = sinf(sidx*M_PI - M_PI_2)*0.5f + 0.5f;
906         const float e_c  = sinf(eidx*M_PI - M_PI_2)*0.5f + 0.5f;
907         s->imdct_win[i]  = s_c / ((s_c * s_c) + (e_c * e_c));
908     }
909
910     /* Allocation curve */
911     alloc_c_len = FF_ARRAY_ELEMS(at9_tab_b_dist);
912     for (int i = 1; i <= alloc_c_len; i++)
913         for (int j = 0; j < i; j++)
914             s->alloc_curve[i - 1][j] = at9_tab_b_dist[(j * alloc_c_len) / i];
915
916     /* Unsigned scalefactor VLCs */
917     for (int i = 1; i < 7; i++) {
918         const HuffmanCodebook *hf = &at9_huffman_sf_unsigned[i];
919
920         init_vlc(&s->sf_vlc[0][i], 9, hf->size, hf->bits, 1, 1, hf->codes,
921                  2, 2, 0);
922     }
923
924     /* Signed scalefactor VLCs */
925     for (int i = 2; i < 6; i++) {
926         const HuffmanCodebook *hf = &at9_huffman_sf_signed[i];
927
928         int nums = hf->size;
929         int16_t sym[32];
930         for (int j = 0; j < nums; j++)
931             sym[j] = sign_extend(j, hf->value_bits);
932
933         ff_init_vlc_sparse(&s->sf_vlc[1][i], 9, hf->size, hf->bits, 1, 1,
934                            hf->codes, 2, 2, sym, sizeof(*sym), sizeof(*sym), 0);
935     }
936
937     /* Coefficient VLCs */
938     for (int i = 0; i < 2; i++) {
939         for (int j = 0; j < 8; j++) {
940             for (int k = 0; k < 4; k++) {
941                 const HuffmanCodebook *hf = &at9_huffman_coeffs[i][j][k];
942                 init_vlc(&s->coeff_vlc[i][j][k], 9, hf->size, hf->bits, 1, 1,
943                          hf->codes, 2, 2, 0);
944             }
945         }
946     }
947
948     return 0;
949 }
950
951 AVCodec ff_atrac9_decoder = {
952     .name           = "atrac9",
953     .long_name      = NULL_IF_CONFIG_SMALL("ATRAC9 (Adaptive TRansform Acoustic Coding 9)"),
954     .type           = AVMEDIA_TYPE_AUDIO,
955     .id             = AV_CODEC_ID_ATRAC9,
956     .priv_data_size = sizeof(ATRAC9Context),
957     .init           = atrac9_decode_init,
958     .close          = atrac9_decode_close,
959     .decode         = atrac9_decode_frame,
960     .flush          = atrac9_decode_flush,
961     .caps_internal  = FF_CODEC_CAP_INIT_THREADSAFE | FF_CODEC_CAP_INIT_CLEANUP,
962     .capabilities   = AV_CODEC_CAP_SUBFRAMES | AV_CODEC_CAP_DR1,
963 };