ADOFAIRenderer/Util/Audio/Clip/AudioClipSource.cs

using System;
using System.Collections.Generic;
using System.Threading;
using MelonLoader;

namespace ADOFAIRenderer.Util.Audio.Clip;

/// <summary>
/// Read position registered on an <see cref="AudioClipSource"/> by one consumer (a playing voice or the spectrum generator).
/// For a windowed source, the chunk containing <see cref="Position"/> and the following chunks
/// (<see cref="AudioClipSource.SlotCount"/> chunks in total) are kept loaded.
/// </summary>
internal sealed class AudioCursor
{
    private readonly AudioClipSource _source;

    /// <summary>Backing value of <see cref="Position"/>, accessed with Interlocked (64-bit reads are not atomic on every platform).</summary>
    private long _position;

    private volatile bool _loop;

    /// <summary>Created by <see cref="AudioClipSource.AddCursor"/>.</summary>
    internal AudioCursor(AudioClipSource source, long position, bool loop)
    {
        _source = source;
        _position = position;
        _loop = loop;
    }

    /// <summary>
    /// Frame in the source at the output sample rate (0 ~ TotalFrames - 1).
    /// Written by the consumer that reads the source; read by the decoder thread.
    /// Moving into another chunk wakes the decoder thread.
    /// </summary>
    public long Position
    {
        get => Interlocked.Read(ref _position);
        set
        {
            var previous = Interlocked.Exchange(ref _position, value);
            if (previous / _source.FramesPerChunk != value / _source.FramesPerChunk)
                _source.OnCursorMoved();
        }
    }

    /// <summary>
    /// When true, the chunk after the last chunk is chunk 0 (the window wraps to the clip start), and a position past the end
    /// wraps into the clip. Setting it wakes the decoder thread.
    /// </summary>
    public bool Loop
    {
        get => _loop;
        set
        {
            _loop = value;
            _source.OnCursorMoved();
        }
    }
}

/// <summary>
/// One loaded chunk of an <see cref="AudioClipSource"/>, held by a consumer between
/// <see cref="AudioClipSource.TryAcquire"/> and <see cref="AudioClipSource.Release"/>.
/// The chunk is not freed while held.
/// </summary>
internal readonly struct ChunkView
{
    /// <summary>Interleaved samples. Frame f of the source is at (f - StartFrame) * Channels.</summary>
    public readonly float[] Data;

    /// <summary>Source frame of the first frame in <see cref="Data"/>.</summary>
    public readonly long StartFrame;

    /// <summary>Valid frames in this chunk. Frames past the end of the source are not included.</summary>
    public readonly int FrameCount;

    /// <summary>Channel count of the stored samples (same as the source clip).</summary>
    public readonly int Channels;

    /// <summary>Index of the chunk in its source (used by <see cref="AudioClipSource.Release"/>).</summary>
    public readonly int ChunkIndex;

    /// <summary>Created by <see cref="AudioClipSource.TryAcquire"/>; default(ChunkView) is the invalid view.</summary>
    internal ChunkView(float[] data, long startFrame, int frameCount, int channels, int chunkIndex)
    {
        Data = data;
        StartFrame = startFrame;
        FrameCount = frameCount;
        Channels = channels;
        ChunkIndex = chunkIndex;
    }

    /// <summary>True for a view that holds a chunk; false for default(ChunkView).</summary>
    public bool IsValid => Data != null;

    /// <summary>Source frame after the last valid frame (exclusive end).</summary>
    public long EndFrame => StartFrame + FrameCount;

    /// <summary>True when source frame <paramref name="frame"/> is one of the valid frames of this chunk.</summary>
    public bool Contains(long frame) => frame >= StartFrame && frame < StartFrame + FrameCount;

    /// <summary>
    /// Sample of <paramref name="frame"/> mapped to output channel <paramref name="outChannel"/> of <paramref name="outChannels"/>.
    /// Mapping: mono output averages all source channels; a mono source is copied to every output channel;
    /// otherwise output channel c takes source channel c, and output channels beyond the source take source channel 0.
    /// <paramref name="frame"/> must be inside this chunk.
    /// </summary>
    public float GetSample(long frame, int outChannel, int outChannels)
    {
        var baseIndex = (int)(frame - StartFrame) * Channels;
        if (Channels == 1)
            return Data[baseIndex];

        if (outChannels == 1)
        {
            var sum = 0f;
            for (var c = 0; c < Channels; c++)
                sum += Data[baseIndex + c];
            return sum / Channels;
        }

        return Data[baseIndex + (outChannel < Channels ? outChannel : 0)];
    }
}

