VRCycling/Assets/Scripts/SicknessReduction/Visual/Rendering/BokehRenderPass.cs

using System;
using System.Runtime.CompilerServices;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.InputSystem.Interactions;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Universal;
using UnityEngine.Rendering.Universal.Internal;

namespace SicknessReduction.Visual.Rendering
{
    internal static class ShaderConstants
    {
        public static readonly int _FullCoCTexture = Shader.PropertyToID("_FullCoCTexture");
        public static readonly int _DofTexture = Shader.PropertyToID("_DofTexture");
        public static readonly int _CoCParams = Shader.PropertyToID("_CoCParams");
        public static readonly int _BokehKernel = Shader.PropertyToID("_BokehKernel");
        public static readonly int _PongTexture = Shader.PropertyToID("_PongTexture");
        public static readonly int _PingTexture = Shader.PropertyToID("_PingTexture");
    }

    class BokehRenderPass : ScriptableRenderPass
    {
        // used to label this pass in Unity's Frame Debug utility
        string profilerTag;

        private Material material;
        private Material uberMaterial;
        private BokehFeature.BokehFeatureSettings bokehSettings;

        private Vector4[] m_BokehKernel;
        private bool kernelPrepared;
        RenderTextureDescriptor m_Descriptor;

        RenderTargetIdentifier cameraColorTargetIdent;
        RenderTargetHandle target;
        private RenderTargetHandle tmp;

        public BokehRenderPass(string profilerTag,
            BokehFeature.BokehFeatureSettings settings)
        {
            this.profilerTag = profilerTag;
            renderPassEvent = settings.WhenToInsert;
            bokehSettings = settings;

            material = CoreUtils.CreateEngineMaterial(
                Shader.Find("Hidden/Universal Render Pipeline/BokehDepthOfField"));

            Debug.Log("Material created: " + material?.shader?.name ?? "Null");
        }

        public void Setup(in RenderTextureDescriptor baseDescriptor, RenderTargetIdentifier target)
        {
            //m_Descriptor = baseDescriptor;
            cameraColorTargetIdent = target;
        }

        // called each frame before Execute, use it to set up things the pass will need
        public override void Configure(CommandBuffer cmd, RenderTextureDescriptor cameraTextureDescriptor)
        {
            int wh = m_Descriptor.width / 2;
            int hh = m_Descriptor.height / 2;
            /*cmd.GetTemporaryRT(ShaderConstants._FullCoCTexture,
                cameraTextureDescriptor);
            cmd.GetTemporaryRT(ShaderConstants._PingTexture,
                cameraTextureDescriptor);
            cmd.GetTemporaryRT(ShaderConstants._PongTexture,
                cameraTextureDescriptor);
            cmd.GetTemporaryRT(target.id, cameraTextureDescriptor);*/
            m_Descriptor = cameraTextureDescriptor;

            // Temporary textures
            cmd.GetTemporaryRT(ShaderConstants._FullCoCTexture,
                GetStereoCompatibleDescriptor(m_Descriptor.width, m_Descriptor.height, GraphicsFormat.R8_UNorm),
                FilterMode.Bilinear);
            cmd.GetTemporaryRT(ShaderConstants._PingTexture,
                GetStereoCompatibleDescriptor(wh, hh, GraphicsFormat.R16G16B16A16_SFloat), FilterMode.Bilinear);
            cmd.GetTemporaryRT(ShaderConstants._PongTexture,
                GetStereoCompatibleDescriptor(wh, hh, GraphicsFormat.R16G16B16A16_SFloat), FilterMode.Bilinear);
            cmd.GetTemporaryRT(target.id, cameraTextureDescriptor);
            cmd.GetTemporaryRT(tmp.id, cameraTextureDescriptor);
        }


