From 988ae6c88f0e5e9f26568825c8daa82a176bb2e1 Mon Sep 17 00:00:00 2001
From: Paul B Mahol <onemda@gmail.com>
Date: Tue, 27 Jun 2023 19:54:25 +0200
Subject: [PATCH 1/2] avcodec: add OSQ audio decoder

Signed-off-by: Paul B Mahol <onemda@gmail.com>
---
 libavcodec/Makefile     |   1 +
 libavcodec/allcodecs.c  |   1 +
 libavcodec/codec_desc.c |   7 +
 libavcodec/codec_id.h   |   1 +
 libavcodec/osq.c        | 435 ++++++++++++++++++++++++++++++++++++++++
 5 files changed, 445 insertions(+)
 create mode 100644 libavcodec/osq.c

diff --git a/libavcodec/Makefile b/libavcodec/Makefile
index 3c16b51462..f3afda61d4 100644
--- a/libavcodec/Makefile
+++ b/libavcodec/Makefile
@@ -579,6 +579,7 @@ OBJS-$(CONFIG_OPUS_DECODER)            += opusdec.o opusdec_celt.o opus_celt.o \
                                           opusdsp.o opus_parse.o opus_rc.o
 OBJS-$(CONFIG_OPUS_ENCODER)            += opusenc.o opusenc_psy.o opus_celt.o \
                                           opus_pvq.o opus_rc.o opustab.o
+OBJS-$(CONFIG_OSQ_DECODER)             += osq.o
 OBJS-$(CONFIG_PAF_AUDIO_DECODER)       += pafaudio.o
 OBJS-$(CONFIG_PAF_VIDEO_DECODER)       += pafvideo.o
 OBJS-$(CONFIG_PAM_DECODER)             += pnmdec.o pnm.o
diff --git a/libavcodec/allcodecs.c b/libavcodec/allcodecs.c
index 8775d15a4f..6e95ca5636 100644
--- a/libavcodec/allcodecs.c
+++ b/libavcodec/allcodecs.c
@@ -516,6 +516,7 @@ extern const FFCodec ff_nellymoser_decoder;
 extern const FFCodec ff_on2avc_decoder;
 extern const FFCodec ff_opus_encoder;
 extern const FFCodec ff_opus_decoder;
+extern const FFCodec ff_osq_decoder;
 extern const FFCodec ff_paf_audio_decoder;
 extern const FFCodec ff_qcelp_decoder;
 extern const FFCodec ff_qdm2_decoder;
diff --git a/libavcodec/codec_desc.c b/libavcodec/codec_desc.c
index 4406dd8318..f556bb94d5 100644
--- a/libavcodec/codec_desc.c
+++ b/libavcodec/codec_desc.c
@@ -3413,6 +3413,13 @@ static const AVCodecDescriptor codec_descriptors[] = {
         .long_name = NULL_IF_CONFIG_SMALL("AC-4"),
         .props     = AV_CODEC_PROP_LOSSY,
     },
