/**

 * MojoShader; generate shader programs from bytecode of compiled

 *  Direct3D shaders.

 *

 * Please see the file LICENSE.txt in the source's root directory.

 *

 *  This file written by Ryan C. Gordon.

 */

#define __MOJOSHADER_INTERNAL__ 1

#include "mojoshader_profile.h"

#pragma GCC visibility push(hidden)

#if SUPPORT_PROFILE_SPIRV

#include "spirv/spirv.h"

#include "spirv/GLSL.std.450.h"

#include <float.h>

static const int SPV_NO_SWIZZLE = 0xE4; // 0xE4 == 11100100 ... 0 1 2 3. No swizzle.

#define EMIT_SPIRV_OPCODE_UNIMPLEMENTED_FUNC(op) \

    void emit_SPIRV_##op(Context *ctx) { \

        fail(ctx, #op " unimplemented in spirv profile"); \

    }

typedef struct SpirvTexm3x3SetupResult

{

    // vec4 load results

    uint32 id_dst_pad0;

    uint32 id_dst_pad1;

    uint32 id_dst;

    // float dot results

    uint32 id_res_x;

    uint32 id_res_y;

    uint32 id_res_z;

} SpirvTexm3x3SetupResult;

static const char *spv_get_uniform_array_varname(Context *ctx,

                                                   const RegisterType regtype,

                                                   char *buf, const size_t len)

{

    const char *shadertype = ctx->shader_type_str;

    const char *type = "";

    switch (regtype)

    {

        case REG_TYPE_CONST: type = "vec4"; break;

        case REG_TYPE_CONSTINT: type = "ivec4"; break;

        case REG_TYPE_CONSTBOOL: type = "bool"; break;

        default: fail(ctx, "BUG: used a uniform we don't know how to define.");

    } // switch

    snprintf(buf, len, "%s_uniforms_%s", shadertype, type);

    return buf;

} // spv_get_uniform_array_varname

static uint32 spv_bumpid(Context *ctx)

{

    return (ctx->spirv.idmax += 1);

} // spv_bumpid

static RegisterList *spv_getreg(Context *ctx, const RegisterType regtype, const int regnum)

{

    RegisterList *r = reglist_find(&ctx->used_registers, regtype, regnum);

    if (!r)

    {

        failf(ctx, "register not found rt=%d, rn=%d", regtype, regnum);

        return NULL;

    } // if

    return r;

} // spv_getreg

static void spv_componentlist_free(Context *ctx, ComponentList *cl)

{

    ComponentList *next;

    while (cl)

    {

        next = cl->next;

        Free(ctx, cl);

        cl = next;

    } // while

} // spv_componentlist_free

static ComponentList *spv_componentlist_alloc(Context *ctx)

{

    ComponentList *ret = (ComponentList *) Malloc(ctx, sizeof(ComponentList));

    if (!ret) return NULL;

    ret->id = 0;

    ret->v.i = 0;

    ret->next = NULL;

    return ret;

} // spv_componentlist_alloc

static const char *get_SPIRV_varname_in_buf(Context *ctx, const RegisterType rt,

                                           const int regnum, char *buf,

                                           const size_t buflen)

{

    // turns out these are identical at the moment.

    return get_D3D_varname_in_buf(ctx, rt, regnum, buf, buflen);

} // get_SPIRV_varname_in_buf

const char *get_SPIRV_varname(Context *ctx, const RegisterType rt,

                                    const int regnum)

{

    // turns out these are identical at the moment.

    return get_D3D_varname(ctx, rt, regnum);

} // get_SPIRV_varname

static inline const char *get_SPIRV_const_array_varname_in_buf(Context *ctx,

                                                const int base, const int size,

                                                char *buf, const size_t buflen)

{

    snprintf(buf, buflen, "c_array_%d_%d", base, size);

    return buf;

} // get_SPIRV_const_array_varname_in_buf

const char *get_SPIRV_const_array_varname(Context *ctx, int base, int size)

