ADOFAIRenderer/Util/FrameReader.cs

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Linq;
using ADOFAIRenderer.Core;
using ADOFAIRenderer.UI;
using FFmpeg.AutoGen;
using MelonLoader;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;
using UnityEngine.Rendering;

namespace ADOFAIRenderer.Util
{
    /// <summary>
    /// GPU capture of rendered frames for the encoder. For each requested frame it optionally converts the RGBA texture to
    /// the encoder's YUV 4:2:0 layout with a compute shader (ColorConversion.compute from the mod's asset bundle, kernels
    /// CSMain_NV12 / CSMain_YUV420P), then reads the result back asynchronously (AsyncGPUReadback.RequestIntoNativeArray)
    /// into a persistent native buffer and adds it to RenderManager.ImageQueue under its video frame index.
    /// When the shader or the kernel for the pixel format is unavailable, the RGBA texture is read back unchanged and the
    /// encoder converts it (IsColorConverted = false).
    /// <para>
    /// Each request uses a <see cref="ReadbackRequestState"/> from a pool (16 created up front, more allocated when the pool
    /// is empty). A state stays out of the pool from the request until the render thread has sent its frame
    /// (RenderManager.SendVideoFrame, then <see cref="RecycleState"/>) or the frame was dropped.
    /// </para>
    /// <para>
    /// Threads: requests and readback completions run on the main thread; <see cref="RecycleState"/> is also called by
    /// the render thread. The pool and the disposed flag are guarded by _poolLock. One instance per render, created in
    /// RenderManager.SetupPlay with the pixel format of the opened encoder.
    /// </para>
    /// </summary>
    internal class FrameReader : IDisposable
    {
        // Property IDs of the compute shader parameters: source texture, output buffer and frame size in pixels.
        private static readonly int Height = Shader.PropertyToID("Height");
        private static readonly int Width = Shader.PropertyToID("Width");
        private static readonly int DestinationBuffer = Shader.PropertyToID("DestinationBuffer");
        private static readonly int SourceTexture = Shader.PropertyToID("SourceTexture");
        /// <summary>Number of readback states created by InitializePool.</summary>
        private const int InitialBufferSize = 16;
        
        /// <summary>States that are free for a new request. Guarded by _poolLock.</summary>
        private readonly Stack<ReadbackRequestState> _statePool = new();
        /// <summary>Every state created by this reader (free, in flight or holding a queued frame), for Dispose. Guarded by _poolLock.</summary>
        private readonly List<ReadbackRequestState> _allocatedStates = new();
        /// <summary>Lock for _statePool, _allocatedStates, _isInitialized, _isDisposed and ReadbackRequestState.IsInFlight.</summary>
        private readonly object _poolLock = new();
        /// <summary>True after InitializePool; false again after Dispose (RecycleState then disposes returned states).</summary>
        private bool _isInitialized;
        /// <summary>Set by Dispose. Completions and recycles that arrive afterwards dispose their state.</summary>
        private volatile bool _isDisposed;

        /// <summary>Color conversion compute shader from the asset bundle (ModAssets.ColorConversionShader); null when not loaded.</summary>
        private readonly ComputeShader? _computeShader;
        /// <summary>Kernel index for the target pixel format; -1 when the format has no kernel or the kernel was not found.</summary>
        private readonly int _kernelIndex = -1;
        /// <summary>True when frames are converted to YUV on the GPU (shader and kernel available).</summary>
        private readonly bool _useGPUConversion;
        /// <summary>Pixel format of the read-back data: the target format with GPU conversion, AV_PIX_FMT_RGBA otherwise.</summary>
        private readonly AVPixelFormat _expectedFormat;

        /// <summary>
        /// Selects the conversion kernel for <paramref name="targetFormat"/> and fills the state pool.
        /// </summary>
        /// <param name="targetFormat">
        /// Pixel format of the encoder input (FFmpegNativeEncoder.SelectedPixelFormat). NV12 and YUV420P use GPU conversion;
        /// any other format falls back to RGBA readback.
        /// </param>
        public FrameReader(AVPixelFormat targetFormat)
        {
            _computeShader = ModAssets.ColorConversionShader;
            if (_computeShader)
            {
                var kernelName = targetFormat switch
                {
                    AVPixelFormat.AV_PIX_FMT_NV12 => "CSMain_NV12",
                    AVPixelFormat.AV_PIX_FMT_YUV420P => "CSMain_YUV420P",
                    _ => null
                };

                if (kernelName != null)
                {
                    try { _kernelIndex = _computeShader.FindKernel(kernelName); }
                    catch { _kernelIndex = -1; }
                }
            }

            _useGPUConversion = _kernelIndex >= 0;
            _expectedFormat = _useGPUConversion ? targetFormat : AVPixelFormat.AV_PIX_FMT_RGBA;

            InitializePool();
        }

