acdream/src/AcDream.App/Rendering/Wb/GLSLShader.cs

using Chorizite.Core.Render;
using AcDream.App.Rendering;
using Microsoft.Extensions.Logging;
using Silk.NET.OpenGL;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Numerics;
using System.Reflection;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;
using System.Xml.Linq;

namespace AcDream.App.Rendering.Wb {

    public unsafe class GLSLShader : BaseShader, IDisposable {
        private OpenGLGraphicsDevice _device;
        private Dictionary<string, int> _uniformLocations = [];
        private Dictionary<int, object> _uniformValues = [];
        private readonly object _lock = new();
        private GL GL => _device.GL;
        public uint Program { get; protected set; }

        public bool HasUniform(string name) {
            lock (_lock) {
                return GetUniformLocation(Program, name) != -1;
            }
        }

        public GLSLShader(OpenGLGraphicsDevice device, string name, string vertSource, string fragSource, ILogger log) : base(name, vertSource, fragSource, log) {
            _device = device;

            Load(vertSource, fragSource);
        }

        public GLSLShader(OpenGLGraphicsDevice device, string name, string shaderDirectory, ILogger log) : base(name, shaderDirectory, log) {
            _device = device;

            Load();
        }

        public override void Dispose() {
            Unload();
            base.Dispose();
        }

        private int GetUniformLocation(uint program, string name) {
            lock (_lock) {
                if (!_uniformLocations.ContainsKey(name)) {
                    _uniformLocations.Add(name, GL.GetUniformLocation(program, name));
                }
                return _uniformLocations[name];
            }
        }

        public override void SetUniform(string location, Matrix4x4 m) {
            lock (_lock) {
                int loc = GetUniformLocation(Program, location);
                if (loc == -1) return;

                if (_uniformValues.TryGetValue(loc, out var val) && val is Matrix4x4 mCached && mCached == m) {
                    return;
                }
                _uniformValues[loc] = m;

                GL.UniformMatrix4(loc, 1, false, (float*)&m);
            }
        }

        public override void SetUniform(string location, int v) {
            lock (_lock) {
                int loc = GetUniformLocation((uint)Program, location);
                if (loc == -1) return;

                if (_uniformValues.TryGetValue(loc, out var val) && val is int vCached && vCached == v) {
                    return;
                }
                _uniformValues[loc] = v;

                GL.Uniform1(loc, v);
            }
        }

        public override void SetUniform(string location, Vector2 vec) {
            lock (_lock) {
                int loc = GetUniformLocation((uint)Program, location);
                if (loc == -1) return;

                if (_uniformValues.TryGetValue(loc, out var val) && val is Vector2 vCached && vCached == vec) {
                    return;
                }
                _uniformValues[loc] = vec;

                GL.Uniform2(loc, vec);
            }
        }

        public override void SetUniform(string location, Vector3 vec) {
            lock (_lock) {
                int loc = GetUniformLocation((uint)Program, location);
                if (loc == -1) return;

                if (_uniformValues.TryGetValue(loc, out var val) && val is Vector3 vCached && vCached == vec) {
                    return;
                }
                _uniformValues[loc] = vec;

                GL.Uniform3(loc, vec);
            }
        }


        public override void SetUniform(string location, Vector3[] vecs) {
            lock (_lock) {
                int loc = GetUniformLocation((uint)Program, location);
                if (loc == -1) return;

                fixed (float* v = &vecs[0].X) {
                    GL.Uniform3(loc, (uint)vecs.Length, v);
                }
            }
        }

        public override void SetUniform(string location, Vector4 vec) {
            lock (_lock) {
                int loc = GetUniformLocation((uint)Program, location);
                if (loc == -1) return;

                if (_uniformValues.TryGetValue(loc, out var val) && val is Vector4 vCached && vCached == vec) {
                    return;
                }
                _uniformValues[loc] = vec;

                GL.Uniform4(loc, vec);
            }
        }