{

    char buf[64];

    get_SPIRV_const_array_varname_in_buf(ctx, base, size, buf, sizeof (buf));

    return StrDup(ctx, buf);

} // get_SPIRV_const_array_varname

static uint32 spv_get_uniform_array_id(Context *ctx, const RegisterType regtype)

{

    uint32 id;

    switch (regtype)

    {

        case REG_TYPE_CONST:

            id = ctx->spirv.uniform_arrays.idvec4;

            if (id == 0)

            {

                id = spv_bumpid(ctx);

                ctx->spirv.uniform_arrays.idvec4 = id;

            } // if

            break;

        case REG_TYPE_CONSTINT:

            id = ctx->spirv.uniform_arrays.idivec4;

            if (id == 0)

            {

                id = spv_bumpid(ctx);

                ctx->spirv.uniform_arrays.idivec4 = id;

            } // if

            break;

        case REG_TYPE_CONSTBOOL:

            id = ctx->spirv.uniform_arrays.idbool;

            if (id == 0)

            {

                id = spv_bumpid(ctx);

                ctx->spirv.uniform_arrays.idbool = id;

            } // if

            break;

        default:

            fail(ctx, "Unexpected register type used to access uniform array.");

            id = 0;

    } // switch

    return id;

} // spv_get_uniform_array_id

static void spv_emit_part_va(Context* ctx, uint32 word_count, uint32 argc, SpvOp op, va_list args)

{

    assert(ctx->output != NULL);

    if (isfail(ctx))

        return;  // we failed previously, don't go on...

    uint32 word = op | (word_count << 16);

    buffer_append(ctx->output, &word, sizeof(word));

    while (--argc)

    {

        word = va_arg(args, uint32);

        buffer_append(ctx->output, &word, sizeof(word));

    } // while

} // spv_emit_part_va

static void spv_emit_part(Context* ctx, uint32 word_count, uint32 argc, SpvOp op, ...)

{

    va_list args;

    va_start(args, op);

    spv_emit_part_va(ctx, word_count, argc, op, args);

    va_end(args);

} // spv_emit_part

static void spv_emit(Context *ctx, uint32 word_count, SpvOp op, ...)

{

    va_list args;

    va_start(args, op);

    spv_emit_part_va(ctx, word_count, word_count, op, args);

    va_end(args);

} // spv_emit

static void spv_emit_word(Context *ctx, uint32 word)

{

    assert(ctx->output != NULL);

    if (isfail(ctx))

        return;  // we failed previously, don't go on...

    buffer_append(ctx->output, &word, sizeof(word));

} // spv_emit_word

static void spv_emit_str(Context *ctx, const char *str)

{

    size_t len;

    uint32 trail;

    assert(ctx->output != NULL);

    if (isfail(ctx))

        return;  // we failed previously, don't go on...

    if (str == NULL)

        return spv_emit_word(ctx, 0);

    len = strlen(str) + 1;

    buffer_append(ctx->output, str, len);

    len = len % 4;

    if (len)

    {

        trail = 0;

        buffer_append(ctx->output, &trail, 4 - len);

    } // if

} // spv_emit_str

// get the word count of a string

static uint32 spv_strlen(const char *str)

{

    size_t len = strlen(str);

    return (uint32) ((len / 4) + 1);

} // spv_strlen

// emits an OpName straight into ctx->globals

static void spv_output_name(Context *ctx, uint32 id, const char *str)

{

    if (isfail(ctx))

        return;  // we failed previously, don't go on...

    push_output(ctx, &ctx->globals);

    spv_emit_part(ctx, 2 + spv_strlen(str), 2, SpvOpName, id);

    spv_emit_str(ctx, str);

    pop_output(ctx);

} // spv_output_name

// emit an OpName instruction to identify a register

static void spv_output_regname(Context *ctx, uint32 id, RegisterType regtype, int regnum)

