nellymoserenc: improve error checking in encode_init()
[ffmpeg.git] / libavcodec / nellymoserenc.c
1 /*
2  * Nellymoser encoder
3  * This code is developed as part of Google Summer of Code 2008 Program.
4  *
5  * Copyright (c) 2008 Bartlomiej Wolowiec
6  *
7  * This file is part of Libav.
8  *
9  * Libav is free software; you can redistribute it and/or
10  * modify it under the terms of the GNU Lesser General Public
11  * License as published by the Free Software Foundation; either
12  * version 2.1 of the License, or (at your option) any later version.
13  *
14  * Libav is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
17  * Lesser General Public License for more details.
18  *
19  * You should have received a copy of the GNU Lesser General Public
20  * License along with Libav; if not, write to the Free Software
21  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
22  */
23
24 /**
25  * @file
26  * Nellymoser encoder
27  * by Bartlomiej Wolowiec
28  *
29  * Generic codec information: libavcodec/nellymoserdec.c
30  *
31  * Some information also from: http://samples.libav.org/A-codecs/Nelly_Moser/ASAO/ASAO.zip
32  *                             (Copyright Joseph Artsimovich and UAB "DKD")
33  *
34  * for more information about nellymoser format, visit:
35  * http://wiki.multimedia.cx/index.php?title=Nellymoser
36  */
37
38 #include "libavutil/mathematics.h"
39 #include "nellymoser.h"
40 #include "avcodec.h"
41 #include "dsputil.h"
42 #include "fft.h"
43 #include "sinewin.h"
44
45 #define BITSTREAM_WRITER_LE
46 #include "put_bits.h"
47
48 #define POW_TABLE_SIZE (1<<11)
49 #define POW_TABLE_OFFSET 3
50 #define OPT_SIZE ((1<<15) + 3000)
51
52 typedef struct NellyMoserEncodeContext {
53     AVCodecContext  *avctx;
54     int             last_frame;
55     int             bufsel;
56     int             have_saved;
57     DSPContext      dsp;
58     FFTContext      mdct_ctx;
59     DECLARE_ALIGNED(32, float, mdct_out)[NELLY_SAMPLES];
60     DECLARE_ALIGNED(32, float, in_buff)[NELLY_SAMPLES];
61     DECLARE_ALIGNED(32, float, buf)[2][3 * NELLY_BUF_LEN];     ///< sample buffer
62     float           (*opt )[NELLY_BANDS];
63     uint8_t         (*path)[NELLY_BANDS];
64 } NellyMoserEncodeContext;
65
66 static float pow_table[POW_TABLE_SIZE];     ///< -pow(2, -i / 2048.0 - 3.0);
67
68 static const uint8_t sf_lut[96] = {
69      0,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  3,  3,  3,  4,  4,
70      5,  5,  5,  6,  7,  7,  8,  8,  9, 10, 11, 11, 12, 13, 13, 14,
71     15, 15, 16, 17, 17, 18, 19, 19, 20, 21, 22, 22, 23, 24, 25, 26,
72     27, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 37, 38, 39, 40,
73     41, 41, 42, 43, 44, 45, 45, 46, 47, 48, 49, 50, 51, 52, 52, 53,
74     54, 55, 55, 56, 57, 57, 58, 59, 59, 60, 60, 60, 61, 61, 61, 62,
75 };
76
77 static const uint8_t sf_delta_lut[78] = {
78      0,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  3,  3,  3,  4,  4,
79      4,  5,  5,  5,  6,  6,  7,  7,  8,  8,  9, 10, 10, 11, 11, 12,
80     13, 13, 14, 15, 16, 17, 17, 18, 19, 19, 20, 21, 21, 22, 22, 23,
81     23, 24, 24, 25, 25, 25, 26, 26, 26, 26, 27, 27, 27, 27, 27, 28,
82     28, 28, 28, 28, 28, 29, 29, 29, 29, 29, 29, 29, 29, 30,
83 };
84
85 static const uint8_t quant_lut[230] = {
86      0,
87
88      0,  1,  2,
89
90      0,  1,  2,  3,  4,  5,  6,
91
92      0,  1,  1,  2,  2,  3,  3,  4,  5,  6,  7,  8,  9, 10, 11, 11,
93     12, 13, 13, 13, 14,
94
95      0,  1,  1,  2,  2,  2,  3,  3,  4,  4,  5,  5,  6,  6,  7,  8,
96      8,  9, 10, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
97     22, 23, 23, 24, 24, 25, 25, 26, 26, 27, 27, 28, 28, 29, 29, 29,
98     30,
99
100      0,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  2,  3,  3,  3,  3,
101      4,  4,  4,  5,  5,  5,  6,  6,  7,  7,  7,  8,  8,  9,  9,  9,
102     10, 10, 11, 11, 11, 12, 12, 13, 13, 13, 13, 14, 14, 14, 15, 15,
103     15, 15, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 20, 20, 20,
104     21, 21, 22, 22, 23, 23, 24, 25, 26, 26, 27, 28, 29, 30, 31, 32,
105     33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 42, 43, 44, 44, 45, 45,
106     46, 47, 47, 48, 48, 49, 49, 50, 50, 50, 51, 51, 51, 52, 52, 52,
107     53, 53, 53, 54, 54, 54, 55, 55, 55, 56, 56, 56, 57, 57, 57, 57,
108     58, 58, 58, 58, 59, 59, 59, 59, 60, 60, 60, 60, 60, 61, 61, 61,
109     61, 61, 61, 61, 62,
110 };
111
112 static const float quant_lut_mul[7] = { 0.0,  0.0,  2.0,  2.0,  5.0, 12.0,  36.6 };
113 static const float quant_lut_add[7] = { 0.0,  0.0,  2.0,  7.0, 21.0, 56.0, 157.0 };
114 static const uint8_t quant_lut_offset[8] = { 0, 0, 1, 4, 11, 32, 81, 230 };
115
116 static void apply_mdct(NellyMoserEncodeContext *s)
117 {
118     s->dsp.vector_fmul(s->in_buff, s->buf[s->bufsel], ff_sine_128, NELLY_BUF_LEN);
119     s->dsp.vector_fmul_reverse(s->in_buff + NELLY_BUF_LEN, s->buf[s->bufsel] + NELLY_BUF_LEN, ff_sine_128,
120                                NELLY_BUF_LEN);
121     s->mdct_ctx.mdct_calc(&s->mdct_ctx, s->mdct_out, s->in_buff);
122
123     s->dsp.vector_fmul(s->buf[s->bufsel] + NELLY_BUF_LEN, s->buf[s->bufsel] + NELLY_BUF_LEN,
124                        ff_sine_128, NELLY_BUF_LEN);
125     s->dsp.vector_fmul_reverse(s->buf[s->bufsel] + 2 * NELLY_BUF_LEN, s->buf[1 - s->bufsel], ff_sine_128,
126                                NELLY_BUF_LEN);
127     s->mdct_ctx.mdct_calc(&s->mdct_ctx, s->mdct_out + NELLY_BUF_LEN, s->buf[s->bufsel] + NELLY_BUF_LEN);
128 }
129
130 static av_cold int encode_end(AVCodecContext *avctx)
131 {
132     NellyMoserEncodeContext *s = avctx->priv_data;
133
134     ff_mdct_end(&s->mdct_ctx);
135
136     if (s->avctx->trellis) {
137         av_free(s->opt);
138         av_free(s->path);
139     }
140
141     return 0;
142 }
143
144 static av_cold int encode_init(AVCodecContext *avctx)
145 {
146     NellyMoserEncodeContext *s = avctx->priv_data;
147     int i, ret;
148
149     if (avctx->channels != 1) {
150         av_log(avctx, AV_LOG_ERROR, "Nellymoser supports only 1 channel\n");
151         return AVERROR(EINVAL);
152     }
153
154     if (avctx->sample_rate != 8000 && avctx->sample_rate != 16000 &&
155         avctx->sample_rate != 11025 &&
156         avctx->sample_rate != 22050 && avctx->sample_rate != 44100 &&
157         avctx->strict_std_compliance >= FF_COMPLIANCE_NORMAL) {
158         av_log(avctx, AV_LOG_ERROR, "Nellymoser works only with 8000, 16000, 11025, 22050 and 44100 sample rate\n");
159         return AVERROR(EINVAL);
160     }
161
162     avctx->frame_size = NELLY_SAMPLES;
163     s->avctx = avctx;
164     if ((ret = ff_mdct_init(&s->mdct_ctx, 8, 0, 32768.0)) < 0)
165         goto error;
166     ff_dsputil_init(&s->dsp, avctx);
167
168     /* Generate overlap window */
169     ff_sine_window_init(ff_sine_128, 128);
170     for (i = 0; i < POW_TABLE_SIZE; i++)
171         pow_table[i] = -pow(2, -i / 2048.0 - 3.0 + POW_TABLE_OFFSET);
172
173     if (s->avctx->trellis) {
174         s->opt  = av_malloc(NELLY_BANDS * OPT_SIZE * sizeof(float  ));
175         s->path = av_malloc(NELLY_BANDS * OPT_SIZE * sizeof(uint8_t));
176         if (!s->opt || !s->path) {
177             ret = AVERROR(ENOMEM);
178             goto error;
179         }
180     }
181
182     return 0;
183 error:
184     encode_end(avctx);
185     return ret;
186 }
187
188 #define find_best(val, table, LUT, LUT_add, LUT_size) \
189     best_idx = \
190         LUT[av_clip ((lrintf(val) >> 8) + LUT_add, 0, LUT_size - 1)]; \
191     if (fabs(val - table[best_idx]) > fabs(val - table[best_idx + 1])) \
192         best_idx++;
193
194 static void get_exponent_greedy(NellyMoserEncodeContext *s, float *cand, int *idx_table)
195 {
196     int band, best_idx, power_idx = 0;
197     float power_candidate;
198
199     //base exponent
200     find_best(cand[0], ff_nelly_init_table, sf_lut, -20, 96);
201     idx_table[0] = best_idx;
202     power_idx = ff_nelly_init_table[best_idx];
203
204     for (band = 1; band < NELLY_BANDS; band++) {
205         power_candidate = cand[band] - power_idx;
206         find_best(power_candidate, ff_nelly_delta_table, sf_delta_lut, 37, 78);
207         idx_table[band] = best_idx;
208         power_idx += ff_nelly_delta_table[best_idx];
209     }
210 }
211
212 static inline float distance(float x, float y, int band)
213 {
214     //return pow(fabs(x-y), 2.0);
215     float tmp = x - y;
216     return tmp * tmp;
217 }
218
219 static void get_exponent_dynamic(NellyMoserEncodeContext *s, float *cand, int *idx_table)
220 {
221     int i, j, band, best_idx;
222     float power_candidate, best_val;
223
224     float  (*opt )[NELLY_BANDS] = s->opt ;
225     uint8_t(*path)[NELLY_BANDS] = s->path;
226
227     for (i = 0; i < NELLY_BANDS * OPT_SIZE; i++) {
228         opt[0][i] = INFINITY;
229     }
230
231     for (i = 0; i < 64; i++) {
232         opt[0][ff_nelly_init_table[i]] = distance(cand[0], ff_nelly_init_table[i], 0);
233         path[0][ff_nelly_init_table[i]] = i;
234     }
235
236     for (band = 1; band < NELLY_BANDS; band++) {
237         int q, c = 0;
238         float tmp;
239         int idx_min, idx_max, idx;
240         power_candidate = cand[band];
241         for (q = 1000; !c && q < OPT_SIZE; q <<= 2) {
242             idx_min = FFMAX(0, cand[band] - q);
243             idx_max = FFMIN(OPT_SIZE, cand[band - 1] + q);
244             for (i = FFMAX(0, cand[band - 1] - q); i < FFMIN(OPT_SIZE, cand[band - 1] + q); i++) {
245                 if ( isinf(opt[band - 1][i]) )
246                     continue;
247                 for (j = 0; j < 32; j++) {
248                     idx = i + ff_nelly_delta_table[j];
249                     if (idx > idx_max)
250                         break;
251                     if (idx >= idx_min) {
252                         tmp = opt[band - 1][i] + distance(idx, power_candidate, band);
253                         if (opt[band][idx] > tmp) {
254                             opt[band][idx] = tmp;
255                             path[band][idx] = j;
256                             c = 1;
257                         }
258                     }
259                 }
260             }
261         }
262         assert(c); //FIXME
263     }
264
265     best_val = INFINITY;
266     best_idx = -1;
267     band = NELLY_BANDS - 1;
268     for (i = 0; i < OPT_SIZE; i++) {
269         if (best_val > opt[band][i]) {
270             best_val = opt[band][i];
271             best_idx = i;
272         }
273     }
274     for (band = NELLY_BANDS - 1; band >= 0; band--) {
275         idx_table[band] = path[band][best_idx];
276         if (band) {
277             best_idx -= ff_nelly_delta_table[path[band][best_idx]];
278         }
279     }
280 }
281
282 /**
283  * Encode NELLY_SAMPLES samples. It assumes, that samples contains 3 * NELLY_BUF_LEN values
284  *  @param s               encoder context
285  *  @param output          output buffer
286  *  @param output_size     size of output buffer
287  */
288 static void encode_block(NellyMoserEncodeContext *s, unsigned char *output, int output_size)
289 {
290     PutBitContext pb;
291     int i, j, band, block, best_idx, power_idx = 0;
292     float power_val, coeff, coeff_sum;
293     float pows[NELLY_FILL_LEN];
294     int bits[NELLY_BUF_LEN], idx_table[NELLY_BANDS];
295     float cand[NELLY_BANDS];
296
297     apply_mdct(s);
298
299     init_put_bits(&pb, output, output_size * 8);
300
301     i = 0;
302     for (band = 0; band < NELLY_BANDS; band++) {
303         coeff_sum = 0;
304         for (j = 0; j < ff_nelly_band_sizes_table[band]; i++, j++) {
305             coeff_sum += s->mdct_out[i                ] * s->mdct_out[i                ]
306                        + s->mdct_out[i + NELLY_BUF_LEN] * s->mdct_out[i + NELLY_BUF_LEN];
307         }
308         cand[band] =
309             log(FFMAX(1.0, coeff_sum / (ff_nelly_band_sizes_table[band] << 7))) * 1024.0 / M_LN2;
310     }
311
312     if (s->avctx->trellis) {
313         get_exponent_dynamic(s, cand, idx_table);
314     } else {
315         get_exponent_greedy(s, cand, idx_table);
316     }
317
318     i = 0;
319     for (band = 0; band < NELLY_BANDS; band++) {
320         if (band) {
321             power_idx += ff_nelly_delta_table[idx_table[band]];
322             put_bits(&pb, 5, idx_table[band]);
323         } else {
324             power_idx = ff_nelly_init_table[idx_table[0]];
325             put_bits(&pb, 6, idx_table[0]);
326         }
327         power_val = pow_table[power_idx & 0x7FF] / (1 << ((power_idx >> 11) + POW_TABLE_OFFSET));
328         for (j = 0; j < ff_nelly_band_sizes_table[band]; i++, j++) {
329             s->mdct_out[i] *= power_val;
330             s->mdct_out[i + NELLY_BUF_LEN] *= power_val;
331             pows[i] = power_idx;
332         }
333     }
334
335     ff_nelly_get_sample_bits(pows, bits);
336
337     for (block = 0; block < 2; block++) {
338         for (i = 0; i < NELLY_FILL_LEN; i++) {
339             if (bits[i] > 0) {
340                 const float *table = ff_nelly_dequantization_table + (1 << bits[i]) - 1;
341                 coeff = s->mdct_out[block * NELLY_BUF_LEN + i];
342                 best_idx =
343                     quant_lut[av_clip (
344                             coeff * quant_lut_mul[bits[i]] + quant_lut_add[bits[i]],
345                             quant_lut_offset[bits[i]],
346                             quant_lut_offset[bits[i]+1] - 1
347                             )];
348                 if (fabs(coeff - table[best_idx]) > fabs(coeff - table[best_idx + 1]))
349                     best_idx++;
350
351                 put_bits(&pb, bits[i], best_idx);
352             }
353         }
354         if (!block)
355             put_bits(&pb, NELLY_HEADER_BITS + NELLY_DETAIL_BITS - put_bits_count(&pb), 0);
356     }
357
358     flush_put_bits(&pb);
359 }
360
361 static int encode_frame(AVCodecContext *avctx, uint8_t *frame, int buf_size, void *data)
362 {
363     NellyMoserEncodeContext *s = avctx->priv_data;
364     const float *samples = data;
365     int i;
366
367     if (s->last_frame)
368         return 0;
369
370     if (data) {
371         memcpy(s->buf[s->bufsel], samples, avctx->frame_size * sizeof(*samples));
372         for (i = avctx->frame_size; i < NELLY_SAMPLES; i++) {
373             s->buf[s->bufsel][i] = 0;
374         }
375         s->bufsel = 1 - s->bufsel;
376         if (!s->have_saved) {
377             s->have_saved = 1;
378             return 0;
379         }
380     } else {
381         memset(s->buf[s->bufsel], 0, sizeof(s->buf[0][0]) * NELLY_BUF_LEN);
382         s->bufsel = 1 - s->bufsel;
383         s->last_frame = 1;
384     }
385
386     if (s->have_saved) {
387         encode_block(s, frame, buf_size);
388         return NELLY_BLOCK_LEN;
389     }
390     return 0;
391 }
392
393 AVCodec ff_nellymoser_encoder = {
394     .name = "nellymoser",
395     .type = AVMEDIA_TYPE_AUDIO,
396     .id = CODEC_ID_NELLYMOSER,
397     .priv_data_size = sizeof(NellyMoserEncodeContext),
398     .init = encode_init,
399     .encode = encode_frame,
400     .close = encode_end,
401     .capabilities = CODEC_CAP_SMALL_LAST_FRAME | CODEC_CAP_DELAY,
402     .long_name = NULL_IF_CONFIG_SMALL("Nellymoser Asao"),
403     .sample_fmts = (const enum AVSampleFormat[]){AV_SAMPLE_FMT_FLT,AV_SAMPLE_FMT_NONE},
404 };