Extend WavPack demuxer and decoder to support >2 channel audio

Signed-off-by: Janne Grunau <janne-ffmpeg@jannau.net>

1 files changed, 216 insertions, 50 deletions

diff --git a/libavcodec/wavpack.c b/libavcodec/wavpack.c
index f78492c924..866cc59879 100644
--- a/libavcodec/wavpack.c
+++ b/libavcodec/wavpack.c
@@ -1,6 +1,6 @@
 /*
  * WavPack lossless audio decoder
- * Copyright (c) 2006 Konstantin Shishkov
+ * Copyright (c) 2006,2011 Konstantin Shishkov
  *
  * This file is part of FFmpeg.
  *
@@ -91,7 +91,7 @@ typedef struct WvChannel {
     int bitrate_acc, bitrate_delta;
 } WvChannel;
 
-typedef struct WavpackContext {
+typedef struct WavpackFrameContext {
     AVCodecContext *avctx;
     int frame_flags;
     int stereo, stereo_in;
@@ -118,6 +118,22 @@ typedef struct WavpackContext {
     int max_samples;
     int pos;
     SavedContext sc, extra_sc;
+} WavpackFrameContext;
+
+#define WV_MAX_FRAME_DECODERS 14
+
+typedef struct WavpackContext {
+    AVCodecContext *avctx;
+
+    WavpackFrameContext *fdec[WV_MAX_FRAME_DECODERS];
+    int fdec_num;
+
+    int multichannel;
+    int mkv_mode;
+    int block;
+    int samples;
+    int samples_left;
+    int ch_offset;
 } WavpackContext;
 
 // exponent table copied from WavPack source
@@ -224,7 +240,7 @@ static av_always_inline int get_tail(GetBitContext *gb, int k)
     return res;
 }
 
-static void update_error_limit(WavpackContext *ctx)
+static void update_error_limit(WavpackFrameContext *ctx)
 {
     int i, br[2], sl[2];
 
@@ -258,7 +274,7 @@ static void update_error_limit(WavpackContext *ctx)
     }
 }
 
-static int wv_get_value(WavpackContext *ctx, GetBitContext *gb, int channel, int *last)
+static int wv_get_value(WavpackFrameContext *ctx, GetBitContext *gb, int channel, int *last)
 {
     int t, t2;
     int sign, base, add, ret;
@@ -361,7 +377,7 @@ static int wv_get_value(WavpackContext *ctx, GetBitContext *gb, int channel, int
     return sign ? ~ret : ret;
 }
 
-static inline int wv_get_value_integer(WavpackContext *s, uint32_t *crc, int S)
+static inline int wv_get_value_integer(WavpackFrameContext *s, uint32_t *crc, int S)
 {
     int bit;
 
@@ -377,7 +393,7 @@ static inline int wv_get_value_integer(WavpackContext *s, uint32_t *crc, int S)
     return (((S + bit) << s->shift) - bit) << s->post_shift;
 }
 
-static float wv_get_value_float(WavpackContext *s, uint32_t *crc, int S)
+static float wv_get_value_float(WavpackFrameContext *s, uint32_t *crc, int S)
 {
     union {
         float    f;
@@ -450,13 +466,13 @@ static float wv_get_value_float(WavpackContext *s, uint32_t *crc, int S)
     return value.f;
 }
 
-static void wv_reset_saved_context(WavpackContext *s)
+static void wv_reset_saved_context(WavpackFrameContext *s)
 {
     s->pos = 0;
     s->sc.crc = s->extra_sc.crc = 0xFFFFFFFF;
 }
 
-static inline int wv_unpack_stereo(WavpackContext *s, GetBitContext *gb, void *dst, const int type)
+static inline int wv_unpack_stereo(WavpackFrameContext *s, GetBitContext *gb, void *dst, const int type)
 {
     int i, j, count = 0;
     int last, t;
@@ -467,6 +483,7 @@ static inline int wv_unpack_stereo(WavpackContext *s, GetBitContext *gb, void *d
     int16_t *dst16 = dst;
     int32_t *dst32 = dst;
     float   *dstfl = dst;
+    const int channel_pad = s->avctx->channels - 2;
 
     if(s->samples_left == s->samples)
         s->one = s->zero = s->zeroes = 0;
@@ -549,12 +566,15 @@ static inline int wv_unpack_stereo(WavpackContext *s, GetBitContext *gb, void *d
         if(type == AV_SAMPLE_FMT_FLT){
             *dstfl++ = wv_get_value_float(s, &crc_extra_bits, L);
             *dstfl++ = wv_get_value_float(s, &crc_extra_bits, R);
+            dstfl += channel_pad;
         } else if(type == AV_SAMPLE_FMT_S32){
             *dst32++ = wv_get_value_integer(s, &crc_extra_bits, L);
             *dst32++ = wv_get_value_integer(s, &crc_extra_bits, R);
+            dst32 += channel_pad;
         } else {
             *dst16++ = wv_get_value_integer(s, &crc_extra_bits, L);
             *dst16++ = wv_get_value_integer(s, &crc_extra_bits, R);
+            dst16 += channel_pad;
         }
         count++;
     }while(!last && count < s->max_samples);
@@ -582,7 +602,7 @@ static inline int wv_unpack_stereo(WavpackContext *s, GetBitContext *gb, void *d
     return count * 2;
 }
 
-static inline int wv_unpack_mono(WavpackContext *s, GetBitContext *gb, void *dst, const int type)
+static inline int wv_unpack_mono(WavpackFrameContext *s, GetBitContext *gb, void *dst, const int type)
 {
     int i, j, count = 0;
     int last, t;
@@ -593,6 +613,7 @@ static inline int wv_unpack_mono(WavpackContext *s, GetBitContext *gb, void *dst
     int16_t *dst16 = dst;
     int32_t *dst32 = dst;
     float   *dstfl = dst;
+    const int channel_stride = s->avctx->channels;
 
     if(s->samples_left == s->samples)
         s->one = s->zero = s->zeroes = 0;
@@ -623,12 +644,16 @@ static inline int wv_unpack_mono(WavpackContext *s, GetBitContext *gb, void *dst
         pos = (pos + 1) & 7;
         crc = crc * 3 + S;
 
-        if(type == AV_SAMPLE_FMT_FLT)
-            *dstfl++ = wv_get_value_float(s, &crc_extra_bits, S);
-        else if(type == AV_SAMPLE_FMT_S32)
-            *dst32++ = wv_get_value_integer(s, &crc_extra_bits, S);
-        else
-            *dst16++ = wv_get_value_integer(s, &crc_extra_bits, S);
+        if(type == AV_SAMPLE_FMT_FLT){
+            *dstfl = wv_get_value_float(s, &crc_extra_bits, S);
+            dstfl += channel_stride;
+        }else if(type == AV_SAMPLE_FMT_S32){
+            *dst32 = wv_get_value_integer(s, &crc_extra_bits, S);
+            dst32 += channel_stride;
+        }else{
+            *dst16 = wv_get_value_integer(s, &crc_extra_bits, S);
+            dst16 += channel_stride;
+        }
         count++;
     }while(!last && count < s->max_samples);
 
@@ -655,47 +680,92 @@ static inline int wv_unpack_mono(WavpackContext *s, GetBitContext *gb, void *dst
     return count;
 }
 
+static av_cold int wv_alloc_frame_context(WavpackContext *c)
+{
+
+    if(c->fdec_num == WV_MAX_FRAME_DECODERS)
+        return -1;
+
+    c->fdec[c->fdec_num] = av_mallocz(sizeof(**c->fdec));
+    if(!c->fdec[c->fdec_num])
+        return -1;
+    c->fdec_num++;
+    c->fdec[c->fdec_num - 1]->avctx = c->avctx;
+    wv_reset_saved_context(c->fdec[c->fdec_num - 1]);
+
+    return 0;
+}
+
 static av_cold int wavpack_decode_init(AVCodecContext *avctx)
 {
     WavpackContext *s = avctx->priv_data;
 
     s->avctx = avctx;
-    if (avctx->channels > 2) {
-        av_log(avctx, AV_LOG_ERROR, "Multichannel WavPack is not supported yet.\n");
-        return -1;
-    }
-    s->stereo = (avctx->channels == 2);
     if(avctx->bits_per_coded_sample <= 16)
         avctx->sample_fmt = AV_SAMPLE_FMT_S16;
     else
         avctx->sample_fmt = AV_SAMPLE_FMT_S32;
+    if(avctx->channels <= 2 && !avctx->channel_layout)
     avctx->channel_layout = (avctx->channels==2) ? CH_LAYOUT_STEREO : CH_LAYOUT_MONO;
 
-    wv_reset_saved_context(s);
+    s->multichannel = avctx->channels > 2;
+    /* lavf demuxer does not provide extradata, Matroska stores 0x403
+       there, use this to detect decoding mode for multichannel */
+    s->mkv_mode = 0;
+    if(s->multichannel && avctx->extradata && avctx->extradata_size == 2){
+        int ver = AV_RL16(avctx->extradata);
+        if(ver >= 0x402 && ver <= 0x410)
+            s->mkv_mode = 1;
+    }
+
+    s->fdec_num = 0;
 
     return 0;
 }
 
