ADOFAIRenderer/Core/FrameHookPipeline.cs

using System;
using System.Collections.Generic;
using ADOFAIRenderer.API;
using MelonLoader;
using UnityEngine;

namespace ADOFAIRenderer.Core;

/// <summary>
/// Runs the frame hooks (session and global, see RenderSession) on a captured frame before it is read back.
/// Owns two RenderTextures (ping and pong) with the size and format of the source frame; each hook pass reads the
/// current frame and writes the other buffer (see FrameHookContext). One instance per render, created in
/// RenderManager.SetupPlay and disposed when the render ends. Main thread only.
/// </summary>
internal class FrameHookPipeline : IDisposable
{
    /// <summary>First ping-pong buffer. Created by EnsureBuffers, released by ReleaseBuffers.</summary>
    private RenderTexture? _pingRT;

    /// <summary>Second ping-pong buffer. Created by EnsureBuffers, released by ReleaseBuffers.</summary>
    private RenderTexture? _pongRT;

    /// <summary>
    /// Buffer returned as the output of the previous Process call (ping or pong), or null when the previous frame was
    /// not modified by any hook. Its readback may still be in flight, so the next frame starts writing the other buffer.
    /// </summary>
    private RenderTexture? _lastOutputRT;

    /// <summary>Context object passed to the hooks; reset and reused for every frame (no allocation per frame).</summary>
    private FrameHookContext? _reusableContext;

    /// <summary>Width of the current buffers in pixels (0 when no buffers exist).</summary>
    private int _width;

    /// <summary>Height of the current buffers in pixels (0 when no buffers exist).</summary>
    private int _height;

    /// <summary>Format of the source texture the current buffers were created for.</summary>
    private RenderTextureFormat _format;

    /// <summary>
    /// Runs <paramref name="hooks"/> in order on <paramref name="sourceTexture"/> and returns the resulting texture.
    /// An exception thrown by a hook is logged and the remaining hooks still run.
    /// RenderTexture.active is restored after the hooks.
    /// </summary>
    /// <param name="sourceTexture">Captured game frame. Hooks do not write into it (they write the ping-pong buffers).</param>
    /// <param name="frameIndex">Video frame index (0-based) of this frame.</param>
    /// <param name="renderFrameIndex">
    /// Render frame index (0-based, RenderManager.frameCountForVideo): counts every game frame since capture started,
    /// including game frames that were not captured because they fall in an already captured video frame.
    /// </param>
    /// <param name="videoTime">Time of this video frame in seconds from the start of the video.</param>
    /// <param name="renderFps">Render frame rate (game frames per second).</param>
    /// <param name="videoFps">Effective video frame rate (min(VideoFPS, RenderFPS)).</param>
    /// <param name="hooks">Hooks to run, in order. Session hooks come before global hooks.</param>
    /// <returns>
    /// <paramref name="sourceTexture"/> when there are no hooks, the source is not created, the buffers could not be
    /// created, or no hook modified the frame; otherwise the texture holding the final frame (a ping-pong buffer, or a
    /// texture a hook passed to FrameHookContext.ReplaceCurrent).
    /// </returns>
    public RenderTexture? Process(
        RenderTexture sourceTexture,
        long frameIndex,
        long renderFrameIndex,
        double videoTime,
        double renderFps,
        double videoFps,
        IReadOnlyList<Action<FrameHookContext>> hooks)
    {
        if (hooks.Count == 0 || !sourceTexture || !sourceTexture.IsCreated())
            return sourceTexture;

        EnsureBuffers(sourceTexture);

        if (!_pingRT || !_pongRT || !_pingRT.IsCreated() || !_pongRT.IsCreated())
            return sourceTexture;

        var prevActive = RenderTexture.active;
        try
        {
            // Inter-frame ping-pong safety: start with the buffer that was NOT used as the output
            // of the previous frame, ensuring the buffer currently being read back is not immediately overwritten.
            var firstDestination = _lastOutputRT == _pingRT ? _pongRT : _pingRT;

            if (_reusableContext == null)
            {
                _reusableContext = new FrameHookContext(
                    sourceTexture,
                    _pingRT,
                    _pongRT,
                    firstDestination,
                    frameIndex,
                    renderFrameIndex,
                    videoTime,
                    renderFps,
                    videoFps
                );
            }
            else
            {
                _reusableContext.Reset(
                    sourceTexture,
                    _pingRT,
                    _pongRT,
                    firstDestination,
                    frameIndex,
                    renderFrameIndex,
                    videoTime,
                    renderFps,
                    videoFps
                );
            }

            for (var i = 0; i < hooks.Count; i++)
            {
                try
                {
                    hooks[i](_reusableContext);
                }
                catch (Exception ex)
                {
                    Melon<RenderMod>.Logger.Error($"[FrameHook] Error executing frame hook #{i} on frame {frameIndex}: {ex}");
                }
            }

            // Current is the source when nothing was modified. A texture set by ReplaceCurrent is stored in
            // _lastOutputRT too; it then matches neither buffer and the next frame starts with ping.
            var result = _reusableContext.Current;
            _lastOutputRT = _reusableContext.HasModified ? result : null;
            return result;
        }
        finally
        {
            RenderTexture.active = prevActive;
        }
    }

