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

using System;
using System.Collections.Generic;
using System.Numerics;
using AcDream.Core.Vfx;
using DatReaderWriter;
using DatReaderWriter.DBObjs;
using DatReaderWriter.Enums;
using Silk.NET.OpenGL;

namespace AcDream.App.Rendering;

/// <summary>
/// Instanced renderer for retail particle emitters.
/// </summary>
public sealed unsafe class ParticleRenderer : IDisposable
{
    private readonly record struct BatchKey(uint TextureHandle, bool UseTexture, bool Additive);
    private readonly record struct ParticleDraw(BatchKey Key, ParticleInstance Instance);

    private readonly struct ParticleInstance
    {
        public readonly Vector3 Position;
        public readonly Vector3 AxisX;
        public readonly Vector3 AxisY;
        public readonly uint ColorArgb;
        public readonly float DistanceSq;

        public ParticleInstance(Vector3 position, Vector3 axisX, Vector3 axisY, uint colorArgb, float distanceSq)
        {
            Position = position;
            AxisX = axisX;
            AxisY = axisY;
            ColorArgb = colorArgb;
            DistanceSq = distanceSq;
        }
    }

    private readonly GL _gl;
    private readonly Shader _shader;
    private readonly TextureCache? _textures;
    private readonly DatCollection? _dats;
    private readonly Dictionary<uint, ParticleGfxInfo> _particleGfxInfoByGfxObj = new();

    private readonly uint _quadVao;
    private readonly uint _quadVbo;
    private readonly uint _quadEbo;
    private readonly uint _instanceVbo;

    private float[] _instanceScratch = new float[256 * 16];

    public ParticleRenderer(GL gl, string shadersDir, TextureCache? textures = null, DatCollection? dats = null)
    {
        _gl = gl ?? throw new ArgumentNullException(nameof(gl));
        _textures = textures;
        _dats = dats;
        _shader = new Shader(_gl,
            System.IO.Path.Combine(shadersDir, "particle.vert"),
            System.IO.Path.Combine(shadersDir, "particle.frag"));

        float[] quadVerts =
        {
            -0.5f, -0.5f, 0f, 0f,
             0.5f, -0.5f, 1f, 0f,
             0.5f,  0.5f, 1f, 1f,
            -0.5f,  0.5f, 0f, 1f,
        };
        uint[] quadIdx = { 0, 1, 2, 0, 2, 3 };

        _quadVao = _gl.GenVertexArray();
        _gl.BindVertexArray(_quadVao);

        _quadVbo = _gl.GenBuffer();
        _gl.BindBuffer(BufferTargetARB.ArrayBuffer, _quadVbo);
        fixed (void* p = quadVerts)
        {
            _gl.BufferData(
                BufferTargetARB.ArrayBuffer,
                (nuint)(quadVerts.Length * sizeof(float)),
                p,
                BufferUsageARB.StaticDraw);
        }

        _gl.EnableVertexAttribArray(0);
        _gl.VertexAttribPointer(0, 2, VertexAttribPointerType.Float, false, 4 * sizeof(float), (void*)0);
        _gl.EnableVertexAttribArray(1);
        _gl.VertexAttribPointer(1, 2, VertexAttribPointerType.Float, false, 4 * sizeof(float), (void*)(2 * sizeof(float)));

        _quadEbo = _gl.GenBuffer();
        _gl.BindBuffer(BufferTargetARB.ElementArrayBuffer, _quadEbo);
        fixed (void* p = quadIdx)
        {
            _gl.BufferData(
                BufferTargetARB.ElementArrayBuffer,
                (nuint)(quadIdx.Length * sizeof(uint)),
                p,
                BufferUsageARB.StaticDraw);
        }

        _instanceVbo = _gl.GenBuffer();
        _gl.BindBuffer(BufferTargetARB.ArrayBuffer, _instanceVbo);
        _gl.BufferData(BufferTargetARB.ArrayBuffer, (nuint)(256 * 16 * sizeof(float)), (void*)0, BufferUsageARB.DynamicDraw);

        _gl.EnableVertexAttribArray(2);
        _gl.VertexAttribPointer(2, 4, VertexAttribPointerType.Float, false, 16 * sizeof(float), (void*)0);
        _gl.VertexAttribDivisor(2, 1);
        _gl.EnableVertexAttribArray(3);
        _gl.VertexAttribPointer(3, 4, VertexAttribPointerType.Float, false, 16 * sizeof(float), (void*)(4 * sizeof(float)));
        _gl.VertexAttribDivisor(3, 1);
        _gl.EnableVertexAttribArray(4);
        _gl.VertexAttribPointer(4, 4, VertexAttribPointerType.Float, false, 16 * sizeof(float), (void*)(8 * sizeof(float)));
        _gl.VertexAttribDivisor(4, 1);
        _gl.EnableVertexAttribArray(5);
        _gl.VertexAttribPointer(5, 4, VertexAttribPointerType.Float, false, 16 * sizeof(float), (void*)(12 * sizeof(float)));
        _gl.VertexAttribDivisor(5, 1);

        _gl.BindVertexArray(0);
    }