-static int wavpack_decode_frame(AVCodecContext *avctx,
-                            void *data, int *data_size,
-                            AVPacket *avpkt)
+static av_cold int wavpack_decode_end(AVCodecContext *avctx)
 {
-    const uint8_t *buf = avpkt->data;
-    int buf_size = avpkt->size;
     WavpackContext *s = avctx->priv_data;
+    int i;
+
+    for(i = 0; i < s->fdec_num; i++)
+        av_freep(&s->fdec[i]);
+    s->fdec_num = 0;
+
+    return 0;
+}
+
+static int wavpack_decode_block(AVCodecContext *avctx, int block_no,
+                                void *data, int *data_size,
+                                const uint8_t *buf, int buf_size)
+{
+    WavpackContext *wc = avctx->priv_data;
+    WavpackFrameContext *s;
     void *samples = data;
     int samplecount;
     int got_terms = 0, got_weights = 0, got_samples = 0, got_entropy = 0, got_bs = 0, got_float = 0;
     int got_hybrid = 0;
+    const uint8_t* orig_buf = buf;
     const uint8_t* buf_end = buf + buf_size;
     int i, j, id, size, ssize, weights, t;
-    int bpp;
+    int bpp, chan, chmask;
 
     if (buf_size == 0){
         *data_size = 0;
         return 0;
     }
 
+    if(block_no >= wc->fdec_num && wv_alloc_frame_context(wc) < 0){
+        av_log(avctx, AV_LOG_ERROR, "Error creating frame decode context\n");
+        return -1;
+    }
+
+    s = wc->fdec[block_no];
+    if(!s){
+        av_log(avctx, AV_LOG_ERROR, "Context for block %d is not present\n", block_no);
+        return -1;
+    }
+
     if(!s->samples_left){
         memset(s->decorr, 0, MAX_TERMS * sizeof(Decorr));
         memset(s->ch, 0, sizeof(s->ch));
@@ -704,11 +774,15 @@ static int wavpack_decode_frame(AVCodecContext *avctx,
         s->got_extra_bits = 0;
     }
 
+    if(!wc->mkv_mode){
     s->samples = AV_RL32(buf); buf += 4;
     if(!s->samples){
         *data_size = 0;
         return buf_size;
     }
+    }else{
+        s->samples = wc->samples;
+    }
     s->frame_flags = AV_RL32(buf); buf += 4;
     if(s->frame_flags&0x80){
         bpp = sizeof(float);
@@ -720,12 +794,19 @@ static int wavpack_decode_frame(AVCodecContext *avctx,
         bpp = 4;
         avctx->sample_fmt = AV_SAMPLE_FMT_S32;
     }
+    samples = (uint8_t*)samples + bpp * wc->ch_offset;
+
+    s->stereo = !(s->frame_flags & WV_MONO);
     s->stereo_in = (s->frame_flags & WV_FALSE_STEREO) ? 0 : s->stereo;
     s->joint = s->frame_flags & WV_JOINT_STEREO;
     s->hybrid = s->frame_flags & WV_HYBRID_MODE;
     s->hybrid_bitrate = s->frame_flags & WV_HYBRID_BITRATE;
     s->post_shift = 8 * (bpp-1-(s->frame_flags&0x03)) + ((s->frame_flags >> 13) & 0x1f);
     s->CRC = AV_RL32(buf); buf += 4;
+    if(wc->mkv_mode)
+        buf += 4; //skip block size;
+
+    wc->ch_offset += 1 + s->stereo;
 
     s->max_samples = *data_size / (bpp * avctx->channels);
     s->max_samples = FFMIN(s->max_samples, s->samples);
@@ -899,7 +980,7 @@ static int wavpack_decode_frame(AVCodecContext *avctx,
             got_float = 1;
             break;
         case WP_ID_DATA:
-            s->sc.offset = buf - avpkt->data;
+            s->sc.offset = buf - orig_buf;
             s->sc.size   = size * 8;
             init_get_bits(&s->gb, buf, size * 8);
             s->data_size = size * 8;
@@ -912,13 +993,50 @@ static int wavpack_decode_frame(AVCodecContext *avctx,
                 buf += size;
                 continue;
             }
-            s->extra_sc.offset = buf - avpkt->data;
+            s->extra_sc.offset = buf - orig_buf;
             s->extra_sc.size   = size * 8;
             init_get_bits(&s->gb_extra_bits, buf, size * 8);
             s->crc_extra_bits = get_bits_long(&s->gb_extra_bits, 32);
             buf += size;
             s->got_extra_bits = 1;
             break;
+        case WP_ID_CHANINFO:
+            if(size <= 1){
+                av_log(avctx, AV_LOG_ERROR, "Insufficient channel information\n");
+                return -1;
+            }
+            chan = *buf++;
+            switch(size - 2){
+            case 0:
+                chmask = *buf;
+                break;
+            case 1:
+                chmask = AV_RL16(buf);
+                break;
+            case 2:
+                chmask = AV_RL24(buf);
+                break;
+            case 3:
+                chmask = AV_RL32(buf);
+                break;
+            case 5:
+                chan |= (buf[1] & 0xF) << 8;
+                chmask = AV_RL24(buf + 2);
+                break;
+            default:
+                av_log(avctx, AV_LOG_ERROR, "Invalid channel info size %d\n", size);
+                chan = avctx->channels;
+                chmask = avctx->channel_layout;
+            }
+            if(chan != avctx->channels){
+                av_log(avctx, AV_LOG_ERROR, "Block reports total %d channels, decoder believes it's %d channels\n",
+                       chan, avctx->channels);
+                return -1;
+            }
+            if(!avctx->channel_layout)
+                avctx->channel_layout = chmask;
+            buf += size - 1;
+            break;
         default:
             buf += size;
         }
@@ -963,10 +1081,10 @@ static int wavpack_decode_frame(AVCodecContext *avctx,
         }
         s->samples_left = s->samples;
     }else{
-        init_get_bits(&s->gb, avpkt->data + s->sc.offset, s->sc.size);
+        init_get_bits(&s->gb, orig_buf + s->sc.offset, s->sc.size);
         skip_bits_long(&s->gb, s->sc.bits_used);
         if(s->got_extra_bits){
-            init_get_bits(&s->gb_extra_bits, avpkt->data + s->extra_sc.offset,
+            init_get_bits(&s->gb_extra_bits, orig_buf + s->extra_sc.offset,
                           s->extra_sc.size);
             skip_bits_long(&s->gb_extra_bits, s->extra_sc.bits_used);
         }
@@ -979,8 +1097,10 @@ static int wavpack_decode_frame(AVCodecContext *avctx,
             samplecount = wv_unpack_stereo(s, &s->gb, samples, AV_SAMPLE_FMT_S32);
         else
             samplecount = wv_unpack_stereo(s, &s->gb, samples, AV_SAMPLE_FMT_FLT);
-
+        samplecount >>= 1;
     }else{
+        const int channel_stride = avctx->channels;
+
         if(avctx->sample_fmt == AV_SAMPLE_FMT_S16)
             samplecount = wv_unpack_mono(s, &s->gb, samples, AV_SAMPLE_FMT_S16);
         else if(avctx->sample_fmt == AV_SAMPLE_FMT_S32)
@@ -989,37 +1109,83 @@ static int wavpack_decode_frame(AVCodecContext *avctx,
             samplecount = wv_unpack_mono(s, &s->gb, samples, AV_SAMPLE_FMT_FLT);
 
         if(s->stereo && avctx->sample_fmt == AV_SAMPLE_FMT_S16){
-            int16_t *dst = (int16_t*)samples + samplecount * 2;
-            int16_t *src = (int16_t*)samples + samplecount;
+            int16_t *dst = (int16_t*)samples + 1;
+            int16_t *src = (int16_t*)samples;
             int cnt = samplecount;
             while(cnt--){
-                *--dst = *--src;
-                *--dst = *src;
+                *dst = *src;
+                src += channel_stride;
+                dst += channel_stride;
             }
-            samplecount *= 2;
         }else if(s->stereo && avctx->sample_fmt == AV_SAMPLE_FMT_S32){
-            int32_t *dst = (int32_t*)samples + samplecount * 2;
-            int32_t *src = (int32_t*)samples + samplecount;
+            int32_t *dst = (int32_t*)samples + 1;
+            int32_t *src = (int32_t*)samples;
             int cnt = samplecount;
             while(cnt--){
-                *--dst = *--src;
-                *--dst = *src;
+                *dst = *src;
+                src += channel_stride;
+                dst += channel_stride;
             }
-            samplecount *= 2;
         }else if(s->stereo){
-            float *dst = (float*)samples + samplecount * 2;
-            float *src = (float*)samples + samplecount;
+            float *dst = (float*)samples + 1;
+            float *src = (float*)samples;
             int cnt = samplecount;
             while(cnt--){
-                *--dst = *--src;
-                *--dst = *src;
+                *dst = *src;
+                src += channel_stride;
+                dst += channel_stride;
+            }
+        }
+    }
+
+    wc->samples_left = s->samples_left;
+
+    return samplecount * bpp;
+}
+
+static int wavpack_decode_frame(AVCodecContext *avctx,
+                            void *data, int *data_size,
+                            AVPacket *avpkt)
+{
+    WavpackContext *s = avctx->priv_data;
+    const uint8_t *buf = avpkt->data;
+    int buf_size = avpkt->size;
+    int frame_size;
+    int samplecount = 0;
+
+    s->block = 0;
+    s->samples_left = 0;
+    s->ch_offset = 0;
+
+    if(s->mkv_mode){
+        s->samples = AV_RL32(buf); buf += 4;
+    }
+    while(buf_size > 0){
+        if(!s->multichannel){
+            frame_size = buf_size;
+        }else{
+            if(!s->mkv_mode){
+                frame_size = AV_RL32(buf) - 12; buf += 4; buf_size -= 4;
+            }else{
+                if(buf_size < 12) //MKV files can have zero flags after last block
+                    break;
+                frame_size = AV_RL32(buf + 8) + 12;
             }
-            samplecount *= 2;
         }
+        if(frame_size < 0 || frame_size > buf_size){
+            av_log(avctx, AV_LOG_ERROR, "Block %d has invalid size (size %d vs. %d bytes left)\n",
+                   s->block, frame_size, buf_size);
+            return -1;
+        }
+        if((samplecount = wavpack_decode_block(avctx, s->block, data,
+                                               data_size, buf, frame_size)) < 0)
+            return -1;
+        s->block++;
+        buf += frame_size; buf_size -= frame_size;
     }
-    *data_size = samplecount * bpp;
+    *data_size = samplecount * avctx->channels;
 
-    return s->samples_left > 0 ? 0 : buf_size;
+    return s->samples_left > 0 ? 0 : avpkt->size;
 }
 
 AVCodec wavpack_decoder = {
@@ -1029,7 +1195,7 @@ AVCodec wavpack_decoder = {
     sizeof(WavpackContext),
     wavpack_decode_init,
     NULL,
-    NULL,
+    wavpack_decode_end,
     wavpack_decode_frame,
     .capabilities = CODEC_CAP_SUBFRAMES,
     .long_name = NULL_IF_CONFIG_SMALL("WavPack"),