{

    char varname[64];

    snprintf(varname, sizeof(varname), "%s_", ctx->shader_type_str);

    size_t offset = strlen(varname);

    get_SPIRV_varname_in_buf(ctx, regtype, regnum, varname + offset, sizeof(varname) - offset);

    spv_output_name(ctx, id, varname);

} // spv_output_regname

// emits an OpDecorate BuiltIn straight into ctx->helpers

static void spv_output_builtin(Context *ctx, uint32 id, SpvBuiltIn builtin)

{

    if (isfail(ctx))

        return;  // we failed previously, don't go on...

    push_output(ctx, &ctx->helpers);

    spv_emit(ctx, 4, SpvOpDecorate, id, SpvDecorationBuiltIn, builtin);

    pop_output(ctx);

} // spv_output_builtin

static uint32 spv_output_location(Context *ctx, uint32 id, uint32 loc)

{

    push_output(ctx, &ctx->helpers);

    spv_emit(ctx, 4, SpvOpDecorate, id, SpvDecorationLocation, loc);

    pop_output(ctx);

    return (buffer_size(ctx->helpers) >> 2) - 1;

} // spv_output_location

static void spv_output_sampler_binding(Context *ctx, uint32 id, uint32 binding)

{

    if (isfail(ctx))

        return;

    uint32 set = 0;

    if (ctx->spirv.mode == SPIRV_MODE_VK)

    {

        set = shader_is_vertex(ctx) ? MOJOSHADER_SPIRV_VS_SAMPLER_SET

                                    : MOJOSHADER_SPIRV_PS_SAMPLER_SET;

    } // if

    push_output(ctx, &ctx->helpers);

    spv_emit(ctx, 4, SpvOpDecorate, id, SpvDecorationDescriptorSet, set);

    spv_emit(ctx, 4, SpvOpDecorate, id, SpvDecorationBinding, binding);

    pop_output(ctx);

} // spv_output_sampler_binding

static SpirvTypeIdx spv_change_base_type_vec_dim(SpirvTypeIdx sti, uint32 dim)

{

    uint32 dimSub1 = dim - 1;

    assert(STI_CORE_START_ <= sti && sti < STI_CORE_END_);

    assert(dimSub1 < 4);

    SpirvTypeIdx sti_base = (SpirvTypeIdx)(sti & ~0x3);

    SpirvTypeIdx sti_new = (SpirvTypeIdx)(sti_base | dimSub1);

    return sti_new;

} // spv_change_base_type_vec_dim

static uint32 spv_get_type(Context *ctx, SpirvTypeIdx tidx)