/// <summary>
/// How an <see cref="AudioClipSource"/> keeps its chunks.
/// </summary>
internal enum ChunkPolicy
{
    /// <summary>
    /// Only the chunks required by registered cursors (current chunk and the next SlotCount - 1 chunks per cursor),
    /// chunks requested by a waiting consumer, and chunks currently held are kept. Other chunks are freed and reloaded on demand.
    /// Used for readers that can seek (files decoded with FFmpeg) when the clip is longer than one chunk.
    /// </summary>
    Windowed,

    /// <summary>
    /// All chunks are loaded once, in order from chunk 0, starting when the source is created, and kept until the store is disposed.
    /// Used for forward-only readers (clip data copied from Unity, released as it is consumed) and for clips of one chunk or less.
    /// </summary>
    Resident
}

/// <summary>
/// Decoded audio of one clip, resampled to the output sample rate and stored in fixed-size chunks.
/// Channel count stays as in the clip.
///
/// Chunks: chunk k covers output frames [k * FramesPerChunk, (k + 1) * FramesPerChunk), where
/// <see cref="FramesPerChunk"/> = <see cref="ChunkFloats"/> / Channels (rounded down; a frame never spans two chunks).
/// Windowed chunks use pooled arrays of ChunkFloats floats; resident chunks use arrays of their exact length.
/// See <see cref="ChunkPolicy"/> for which chunks are kept.
///
/// Threads: cursors and TryAcquire/Release may be used from any thread.
/// Loading and freeing run only on the store's decoder thread (<see cref="AudioClipStore"/>), through
/// <see cref="FreeUnused"/>, <see cref="TryGetNextLoad"/> and <see cref="LoadChunk"/>.
/// </summary>
internal sealed class AudioClipSource
{
    /// <summary>Floats per chunk array: 2^22 floats = 16 MiB.</summary>
    public const int ChunkFloats = 1 << 22;

    /// <summary>Chunks kept per cursor of a windowed source: the chunk containing the cursor and the next SlotCount - 1 chunks.</summary>
    public const int SlotCount = 3;

    /// <summary>Load priority of a chunk a consumer is waiting for (lowest value is loaded first).</summary>
    internal const int PriorityDemanded = 0;

    /// <summary>Load priority of slot s of a cursor window is PriorityWindow + s.</summary>
    internal const int PriorityWindow = 1;

    /// <summary>Load priority of resident chunks nobody is waiting for.</summary>
    internal const int PriorityResident = PriorityWindow + SlotCount;

    /// <summary>
    /// Most decode streams (reader + resampler) a windowed source opens. A load continues the stream positioned at the chunk start,
    /// otherwise opens another stream up to this count, otherwise restarts the least recently used stream.
    /// Each stream holds one open reader (file handle and decoder state, measured 0.3 ~ 2.7 MiB for ogg/flac/MP3 with the clip harness).
    /// </summary>
    public const int MaxDecodeStreams = 4;

    /// <summary>Decode stream limit of this source (<see cref="MaxDecodeStreams"/> unless changed for tests).</summary>
    internal int DecodeStreamLimit { get; set; } = MaxDecodeStreams;

    private readonly AudioClipStore _store;

    /// <summary>Reader passed to the constructor (the first decode stream); disposed again in <see cref="DisposeInternal"/>.</summary>
    private readonly IPCMReader _firstReader;

    /// <summary>Opens another reader for an additional decode stream; null when only the first reader can be used.</summary>
    private readonly Func<IPCMReader?>? _openReader;

    // Guarded by _sync: chunk table, hold counts, cursors, demand set, loaded set.
    private readonly object _sync = new();

