Add FFmpeg Audio File Reader

Util/Audio/Clip/FFmpegPCMReader.cs

@@ -0,0 +1,327 @@
+using System;
+using System.Runtime.InteropServices;
+using FFmpeg.AutoGen;
+
+namespace ADOFAIRenderer.Util.Audio.Clip;
+
+/// <summary>
+/// <see cref="IPCMReader"/> that decodes the best audio stream of a file with libavformat/libavcodec.
+/// Output keeps the stream's sample rate and channel count; libswresample converts only the sample format to interleaved float.
+///
+/// Frame positions: frame 0 is the first sample the decoder outputs for the stream's start time.
+/// The position of the first frame decoded after opening or seeking is taken from its timestamp
+/// (best_effort_timestamp - stream start_time, in samples); following frames are counted sequentially.
+/// Seek moves the demuxer to <see cref="SeekMarginSeconds"/> before the target and discards decoded samples before the target.
+///
+/// <see cref="TotalFrames"/> comes from the stream duration (or the container duration) and may differ from the decoded length;
+/// reads stop at whichever ends first.
+/// Reference: https://ffmpeg.org/doxygen/trunk/group__lavc__decoding.html (send/receive API)
+/// Reference: https://ffmpeg.org/doxygen/trunk/group__lavf__decoding.html (av_read_frame, av_seek_frame)
+/// </summary>
+internal sealed unsafe class FFmpegPCMReader : IPCMReader
+{
+    /// <summary>
+    /// Time the demuxer is moved before a seek target, in seconds. Decoders that depend on previous packets
+    /// (for example the MP3 bit reservoir) produce settled output after this margin; the samples in it are discarded.
+    /// </summary>
+    private const double SeekMarginSeconds = 0.25;
+
+    private AVFormatContext* _format;
+    private AVCodecContext* _codec;
+    private SwrContext* _swr;
+    private AVPacket* _packet;
+    private AVFrame* _frame;
+    private readonly int _streamIndex;
+    private readonly AVRational _timeBase;
+    private readonly long _startPts;
+
+    // Decoded samples not yet returned: frames [_pendingStart, _pendingStart + _pendingCount) at _pending[_pendingOffset * Channels].
+    private float[] _pending = Array.Empty<float>();
+    private int _pendingOffset;
+    private int _pendingCount;
+    private long _pendingStart;
+
+    private long _nextDecodedFrame; // Frame index of the next decoded frame when counting sequentially
+    private bool _syncToTimestamp;  // True until the first frame after open/seek sets _nextDecodedFrame from its timestamp
+    private bool _flushSent;
+    private bool _endOfStream;
+
+    public string Path { get; }
+    public int Channels { get; }
+    public int SampleRate { get; }
+    public long TotalFrames { get; }
+    public bool IsForwardOnly => false;
+    public long Position { get; private set; }
+
+    /// <summary>Number of demuxer seeks performed (for diagnostics).</summary>
+    public int SeekCount { get; private set; }
+
+    public FFmpegPCMReader(string path)
+    {
+        Path = path;
+        try
+        {
+            AVFormatContext* format = null;
+            var ret = ffmpeg.avformat_open_input(&format, path, null, null);
+            if (ret < 0) throw new InvalidOperationException($"Failed to open audio file '{path}': {ErrorString(ret)}");
+            _format = format;
+
+            ret = ffmpeg.avformat_find_stream_info(_format, null);
+            if (ret < 0) throw new InvalidOperationException($"Failed to read stream info of '{path}': {ErrorString(ret)}");
+
+            AVCodec* decoder = null;
+            _streamIndex = ffmpeg.av_find_best_stream(_format, AVMediaType.AVMEDIA_TYPE_AUDIO, -1, -1, &decoder, 0);
+            if (_streamIndex < 0 || decoder == null) throw new InvalidOperationException($"No decodable audio stream in '{path}': {ErrorString(_streamIndex)}");
+
+            for (var i = 0; i < (int)_format->nb_streams; i++)
+            {
+                if (i != _streamIndex) _format->streams[i]->discard = AVDiscard.AVDISCARD_ALL;
+            }
+
+            var stream = _format->streams[_streamIndex];
+            _timeBase = stream->time_base;
+            _startPts = stream->start_time != ffmpeg.AV_NOPTS_VALUE ? stream->start_time : 0;
+
+            _codec = ffmpeg.avcodec_alloc_context3(decoder);
+            if (_codec == null) throw new InvalidOperationException("avcodec_alloc_context3 failed.");
+            ret = ffmpeg.avcodec_parameters_to_context(_codec, stream->codecpar);
+            if (ret < 0) throw new InvalidOperationException($"avcodec_parameters_to_context failed: {ErrorString(ret)}");
+            _codec->pkt_timebase = _timeBase;
+            ret = ffmpeg.avcodec_open2(_codec, decoder, null);
+            if (ret < 0) throw new InvalidOperationException($"Failed to open audio decoder for '{path}': {ErrorString(ret)}");
+
+            Channels = _codec->ch_layout.nb_channels;
+            SampleRate = _codec->sample_rate;
+            if (Channels <= 0 || SampleRate <= 0) throw new InvalidOperationException($"Invalid audio format in '{path}': {Channels}ch {SampleRate}Hz.");
+
+            TotalFrames = EstimateTotalFrames(stream);
+
+            SwrContext* swr = null;
+            ret = ffmpeg.swr_alloc_set_opts2(&swr,
+                &_codec->ch_layout, AVSampleFormat.AV_SAMPLE_FMT_FLT, SampleRate,
+                &_codec->ch_layout, _codec->sample_fmt, SampleRate, 0, null);
+            if (ret < 0 || swr == null) throw new InvalidOperationException($"swr_alloc_set_opts2 failed: {ErrorString(ret)}");
+            _swr = swr;
+            ret = ffmpeg.swr_init(_swr);
+            if (ret < 0) throw new InvalidOperationException($"swr_init failed: {ErrorString(ret)}");
+
+            _packet = ffmpeg.av_packet_alloc();
+            _frame = ffmpeg.av_frame_alloc();
+            if (_packet == null || _frame == null) throw new InvalidOperationException("Failed to allocate packet/frame.");
+
+            _syncToTimestamp = true;
+        }
+        catch
+        {
+            Dispose();
+            throw;
+        }
+    }
+
+    private long EstimateTotalFrames(AVStream* stream)
+    {
+        if (stream->duration != ffmpeg.AV_NOPTS_VALUE && stream->duration > 0)
+            return ffmpeg.av_rescale_q(stream->duration, _timeBase, new AVRational { num = 1, den = SampleRate });
+        if (_format->duration != ffmpeg.AV_NOPTS_VALUE && _format->duration > 0)
+            return ffmpeg.av_rescale(_format->duration, SampleRate, ffmpeg.AV_TIME_BASE);
+        throw new InvalidOperationException($"Audio file '{Path}' has no duration.");
+    }
+
+    public void Seek(long frame)
+    {
+        if (frame < 0 || frame > TotalFrames) throw new ArgumentOutOfRangeException(nameof(frame));
+        if (frame == Position) return;
+
+        var marginFrames = (long)(SeekMarginSeconds * SampleRate);
+
+        // Target inside the pending block or less than the seek margin after it: decode forward and discard.
+        if (!_syncToTimestamp && frame >= _pendingStart + _pendingOffset
+            && frame <= _pendingStart + _pendingOffset + _pendingCount + marginFrames)
+        {
+            Position = frame;
+            return;
+        }
+
+        var targetTs = _startPts + ffmpeg.av_rescale_q(Math.Max(0, frame - marginFrames), new AVRational { num = 1, den = SampleRate }, _timeBase);
+        var ret = ffmpeg.av_seek_frame(_format, _streamIndex, targetTs, ffmpeg.AVSEEK_FLAG_BACKWARD);
+        if (ret < 0)
+            ret = ffmpeg.av_seek_frame(_format, _streamIndex, _startPts, ffmpeg.AVSEEK_FLAG_BACKWARD);
+        if (ret < 0)
+            throw new InvalidOperationException($"Failed to seek '{Path}' to frame {frame}: {ErrorString(ret)}");
+
+        ffmpeg.avcodec_flush_buffers(_codec);
+        SeekCount++;
+        _pendingCount = 0;
+        _pendingOffset = 0;
+        _syncToTimestamp = true;
+        _flushSent = false;
+        _endOfStream = false;
+        Position = frame;
+    }
+
+    public int Read(float[] buffer, int frameCount)
+    {
+        var limit = (int)Math.Min(frameCount, TotalFrames - Position);
+        var written = 0;
+        while (written < limit)
+        {
+            if (_pendingCount == 0 || Position >= _pendingStart + _pendingOffset + _pendingCount)
+            {
+                if (!DecodeNextFrame()) break;
+                continue;
+            }
+
+            var pendingPosition = _pendingStart + _pendingOffset;
+            if (Position > pendingPosition)
+            {
+                // Discard decoded samples before the read position (after a seek).
+                var skip = (int)Math.Min(Position - pendingPosition, _pendingCount);
+                _pendingOffset += skip;
+                _pendingCount -= skip;
+                continue;
+            }
+            if (Position < pendingPosition)
+            {
+                // A gap in the decoded timestamps: fill with silence up to the next decoded sample.
+                var gap = (int)Math.Min(pendingPosition - Position, limit - written);
+                Array.Clear(buffer, written * Channels, gap * Channels);
+                written += gap;
+                Position += gap;
+                continue;
+            }
+
+            var count = Math.Min(_pendingCount, limit - written);
+            Array.Copy(_pending, _pendingOffset * Channels, buffer, written * Channels, count * Channels);
+            _pendingOffset += count;
+            _pendingCount -= count;
+            written += count;
+            Position += count;
+        }
+        return written;
+    }
+
+    /// <summary>Decodes the next audio frame into the pending buffer. Returns false at the end of the stream.</summary>
+    private bool DecodeNextFrame()
+    {
+        while (!_endOfStream)
+        {
+            var ret = ffmpeg.avcodec_receive_frame(_codec, _frame);
+            if (ret == 0)
+            {
+                StoreFrame();
+                ffmpeg.av_frame_unref(_frame);
+                return true;
+            }
+            if (ret == ffmpeg.AVERROR_EOF)
+            {
+                _endOfStream = true;
+                break;
+            }
+            if (ret != ffmpeg.AVERROR(ffmpeg.EAGAIN))
+                throw new InvalidOperationException($"Audio decoding failed in '{Path}': {ErrorString(ret)}");
+
+            if (_flushSent)
+            {
+                _endOfStream = true;
+                break;
+            }
+
+            ret = ffmpeg.av_read_frame(_format, _packet);
+            if (ret < 0)
+            {
+                if (ret != ffmpeg.AVERROR_EOF)
+                    throw new InvalidOperationException($"Failed to read '{Path}': {ErrorString(ret)}");
+                ffmpeg.avcodec_send_packet(_codec, null);
+                _flushSent = true;
+                continue;
+            }
+
+            try
+            {
+                if (_packet->stream_index != _streamIndex) continue;
+                ret = ffmpeg.avcodec_send_packet(_codec, _packet);
+                // A corrupt packet is skipped; decoding continues with the next one.
+                if (ret < 0 && ret != ffmpeg.AVERROR_INVALIDDATA && ret != ffmpeg.AVERROR(ffmpeg.EAGAIN))
+                    throw new InvalidOperationException($"Failed to send audio packet of '{Path}': {ErrorString(ret)}");
+            }
+            finally
+            {
+                ffmpeg.av_packet_unref(_packet);
+            }
+        }
+        return false;
+    }
+
+    private void StoreFrame()
+    {
+        var samples = _frame->nb_samples;
+        if (samples <= 0) return;
+
+        if (_syncToTimestamp)
+        {
+            var pts = _frame->best_effort_timestamp;
+            _nextDecodedFrame = pts != ffmpeg.AV_NOPTS_VALUE
+                ? ffmpeg.av_rescale_q(pts - _startPts, _timeBase, new AVRational { num = 1, den = SampleRate })
+                : Position;
+            _syncToTimestamp = false;
+        }
+
+        if (_pending.Length < samples * Channels)
+            _pending = new float[samples * Channels];
+
+        fixed (float* pPending = _pending)
+        {
+            var outPlane = (byte*)pPending;
+            var converted = ffmpeg.swr_convert(_swr, &outPlane, samples, _frame->extended_data, samples);
+            if (converted < 0)
+                throw new InvalidOperationException($"Sample format conversion failed in '{Path}': {ErrorString(converted)}");
+            samples = converted;
+        }
+
+        _pendingStart = _nextDecodedFrame;
+        _pendingOffset = 0;
+        _pendingCount = samples;
+        _nextDecodedFrame += samples;
+    }
+
+    public void Dispose()
+    {
+        if (_frame != null)
+        {
+            var frame = _frame;
+            ffmpeg.av_frame_free(&frame);
+            _frame = null;
+        }
+        if (_packet != null)
+        {
+            var packet = _packet;
+            ffmpeg.av_packet_free(&packet);
+            _packet = null;
+        }
+        if (_swr != null)
+        {
+            var swr = _swr;
+            ffmpeg.swr_free(&swr);
+            _swr = null;
+        }
+        if (_codec != null)
+        {
+            var codec = _codec;
+            ffmpeg.avcodec_free_context(&codec);
+            _codec = null;
+        }
+        if (_format != null)
+        {
+            var format = _format;
+            ffmpeg.avformat_close_input(&format);
+            _format = null;
+        }
+    }
+
+    private static string ErrorString(int error)
+    {
+        var buffer = stackalloc byte[1024];
+        ffmpeg.av_strerror(error, buffer, 1024);
+        return $"{Marshal.PtrToStringAnsi((IntPtr)buffer)} ({error})";
+    }
+}