#!amber
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

SHADER vertex vert_shader PASSTHROUGH

SHADER fragment frag_shader GLSL TARGET_ENV spv1.3
#version 460
#extension GL_EXT_demote_to_helper_invocation : require
#extension GL_KHR_shader_subgroup_quad : require

layout(binding = 0) readonly buffer Block0
{
    float alpha[];
};

layout(binding = 1) buffer Block1
{
    uint atomics[];
};

layout(location = 0) out vec4 color;

float build_alpha_shuffle(float v)
{
    v = (helperInvocationEXT() ? 8.0 : roundEven(v));
    vec4 helpers;
    helpers.x = subgroupQuadBroadcast(v, 0u);
    helpers.y = subgroupQuadBroadcast(v, 1u);
    helpers.z = subgroupQuadBroadcast(v, 2u);
    helpers.w = subgroupQuadBroadcast(v, 3u);

    return dot(helpers, vec4(1, 10, 100, 1000));
}

void main()
{
    ivec2 coord = ivec2(gl_FragCoord.xy);
    int linear_coord = coord.y * 2 + coord.x;

    float alpha_value = alpha[linear_coord];
    float mask0 = build_alpha_shuffle(alpha_value);

    // Lane 1 and 2 should be nuked by this.
    if (fract(alpha_value) < 0.5)
    {
        demote;
    }

    float mask1 = build_alpha_shuffle(alpha_value);
    uint last_value = 0u;

    last_value = atomicAdd(atomics[linear_coord], 101u);

    if (linear_coord == 3 || last_value > 1000)
    {
        demote;
    }

    float mask2 = build_alpha_shuffle(alpha_value);

    color = vec4(1.0, mask0, mask1, mask2);
}
END

SHADER compute comp_shader GLSL
#version 460
layout(local_size_x=1,local_size_y=1) in;
uniform layout(set=0, binding=0, rgba32f) image2D resultImage;

layout(set = 0, binding = 1) buffer block0
{
    vec4 results[];
};

layout(set = 0, binding = 2) readonly buffer Block1
{
    uvec4 atomics;
};

void main()
{
    vec4 colorValues = imageLoad(resultImage, ivec2(0, 0));

    results[0] = vec4(0.0, .0, 0.0, 0.0);
    results[1] = imageLoad(resultImage, ivec2(1, 0));
    results[2] = imageLoad(resultImage, ivec2(0, 1));
    results[3] = imageLoad(resultImage, ivec2(1, 1));

    // We don't know if the invocations are helpers or not at the start
    // and therefore all the possible outcomes are allowed.
    for (int x = 0; x < 2; x++)
    {
        for (int y = 0; y < 2; y++)
        {
            for (int z = 0; z < 2; z++)
            {
                for (int w = 0; w < 2; w++)
                {
                    vec4 testVec = vec4(x * 1.0, y * 2.0, z * 3.0, w * 4.0);
                    uvec4 atomicVec = uvec4(0, 0, 0, 0);

                    for (int j = 0; j < 4; j++)
                    {
                        if (testVec[j] > 0.0)
                        {
                            testVec[j] = 8.0;
                            atomicVec[j] = 0;
                        }
                        else
                        {
                            testVec[j] = j + 1.0;
                            atomicVec[j] = 101;
                        }
                    }

                    float mask0 = dot(testVec, vec4(1, 10, 100, 1000));
                    float mask1 = dot(vec4(testVec.x, 8.0, 8.0, testVec.w), vec4(1, 10, 100, 1000));
                    float mask2 = dot(vec4(testVec.x, 8.0, 8.0, 8.0), vec4(1, 10, 100, 1000));

                    atomicVec[1] = 0;
                    atomicVec[2] = 0;

                    if (colorValues.x == 1.0 && atomics.x == atomicVec.x
                    && colorValues.y == mask0 && atomics.y == atomicVec.y
                    && colorValues.z == mask1 && atomics.z == atomicVec.z
                    && colorValues.w == mask2 && atomics.w == atomicVec.w)
                    {
                        results[0] = vec4(1.0, 1.0, 1.0, 1.0);
                    }
                }
            }
        }
    }
}
END

BUFFER alpha_keys DATA_TYPE float DATA
    0.75
    2.25
    3.25
    3.75
END

BUFFER atomics DATA_TYPE uint32 DATA
    0
    0
    0
    0
END

BUFFER results DATA_TYPE float SIZE 16 FILL 0.0
BUFFER ref_buffer DATA_TYPE float SIZE 16 FILL 1.0

BUFFER framebuffer FORMAT R32G32B32A32_SFLOAT

PIPELINE graphics myPipeline
    ATTACH vert_shader
    ATTACH frag_shader

    FRAMEBUFFER_SIZE 2 2

    BIND BUFFER framebuffer AS color LOCATION 0
    BIND BUFFER alpha_keys AS storage DESCRIPTOR_SET 0 BINDING 0
    BIND BUFFER atomics AS storage DESCRIPTOR_SET 0 BINDING 1
END

PIPELINE compute verifyPipeline
    ATTACH comp_shader
    FRAMEBUFFER_SIZE 2 2
    BIND BUFFER framebuffer AS storage_image DESCRIPTOR_SET 0 BINDING 0
    BIND BUFFER results AS storage DESCRIPTOR_SET 0 BINDING 1
    BIND BUFFER atomics AS storage DESCRIPTOR_SET 0 BINDING 2
END

CLEAR_COLOR myPipeline 255 255 255 255
CLEAR myPipeline

RUN myPipeline DRAW_RECT POS 0 0 SIZE 2 2
RUN verifyPipeline 1 1 1

EXPECT results EQ_BUFFER ref_buffer