acdream/src/AcDream.App/Rendering/RetailAlphaQueue.cs

using System;
using System.Collections.Generic;
using System.Numerics;
using System.Runtime.ExceptionServices;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using AcDream.App.Rendering.Residency;

namespace AcDream.App.Rendering;

internal static class RetailAlphaOrdering
{
    /// <summary>
    /// Retail <c>CPhysicsPart::UpdateViewerDistance</c> 0x0050E030:
    /// transform the GfxObj's authored sort center through the part draw frame,
    /// then store its Euclidean distance from the viewer as <c>CYpt</c>.
    /// </summary>
    public static float ComputeViewerDistance(
        Vector3 localSortCenter,
        Matrix4x4 model,
        Vector3 cameraWorldPosition)
        => Vector3.Distance(Vector3.Transform(localSortCenter, model), cameraWorldPosition);
}

/// <summary>
/// Renderer-owned half of retail's delayed alpha list. A source keeps the
/// immutable payload addressed by each token until the queue flushes, then
/// releases the payload when <see cref="ResetAlphaSubmissions"/> is called.
/// </summary>
internal interface IRetailAlphaDrawSource
{
    /// <summary>Uploads this source's complete payload for the current sorted
    /// alpha scope exactly once. Tokens arrive in final far-to-near order.</summary>
    void PrepareAlphaDraws(ReadOnlySpan<int> tokens);

    /// <summary>Draws a contiguous range from the payload prepared above.</summary>
    void DrawPreparedAlphaBatch(int firstPreparedDraw, int drawCount);

    void ResetAlphaSubmissions();
}

internal readonly record struct RetailAlphaSubmission(
    IRetailAlphaDrawSource Source,
    int Token,
    float ViewerDistance);

/// <summary>
/// Frame-scoped port of retail's shared <c>D3DPolyRender</c> alpha list.
/// Scene-particle parts and ordinary GfxObj parts enter one stable,
/// far-to-near stream instead of being composited in renderer-local passes.
///
/// Retail oracle:
/// <list type="bullet">
/// <item><c>CShadowPart::insertion_sort</c> 0x006B5130 orders all cell parts by
/// <c>CPhysicsPart::CYpt</c>.</item>
/// <item><c>D3DPolyRender::AddMeshToAlphaList</c> 0x0059C230 appends delayed
/// surface subsets in that established order.</item>
/// <item><c>D3DPolyRender::FlushAlphaList</c> 0x0059D2E0 drains without another
/// material/renderer sort.</item>
/// </list>
/// </summary>
internal interface IWorldSceneAlphaFrame
{
    void BeginFrame();

    void EndFrame();

    void AbortFrame();
}

internal sealed class RetailAlphaQueue : IWorldSceneAlphaFrame
{
    private const int MinimumSubmissionCapacity = 256;
    private const int SubmissionGrowthQuantum = 256;
    private const int MinimumSourceCapacity = 4;
    private readonly List<RetailAlphaSubmission> _submissions = new(256);
    private readonly List<IRetailAlphaDrawSource> _sources = new(4);
    private int[] _tokenScratch = new int[256];
    private int[] _sourceDrawOffsets = new int[4];
    private RetailAlphaSubmission[] _sortScratch = new RetailAlphaSubmission[256];
    private readonly int[] _radixOffsets = new int[256];
    private readonly RetainedScratchCapacityPolicy _scratchPolicy;
    private readonly long _scratchBudgetBytes;

    internal RetailAlphaQueue(long? scratchBudgetBytes = null)
    {
        long budget = scratchBudgetBytes
            ?? AlphaScratchBudgetProfile.Create(
                ResidencyBudgetOptions.Default.AlphaScratchBytes).QueueBytes;
        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(budget);
        _scratchBudgetBytes = budget;
        _scratchPolicy = new RetainedScratchCapacityPolicy(Math.Max(
            1,
            budget - (long)_radixOffsets.Length * sizeof(int)));
    }

    public bool IsCollecting { get; private set; }
    internal int PendingCount => _submissions.Count;
    internal long ScratchBudgetBytes => _scratchBudgetBytes;
    internal long RetainedScratchBytes => checked(
        (long)_submissions.Capacity
            * Unsafe.SizeOf<RetailAlphaSubmission>()
        + (long)_sortScratch.Length
            * Unsafe.SizeOf<RetailAlphaSubmission>()
        + (long)_tokenScratch.Length * sizeof(int)
        + (long)_sources.Capacity * IntPtr.Size
        + (long)_sourceDrawOffsets.Length * sizeof(int)
        + (long)_radixOffsets.Length * sizeof(int));

    public void BeginFrame()
    {
        if (IsCollecting)
            throw new InvalidOperationException("Retail alpha frame is already active.");
        if (_submissions.Count != 0 || _sources.Count != 0)
            throw new InvalidOperationException("Retail alpha queue retained payload outside a frame.");

        IsCollecting = true;
    }