+    {
+        .id        = AV_CODEC_ID_OSQ,
+        .type      = AVMEDIA_TYPE_AUDIO,
+        .name      = "osq",
+        .long_name = NULL_IF_CONFIG_SMALL("OSQ (Original Sound Quality)"),
+        .props     = AV_CODEC_PROP_INTRA_ONLY | AV_CODEC_PROP_LOSSLESS,
+    },
 
     /* subtitle codecs */
     {
diff --git a/libavcodec/codec_id.h b/libavcodec/codec_id.h
index a5a0cb8525..29b410b8d3 100644
--- a/libavcodec/codec_id.h
+++ b/libavcodec/codec_id.h
@@ -543,6 +543,7 @@ enum AVCodecID {
     AV_CODEC_ID_WAVARC,
     AV_CODEC_ID_RKA,
     AV_CODEC_ID_AC4,
+    AV_CODEC_ID_OSQ,
 
     /* subtitle codecs */
     AV_CODEC_ID_FIRST_SUBTITLE = 0x17000,          ///< A dummy ID pointing at the start of subtitle codecs.
diff --git a/libavcodec/osq.c b/libavcodec/osq.c
new file mode 100644
index 0000000000..b6dc5c1bb4
--- /dev/null
+++ b/libavcodec/osq.c
@@ -0,0 +1,435 @@
+/*
+ * OSQ audio decoder
+ * Copyright (c) 2023 Paul B Mahol
+ *
+ * This file is part of FFmpeg.
+ *
+ * FFmpeg is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * FFmpeg is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with FFmpeg; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#define ASSERT_LEVEL 5
+#include "libavutil/avassert.h"
+#include "libavutil/internal.h"
+#include "libavutil/intreadwrite.h"
+#include "avcodec.h"
+#include "codec_internal.h"
+#include "decode.h"
+#include "internal.h"
+#define BITSTREAM_READER_LE
+#include "get_bits.h"
+#include "unary.h"
+
+typedef struct OSQChannel {
+    int prediction;
+    int coding_mode;
+    int residue_parameter;
+    int residue_bits;
+} OSQChannel;
+
+typedef struct OSQContext {
+    GetBitContext gb;
+    OSQChannel ch[2];
+
+    uint8_t *bitstream;
+    int64_t max_framesize;
+    int bitstream_size;
+
+    int frame_samples;
+    int64_t nb_samples;
+
+    int32_t *decode_buffer[2];
+
+    AVPacket *pkt;
+    int pkt_offset;
+} OSQContext;
+
+static av_cold int osq_close(AVCodecContext *avctx)
+{
+    OSQContext *s = avctx->priv_data;
+
+    av_freep(&s->bitstream);
+    s->bitstream_size = 0;
+
+    for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++)
+        av_freep(&s->decode_buffer[ch]);
+
+    return 0;
+}
+
+static av_cold int osq_init(AVCodecContext *avctx)
+{
+    OSQContext *s = avctx->priv_data;
+
+    if (avctx->extradata_size < 48)
+        return AVERROR(EINVAL);
+
+    if (avctx->extradata[0] != 1) {
+        av_log(avctx, AV_LOG_ERROR, "Unsupported version.\n");
+        return AVERROR_INVALIDDATA;
+    }
+
+    avctx->sample_rate = AV_RL32(avctx->extradata + 4);
+    if (avctx->sample_rate < 1)
+        return AVERROR_INVALIDDATA;
+
+    avctx->ch_layout.nb_channels = avctx->extradata[3];
+    if (avctx->ch_layout.nb_channels < 1)
+        return AVERROR_INVALIDDATA;
+
+    switch (avctx->extradata[2]) {
+    case  8: avctx->sample_fmt = AV_SAMPLE_FMT_U8P; break;
+    case 16: avctx->sample_fmt = AV_SAMPLE_FMT_S16P; break;
+    case 20:
+    case 24:
+    case 28:
+    case 32: avctx->sample_fmt = AV_SAMPLE_FMT_S32P; break;
+    default: return AVERROR_INVALIDDATA;
+    }
+
+    s->nb_samples = AV_RL64(avctx->extradata + 16);
+    s->frame_samples = AV_RL16(avctx->extradata + 8);
+    s->max_framesize = (s->frame_samples * 16 + 1024) * avctx->ch_layout.nb_channels;
+
+    s->bitstream = av_calloc(s->max_framesize + AV_INPUT_BUFFER_PADDING_SIZE, sizeof(*s->bitstream));
+    if (!s->bitstream)
+        return AVERROR(ENOMEM);
+
+    for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++) {
+        s->decode_buffer[ch] = av_calloc(s->frame_samples + 4,
+                                         sizeof(*s->decode_buffer[ch]));
+        if (!s->decode_buffer[ch])
+            return AVERROR(ENOMEM);
+    }
+
+    s->pkt = avctx->internal->in_pkt;
+
+    return 0;
+}
+
+static uint32_t get_urice(GetBitContext *gb, int k)
+{
+    uint32_t z, x, b;
+
+    x = get_unary(gb, 1, 512);
+    b = get_bits_long(gb, k & 31);
+    z = b | x << (k & 31);
+
+    return z;
+}
+
+static int32_t get_srice(GetBitContext *gb, int x)
+{
+    int32_t y = get_urice(gb, x);
+    return get_bits1(gb) ? -y : y;
+}
+
+static int osq_channel_parameters(AVCodecContext *avctx, int ch)
+{
+    OSQContext *s = avctx->priv_data;
+    OSQChannel *cb = &s->ch[ch];
+    GetBitContext *gb = &s->gb;
+
+    cb->prediction = get_urice(gb, 5);
+    cb->coding_mode = get_urice(gb, 3);
+    if (cb->prediction >= 15)
+        return AVERROR_INVALIDDATA;
+    if (cb->coding_mode > 0 && cb->coding_mode < 3) {
+        cb->residue_parameter = get_urice(gb, 4);
+        if (!cb->residue_parameter || cb->residue_parameter >= 31)
+            return AVERROR_INVALIDDATA;
+    } else if (cb->coding_mode == 3) {
+        cb->residue_bits = get_urice(gb, 4);
+        if (!cb->residue_bits || cb->residue_bits >= 31)
+            return AVERROR_INVALIDDATA;
+    } else if (cb->coding_mode) {
+        return AVERROR_INVALIDDATA;
+    }
+
+    return 0;
+}
+
+static int osq_decode_block(AVCodecContext *avctx, AVFrame *frame)
+{
+    OSQContext *s = avctx->priv_data;
+    int ret, decorrelate, downsample;
+    GetBitContext *gb = &s->gb;
+
+    skip_bits1(gb);
+    decorrelate = get_bits1(gb);
+    downsample = get_bits1(gb);
+
+    for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++) {
+        if ((ret = osq_channel_parameters(avctx, ch)) < 0) {
+            av_log(avctx, AV_LOG_ERROR, "invalid channel parameters\n");
+            return ret;
+        }
+    }
+
+    for (int n = 0; n < frame->nb_samples; n++) {
+        for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++) {
+            OSQChannel *cb = &s->ch[ch];
+            int32_t *dst = s->decode_buffer[ch] + 4;
+
+            if (!cb->coding_mode) {
+                dst[n] = 0;
+            } else if (cb->coding_mode == 3) {
+                dst[n] = get_sbits_long(gb, cb->residue_bits);
+            } else {
+                dst[n] = get_srice(gb, cb->residue_parameter);
+            }
+
+            if (get_bits_left(gb) < 0) {
+                av_log(avctx, AV_LOG_ERROR, "overread!\n");
+                return AVERROR_INVALIDDATA;
+            }
+        }
+    }
+
+    align_get_bits(gb);
+
+#define A (n-1)
+#define B (n-2)
+#define C (n-3)
+#define D (n-4)
+#define P2 (dst[A] * 2LL - dst[B])
+#define P3 ((dst[A] * 1LL - dst[B]) * 3LL + dst[C])
+
+    for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++) {
+        int32_t pred = 0, *dst = s->decode_buffer[ch] + 4;
+        const int nb_samples = frame->nb_samples;
+        OSQChannel *cb = &s->ch[ch];
+
+        av_log(avctx, AV_LOG_DEBUG, "prediction: %d\n", cb->prediction);
+        switch (cb->prediction) {
+        case 0:
+            break;
+        case 1:
+            for (int n = 0; n < nb_samples; n++)
+                dst[n] += dst[A];
+            break;
+        case 2:
+            for (int n = 0; n < nb_samples; n++) {
+                int32_t p = pred;
+                pred = dst[n] / 2;
+                dst[n] += dst[A] + p;
+            }
+            break;
+        case 3:
+            for (int n = 0; n < nb_samples; n++)
+                dst[n] += P2;
+            break;
+        case 4:
+            for (int n = 0; n < nb_samples; n++) {
+                int32_t p = pred;
+                pred = dst[n] / 2;
+                dst[n] += P2 + p;
+            }
+            break;
+        case 5:
+            for (int n = 0; n < nb_samples; n++)
+                dst[n] += P3;
+            break;
+        case 6:
+            for (int n = 0; n < nb_samples; n++) {
+                int32_t p = pred;
+                pred = dst[n] / 2;
+                dst[n] += P3 + p;
+            }
+            break;
+        case 7:
+            for (int n = 0; n < nb_samples; n++) {
+                int32_t p = pred;
+                pred = dst[n] / 2;
+                dst[n] += (P2 + P3) / 2 + p;
+            }
+            break;
+        case 8:
+            for (int n = 0; n < nb_samples; n++)
+                dst[n] += (P2 * 1LL + P3) / 2;
+            break;
+        case 9:
+            for (int n = 0; n < nb_samples; n++) {
+                int32_t p = pred;
+                pred = dst[n] / 2;
+                dst[n] += (P2 * 2 + P3) / 3 + p;
+            }
+            break;
+        case 10:
+            for (int n = 0; n < nb_samples; n++) {
+                int32_t p = pred;
+                pred = dst[n] / 2;
+                dst[n] += (P2 + P3 * 2) / 3 + p;
+            }
+            break;
+        case 11:
+            for (int n = 0; n < nb_samples; n++)
+                dst[n] += (dst[A] * 1LL + dst[B]) / 2;
+            break;
+        case 12:
+            for (int n = 0; n < nb_samples; n++)
+                dst[n] += dst[B];
+            break;
+        case 13:
+            for (int n = 0; n < nb_samples; n++)
+                dst[n] += (dst[D] * 1LL + dst[B]) / 2;
+            break;
+        case 14:
+            for (int n = 0; n < nb_samples; n++) {
+                int32_t p = pred;
+                pred = dst[n] / 2;
+                dst[n] += (P2 + dst[A]) / 2 + p;
+            }
+            break;
+        default:
+            return AVERROR_INVALIDDATA;
+        }
+
+        memcpy(s->decode_buffer[ch], s->decode_buffer[ch] + frame->nb_samples, 4 * sizeof(*s->decode_buffer[0]));
+
+        if (downsample) {
+            int32_t *dst = s->decode_buffer[ch];
+
+            for (int n = 0; n < frame->nb_samples + 4; n++)
+                dst[n] = (dst[n] / 256) * 256;
+        }
+    }
+
+    if (decorrelate && avctx->ch_layout.nb_channels == 2) {
+        int32_t *l = s->decode_buffer[0] + 4;
+        int32_t *r = s->decode_buffer[1] + 4;
+
+        for (int n = 0; n < frame->nb_samples; n++) {
+            int64_t L = l[n];
+            int64_t R = r[n];
+
+            l[n] = L - R;
+            r[n] = L + R;
+        }
+    }
+
+    switch (avctx->sample_fmt) {
+    case AV_SAMPLE_FMT_U8P:
+        for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++) {
+            uint8_t *dst = (uint8_t *)frame->extended_data[ch];
+            int32_t *src = s->decode_buffer[ch] + 4;
+
+            for (int n = 0; n < frame->nb_samples; n++)
+                dst[n] = av_clip_uint8(src[n] + 0x80);
+        }
+        break;
+    case AV_SAMPLE_FMT_S16P:
+        for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++) {
+            int16_t *dst = (int16_t *)frame->extended_data[ch];
+            int32_t *src = s->decode_buffer[ch] + 4;
+
+            for (int n = 0; n < frame->nb_samples; n++)
+                dst[n] = av_clip_int16(src[n]);
+        }
+        break;
+    case AV_SAMPLE_FMT_S32P:
+        for (int ch = 0; ch < avctx->ch_layout.nb_channels; ch++) {
+            int32_t *dst = (int32_t *)frame->extended_data[ch];
+            int32_t *src = s->decode_buffer[ch] + 4;
+
+            for (int n = 0; n < frame->nb_samples; n++)
+                dst[n] = src[n];
+        }
+        break;
+    default:
+        return AVERROR_BUG;
+    }
+
+    return 0;
+}
+
+static int osq_receive_frame(AVCodecContext *avctx, AVFrame *frame)
+{
+    OSQContext *s = avctx->priv_data;
+    GetBitContext *gb = &s->gb;
+    int ret, n;
+
+    while (s->bitstream_size < s->max_framesize) {
+        int size;
+
+        if (!s->pkt->data) {
+            ret = ff_decode_get_packet(avctx, s->pkt);
+            if (ret == AVERROR_EOF && s->bitstream_size > 0)
+                break;
+            if (ret < 0)
+                return ret;
+        }
+
+        size = FFMIN(s->pkt->size - s->pkt_offset, s->max_framesize - s->bitstream_size);
+        memcpy(s->bitstream + s->bitstream_size, s->pkt->data + s->pkt_offset, size);
+        s->bitstream_size += size;
+        s->pkt_offset += size;
+
+        if (s->pkt_offset == s->pkt->size) {
+            av_packet_unref(s->pkt);
+            s->pkt_offset = 0;
+        }
+    }
+
+    frame->nb_samples = FFMIN(s->frame_samples, s->nb_samples);
+    if (frame->nb_samples <= 0)
+        return AVERROR_EOF;
+
+    if ((ret = ff_get_buffer(avctx, frame, 0)) < 0)
+        goto fail;
+
+    if ((ret = init_get_bits8(gb, s->bitstream, s->bitstream_size)) < 0)
+        goto fail;
+
+    if ((ret = osq_decode_block(avctx, frame)) < 0)
+        goto fail;
+
+    s->nb_samples -= frame->nb_samples;
+
+    n = get_bits_count(gb) / 8;
+    if (n > s->bitstream_size) {
+        ret = AVERROR_INVALIDDATA;
+        goto fail;
+    }
+
+    memmove(s->bitstream, &s->bitstream[n], s->bitstream_size - n);
+    s->bitstream_size -= n;
+
+    return 0;
+
+fail:
+    s->bitstream_size = 0;
+    s->pkt_offset = 0;
+    av_packet_unref(s->pkt);
+
+    return ret;
+}
+
+const FFCodec ff_osq_decoder = {
+    .p.name           = "osq",
+    CODEC_LONG_NAME("OSQ (Original Sound Quality)"),
+    .p.type           = AVMEDIA_TYPE_AUDIO,
+    .p.id             = AV_CODEC_ID_OSQ,
+    .priv_data_size   = sizeof(OSQContext),
+    .init             = osq_init,
+    FF_CODEC_RECEIVE_FRAME_CB(osq_receive_frame),
+    .close            = osq_close,
+    .p.capabilities   = AV_CODEC_CAP_CHANNEL_CONF |
+                        AV_CODEC_CAP_DR1,
+    .caps_internal    = FF_CODEC_CAP_INIT_CLEANUP,
+    .p.sample_fmts    = (const enum AVSampleFormat[]) { AV_SAMPLE_FMT_U8P,
+                                                        AV_SAMPLE_FMT_S16P,
+                                                        AV_SAMPLE_FMT_S32P,
+                                                        AV_SAMPLE_FMT_NONE },
+};
-- 
2.39.1