        /// <summary>Creates InitialBufferSize states once. Their buffers are allocated on first use (ReadbackRequestState.Setup).</summary>
        private void InitializePool()
        {
            lock (_poolLock)
            {
                if (_isInitialized) return;
                for (var i = 0; i < InitialBufferSize; i++)
                {
                    var state = new ReadbackRequestState(this);
                    _statePool.Push(state);
                    _allocatedStates.Add(state);
                }

                _isInitialized = true;
            }
        }

        /// <summary>
        /// Bytes of the readback buffer of one frame: Width x Height x 3 / 2 when the GPU conversion to NV12 / YUV420P
        /// is used, Width x Height x 4 (RGBA) otherwise. The same size as ReadbackRequestState.Setup uses.
        /// </summary>
        public long GetFrameBufferBytes(int width, int height)
        {
            var pixels = (long)width * height;
            return _useGPUConversion ? pixels * 3 / 2 : pixels * 4;
        }

        /// <summary>
        /// Requests the capture of <paramref name="texture"/> as video frame <paramref name="currentFrame"/>. Does nothing
        /// outside the Render stage. Main thread.
        /// With GPU conversion, the compute kernel writes the YUV data of the texture into the state's ComputeBuffer and that
        /// buffer is read back; otherwise the texture itself is read back as RGBA. The buffers are sized from the texture
        /// dimensions of this call. The result arrives in OnReadbackComplete.
        /// </summary>
        /// <param name="texture">Frame to capture (game frame or the output of the frame hooks).</param>
        /// <param name="currentFrame">Video frame index; the key of the frame in ImageQueue.</param>
        public void RequestFrameCapture(Texture texture, long currentFrame)
        {
            if (!RenderManager.IsCapturingActive) return;
            ReadbackRequestState state;

            lock (_poolLock)
            {
                if (_statePool.Count == 0)
                {
                    state = new ReadbackRequestState(this);
                    _allocatedStates.Add(state);
                }
                else
                {
                    state = _statePool.Pop();
                }
                
                state.IsInFlight = true;
            }

            var bufferSize = texture.width * texture.height * 4; // RGBA = 4 channel
            state.Setup(currentFrame, bufferSize, _useGPUConversion, _expectedFormat, texture.width, texture.height);

            if (_useGPUConversion && state.GPUBuffer != null && _computeShader)
            {
                _computeShader.SetTexture(_kernelIndex, SourceTexture, texture);
                _computeShader.SetBuffer(_kernelIndex, DestinationBuffer, state.GPUBuffer);
                _computeShader.SetInt(Width, texture.width);
                _computeShader.SetInt(Height, texture.height);

                // Thread group counts cover the frame in tiles of 32 x 16 pixels (NV12) or 64 x 16 pixels (YUV420P), rounded up.
            // The tile size is the area one thread group of the ColorConversion.compute kernels processes.
                var threadGroupsX = _expectedFormat == AVPixelFormat.AV_PIX_FMT_NV12
                    ? (texture.width + 31) / 32
                    : (texture.width + 63) / 64;
                var threadGroupsY = (texture.height + 15) / 16;

                _computeShader.Dispatch(_kernelIndex, threadGroupsX, threadGroupsY, 1);

                // The native buffer must not be disposed until the request completes; Dispose skips states that are in flight.
                // Reference: https://docs.unity3d.com/ScriptReference/Rendering.AsyncGPUReadback.RequestIntoNativeArray.html
                AsyncGPUReadback.RequestIntoNativeArray(
                    ref state.NativeBuffer, state.GPUBuffer, state.CachedCallback
                );
                /*AsyncGPUReadback.Request(
                    state.GPUBuffer,
                    state.CachedCallback
                );*/
            }
            else
            {
                AsyncGPUReadback.RequestIntoNativeArray(
                    ref state.NativeBuffer, texture, 0, state.CachedCallback
                );
            }
        }