{

    assert(((uint32)tidx) < ((uint32)STI_LENGTH_));

    uint32 tid = ctx->spirv.tid[tidx];

    if (tid)

        return tid;

    push_output(ctx, &ctx->mainline_intro);

    if (STI_CORE_START_ <= tidx && tidx < STI_CORE_END_)

    {

        uint32 dim = tidx & 0x3;

        SpirvType type = (SpirvType)((tidx >> 2) & 0x3);

        if (dim)

        {

            uint32 tid_base = spv_get_type(ctx, (SpirvTypeIdx)(tidx - dim));

            tid = spv_bumpid(ctx);

            spv_emit(ctx, 4, SpvOpTypeVector, tid, tid_base, dim + 1);

        } // if

        else

        {

            tid = spv_bumpid(ctx);

            switch (type)

            {

                case ST_FLOAT: spv_emit(ctx, 3, SpvOpTypeFloat, tid, 32); break;

                case ST_SINT: spv_emit(ctx, 4, SpvOpTypeInt, tid, 32, 1); break;

                case ST_UINT: spv_emit(ctx, 4, SpvOpTypeInt, tid, 32, 0); break;

                case ST_BOOL: spv_emit(ctx, 2, SpvOpTypeBool, tid); break;

                default: assert(!"Unexpected value of SpirvType."); break;

            } // switch

        } // else

    } // if

    else if (STI_IMAGE2D <= tidx && tidx <= STI_IMAGECUBE)

    {

        static const SpvDim dim_table[] = {SpvDim2D, SpvDim3D, SpvDimCube};

        SpvDim dim = dim_table[tidx - STI_IMAGE2D];

        uint32 tid_float = spv_get_type(ctx, STI_FLOAT);

        uint32 id_image = spv_bumpid(ctx);

        tid = spv_bumpid(ctx);

        spv_emit(ctx, 9, SpvOpTypeImage, id_image, tid_float, dim, 0, 0, 0, 1, SpvImageFormatUnknown);

        spv_emit(ctx, 3, SpvOpTypeSampledImage, tid, id_image);

    } // else if

    else if (tidx == STI_VOID)

    {

        tid = spv_bumpid(ctx);

        spv_emit(ctx, 2, SpvOpTypeVoid, tid);

    } // else if

    else if (tidx == STI_FUNC_VOID)

    {

        uint32 tid_void = spv_get_type(ctx, STI_VOID);

        tid = spv_bumpid(ctx);

        spv_emit(ctx, 3, SpvOpTypeFunction, tid, tid_void);

    } // else if

    else if (tidx == STI_FUNC_LIT)

    {

        uint32 tid_vec4 = spv_get_type(ctx, STI_VEC4);

        tid = spv_bumpid(ctx);

        spv_emit(ctx, 3 + 1, SpvOpTypeFunction, tid, tid_vec4, tid_vec4);

    } // else if

    else if (STI_PTR_START_ <= tidx && tidx < STI_PTR_END_)

    {

        uint32 dim = (tidx & (1 << 4)) ? 3 : 0;

        SpirvType type = (SpirvType)((tidx >> 2) & 0x3);

        uint32 tid_base = spv_get_type(ctx, (SpirvTypeIdx)((1 << 4) | (type << 2) | dim));

        static const SpvStorageClass sc_map[] = {

            SpvStorageClassInput,

            SpvStorageClassInput,

            SpvStorageClassOutput,

            SpvStorageClassOutput,

            SpvStorageClassPrivate,

            SpvStorageClassPrivate,

            SpvStorageClassUniformConstant,

            SpvStorageClassUniform,

        };

        SpvStorageClass sc = sc_map[((tidx & 0x3) << 1) | (ctx->spirv.mode == SPIRV_MODE_VK)];

        tid = spv_bumpid(ctx);

        spv_emit(ctx, 4, SpvOpTypePointer, tid, sc, tid_base);

    } // else if

    else if (STI_PTR_IMAGE2D <= tidx && tidx <= STI_PTR_IMAGECUBE)

    {

        uint32 tid_image = spv_get_type(ctx, (SpirvTypeIdx)(tidx - (STI_PTR_IMAGE2D - STI_IMAGE2D)));

        tid = spv_bumpid(ctx);

        spv_emit(ctx, 4, SpvOpTypePointer, tid, SpvStorageClassUniformConstant, tid_image);

    } // else if

    else

        assert(!"Unexpected value of type index.");

    pop_output(ctx);

    ctx->spirv.tid[tidx] = tid;

    return tid;

} // spv_get_type

static uint32 spv_gettrue(Context *ctx)

{

    if (ctx->spirv.idtrue)

        return ctx->spirv.idtrue;

    uint32 tid_bool = spv_get_type(ctx, STI_BOOL);

    uint32 id = spv_bumpid(ctx);

    push_output(ctx, &ctx->mainline_intro);

    spv_emit(ctx, 3, SpvOpConstantTrue, tid_bool, id);

    pop_output(ctx);

    return ctx->spirv.idtrue = id;

} // spv_gettrue

static uint32 spv_getfalse(Context *ctx)