        // Execute is called for every eligible camera every frame. It's not called at the moment that
        // rendering is actually taking place, so don't directly execute rendering commands here.
        // Instead use the methods on ScriptableRenderContext to set up instructions.
        // RenderingData provides a bunch of (not very well documented) information about the scene
        // and what's being rendered.
        public override void Execute(ScriptableRenderContext context, ref RenderingData renderingData)
        {
            // fetch a command buffer to use
            CommandBuffer cmd = CommandBufferPool.Get(profilerTag);
            cmd.Clear();

            int wh = m_Descriptor.width / 2;
            int hh = m_Descriptor.height / 2;

            // "A Lens and Aperture Camera Model for Synthetic Image Generation" [Potmesil81]
            float maxCoC = bokehSettings.maxCoc;
            float P = bokehSettings.focusDistance;
            float maxRadius = GetMaxBokehRadiusInPixels(m_Descriptor.height);
            float rcpAspect = 1f / (wh / (float) hh);

            cmd.SetGlobalVector(ShaderConstants._CoCParams,
                new Vector4(P, maxCoC, maxRadius, rcpAspect));

            // Prepare the bokeh kernel constant buffer
            if (!kernelPrepared)
            {
                PrepareBokehKernel();
                kernelPrepared = true;
            }

            cmd.SetGlobalVectorArray(ShaderConstants._BokehKernel, m_BokehKernel);

            // Compute CoC
            cmd.Blit(cameraColorTargetIdent, ShaderConstants._FullCoCTexture, material, 0);
            cmd.SetGlobalTexture(ShaderConstants._FullCoCTexture,
                ShaderConstants._FullCoCTexture);

            // Downscale & prefilter color + coc
            cmd.Blit(cameraColorTargetIdent, ShaderConstants._PingTexture, material, 1);

            // Bokeh blur
            cmd.Blit(ShaderConstants._PingTexture, ShaderConstants._PongTexture,
                material, 2);

            // Post-filtering
            cmd.Blit(ShaderConstants._PongTexture,
                BlitDstDiscardContent(cmd, ShaderConstants._PingTexture), material, 3);

            // Composite
            cmd.SetGlobalTexture(ShaderConstants._DofTexture,
                ShaderConstants._PingTexture);
            cmd.Blit(cameraColorTargetIdent, BlitDstDiscardContent(cmd, target.Identifier()), material, 4);

            //cmd.Blit(tmp.Identifier(), target.Identifier(), material, 5);
            //cmd.SetGlobalTexture("_BlitTex", cameraColorTargetIdent);
            //cmd.Blit(target.Identifier(), cameraColorTargetIdent, uberMaterial);


            cmd.Blit(target.Identifier(), cameraColorTargetIdent);

            context.ExecuteCommandBuffer(cmd);

            // tidy up after ourselves
            cmd.Clear();
            CommandBufferPool.Release(cmd);
        }

        // called after Execute, use it to clean up anything allocated in Configure
        public override void FrameCleanup(CommandBuffer cmd)
        {
            // Cleanup
            cmd.ReleaseTemporaryRT(ShaderConstants._FullCoCTexture);
            cmd.ReleaseTemporaryRT(ShaderConstants._PingTexture);
            cmd.ReleaseTemporaryRT(ShaderConstants._PongTexture);
            cmd.ReleaseTemporaryRT(target.id);
            cmd.ReleaseTemporaryRT(tmp.id);
        }

        #region Depth Of Field

        void PrepareBokehKernel()
        {
            const int kRings = 4;
            const int kPointsPerRing = 7;

            // Check the existing array
            if (m_BokehKernel == null)
                m_BokehKernel = new Vector4[42];

            // Fill in sample points (concentric circles transformed to rotated N-Gon)
            int idx = 0;
            float bladeCount = 5f;
            float curvature = 0f;
            float rotation = 10f * Mathf.Deg2Rad;
            const float PI = Mathf.PI;
            const float TWO_PI = Mathf.PI * 2f;

            for (int ring = 1; ring < kRings; ring++)
            {
                float bias = 1f / kPointsPerRing;
                float radius = (ring + bias) / (kRings - 1f + bias);
                int points = ring * kPointsPerRing;

                for (int point = 0; point < points; point++)
                {
                    // Angle on ring
                    float phi = 2f * PI * point / points;

                    // Transform to rotated N-Gon
                    // Adapted from "CryEngine 3 Graphics Gems" [Sousa13]
                    float nt = Mathf.Cos(PI / bladeCount);
                    float dt = Mathf.Cos(phi - (TWO_PI / bladeCount) *
                        Mathf.Floor((bladeCount * phi + Mathf.PI) / TWO_PI));
                    float r = radius * Mathf.Pow(nt / dt, curvature);
                    float u = r * Mathf.Cos(phi - rotation);
                    float v = r * Mathf.Sin(phi - rotation);

                    m_BokehKernel[idx] = new Vector4(u, v);
                    idx++;
                }
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        static float GetMaxBokehRadiusInPixels(float viewportHeight)
        {
            // Estimate the maximum radius of bokeh (empirically derived from the ring count)
            const float kRadiusInPixels = 14f;
            return Mathf.Min(0.05f, kRadiusInPixels / viewportHeight);
        }

        RenderTextureDescriptor GetStereoCompatibleDescriptor()
            => GetStereoCompatibleDescriptor(m_Descriptor.width, m_Descriptor.height, m_Descriptor.graphicsFormat,
                m_Descriptor.depthBufferBits);

        RenderTextureDescriptor GetStereoCompatibleDescriptor(int width, int height, GraphicsFormat format,
            int depthBufferBits = 0)
        {
            // Inherit the VR setup from the camera descriptor
            var desc = m_Descriptor;
            desc.depthBufferBits = depthBufferBits;
            desc.msaaSamples = 1;
            desc.width = width;
            desc.height = height;
            desc.graphicsFormat = format;
            return desc;
        }

        private BuiltinRenderTextureType BlitDstDiscardContent(CommandBuffer cmd, RenderTargetIdentifier rt)
        {
            // We set depth to DontCare because rt might be the source of PostProcessing used as a temporary target
            // Source typically comes with a depth buffer and right now we don't have a way to only bind the color attachment of a RenderTargetIdentifier
            cmd.SetRenderTarget(new RenderTargetIdentifier(rt, 0, CubemapFace.Unknown, -1),
                RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store,
                RenderBufferLoadAction.DontCare, RenderBufferStoreAction.DontCare);
            return BuiltinRenderTextureType.CurrentActive;
        }

        #endregion
    }
}