    /// <summary>
    /// Creates the ping and pong buffers when they do not exist, were released, or the source size or format changed.
    /// The buffers copy the source descriptor without depth, MSAA and mipmaps, and enable random write (UAV) when the
    /// GPU supports compute shaders and random write for that format, so compute shaders can write them.
    /// When creating them from the descriptor throws, ARGB32 buffers of the same size are created.
    /// </summary>
    private void EnsureBuffers(RenderTexture sourceTexture)
    {
        var width = sourceTexture.width;
        var height = sourceTexture.height;
        var format = sourceTexture.format;

        if (_pingRT && _pongRT &&
            _width == width && _height == height && _format == format &&
            _pingRT.IsCreated() && _pongRT.IsCreated())
        {
            return;
        }

        ReleaseBuffers();

        _width = width;
        _height = height;
        _format = format;

        try
        {
            var desc = sourceTexture.descriptor;
            desc.depthBufferBits = 0;
            desc.msaaSamples = 1;
            desc.useMipMap = false;
            desc.autoGenerateMips = false;

            var supportsUav = SystemInfo.supportsComputeShaders &&
                              SystemInfo.SupportsRandomWriteOnRenderTextureFormat(desc.colorFormat);
            desc.enableRandomWrite = supportsUav;

            _pingRT = new RenderTexture(desc)
            {
                name = "ADOFAIRenderer_PingBuffer",
                filterMode = sourceTexture.filterMode,
                wrapMode = sourceTexture.wrapMode
            };
            _pingRT.Create();

            _pongRT = new RenderTexture(desc)
            {
                name = "ADOFAIRenderer_PongBuffer",
                filterMode = sourceTexture.filterMode,
                wrapMode = sourceTexture.wrapMode
            };
            _pongRT.Create();
        }
        catch (Exception ex)
        {
            ReleaseBuffers();
            Melon<RenderMod>.Logger.Warning($"[FrameHook] Failed to create descriptor-based PingPong RTs: {ex.Message}. Falling back to standard ARGB32.");
            _pingRT = new RenderTexture(width, height, 0, RenderTextureFormat.ARGB32)
            {
                name = "ADOFAIRenderer_PingBuffer",
                filterMode = sourceTexture.filterMode,
                wrapMode = sourceTexture.wrapMode
            };
            _pingRT.Create();

            _pongRT = new RenderTexture(width, height, 0, RenderTextureFormat.ARGB32)
            {
                name = "ADOFAIRenderer_PongBuffer",
                filterMode = sourceTexture.filterMode,
                wrapMode = sourceTexture.wrapMode
            };
            _pongRT.Create();
        }
    }

    /// <summary>
    /// Releases and destroys both buffers and drops the reusable context. The next Process call creates them again.
    /// </summary>
    public void ReleaseBuffers()
    {
        if (_pingRT)
        {
            if (_pingRT.IsCreated()) _pingRT.Release();
            UnityEngine.Object.Destroy(_pingRT);
            _pingRT = null;
        }

        if (_pongRT)
        {
            if (_pongRT.IsCreated()) _pongRT.Release();
            UnityEngine.Object.Destroy(_pongRT);
            _pongRT = null;
        }

        _width = 0;
        _height = 0;
        _lastOutputRT = null;
        _reusableContext = null;
    }

    /// <summary>Same as <see cref="ReleaseBuffers"/>. Main thread (Object.Destroy is used).</summary>
    public void Dispose()
    {
        ReleaseBuffers();
    }
}