{

    if (ctx->spirv.idfalse)

        return ctx->spirv.idfalse;

    uint32 tid_bool = spv_get_type(ctx, STI_BOOL);

    uint32 id = spv_bumpid(ctx);

    push_output(ctx, &ctx->mainline_intro);

    spv_emit(ctx, 3, SpvOpConstantFalse, tid_bool, id);

    pop_output(ctx);

    return ctx->spirv.idfalse = id;

} // spv_getfalse

static uint32 spv_getext(Context *ctx)

{

    if (ctx->spirv.idext)

        return ctx->spirv.idext;

    return ctx->spirv.idext = spv_bumpid(ctx);

} // spv_getext

static uint32 spv_output_scalar(Context *ctx, ComponentList *cl,

                             MOJOSHADER_attributeType type)

{

    uint32 idret, idtype;

    if (type == MOJOSHADER_ATTRIBUTE_FLOAT)

        idtype = spv_get_type(ctx, STI_FLOAT);

    else if (type == MOJOSHADER_ATTRIBUTE_INT)

        idtype = spv_get_type(ctx, STI_INT);

    else if (type == MOJOSHADER_ATTRIBUTE_UINT)

        idtype = spv_get_type(ctx, STI_UINT);

    else

    {

        failf(ctx, "spv_output_scalar: invalid attribute type %d", type);

        return 0;

    } // else

    idret = spv_bumpid(ctx);

    push_output(ctx, &ctx->mainline_intro);

    spv_emit(ctx, 4, SpvOpConstant, idtype, idret, cl->v.u);

    pop_output(ctx);

    return idret;

} // spv_output_scalar

// The spv_getscalar* functions retrieve the result id of an OpConstant

// instruction with the corresponding value v, or generate a new one.

static uint32 spv_getscalarf(Context *ctx, float v)

{

    ComponentList *prev = &(ctx->spirv.cl.f), *cl = ctx->spirv.cl.f.next;

    while (cl)

    {

        if (v == cl->v.f)

            return cl->id;

        else if (v < cl->v.f)

            break;

        prev = cl;

        cl = cl->next;

    } // while

    cl = spv_componentlist_alloc(ctx);

    cl->next = prev->next;

    prev->next = cl;

    cl->v.f = v;

    cl->id = spv_output_scalar(ctx, cl, MOJOSHADER_ATTRIBUTE_FLOAT);

    return cl->id;

} // spv_getscalarf

static uint32 spv_getscalari(Context *ctx, int v)

{

    ComponentList *prev = &(ctx->spirv.cl.i), *cl = ctx->spirv.cl.i.next;

    while (cl)

    {

        if (v == cl->v.i)

            return cl->id;

        else if (v < cl->v.i)

            break;

        prev = cl;

        cl = cl->next;

    } // while

    cl = spv_componentlist_alloc(ctx);

    cl->next = prev->next;

    prev->next = cl;

    cl->v.i = v;

    cl->id = spv_output_scalar(ctx, cl, MOJOSHADER_ATTRIBUTE_INT);

    return cl->id;

} // spv_getscalari

static uint32 spv_get_constant_composite(Context *ctx, uint32 tid, uint32* cache, float scalar)

{

    uint32 i;

    assert(tid != 0);

    uint32 dim =

       (tid == ctx->spirv.tid[STI_VEC4]) ? 4 :

       (tid == ctx->spirv.tid[STI_VEC3]) ? 3 :

       (tid == ctx->spirv.tid[STI_VEC2]) ? 2 : 1;

    uint32 id = cache[dim - 1];

    if (id)

        return id;

    uint32 sid = spv_getscalarf(ctx, scalar);

    if (dim == 1)

    {

        cache[0] = sid;

        return sid;

    } // if

    id = spv_bumpid(ctx);

    push_output(ctx, &ctx->mainline_intro);

    spv_emit_part(ctx, 3 + dim, 3, SpvOpConstantComposite, tid, id);

    for (i = 0; i < dim; i++)

        spv_emit_word(ctx, sid);

    pop_output(ctx);

    cache[dim - 1] = id;

    return id;

} // spv_get_constant_composite