    /// <summary>Sample array of each chunk, null while not loaded.</summary>
    private readonly float[]?[] _chunks;

    /// <summary>Number of consumers holding each chunk (TryAcquire minus Release); a held chunk is not freed.</summary>
    private readonly int[] _holdCounts;

    private readonly List<AudioCursor> _cursors = new();

    /// <summary>Chunks a consumer asked for while they were not loaded; an index is removed when the chunk is acquired.</summary>
    private readonly HashSet<int> _demanded = new();

    /// <summary>Indices of the loaded chunks (non-null entries of <see cref="_chunks"/>).</summary>
    private readonly HashSet<int> _loaded = new();

    private int _residentLoadedCount; // Resident: chunks [0, _residentLoadedCount) are loaded

    // Guarded by _loadLock: decode streams.
    private readonly object _loadLock = new();
    private readonly List<DecodeStream> _streams = new();

    /// <summary>Incremented on every load; stored in <see cref="DecodeStream.LastUse"/> for least-recently-used selection.</summary>
    private long _streamUseCounter;

    /// <summary>True after opening an additional reader failed; no further readers are opened for this source.</summary>
    private bool _openFailed;

    private volatile bool _isFailed;

    /// <summary>Key the source is registered under in the store (the game's audioLib key / AudioClip.name).</summary>
    public string Key { get; }

    /// <summary>Channel count of the stored samples (same as the clip).</summary>
    public int Channels { get; }

    /// <summary>Sample rate of the stored samples in Hz (the output sample rate).</summary>
    public int Frequency { get; }

    /// <summary>Frame count at <see cref="Frequency"/>: ceil(clip frames * Frequency / clip rate).</summary>
    public long TotalFrames { get; }

    /// <summary>Clip length in seconds: clip frames / clip rate.</summary>
    public double Duration { get; }

    /// <summary>Frames per chunk: <see cref="ChunkFloats"/> / <see cref="Channels"/> (rounded down).</summary>
    public int FramesPerChunk { get; }

    /// <summary>Number of chunks: ceil(<see cref="TotalFrames"/> / <see cref="FramesPerChunk"/>); 0 for an empty clip.</summary>
    public int ChunkCount { get; }

    /// <summary>
    /// <see cref="ChunkPolicy.Resident"/> for forward-only readers and clips of at most one chunk, otherwise
    /// <see cref="ChunkPolicy.Windowed"/>. Fixed at construction.
    /// </summary>
    public ChunkPolicy Policy { get; }

    /// <summary>True when <see cref="Policy"/> is <see cref="ChunkPolicy.Resident"/>.</summary>
    public bool IsResident => Policy == ChunkPolicy.Resident;

    /// <summary>True after loading failed. A failed source returns no chunks.</summary>
    public bool IsFailed => _isFailed;

    /// <summary>True when every chunk of a resident source is loaded.</summary>
    public bool IsFullyLoaded
    {
        get
        {
            lock (_sync) return _loaded.Count == ChunkCount;
        }
    }

    /// <summary>Number of chunk loads performed (for diagnostics).</summary>
    public int LoadCount { get; private set; }

    /// <summary>Decode streams opened so far (for diagnostics).</summary>
    public int StreamCount
    {
        get
        {
            lock (_loadLock) return _streams.Count;
        }
    }