    public void Draw(
        ParticleSystem particles,
        ICamera camera,
        Vector3 cameraWorldPos,
        ParticleRenderPass renderPass = ParticleRenderPass.Scene,
        Func<AcDream.Core.Vfx.ParticleEmitter, bool>? emitterFilter = null)
    {
        if (particles is null || camera is null)
            return;

        Matrix4x4.Invert(camera.View, out var invView);
        Vector3 cameraRight = Vector3.Normalize(new Vector3(invView.M11, invView.M12, invView.M13));
        Vector3 cameraUp = Vector3.Normalize(new Vector3(invView.M21, invView.M22, invView.M23));
        var draws = BuildDrawList(particles, cameraWorldPos, renderPass, cameraRight, cameraUp, emitterFilter);
        if (draws.Count == 0)
            return;
        draws.Sort(static (a, b) => b.Instance.DistanceSq.CompareTo(a.Instance.DistanceSq));

        _shader.Use();
        _shader.SetMatrix4("uViewProjection", camera.View * camera.Projection);
        _shader.SetInt("uParticleTexture", 0);

        _gl.Enable(EnableCap.DepthTest);
        _gl.Enable(EnableCap.Blend);
        _gl.DepthMask(false);
        _gl.Disable(EnableCap.CullFace);
        _gl.ActiveTexture(TextureUnit.Texture0);

        var run = new List<ParticleInstance>(64);
        for (int i = 0; i < draws.Count;)
        {
            var key = draws[i].Key;
            run.Clear();
            do
            {
                run.Add(draws[i].Instance);
                i++;
            }
            while (i < draws.Count && draws[i].Key == key);

            _gl.BlendFunc(
                BlendingFactor.SrcAlpha,
                key.Additive ? BlendingFactor.One : BlendingFactor.OneMinusSrcAlpha);
            _shader.SetInt("uUseTexture", key.UseTexture ? 1 : 0);
            _gl.BindTexture(TextureTarget.Texture2D, key.UseTexture ? key.TextureHandle : 0);
            DrawInstances(run);
        }

        _gl.BindTexture(TextureTarget.Texture2D, 0);
        _gl.BindVertexArray(0);
        _gl.DepthMask(true);
        _gl.Disable(EnableCap.Blend);
    }

    private List<ParticleDraw> BuildDrawList(
        ParticleSystem particles,
        Vector3 cameraWorldPos,
        ParticleRenderPass renderPass,
        Vector3 cameraRight,
        Vector3 cameraUp,
        Func<AcDream.Core.Vfx.ParticleEmitter, bool>? emitterFilter)
    {
        var draws = new List<ParticleDraw>(Math.Max(64, particles.ActiveParticleCount));
        foreach (var (em, idx) in particles.EnumerateLive())
        {
            if (em.RenderPass != renderPass)
                continue;
            if (emitterFilter is not null && !emitterFilter(em))
                continue;

            ref var p = ref em.Particles[idx];
            // `p.Position` is already in world coordinates: AttachLocal
            // emitters get their AnchorPos refreshed each frame by the
            // owning subsystem (sky-PES driver, animation tick, etc.) which
            // mirrors retail's live-parent-frame read at
            // ParticleEmitter::UpdateParticles 0x0051d2d4 for is_parent_local=1.
            Vector3 pos = p.Position;
            float distSq = Vector3.DistanceSquared(pos, cameraWorldPos);
            var gfxInfo = ResolveParticleGfxInfo(em.Desc);
            uint texture = gfxInfo.TextureHandle;
            bool useTexture = texture != 0;
            bool additive = gfxInfo.HasMaterial
                ? gfxInfo.Additive
                : (em.Desc.Flags & EmitterFlags.Additive) != 0;
            var key = new BatchKey(texture, useTexture, additive);
            Vector3 axisX;
            Vector3 axisY;
            if (gfxInfo.IsBillboard)
            {
                pos += Vector3.UnitZ * (gfxInfo.CenterOffset.Z * p.Size);
                axisX = cameraRight * (gfxInfo.Size.X * p.Size);
                axisY = cameraUp * (gfxInfo.Size.Y * p.Size);
            }
            else
            {
                Quaternion orientation = ParticleOrientation(em, p);
                pos += Vector3.Transform(gfxInfo.CenterOffset * p.Size, orientation);
                axisX = Vector3.Transform(gfxInfo.AxisX, orientation) * (gfxInfo.Size.X * p.Size);
                axisY = Vector3.Transform(gfxInfo.AxisY, orientation) * (gfxInfo.Size.Y * p.Size);
            }

            draws.Add(new ParticleDraw(key, new ParticleInstance(pos, axisX, axisY, p.ColorArgb, distSq)));
        }

        return draws;
    }

