o3de/Gems/AtomLyIntegration/EditorModeFeedback/Assets/Shaders/EditorModeMask.azsl

/*
 * 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 <Atom/Features/SrgSemantics.azsli>

#include <viewsrg_all.srgi>
#include <scenesrg_all.srgi>

ShaderResourceGroup ObjectSrg : SRG_PerObject
{
    //! Id of this draw object for retrieval of transformation matrices.
    uint m_objectId;

    //! Returns the matrix for transforming points from Object Space to World Space.
    float4x4 GetWorldMatrix()
    {
        return SceneSrg::GetObjectToWorldMatrix(m_objectId);
    }
}

ShaderResourceGroup PassSrg : SRG_PerPass
{
    //! The depth buffer data for existing fragments in the color buffer.
    Texture2D<float> m_existingDepth;

    //! Dimensions of depth buffer texture (xy is width and height, zw is 1/width and 1/height).
    float4 m_existingDepthDimensions;
    
    Sampler LinearSampler
    {
        MinFilter = Linear;
        MagFilter = Linear;
        MipFilter = Linear;
        AddressU = Clamp;
        AddressV = Clamp;
        AddressW = Clamp;
    };
}

struct VSInput
{
    // Base fields (required by the template azsli file)...
    float3 m_position : POSITION;
};

struct VSOutput
{
    // Base fields (required by the template azsli file)...
    // "centroid" is needed for SV_Depth to compile
    precise linear centroid float4 m_position : SV_Position;

    // Fragment depth
    float m_depth : UV0;
};

VSOutput MainVS(VSInput IN)
{
    VSOutput OUT;

    const float3 worldPosition = mul(ObjectSrg::GetWorldMatrix(), float4(IN.m_position, 1.0)).xyz;
    OUT.m_position = mul(ViewSrg::m_viewProjectionMatrix, float4(worldPosition, 1.0));
    
    // The existing depth buffer contents are the result of the resolved MSAA depth buffer so contain depth contributions
    // from surrounding sub-fragments so we need to pull the current fragment towards the viewer slightly in order to reduce
    // z fighting with the sub-fragment depth contributions
    OUT.m_position.z *= 1.01;

    // Fragment depth
    OUT.m_depth = OUT.m_position.z / OUT.m_position.w;
    return OUT;
}

struct PixelOutput
{
    float2 m_mask : SV_Target0;
};

//! Returns the linearized depth of the specified non-linear depth value.
float LinearizeDepth(const float depth)
{
    return abs(((ViewSrg::GetFarZTimesNearZ()) / (ViewSrg::GetFarZMinusNearZ() * depth - ViewSrg::GetFarZ())));
}

PixelOutput MainPS(VSOutput IN)
{
    PixelOutput OUT;

    const float2 screenUV = IN.m_position.xy * PassSrg::m_existingDepthDimensions.zw;
    const float existingLinearDepth = LinearizeDepth(PassSrg::m_existingDepth.Sample(PassSrg::LinearSampler, screenUV).r);
    const float linearDepth = LinearizeDepth(IN.m_depth);

    // Do not write out any visible fragments to the red channel unless the depth test below passes
    OUT.m_mask.r = 0.0;

    // Write out all visible and occluded fragments to the green channel
    OUT.m_mask.g = 1.0;

    // Performing a equal or greater than depth test is still to flakey so we will instead perform a depth test within a small
    // tolerance range in order for this fragment to pass the test
    if((linearDepth - existingLinearDepth) <= 0.0001)
    {
        // Depth test passed, update the mask with this closest value
        OUT.m_mask.r = 1.0;
    }

    return OUT;
}