    /// <param name="store">Store that owns this source.</param>
    /// <param name="key">Clip name or path the source is registered under.</param>
    /// <param name="reader">First reader; also defines the clip format.</param>
    /// <param name="openReader">Opens another reader of the same clip for an additional decode stream, or null when not possible.</param>
    /// <param name="outputRate">Output sample rate in Hz.</param>
    internal AudioClipSource(AudioClipStore store, string key, IPCMReader reader, Func<IPCMReader?>? openReader, int outputRate)
    {
        if (reader.Channels <= 0) throw new ArgumentException("Reader channel count must be positive.", nameof(reader));
        if (reader.SampleRate <= 0) throw new ArgumentException("Reader sample rate must be positive.", nameof(reader));
        if (reader.TotalFrames < 0) throw new ArgumentException("Reader frame count must not be negative.", nameof(reader));

        _store = store;
        _firstReader = reader;
        _openReader = openReader;
        Key = key;
        Channels = reader.Channels;
        Frequency = outputRate;
        Duration = (double)reader.TotalFrames / reader.SampleRate;
        TotalFrames = CeilDiv(reader.TotalFrames * outputRate, reader.SampleRate);
        FramesPerChunk = ChunkFloats / Channels;
        ChunkCount = (int)CeilDiv(TotalFrames, FramesPerChunk);
        Policy = reader.IsForwardOnly || ChunkCount <= 1 ? ChunkPolicy.Resident : ChunkPolicy.Windowed;

        _chunks = new float[ChunkCount][];
        _holdCounts = new int[ChunkCount];
        _streams.Add(new DecodeStream(reader, outputRate));
    }

    // ── Cursors ──

    /// <summary>Registers a read position. The decoder thread loads the chunks it requires.</summary>
    public AudioCursor AddCursor(long position, bool loop)
    {
        var cursor = new AudioCursor(this, position, loop);
        lock (_sync)
        {
            _cursors.Add(cursor);
        }
        _store.RequestService();
        return cursor;
    }

    /// <summary>Unregisters <paramref name="cursor"/>; its window chunks become freeable on the decoder thread's next pass.</summary>
    public void RemoveCursor(AudioCursor cursor)
    {
        lock (_sync)
        {
            _cursors.Remove(cursor);
        }
        _store.RequestService();
    }

    /// <summary>Wakes the decoder thread after a cursor changed chunk or loop flag (resident sources ignore cursors).</summary>
    internal void OnCursorMoved()
    {
        if (!IsResident) _store.RequestService();
    }

    // ── Consumer access ──

    /// <summary>Chunk index containing <paramref name="frame"/> (0 ~ TotalFrames - 1).</summary>
    public int ChunkIndexOf(long frame) => (int)(frame / FramesPerChunk);

    /// <summary>
    /// Holds the chunk containing <paramref name="frame"/> when it is loaded.
    /// When it is not loaded, records a load request (kept until the chunk is acquired), wakes the decoder thread and returns false.
    /// Every successful call must be paired with <see cref="Release"/>.
    /// </summary>
    public bool TryAcquire(long frame, out ChunkView view)
    {
        view = default;
        if (_isFailed || frame < 0 || frame >= TotalFrames) return false;

        var index = ChunkIndexOf(frame);
        lock (_sync)
        {
            var data = _chunks[index];
            if (data != null)
            {
                _demanded.Remove(index);
                _holdCounts[index]++;
                view = CreateView(index, data);
                return true;
            }
            _demanded.Add(index);
        }
        _store.RequestService();
        return false;
    }

    /// <summary>Ends one hold of the chunk of <paramref name="view"/> taken by <see cref="TryAcquire"/>. An invalid view is ignored.</summary>
    public void Release(in ChunkView view)
    {
        if (!view.IsValid) return;
        lock (_sync)
        {
            if (_holdCounts[view.ChunkIndex] > 0)
                _holdCounts[view.ChunkIndex]--;
        }
    }

    /// <summary>View of chunk <paramref name="index"/>; the last chunk's frame count stops at <see cref="TotalFrames"/>.</summary>
    private ChunkView CreateView(int index, float[] data)
    {
        var start = (long)index * FramesPerChunk;
        var count = (int)Math.Min(FramesPerChunk, TotalFrames - start);
        return new ChunkView(data, start, count, Channels, index);
    }

    // ── Decoder thread ──

    /// <summary>
    /// Windowed: frees loaded chunks that are not required (see <see cref="ChunkPolicy.Windowed"/>) and not held.
    /// Resident: does nothing. Decoder thread only.
    /// </summary>
    internal void FreeUnused()
    {
        if (IsResident) return;

        List<float[]>? freed = null;
        lock (_sync)
        {
            if (_loaded.Count == 0) return;
            List<int>? remove = null;
            foreach (var index in _loaded)
            {
                if (_holdCounts[index] > 0 || IsRequired(index)) continue;
                (remove ??= new List<int>()).Add(index);
            }
            if (remove == null) return;

            foreach (var index in remove)
            {
                (freed ??= new List<float[]>()).Add(_chunks[index]!);
                _chunks[index] = null;
                _loaded.Remove(index);
            }
        }

        foreach (var data in freed!)
            _store.ReturnChunkArray(data);
    }