    private void DrawInstances(List<ParticleInstance> instances)
    {
        if (instances.Count == 0)
            return;

        int needed = instances.Count * 16;
        if (_instanceScratch.Length < needed)
            _instanceScratch = new float[needed + 256 * 16];

        for (int i = 0; i < instances.Count; i++)
        {
            var p = instances[i];
            int o = i * 16;
            _instanceScratch[o + 0] = p.Position.X;
            _instanceScratch[o + 1] = p.Position.Y;
            _instanceScratch[o + 2] = p.Position.Z;
            _instanceScratch[o + 3] = 0f;

            _instanceScratch[o + 4] = p.AxisX.X;
            _instanceScratch[o + 5] = p.AxisX.Y;
            _instanceScratch[o + 6] = p.AxisX.Z;
            _instanceScratch[o + 7] = 0f;

            _instanceScratch[o + 8] = p.AxisY.X;
            _instanceScratch[o + 9] = p.AxisY.Y;
            _instanceScratch[o + 10] = p.AxisY.Z;
            _instanceScratch[o + 11] = 0f;

            _instanceScratch[o + 12] = ((p.ColorArgb >> 16) & 0xFF) / 255f;
            _instanceScratch[o + 13] = ((p.ColorArgb >> 8) & 0xFF) / 255f;
            _instanceScratch[o + 14] = (p.ColorArgb & 0xFF) / 255f;
            _instanceScratch[o + 15] = ((p.ColorArgb >> 24) & 0xFF) / 255f;
        }

        _gl.BindBuffer(BufferTargetARB.ArrayBuffer, _instanceVbo);
        fixed (void* bp = _instanceScratch)
        {
            _gl.BufferData(
                BufferTargetARB.ArrayBuffer,
                (nuint)(instances.Count * 16 * sizeof(float)),
                bp,
                BufferUsageARB.DynamicDraw);
        }

        _gl.BindVertexArray(_quadVao);
        _gl.DrawElementsInstanced(PrimitiveType.Triangles, 6, DrawElementsType.UnsignedInt, (void*)0, (uint)instances.Count);
    }

    private ParticleGfxInfo ResolveParticleGfxInfo(EmitterDesc desc)
    {
        if (_textures is null)
            return ParticleGfxInfo.Default;

        if (desc.TextureSurfaceId != 0)
            return ParticleGfxInfo.Billboard(
                _textures.GetOrUpload(desc.TextureSurfaceId),
                Vector2.One,
                Vector3.Zero,
                additive: (desc.Flags & EmitterFlags.Additive) != 0,
                hasMaterial: false);

        uint gfxObjId = desc.HwGfxObjId != 0 ? desc.HwGfxObjId : desc.GfxObjId;
        if (gfxObjId == 0 || _dats is null)
            return ParticleGfxInfo.Default;

        if (!_particleGfxInfoByGfxObj.TryGetValue(gfxObjId, out var info))
        {
            info = ReadParticleGfxInfo(gfxObjId);
            _particleGfxInfoByGfxObj[gfxObjId] = info;
        }

        return info.TextureHandle != 0 ? info : ParticleGfxInfo.Default;
    }

    private ParticleGfxInfo ReadParticleGfxInfo(uint gfxObjId)
    {
        try
        {
            var gfx = _dats?.Get<GfxObj>(gfxObjId);
            if (gfx is null)
                return ParticleGfxInfo.Default;

            uint surfaceId = gfx.Surfaces.Count > 0 ? gfx.Surfaces[0].DataId : 0u;
            uint texture = surfaceId != 0 && _textures is not null ? _textures.GetOrUpload(surfaceId) : 0u;
            bool additive = false;
            if (surfaceId != 0)
            {
                var surface = _dats?.Get<Surface>(surfaceId);
                additive = surface is not null && surface.Type.HasFlag(SurfaceType.Additive);
            }
            return AuthoredParticleGfxInfo(gfx, texture, additive, surfaceId != 0);
        }
        catch
        {
            return ParticleGfxInfo.Default;
        }
    }

