#version 450
/* Copyright (c) 2020 Themaister
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include "small_types.h"

layout(local_size_x_id = 3) in;

layout(constant_id = 0) const int RDRAM_SIZE = 8 * 1024 * 1024;
const int RDRAM_MASK_8 = RDRAM_SIZE - 1;
const int RDRAM_MASK_16 = RDRAM_MASK_8 >> 1;
const int RDRAM_MASK_32 = RDRAM_MASK_8 >> 2;
layout(constant_id = 1) const int FB_SIZE_LOG2 = 0;
layout(constant_id = 2) const bool COLOR_DEPTH_ALIAS = false;
layout(constant_id = 4) const int NUM_SAMPLES = 1;

layout(push_constant) uniform Registers
{
    uint num_pixels, fb_addr, fb_depth_addr;
} registers;

layout(set = 0, binding = 0) readonly buffer RDRAMSingleSampled8
{
    uint8_t elems[];
} vram8;

layout(set = 0, binding = 0) readonly buffer RDRAMSingleSampled16
{
    uint16_t elems[];
} vram16;

layout(set = 0, binding = 0) readonly buffer RDRAMSingleSampled32
{
    uint elems[];
} vram32;

layout(set = 0, binding = 1) readonly buffer RDRAMHiddenSingleSampled
{
    uint8_t elems[];
} hidden_vram;

layout(set = 0, binding = 2) buffer RDRAMUpscalingReference8
{
    uint8_t elems[];
} vram_reference8;

layout(set = 0, binding = 2) buffer RDRAMUpscalingReference16
{
    uint16_t elems[];
} vram_reference16;

layout(set = 0, binding = 2) buffer RDRAMUpscalingReference32
{
    uint elems[];
} vram_reference32;

layout(set = 0, binding = 3) buffer RDRAMUpscaling8
{
    uint8_t elems[];
} vram_upscaled8;

layout(set = 0, binding = 3) buffer RDRAMUpscaling16
{
    uint16_t elems[];
} vram_upscaled16;

layout(set = 0, binding = 3) buffer RDRAMUpscaling32
{
    uint elems[];
} vram_upscaled32;

layout(set = 0, binding = 4) buffer RDRAMHiddenUpscaling
{
    uint8_t elems[];
} hidden_vram_upscaled;

void update_rdram_8(uint index)
{
    index &= RDRAM_MASK_8;

    uint real_word = uint(vram8.elems[index]);
    uint reference_word = uint(vram_reference8.elems[index]);

    if (real_word != reference_word)
    {
        uint mirrored_index = index ^ 3u;
        uint real_hidden_word = uint(hidden_vram.elems[mirrored_index >> 1u]);
        for (int i = 0; i < NUM_SAMPLES; i++)
        {
            vram_upscaled8.elems[index + i * RDRAM_SIZE] = uint8_t(real_word);
            if ((mirrored_index & 1u) != 0u)
                hidden_vram_upscaled.elems[(mirrored_index >> 1u) + i * (RDRAM_SIZE >> 1)] = uint8_t(real_hidden_word);
        }
        vram_reference8.elems[index] = uint8_t(real_word);
    }
}

void update_rdram_16(uint index)
{
    index &= RDRAM_MASK_16;

    uint real_word = uint(vram16.elems[index]);
    uint reference_word = uint(vram_reference16.elems[index]);

    if (real_word != reference_word)
    {
        uint mirrored_index = index ^ 1u;
        uint real_hidden_word = uint(hidden_vram.elems[mirrored_index]);
        for (int i = 0; i < NUM_SAMPLES; i++)
        {
            vram_upscaled16.elems[index + i * (RDRAM_SIZE >> 1)] = uint16_t(real_word);
            hidden_vram_upscaled.elems[mirrored_index + i * (RDRAM_SIZE >> 1)] = uint8_t(real_hidden_word);
        }
        vram_reference16.elems[index] = uint16_t(real_word);
    }
}

void update_rdram_32(uint index)
{
    index &= RDRAM_MASK_32;

    uint real_word = vram32.elems[index];
    uint reference_word = vram_reference32.elems[index];

    if (real_word != reference_word)
    {
        uint real_hidden_word0 = uint(hidden_vram.elems[2u * index]);
        uint real_hidden_word1 = uint(hidden_vram.elems[2u * index + 1u]);

        for (int i = 0; i < NUM_SAMPLES; i++)
        {
            vram_upscaled32.elems[index + i * (RDRAM_SIZE >> 2)] = real_word;
            hidden_vram_upscaled.elems[2u * index + i * (RDRAM_SIZE >> 1)] = uint8_t(real_hidden_word0);
            hidden_vram_upscaled.elems[2u * index + 1u + i * (RDRAM_SIZE >> 1)] = uint8_t(real_hidden_word1);
        }
        vram_reference32.elems[index] = real_word;
    }
}

void main()
{
    uint index = gl_GlobalInvocationID.x;
    if (index >= registers.num_pixels)
        return;

    uint depth_index = index + registers.fb_depth_addr;
    uint color_index = index + registers.fb_addr;

    switch (FB_SIZE_LOG2)
    {
        case 0:
            update_rdram_8(color_index);
            break;

        case 1:
            update_rdram_16(color_index);
            break;

        case 2:
            update_rdram_32(color_index);
            break;
    }

    if (!COLOR_DEPTH_ALIAS)
        update_rdram_16(depth_index);
}