    /// <summary>
    /// Next chunk to load and its priority (lower loads first), or false when nothing is missing.
    /// Windowed: demanded chunks, then each cursor's window by slot.
    /// Resident: the next chunk in order; priority is <see cref="PriorityDemanded"/> while a consumer is waiting on this source.
    /// Decoder thread only.
    /// </summary>
    internal bool TryGetNextLoad(out int index, out int priority)
    {
        index = -1;
        priority = int.MaxValue;
        if (_isFailed) return false;

        lock (_sync)
        {
            if (IsResident)
            {
                if (_residentLoadedCount >= ChunkCount) return false;
                index = _residentLoadedCount;
                priority = _demanded.Count > 0 ? PriorityDemanded : PriorityResident;
                return true;
            }

            foreach (var demanded in _demanded)
            {
                if (_chunks[demanded] != null) continue;
                index = demanded;
                priority = PriorityDemanded;
                return true;
            }

            // For each cursor, the first missing chunk of its window; the lowest slot over all cursors wins
            // (ties go to the earlier cursor).
            foreach (var cursor in _cursors)
            {
                if (!TryGetWindowStart(cursor, out var first, out var loop)) continue;
                for (var slot = 0; slot < SlotCount; slot++)
                {
                    if (!TryWindowIndex(first, slot, loop, out var candidate)) break;
                    if (_chunks[candidate] != null) continue;
                    if (PriorityWindow + slot < priority)
                    {
                        index = candidate;
                        priority = PriorityWindow + slot;
                    }
                    break;
                }
            }
            return index >= 0;
        }
    }

    /// <summary>True when a demand or a cursor window covers <paramref name="index"/>. Caller holds _sync.</summary>
    private bool IsRequired(int index)
    {
        if (_demanded.Contains(index)) return true;
        foreach (var cursor in _cursors)
        {
            if (!TryGetWindowStart(cursor, out var first, out var loop)) continue;
            for (var slot = 0; slot < SlotCount; slot++)
            {
                if (!TryWindowIndex(first, slot, loop, out var candidate)) break;
                if (candidate == index) return true;
            }
        }
        return false;
    }

    /// <summary>
    /// First chunk of the cursor's window: the chunk of its position (negative positions count as 0; positions past the end
    /// wrap when looping). False when a non-looping cursor is past the end (no window). Caller holds _sync.
    /// </summary>
    private bool TryGetWindowStart(AudioCursor cursor, out int first, out bool loop)
    {
        var position = cursor.Position;
        loop = cursor.Loop;
        first = 0;
        if (position < 0) position = 0;
        if (position >= TotalFrames)
        {
            if (!loop || TotalFrames == 0) return false;
            position %= TotalFrames;
        }
        first = ChunkIndexOf(position);
        return true;
    }

    /// <summary>Chunk index of window slot <paramref name="slot"/> after <paramref name="first"/>; wraps when looping, false past the end otherwise.</summary>
    private bool TryWindowIndex(int first, int slot, bool loop, out int index)
    {
        index = first + slot;
        if (index < ChunkCount) return true;
        if (!loop) return false;
        index %= ChunkCount;
        return true;
    }