        public override void SetUniform(string location, float v) {
            lock (_lock) {
                int loc = GetUniformLocation((uint)Program, location);
                if (loc == -1) return;

                if (_uniformValues.TryGetValue(loc, out var val) && val is float vCached && vCached == v) {
                    return;
                }
                _uniformValues[loc] = v;

                GL.Uniform1(loc, v);
            }
        }

        public override void SetUniform(string location, float[] vs) {
            lock (_lock) {
                fixed (float* v = &vs[0]) {
                    GL.Uniform1(GetUniformLocation((uint)Program, location), (uint)vs.Length, v);
                }
            }
        }

        public override void Load(string vertShaderSource, string fragShaderSource) {

            if (string.IsNullOrWhiteSpace(vertShaderSource) || string.IsNullOrWhiteSpace(fragShaderSource)) {
                _log.LogError($"Shader {Name} has no source code!");
                throw new InvalidOperationException($"Shader {Name} has no source code.");
            }

            if (_device.HasOpenGL43 && _device.HasBindless) {
                string replacement = "#version 430 core\n#extension GL_ARB_bindless_texture : require";
                vertShaderSource = vertShaderSource.Replace("#version 330 core", replacement);
                fragShaderSource = fragShaderSource.Replace("#version 330 core", replacement);
            }

            var resources = new ResourceCleanupGroup();
            uint prog = 0;
            bool accountingPublished = false;
            try {
                prog = ShaderProgramConstruction.Build(
                    new GlShaderProgramBuildApi(GL),
                    vertShaderSource,
                    fragShaderSource);
                uint ownedProgram = prog;
                var unpublishedProgramRelease = new RetryableGpuResourceRelease(
                    () => GlResourceCommand.DeleteProgram(
                        GL,
                        ownedProgram,
                        $"delete unpublished WB shader program {ownedProgram}"),
                    () => {
                        if (accountingPublished)
                        {
                            GpuMemoryTracker.TrackResourceDeallocation(
                                GpuResourceType.Shader);
                        }
                    });
                resources.Add("WB shader program", unpublishedProgramRelease.Run);

                // Bind SceneData uniform block to point 0 if it exists.
                GlResourceCommand.Execute(GL, $"configure shader {Name} SceneData binding", () => {
                    uint sceneDataIndex = GL.GetUniformBlockIndex(prog, "SceneData");
                    if (sceneDataIndex != uint.MaxValue)
                        GL.UniformBlockBinding(prog, sceneDataIndex, 0);
                });

                GpuMemoryTracker.TrackResourceAllocation(GpuResourceType.Shader);
                accountingPublished = true;
                resources.TransferAll();
            } catch (Exception constructionFailure) {
                _log.LogError(constructionFailure, "Failed to construct shader {ShaderName}", Name);
                resources.RollbackConstructionAndThrow(
                    $"Shader {Name} construction failed and its GL program did not cleanly roll back.",
                    constructionFailure);
            }

            _log.LogTrace($"{(Program != 0 ? "Reloaded" : "Loaded")} shader: {Name}");

            if (Program != 0) {
                Unload();
            }
            _uniformLocations.Clear();
            _uniformValues.Clear();

            Program = prog;
            ProgramId = prog;
            NeedsLoad = false;
            GLHelpers.CheckErrors(GL);
        }

        public override void Bind() {
            lock (_lock) {
                SetActive();
                if (Program != 0) {
                    GL.UseProgram((uint)Program);
                }
            }
        }

        public override void Unbind() {
            lock (_lock) {
                GL.UseProgram(0);
                GLHelpers.CheckErrors(GL);
            }
        }

        protected override void Unload() {
            lock (_lock) {
                if (Program != 0) {
                    var prog = Program;
                    Program = 0;
                    ProgramId = 0;
                    _device.QueueGLAction(gl => {
                        gl.DeleteProgram(prog);
                        GpuMemoryTracker.TrackResourceDeallocation(GpuResourceType.Shader);
                    });
                }
            }
        }
    }
}