    private ParticleGfxInfo AuthoredParticleGfxInfo(GfxObj gfx, uint texture, bool additive, bool hasMaterial)
    {
        if (gfx.VertexArray.Vertices.Count == 0)
            return ParticleGfxInfo.Billboard(texture, Vector2.One, Vector3.Zero, additive, hasMaterial);

        var min = new Vector3(float.PositiveInfinity);
        var max = new Vector3(float.NegativeInfinity);
        foreach (var (_, v) in gfx.VertexArray.Vertices)
        {
            min = Vector3.Min(min, v.Origin);
            max = Vector3.Max(max, v.Origin);
        }

        var size = max - min;
        var center = (min + max) * 0.5f;
        if (IsPointSprite(gfx))
        {
            float sx = FallbackParticleExtent(size.X) * 0.9f;
            float sy = FallbackParticleExtent(size.Z) * 0.9f;
            return ParticleGfxInfo.Billboard(texture, new Vector2(sx, sy), center, additive, hasMaterial);
        }

        Vector3 axisX;
        Vector3 axisY;
        Vector2 planeSize;
        if (size.Y > size.X && size.Y > size.Z)
        {
            if (size.X > size.Z)
            {
                axisX = Vector3.UnitX;
                axisY = Vector3.UnitY;
                planeSize = new Vector2(size.X, size.Y);
            }
            else
            {
                axisX = Vector3.UnitY;
                axisY = Vector3.UnitZ;
                planeSize = new Vector2(size.Y, size.Z);
            }
        }
        else if (size.X > size.Y && size.X > size.Z)
        {
            if (size.Z > size.Y)
            {
                axisX = Vector3.UnitX;
                axisY = Vector3.UnitZ;
                planeSize = new Vector2(size.X, size.Z);
            }
            else
            {
                axisX = Vector3.UnitX;
                axisY = Vector3.UnitY;
                planeSize = new Vector2(size.X, size.Y);
            }
        }
        else
        {
            if (size.X > size.Y)
            {
                axisX = Vector3.UnitX;
                axisY = Vector3.UnitZ;
                planeSize = new Vector2(size.X, size.Z);
            }
            else
            {
                axisX = Vector3.UnitY;
                axisY = Vector3.UnitZ;
                planeSize = new Vector2(size.Y, size.Z);
            }
        }

        planeSize.X = FallbackParticleExtent(planeSize.X);
        planeSize.Y = FallbackParticleExtent(planeSize.Y);
        return new ParticleGfxInfo(texture, planeSize, axisX, axisY, center, false, additive, hasMaterial);
    }

    private bool IsPointSprite(GfxObj gfx)
    {
        if (!gfx.Flags.HasFlag(GfxObjFlags.HasDIDDegrade) || gfx.DIDDegrade == 0 || _dats is null)
            return false;

        try
        {
            var degrade = _dats.Get<GfxObjDegradeInfo>(gfx.DIDDegrade);
            return degrade?.Degrades.Count > 0 && degrade.Degrades[0].DegradeMode == 2;
        }
        catch
        {
            return false;
        }
    }

    private static float FallbackParticleExtent(float value)
        => value > 1e-4f ? Math.Clamp(value, 1e-4f, 10_000f) : 1f;

    private static Quaternion ParticleOrientation(AcDream.Core.Vfx.ParticleEmitter em, Particle p)
    {
        Quaternion orientation = (em.Desc.Flags & EmitterFlags.AttachLocal) != 0
            ? em.AnchorRot
            : p.SpawnRotation;

        if (em.Desc.Type is AcDream.Core.Vfx.ParticleType.ParabolicLVGAGR
            or AcDream.Core.Vfx.ParticleType.ParabolicLVLALR
            or AcDream.Core.Vfx.ParticleType.ParabolicGVGAGR)
        {
            Vector3 angular = p.C * p.Age;
            float radians = angular.Length();
            if (radians > 1e-6f)
                orientation = Quaternion.Normalize(orientation * Quaternion.CreateFromAxisAngle(angular / radians, radians));
        }

        return orientation;
    }

    public void Dispose()
    {
        _gl.DeleteBuffer(_quadVbo);
        _gl.DeleteBuffer(_quadEbo);
        _gl.DeleteBuffer(_instanceVbo);
        _gl.DeleteVertexArray(_quadVao);
        _shader.Dispose();
    }

    private readonly record struct ParticleGfxInfo(
        uint TextureHandle,
        Vector2 Size,
        Vector3 AxisX,
        Vector3 AxisY,
        Vector3 CenterOffset,
        bool IsBillboard,
        bool Additive,
        bool HasMaterial)
    {
        public static ParticleGfxInfo Default { get; } =
            Billboard(0u, Vector2.One, Vector3.Zero, additive: false, hasMaterial: false);

        public static ParticleGfxInfo Billboard(
            uint textureHandle,
            Vector2 size,
            Vector3 centerOffset,
            bool additive,
            bool hasMaterial) =>
            new(textureHandle, size, Vector3.UnitX, Vector3.UnitY, centerOffset, true, additive, hasMaterial);
    }
}