// SPDX-FileCopyrightText: Copyright 2024 shadPS4 Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#version 450

layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

layout(std430, binding = 0) buffer input_buf {
    uint in_data[];
};
layout(std430, binding = 1) buffer output_buf {
    uint out_data[];
};

layout(push_constant) uniform image_info {
    uint num_levels;
    uint pitch;
    uint height;
    uint sizes[16];
} info;

// Inverse morton LUT, small enough to fit into K$
const uint rmort[16] = {
    0x11011000, 0x31213020,
    0x13031202, 0x33233222,
    0x51415040, 0x71617060,
    0x53435242, 0x73637262,

    0x15051404, 0x35253424,
    0x17071606, 0x37273626,
    0x55455444, 0x75657464,
    0x57475646, 0x77677666,
};

#define MICRO_TILE_DIM      (8)

void main() {
    uint block_ofs = 4 * gl_GlobalInvocationID.x;
    uint p0 = in_data[block_ofs + 0];
    uint p1 = in_data[block_ofs + 1];
    uint p2 = in_data[block_ofs + 2];
    uint p3 = in_data[block_ofs + 3];

    uint bit_ofs = 8 * (gl_LocalInvocationID.x % 4);
    uint packed_pos = rmort[gl_LocalInvocationID.x >> 2] >> bit_ofs;
    uint col = bitfieldExtract(packed_pos, 4, 4);
    uint row = bitfieldExtract(packed_pos, 0, 4);

    uint mip = 0;
    for (int m = 0; m < info.num_levels; ++m) {
        mip += (gl_GlobalInvocationID.x * 16) >= info.sizes[m] ? 1 : 0;
    }

    uint tiles_per_pitch = max(((info.pitch >> mip) / MICRO_TILE_DIM), 1u) * 4;
    uint target_tile_x = 4 * gl_WorkGroupID.x % tiles_per_pitch;
    uint target_tile_y = 4 * gl_WorkGroupID.x / tiles_per_pitch;
    uint dw_ofs_x = (target_tile_x * MICRO_TILE_DIM) + 4 * col;
    uint dw_ofs_y = ((target_tile_y * tiles_per_pitch) * 64u) + ((row * tiles_per_pitch) * MICRO_TILE_DIM);
    out_data[dw_ofs_x + dw_ofs_y] = p0;
    out_data[dw_ofs_x + dw_ofs_y + 1] = p1; 
    out_data[dw_ofs_x + dw_ofs_y + 2] = p2; 
    out_data[dw_ofs_x + dw_ofs_y + 3] = p3; 
}
