o3de/Gems/Atom/Feature/Common/Assets/Shaders/Reflections/ReflectionScreenSpaceTrace.azsli

/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#include <scenesrg_all.srgi>
#include <viewsrg_all.srgi>
#include <Atom/RPI/Math.azsli>

float DistanceSqr(float2 v1, float2 v2) 
{ 
    v1 -= v2; 
    return dot(v1, v1); 
}

//
// This is a modified version of the following work:
// http://jcgt.org/published/0003/04/04/paper.pdf
// 
// Copyright (c) 2014, Morgan McGuire and Michael Mara
// All rights reserved.
//
// Released as open source under the BSD 2-Clause License
// http://opensource.org/licenses/BSD-2-Clause
//
// 1. Redistributions of source code must retain the above copyright notice, this list of conditions 
// and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions 
// and the following disclaimer in the documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED 
// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
// PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED 
// TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) 
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING 
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
// POSSIBILITY OF SUCH DAMAGE.
//
// See http://kode80.com/blog/2015/03/11/screen-space-reflections-in-unity-5/index.html for additional
// extensions and optimizations.
//

static const uint MaxSteps = 96;
static const uint BinarySearchSteps = 16;

bool TraceRayScreenSpace(float3 rayStartVS, float3 rayDirectionVS, uint2 dimensions, out float2 hitCoords)
{
    float3 rayEndVS = rayStartVS + rayDirectionVS * PassSrg::m_maxRayDistance;

    // early out if endpoint is behind camera
    if (rayEndVS.z >= 0.0f)
    {
        return false;
    }

    // project into homogeneous clip space
    float4 H0 = mul(ViewSrg::m_projectionMatrix, float4(rayStartVS, 1.0f));
    float4 H1 = mul(ViewSrg::m_projectionMatrix, float4(rayEndVS, 1.0f));
    float k0 = 1.0f / H0.w;
    float k1 = 1.0f / H1.w;
    
    // convert to screen-space endpoints 
    float2 P0 = H0.xy * k0 * 0.5f + 0.5f;
    P0 = float2(P0.x * dimensions.x, (1.0f - P0.y) * dimensions.y);
    float2 P1 = H1.xy * k1 * 0.5f + 0.5f;
    P1 = float2(P1.x * dimensions.x, (1.0f - P1.y) * dimensions.y);

    // the interpolated homogeneous version of the viewspace points  
    float3 Q0 = rayStartVS * k0;
    float3 Q1 = rayEndVS * k1;

    // if the line is degenerate, make it cover at least one pixel
#ifdef PRE_HLSL_2021
    // See /o3de/Gems/Atom/Feature/Common/Assets/ShaderLib/PRE_HLSL_2021.md for details.
    P1 += DistanceSqr(P0, P1) < EPSILON ? float2(0.01f, 0.01f) : 0.0f;
#else
    P1 += select(DistanceSqr(P0, P1) < EPSILON, float2(0.01f, 0.01f), 0.0f);
#endif

    // store all of the start variables in a single float4
    float4 PQK = float4(P0, Q0.z, k0);

    // compute the step amount for each variable (SIMD)
    float4 dPQK = float4(P1, Q1.z, k1);
    dPQK -= PQK;

    bool permute = (abs(dPQK.x) < abs(dPQK.y));
    if (permute)
    {
        PQK.xy = PQK.yx;
        dPQK.xy = dPQK.yx;
    }

    dPQK /= MaxSteps;

    // advance by one step before starting the ray march
    PQK += dPQK;

    // ray march until the expected ray depth is beyond the actual scene depth
    bool foundHit = false;
    float rayZMin = 0.0f;
    float rayZMax = 0.0f;
    float prevRayZMaxEstimate = rayStartVS.z;
    for (uint step = 0; step < MaxSteps; ++step)
    {
        // validate the current screenspace coordinates (stored in PQK.xy)
        hitCoords = permute ? PQK.yx : PQK.xy;
        float4 validate = float4(hitCoords.x >= dimensions.x, hitCoords.y >= dimensions.y, hitCoords.x < 0, hitCoords.y < 0);
        if (any(validate))
        {
            break;
        }
    
        // retrieve scene depth from the linear viewspace depth buffer
        float sceneDepth = -PassSrg::m_depthLinear.Load(uint3(hitCoords, 0)).r;
        if (sceneDepth == 0.0f)
        {
            break;
        }

        // compute expected depth, and swap min/max if necessary
        rayZMin = prevRayZMaxEstimate;
        rayZMax = (dPQK.z * 0.5f + PQK.z) / (dPQK.w * 0.5f + PQK.w);
        prevRayZMaxEstimate = rayZMax;
        if (rayZMin > rayZMax)
        {
            swap(rayZMin, rayZMax);
        }
    
        // a hit occurs when the expected ray depth is beyond the scene depth and within the depth threshold (thickness)
        if ((rayZMin < sceneDepth) && (rayZMax >= sceneDepth - PassSrg::m_maxDepthThreshold))
        {
            foundHit = true;
            break;
        }

        // increase all three variables by the step amount
        PQK += dPQK;
    }

    if (foundHit)
    {
        // binary search refinement on the hit
        // start by moving back one step to just before the hit       
        PQK -= dPQK;

        // the stride reduces the dQKP increment each iteration of the binary search
        float stride = 0.5f;

        // the sign of the stride is changed each iteration to control the step direction
        float strideAndDirection = stride;
        
        for (uint binarySearchStep = 0; binarySearchStep < BinarySearchSteps; ++binarySearchStep)
        {
            dPQK *= strideAndDirection;
            PQK += dPQK;

            // current screen coordinates are stored in PQK.xy
            hitCoords = permute ? PQK.yx : PQK.xy;

            // retrieve linear depth
            float sceneDepth = -PassSrg::m_depthLinear.Load(uint3(hitCoords, 0)).r;

            // compute the expected depth of the ray at this point, by performing the perspective-divide
            // on the current homogenous z-coordinate (in PQK.z) by the current w-coordinate (in PQK.w)
            float rayDepth = PQK.z / PQK.w;

            // determine if the expected ray depth is beyond the scene depth and within the depth tolerance
            bool exceedsSceneDepth = (rayDepth < sceneDepth);
        
            // reduce the stride each iteration
            stride *= 0.5f;

            // move backwards if the ray depth exceeds the scene depth, otherwise move forward
            strideAndDirection = exceedsSceneDepth ? -stride : stride;
        }
    }

    return foundHit;
}