    public void Submit(IRetailAlphaDrawSource source, int token, float viewerDistance)
    {
        ArgumentNullException.ThrowIfNull(source);
        if (!IsCollecting)
            throw new InvalidOperationException("Retail alpha submission requires an active frame.");
        if (token < 0)
            throw new ArgumentOutOfRangeException(nameof(token));

        _submissions.Add(new RetailAlphaSubmission(
            source,
            token,
            NormalizeDistance(viewerDistance)));

        for (int i = 0; i < _sources.Count; i++)
            if (ReferenceEquals(_sources[i], source))
                return;
        _sources.Add(source);
    }

    /// <summary>
    /// Drains the current retail alpha scope and keeps the frame open so the
    /// caller can collect the post-depth-clear scope. Only adjacent entries
    /// from one renderer are handed over as a batch; the queue never groups by
    /// renderer across another entry and therefore cannot alter compositing.
    /// </summary>
    public void Flush()
    {
        if (!IsCollecting)
            throw new InvalidOperationException("Retail alpha flush requires an active frame.");

        Exception? drawFailure = null;
        List<Exception>? resetFailures = null;
        try
        {
            if (_submissions.Count > 0)
            {
                SortRetailOrder();
                EnsureTokenCapacity(_submissions.Count);
                EnsureSourceCapacity(_sources.Count);
                Array.Clear(_sourceDrawOffsets, 0, _sources.Count);

                // Prepare each renderer once for this alpha scope. The filtered
                // token sequence preserves final retail order for that source, so
                // every later adjacent source-run maps to one contiguous prepared
                // range without another GPU upload.
                for (int sourceIndex = 0; sourceIndex < _sources.Count; sourceIndex++)
                {
                    IRetailAlphaDrawSource source = _sources[sourceIndex];
                    int sourceCount = 0;
                    for (int i = 0; i < _submissions.Count; i++)
                    {
                        RetailAlphaSubmission submission = _submissions[i];
                        if (ReferenceEquals(submission.Source, source))
                            _tokenScratch[sourceCount++] = submission.Token;
                    }

                    source.PrepareAlphaDraws(_tokenScratch.AsSpan(0, sourceCount));
                }

                int start = 0;
                while (start < _submissions.Count)
                {
                    IRetailAlphaDrawSource source = _submissions[start].Source;
                    int end = start + 1;
                    while (end < _submissions.Count
                           && ReferenceEquals(_submissions[end].Source, source))
                        end++;

                    int count = end - start;
                    int sourceIndex = FindSourceIndex(source);
                    int firstPreparedDraw = _sourceDrawOffsets[sourceIndex];
                    source.DrawPreparedAlphaBatch(firstPreparedDraw, count);
                    _sourceDrawOffsets[sourceIndex] += count;
                    start = end;
                }
            }
        }
        catch (Exception error)
        {
            drawFailure = error;
        }
        finally
        {
            int submissionCount = _submissions.Count;
            int sourceCount = _sources.Count;
            for (int i = 0; i < _sources.Count; i++)
            {
                try
                {
                    _sources[i].ResetAlphaSubmissions();
                }
                catch (Exception error)
                {
                    (resetFailures ??= []).Add(error);
                }
            }
            _sources.Clear();
            _submissions.Clear();
            ApplyScratchRetention(submissionCount, sourceCount);
        }

        if (drawFailure is not null)
        {
            if (resetFailures is { Count: > 0 })
            {
                resetFailures.Insert(0, drawFailure);
                throw new AggregateException(
                    "Retail alpha drawing failed and its submissions could not be fully reset.",
                    resetFailures);
            }

            ExceptionDispatchInfo.Capture(drawFailure).Throw();
        }

        if (resetFailures is { Count: > 0 })
        {
            throw new AggregateException(
                "Retail alpha submissions could not be fully reset.",
                resetFailures);
        }
    }

    public void EndFrame()
    {
        if (!IsCollecting)
            throw new InvalidOperationException("Retail alpha frame is not active.");

        try
        {
            Flush();
        }
        finally
        {
            IsCollecting = false;
        }
    }

    /// <summary>
    /// Discards an incomplete frame without drawing it. Every source is still
    /// told to release its retained submission payload so the next frame begins
    /// from the same empty invariant as a successful flush.
    /// </summary>
    public void AbortFrame()
    {
        List<Exception>? failures = null;
        try
        {
            for (int i = 0; i < _sources.Count; i++)
            {
                try
                {
                    _sources[i].ResetAlphaSubmissions();
                }
                catch (Exception error)
                {
                    (failures ??= []).Add(error);
                }
            }
        }
        finally
        {
            int submissionCount = _submissions.Count;
            int sourceCount = _sources.Count;
            _sources.Clear();
            _submissions.Clear();
            IsCollecting = false;
            ApplyScratchRetention(submissionCount, sourceCount);
        }

        if (failures is { Count: > 0 })
            throw new AggregateException("Retail alpha frame abort failed.", failures);
    }