    /// <summary>
    /// Loads chunk <paramref name="index"/> if it is not loaded. Resident sources load only the next chunk in order.
    /// Decoder thread only.
    /// </summary>
    /// <returns>True when this call loaded the chunk.</returns>
    internal bool LoadChunk(int index)
    {
        if (_isFailed || index < 0 || index >= ChunkCount) return false;

        lock (_loadLock)
        {
            lock (_sync)
            {
                if (_chunks[index] != null) return false;
                if (IsResident && index != _residentLoadedCount)
                    throw new InvalidOperationException($"Resident clip '{Key}' loads chunks in order: expected {_residentLoadedCount}, got {index}.");
            }

            float[]? data = null;
            try
            {
                var start = (long)index * FramesPerChunk;
                var frames = (int)Math.Min(FramesPerChunk, TotalFrames - start);
                data = IsResident ? _store.AllocateExactArray(frames * Channels) : _store.RentChunkArray();
                var stream = SelectStream(start);
                stream.LastUse = ++_streamUseCounter;
                stream.Fill(start, frames, data);
            }
            // Out of memory fails the whole render (store fatal error); any other error makes only this source silent.
            catch (OutOfMemoryException e)
            {
                ReturnArray(data);
                Fail();
                _store.ReportFatal(this, e);
                return false;
            }
            catch (Exception e)
            {
                ReturnArray(data);
                Fail();
                _store.ReportSourceFailure(this, e);
                return false;
            }

            lock (_sync)
            {
                _chunks[index] = data;
                _loaded.Add(index);
                if (IsResident) _residentLoadedCount++;
            }
            LoadCount++;

            // A resident source never reads again after its last chunk; the reader releases its data now.
            if (IsResident && _residentLoadedCount == ChunkCount)
                DisposeStreams();
        }

        _store.NotifyChunkLoaded();
        return true;
    }

    /// <summary>
    /// Decode stream for a load starting at output frame <paramref name="start"/>:
    /// the stream positioned at <paramref name="start"/>; otherwise a new stream while fewer than <see cref="DecodeStreamLimit"/>
    /// are open (windowed sources only); otherwise the least recently used stream. Caller holds _loadLock.
    /// </summary>
    private DecodeStream SelectStream(long start)
    {
        DecodeStream? oldest = null;
        foreach (var stream in _streams)
        {
            if (stream.NextOutFrame == start) return stream;
            if (oldest == null || stream.LastUse < oldest.LastUse) oldest = stream;
        }

        if (!IsResident && !_openFailed && _openReader != null && _streams.Count < DecodeStreamLimit)
        {
            IPCMReader? reader = null;
            try
            {
                reader = _openReader();
                if (reader is { IsForwardOnly: true })
                {
                    reader.Dispose();
                    reader = null;
                }
                if (reader != null)
                {
                    var stream = new DecodeStream(reader, Frequency);
                    _streams.Add(stream);
                    return stream;
                }
            }
            catch (Exception e)
            {
                reader?.Dispose();
                Melon<RenderMod>.Logger.Warning($"[AudioClipSource] Could not open another decode stream for '{Key}', reusing open streams: {e.Message}");
            }
            _openFailed = true;
        }

        return oldest!;
    }

    /// <summary>Marks the source failed (no more chunks are returned or loaded); a resident source also closes its readers now. Caller holds _loadLock.</summary>
    private void Fail()
    {
        _isFailed = true;
        if (IsResident) DisposeStreams();
    }

    /// <summary>Disposes every decode stream (reader and resampler) and clears the list. Caller holds _loadLock.</summary>
    private void DisposeStreams()
    {
        foreach (var stream in _streams)
            stream.Dispose();
        _streams.Clear();
    }

    /// <summary>Gives a chunk array back to the store: exact arrays (resident) are uncounted, pooled arrays (windowed) are returned to the pool.</summary>
    private void ReturnArray(float[]? data)
    {
        if (data == null) return;
        if (IsResident) _store.FreeExactArray(data);
        else _store.ReturnChunkArray(data);
    }

    /// <summary>Releases all chunk arrays and readers. Called by the store after the decoder thread stopped.</summary>
    internal void DisposeInternal()
    {
        lock (_loadLock)
        {
            lock (_sync)
            {
                foreach (var index in _loaded)
                {
                    ReturnArray(_chunks[index]);
                    _chunks[index] = null;
                }
                _loaded.Clear();
                _cursors.Clear();
                _demanded.Clear();
            }
            DisposeStreams();
            _firstReader.Dispose();
        }
    }

    /// <summary>ceil(a / b) for a ≥ 0, b &gt; 0.</summary>
    private static long CeilDiv(long a, long b) => (a + b - 1) / b;
}