        /// <summary>
        /// Completion of a readback request (through ReadbackRequestState.CachedCallback).
        /// <list type="bullet">
        /// <item>Reader disposed: the state is disposed.</item>
        /// <item>Render aborted: the state is returned to the pool and nothing is queued.</item>
        /// <item>Readback error: a black frame (managed buffer) is queued under the same index, so the render thread does not
        /// wait forever for the missing index, and the state is returned to the pool.</item>
        /// <item>Success: the state's native buffer is queued without copying; the state stays out of the pool until the
        /// render thread has encoded the frame.</item>
        /// </list>
        /// </summary>
        internal void OnReadbackComplete(ReadbackRequestState state, AsyncGPUReadbackRequest request)
        {
            var frame = state.FrameIndex;

            lock (_poolLock)
            {
                state.IsInFlight = false;
                if (_isDisposed)
                {
                    state.Dispose();
                    return;
                }
            }

            if (RenderManager.Instance && RenderManager.Instance.IsRenderAborted)
            {
                RecycleState(state);
                return;
            }

            if (request.hasError)
            {
                if (RenderManager.Instance)
                {
                    RenderManager.Instance.readbackFailureCount++;
                }
                Melon<RenderMod>.Logger.Warning($"[OOM Prevent] GPU Readback failed for frame {frame}. Queue deadlock prevented by sending a black frame.");

                var blackFrame = CreateBlackFrame(state, out var byteSize);
                var capFrame = new RenderManager.CapturedFrame
                {
                    ManagedBuffer = blackFrame, Length = byteSize,
                    IsColorConverted = state.IsGPUConversion
                };
                PushToQueue(frame, capFrame);
                RecycleState(state);
            }
            else
            {
                unsafe
                {
                    /*
                    if (state.IsGPUConversion)
                    {
                        var data = request.GetData<byte>();
                        // Copy async request buffer data to persistent NativeBuffer memory
                        NativeArray<byte>.Copy(data, state.NativeBuffer, data.Length);
                    }
                    */
                    
                    var ptr = (IntPtr)state.NativeBuffer.GetUnsafeReadOnlyPtr();
                    var capFrame = new RenderManager.CapturedFrame
                    {
                        NativePointer = ptr, Length = state.NativeBuffer.Length,
                        ReadbackState = state, IsColorConverted = state.IsGPUConversion
                    };
                    PushToQueue(frame, capFrame);
                }
            }
        }

        /// <summary>
        /// Creates a black frame in the same format as the frames of this reader, from ArrayPool&lt;byte&gt;.Shared
        /// (returned to the pool by the render thread after encoding).
        ///
        ///   GPU conversion (NV12 / YUV420P, full range): Width × Height × 3 / 2 bytes.
        ///     Y plane (Width × Height bytes) = 0, U/V data (the rest) = 128.
        ///     128 is the zero chroma value of the 8-bit full range: the U/V offset 0.5 of RGB2YUV_BT709_Full
        ///     in ColorConversion.compute, which the encoder input of this path uses.
        ///     Both NV12 (interleaved UV) and YUV420P (U plane, then V plane) store all chroma after the Y plane,
        ///     so the same fill applies to both.
        ///   No GPU conversion (RGBA input, converted by the encoder): Width × Height × 4 bytes, all 0.
        /// </summary>
        private static byte[] CreateBlackFrame(ReadbackRequestState state, out int byteSize)
        {
            var lumaSize = state.Width * state.Height;
            byteSize = state.IsGPUConversion ? lumaSize * 3 / 2 : lumaSize * 4;

            var frame = ArrayPool<byte>.Shared.Rent(byteSize);
            if (state.IsGPUConversion)
            {
                Array.Clear(frame, 0, lumaSize);
                frame.AsSpan(lumaSize, byteSize - lumaSize).Fill(128);
            }
            else
            {
                Array.Clear(frame, 0, byteSize);
            }
            return frame;
        }

        /// <summary>
        /// Returns a state to the pool after its frame was sent or dropped, or disposes it when the reader is disposed.
        /// Called by the render thread (after sending) and by the main thread (abort, errors, queue cleanup).
        /// </summary>
        internal void RecycleState(ReadbackRequestState state)
        {
            lock (_poolLock)
            {
                if (_isDisposed || !_isInitialized)
                {
                    state.Dispose();
                    return;
                }
                
                state.Recycle();
                _statePool.Push(state);
            }
        }
        
        /// <summary>
        /// Adds the frame to RenderManager.ImageQueue under <paramref name="frame"/> (replacing a frame with the same index)
        /// and wakes the render thread.
        /// </summary>
        private static void PushToQueue(long frame, RenderManager.CapturedFrame capturedFrame)
        {
            if (!RenderManager.Instance) return;
            lock (RenderManager.Instance.ImageQueue)
            {
                RenderManager.Instance.ImageQueue[frame] = capturedFrame;
            }
            RenderManager.Instance.FrameAvailableEvent.Set();
        }
        
        /// <summary>
        /// Disposes every state that is not in flight and clears the pool. States in flight are disposed when their
        /// readback completes. States still referenced by frames in ImageQueue are disposed too, so callers clear the queue
        /// first (RenderManager.StopFFmpeg / RestoreRenderEnvironment). Main thread.
        /// </summary>
        public void Dispose()
        {
            lock (_poolLock)
            {
                if (_isDisposed) return;
                _isDisposed = true;
                foreach (var state in _allocatedStates.Where(state => !state.IsInFlight))
                {
                    state.Dispose();
                }
                _statePool.Clear();
                _allocatedStates.Clear();
                _isInitialized = false;
            }
        }
    }