    private void SortRetailOrder()
    {
        int count = _submissions.Count;
        if (count <= 1)
            return;

        EnsureSortCapacity(count);
        Span<RetailAlphaSubmission> submissions = CollectionsMarshal.AsSpan(_submissions);
        Span<RetailAlphaSubmission> scratch = _sortScratch.AsSpan(0, count);

        // Positive IEEE-754 float bits sort in the same order as their numeric
        // values. Complementing those bits turns an ascending stable radix pass
        // into retail's descending CYpt order. Four stable byte passes preserve
        // original submission order for equal distances, matching
        // CShadowPart::insertion_sort without List.Sort's O(n log n) interface
        // comparator overhead in particle-heavy views.
        RadixPass(submissions, scratch, shift: 0);
        RadixPass(scratch, submissions, shift: 8);
        RadixPass(submissions, scratch, shift: 16);
        RadixPass(scratch, submissions, shift: 24);
    }

    private static float NormalizeDistance(float value)
        // value > 0 also canonicalizes -0 to +0 so its IEEE sort key joins the
        // same stable equal-distance run as ordinary zero.
        => float.IsFinite(value) && value > 0f ? value : 0f;

    private void RadixPass(
        ReadOnlySpan<RetailAlphaSubmission> source,
        Span<RetailAlphaSubmission> destination,
        int shift)
    {
        Array.Clear(_radixOffsets);
        for (int i = 0; i < source.Length; i++)
        {
            uint key = ~BitConverter.SingleToUInt32Bits(source[i].ViewerDistance);
            _radixOffsets[(int)((key >> shift) & 0xFF)]++;
        }

        int prefix = 0;
        for (int bucket = 0; bucket < _radixOffsets.Length; bucket++)
        {
            int bucketCount = _radixOffsets[bucket];
            _radixOffsets[bucket] = prefix;
            prefix += bucketCount;
        }

        for (int i = 0; i < source.Length; i++)
        {
            RetailAlphaSubmission submission = source[i];
            uint key = ~BitConverter.SingleToUInt32Bits(submission.ViewerDistance);
            int bucket = (int)((key >> shift) & 0xFF);
            destination[_radixOffsets[bucket]++] = submission;
        }
    }

    private void EnsureTokenCapacity(int count)
    {
        if (_tokenScratch.Length >= count)
            return;
        Array.Resize(ref _tokenScratch, count + 256);
    }

    private void EnsureSourceCapacity(int count)
    {
        if (_sourceDrawOffsets.Length >= count)
            return;
        Array.Resize(ref _sourceDrawOffsets, count + 4);
    }

    private void EnsureSortCapacity(int count)
    {
        if (_sortScratch.Length >= count)
            return;
        Array.Resize(ref _sortScratch, count + 256);
    }

    private int FindSourceIndex(IRetailAlphaDrawSource source)
    {
        for (int i = 0; i < _sources.Count; i++)
            if (ReferenceEquals(_sources[i], source))
                return i;
        throw new InvalidOperationException("Retail alpha source was not registered for this frame.");
    }

    private void ApplyScratchRetention(
        int observedSubmissionCount,
        int observedSourceCount)
    {
        int currentCapacity = Math.Max(
            _submissions.Capacity,
            Math.Max(_tokenScratch.Length, _sortScratch.Length));
        int bytesPerSubmission =
            checked(
                2 * Unsafe.SizeOf<RetailAlphaSubmission>()
                + 2 * sizeof(int)
                + IntPtr.Size);
        int targetCapacity = _scratchPolicy.ObserveAndSelectCapacity(
            currentCapacity,
            observedSubmissionCount,
            bytesPerSubmission,
            MinimumSubmissionCapacity,
            SubmissionGrowthQuantum);
        if (targetCapacity < currentCapacity)
        {
            _submissions.Capacity = targetCapacity;
            Array.Resize(ref _tokenScratch, targetCapacity);
            Array.Resize(ref _sortScratch, targetCapacity);

            int sourceTarget = Math.Max(
                MinimumSourceCapacity,
                observedSourceCount == 0
                    ? MinimumSourceCapacity
                    : checked(observedSourceCount * 2));
            sourceTarget = Math.Min(sourceTarget, _sources.Capacity);
            _sources.Capacity = sourceTarget;
            if (_sourceDrawOffsets.Length > sourceTarget)
                Array.Resize(ref _sourceDrawOffsets, sourceTarget);
        }
    }
}