    /// <summary>
    /// One reusable capture slot of <see cref="FrameReader"/>: the persistent native buffer the readback writes into, the
    /// ComputeBuffer the conversion kernel writes (GPU conversion only), and the metadata of the frame it currently holds.
    /// The buffers are kept between frames and reallocated only when the frame size or conversion mode changes.
    /// </summary>
    internal class ReadbackRequestState : IDisposable
    {
        /// <summary>Reader that receives the completion and owns the pool.</summary>
        private readonly FrameReader _owner;
        /// <summary>Video frame index of the current request; -1 while the state is free.</summary>
        public long FrameIndex { get; private set; }
        /// <summary>
        /// Persistent native buffer receiving the readback: Width x Height x 3 / 2 bytes of YUV 4:2:0 with GPU conversion,
        /// Width x Height x 4 bytes of RGBA otherwise. Read by the render thread through CapturedFrame.NativePointer.
        /// </summary>
        public NativeArray<byte> NativeBuffer;
        /// <summary>True between the request and its completion. Written under FrameReader._poolLock.</summary>
        public bool IsInFlight { get; internal set; }

        /// <summary>Raw ComputeBuffer (4-byte elements) the conversion kernel writes; null without GPU conversion.</summary>
        public ComputeBuffer? GPUBuffer;
        /// <summary>True when the current frame is converted to YUV on the GPU.</summary>
        public bool IsGPUConversion { get; private set; }
        /// <summary>Width of the current frame in pixels.</summary>
        public int Width { get; private set; }
        /// <summary>Height of the current frame in pixels.</summary>
        public int Height { get; private set; }

        /// <summary>Completion delegate created once per state, so a request does not allocate a new delegate.</summary>
        public readonly Action<AsyncGPUReadbackRequest> CachedCallback;

        /// <summary>Creates a free state without buffers.</summary>
        public ReadbackRequestState(FrameReader owner)
        {
            _owner = owner;
            CachedCallback = OnComplete;
            FrameIndex = -1;
        }

        /// <summary>
        /// Prepares the state for a request: stores the frame metadata and (re)allocates the buffers when their size differs.
        /// The native buffer holds width x height x 3 / 2 bytes with GPU conversion, otherwise <paramref name="bufferSize"/>.
        /// The ComputeBuffer holds the same number of bytes as 4-byte elements (byte count / 4, rounded down).
        /// </summary>
        /// <param name="frameIndex">Video frame index of the request.</param>
        /// <param name="bufferSize">Readback size in bytes for RGBA readback (width x height x 4).</param>
        /// <param name="isGPUConversion">Whether the kernel converts the frame to YUV.</param>
        /// <param name="pixelFormat">Pixel format of the read-back data (not used by this method).</param>
        /// <param name="width">Frame width in pixels.</param>
        /// <param name="height">Frame height in pixels.</param>
        public void Setup(long frameIndex, int bufferSize, bool isGPUConversion, AVPixelFormat pixelFormat, int width, int height)
        {
            FrameIndex = frameIndex;
            IsGPUConversion = isGPUConversion;
            Width = width; Height = height;

            var readbackSize = isGPUConversion ? (width * height * 3 / 2) : bufferSize;
            if (!NativeBuffer.IsCreated || NativeBuffer.Length != readbackSize)
            {
                if (NativeBuffer.IsCreated) NativeBuffer.Dispose();

                NativeBuffer = new NativeArray<byte>(readbackSize, Allocator.Persistent,
                    NativeArrayOptions.UninitializedMemory);
            }

            if (isGPUConversion)
            {
                var computeBufferSize = readbackSize / 4;
                if (GPUBuffer != null && GPUBuffer.count == computeBufferSize) return;
                
                GPUBuffer?.Dispose();
                GPUBuffer = new ComputeBuffer(computeBufferSize, 4, ComputeBufferType.Raw);
            }
            else
            {
                GPUBuffer?.Dispose();
                GPUBuffer = null;
            }
        }

        /// <summary>Readback completion; forwards to FrameReader.OnReadbackComplete.</summary>
        private void OnComplete(AsyncGPUReadbackRequest request)
        {
            _owner.OnReadbackComplete(this, request);
        }

        /// <summary>Marks the state as free (FrameIndex -1). The buffers are kept for the next request.</summary>
        public void Recycle()
        {
            FrameIndex = -1;
        }

        /// <summary>Releases the native buffer and the ComputeBuffer.</summary>
        public void Dispose()
        {
            if (NativeBuffer.IsCreated) NativeBuffer.Dispose();
            GPUBuffer?.Dispose();
            GPUBuffer = null;
        }
    }
}