/**

 * 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.

 */

#ifdef _WIN32

#define WIN32_LEAN_AND_MEAN 1

#include <windows.h> // Include this early to avoid SDL conflicts

#endif

#define __MOJOSHADER_INTERNAL__ 1

#include "mojoshader_internal.h"

#if SUPPORT_PROFILE_HLSL

#define D3D11_NO_HELPERS

#define CINTERFACE

#define COBJMACROS

#include <d3d11.h>

#ifndef WINAPI_FAMILY_WINRT

#define WINAPI_FAMILY_WINRT 0

#endif

#if WINAPI_FAMILY_WINRT

#include <d3dcompiler.h>

#endif

// D3DCompile optimization can be overzealous and cause very visible bugs,

//  so we disable it when compiling shaders to preserve correctness.

#define D3D_SKIP_OPT (1 << 2)

/* Error state */

static char error_buffer[1024] = { '\0' };

static void set_error(const char *str)

{

    snprintf(error_buffer, sizeof (error_buffer), "%s", str);

} // set_error

static inline void out_of_memory(void)

{

    set_error("out of memory");

} // out_of_memory

/* D3DCompile signature */

typedef HRESULT(WINAPI *PFN_D3DCOMPILE)(

    LPCVOID pSrcData,

    SIZE_T SrcDataSize,

    LPCSTR pSourceName,

    const D3D_SHADER_MACRO *pDefines,

    ID3DInclude *pInclude,

    LPCSTR pEntrypoint,

    LPCSTR pTarget,

    UINT Flags1,

    UINT Flags2,

    ID3DBlob **ppCode,

    ID3DBlob **ppErrorMsgs

);

/* Structs */

typedef struct d3d11ShaderMap

{

    void *val;

    union

    {

        struct

        {

            uint64 layoutHash;

            ID3D10Blob *blob;

        } vertex;

        struct

        {

            MOJOSHADER_d3d11Shader *vshader;

        } pixel;

    };

} d3d11ShaderMap;

typedef struct MOJOSHADER_d3d11Shader

{

    const MOJOSHADER_parseData *parseData;

    uint32 refcount;

    ID3D11Buffer *ubo;

    size_t buflen;

    unsigned char *constantData;

    unsigned int mapCapacity;

    unsigned int numMaps;

    d3d11ShaderMap *shaderMaps;

} MOJOSHADER_d3d11Shader;

// Max entries for each register file type...

#define MAX_REG_FILE_F 8192

#define MAX_REG_FILE_I 2047

#define MAX_REG_FILE_B 2047

typedef struct MOJOSHADER_d3d11Context

{

    // Allocators...

    MOJOSHADER_malloc malloc_fn;

    MOJOSHADER_free free_fn;

    void *malloc_data;

    // The constant register files...

    // !!! FIXME: Man, it kills me how much memory this takes...

    // !!! FIXME:  ... make this dynamically allocated on demand.

    float vs_reg_file_f[MAX_REG_FILE_F * 4];

    int vs_reg_file_i[MAX_REG_FILE_I * 4];

    uint8 vs_reg_file_b[MAX_REG_FILE_B];

    float ps_reg_file_f[MAX_REG_FILE_F * 4];

    int ps_reg_file_i[MAX_REG_FILE_I * 4];

    uint8 ps_reg_file_b[MAX_REG_FILE_B];

    // Pointer to the active ID3D11Device.

    ID3D11Device *device;

    // Pointer to the ID3D11DeviceContext.

    ID3D11DeviceContext *deviceContext;

    // Currently bound vertex and pixel shaders.

    MOJOSHADER_d3d11Shader *vertexShader;

    MOJOSHADER_d3d11Shader *pixelShader;

    int vertexNeedsBound;

    int pixelNeedsBound;

    // D3DCompile function pointer.

    PFN_D3DCOMPILE D3DCompileFunc;

#if !WINAPI_FAMILY_WINRT

    HMODULE d3dcompilerDLL;

#endif

} MOJOSHADER_d3d11Context;

static MOJOSHADER_d3d11Context *ctx = NULL;

/* Uniform buffer utilities */

static inline int next_highest_alignment(int n)

{

    const int align = 16;

    return align * ((n + align - 1) / align);

} // next_highest_alignment

static inline void *get_uniform_buffer(MOJOSHADER_d3d11Shader *shader)

{

    return (shader == NULL || shader->ubo == NULL) ? NULL : shader->ubo;

} // get_uniform_buffer

static void update_uniform_buffer(MOJOSHADER_d3d11Shader *shader)

{

    if (shader == NULL || shader->ubo == NULL)

        return;

    float *regF; int *regI; uint8 *regB;

    if (shader->parseData->shader_type == MOJOSHADER_TYPE_VERTEX)

    {

        regF = ctx->vs_reg_file_f;

        regI = ctx->vs_reg_file_i;

        regB = ctx->vs_reg_file_b;

    } // if

    else

    {

        regF = ctx->ps_reg_file_f;

        regI = ctx->ps_reg_file_i;

        regB = ctx->ps_reg_file_b;

    } // else

    // Update the buffer contents

    int needsUpdate = 0;

    size_t offset = 0;

    for (int i = 0; i < shader->parseData->uniform_count; i++)

    {

        if (shader->parseData->uniforms[i].constant)

            continue;

        int idx = shader->parseData->uniforms[i].index;

        int arrayCount = shader->parseData->uniforms[i].array_count;

        void *src = NULL;

        void *dst = NULL;

        size_t size = arrayCount ? arrayCount : 1;

        switch (shader->parseData->uniforms[i].type)

        {

            case MOJOSHADER_UNIFORM_FLOAT:

                src = &regF[4 * idx];

                dst = shader->constantData + offset;

                size *= 16;

                break;

            case MOJOSHADER_UNIFORM_INT:

                src = &regI[4 * idx];

                dst = shader->constantData + offset;

                size *= 16;

                break;

            case MOJOSHADER_UNIFORM_BOOL:

                src = &regB[idx];

                dst = shader->constantData + offset;

                break;

            default:

                assert(0); // This should never happen.

                break;

        } // switch

        if (memcmp(dst, src, size) != 0)

        {

            memcpy(dst, src, size);

            needsUpdate = 1;

        } // if

        offset += size;

    } // for

    if (needsUpdate)

    {

        // Map the buffer

        D3D11_MAPPED_SUBRESOURCE res;

        ID3D11DeviceContext_Map((ID3D11DeviceContext*) ctx->deviceContext,

                                (ID3D11Resource*) shader->ubo, 0,

                                D3D11_MAP_WRITE_DISCARD, 0, &res);

        // Copy the contents

        memcpy(res.pData, shader->constantData, shader->buflen);

        // Unmap the buffer

        ID3D11DeviceContext_Unmap(

            (ID3D11DeviceContext*) ctx->deviceContext,

            (ID3D11Resource*) shader->ubo,

            0

        );

    } // if

} // update_uniform_buffer

static inline void expand_map(MOJOSHADER_d3d11Shader *shader)

{

    if (shader->numMaps == shader->mapCapacity)

    {

        d3d11ShaderMap *newMap = (d3d11ShaderMap *) ctx->malloc_fn(

            sizeof(d3d11ShaderMap) * shader->mapCapacity * 2,

            ctx->malloc_data

        );

        memcpy(newMap, shader->shaderMaps,

            sizeof(d3d11ShaderMap) * shader->mapCapacity);

        shader->mapCapacity *= 2;

        ctx->free_fn(shader->shaderMaps, ctx->malloc_data);

        shader->shaderMaps = newMap;

        newMap = NULL;

    } // if

} // expand_map

static inline int element_is_uint(DXGI_FORMAT format)

{

    return  format == DXGI_FORMAT_R32G32B32A32_UINT ||

            format == DXGI_FORMAT_R32G32B32_UINT ||

            format == DXGI_FORMAT_R16G16B16A16_UINT ||

            format == DXGI_FORMAT_R32G32_UINT ||

            format == DXGI_FORMAT_R10G10B10A2_UINT ||

            format == DXGI_FORMAT_R8G8B8A8_UINT ||

            format == DXGI_FORMAT_R16G16_UINT ||

            format == DXGI_FORMAT_R32_UINT ||

            format == DXGI_FORMAT_R8G8_UINT ||

            format == DXGI_FORMAT_R16_UINT ||

            format == DXGI_FORMAT_R8_UINT;

} // element_is_uint

static inline int element_is_int(DXGI_FORMAT format)

{

    return  format == DXGI_FORMAT_R32G32B32A32_SINT ||

            format == DXGI_FORMAT_R32G32B32_SINT ||

            format == DXGI_FORMAT_R16G16B16A16_SINT ||

            format == DXGI_FORMAT_R32G32_SINT ||

            format == DXGI_FORMAT_R8G8B8A8_SINT ||

            format == DXGI_FORMAT_R16G16_SINT ||

            format == DXGI_FORMAT_R32_SINT ||

            format == DXGI_FORMAT_R8G8_SINT ||

            format == DXGI_FORMAT_R16_SINT ||

            format == DXGI_FORMAT_R8_SINT;

} // element_is_int

/* Shader Compilation Utilities */

static ID3D11VertexShader *compileVertexShader(MOJOSHADER_d3d11Shader *shader,

                                               const char *src, int src_len,

                                               ID3D10Blob **blob)

{

    const MOJOSHADER_parseData *pd = shader->parseData;

    HRESULT result = ctx->D3DCompileFunc(src, src_len, pd->mainfn,

                                         NULL, NULL, pd->mainfn, "vs_4_0",

                                         D3D_SKIP_OPT, 0, blob, blob);

    if (result < 0)

    {

        set_error((const char *) ID3D10Blob_GetBufferPointer(*blob));

        ID3D10Blob_Release(*blob);

        return NULL;

    } // if

    void *bytecode = ID3D10Blob_GetBufferPointer(*blob);

    int bytecodeLength = ID3D10Blob_GetBufferSize(*blob);

    ID3D11VertexShader *ret = NULL;

    ID3D11Device_CreateVertexShader(ctx->device, bytecode, bytecodeLength,

                                    NULL, &ret);

    return ret;

} // compileVertexShader

static void replaceVarname(const char *find, const char *replace,

                           const char **source)

{

    const char *srcbuf = *source;

    size_t find_len = strlen(find);

    size_t replace_len = strlen(replace);

    #define IS_PARTIAL_TOKEN(token) \

        (isalnum(*(token + find_len)) || isalnum(*(token-1)))

    // How many times does `find` occur in the source buffer?

    int count = 0;

    char *ptr = (char *) strstr(srcbuf, find);

    while (ptr != NULL)

    {

        if (!IS_PARTIAL_TOKEN(ptr))

            count++;

        ptr = strstr(ptr + find_len, find);

    } // while

    // How big should we make the new text buffer?

    size_t oldlen = strlen(srcbuf) + 1;

    size_t newlen = oldlen + (count * (replace_len - find_len));

    // Easy case; just find/replace in the original buffer

    if (newlen == oldlen)

    {

        ptr = (char *) strstr(srcbuf, find);

        while (ptr != NULL)

        {

            if (!IS_PARTIAL_TOKEN(ptr))

                memcpy(ptr, replace, replace_len);

            ptr = strstr(ptr + find_len, find);

        } // while

        return;

    } // if

    // Allocate a new buffer

    char *newbuf = (char *) ctx->malloc_fn(newlen, ctx->malloc_data);

    memset(newbuf, '\0', newlen);

    // Find + replace

    char *prev_ptr = (char *) srcbuf;

    char *curr_ptr = (char *) newbuf;

    ptr = (char*) strstr(srcbuf, find);

    while (ptr != NULL)

    {

        memcpy(curr_ptr, prev_ptr, ptr - prev_ptr);

        curr_ptr += ptr - prev_ptr;

        if (!IS_PARTIAL_TOKEN(ptr))

        {

            memcpy(curr_ptr, replace, replace_len);

            curr_ptr += replace_len;

        } // if

        else

        {

            // Don't accidentally eat partial tokens...

            memcpy(curr_ptr, find, find_len);

            curr_ptr += find_len;

        } // else

        prev_ptr = ptr + find_len;

        ptr = strstr(prev_ptr, find);

    } // while

    #undef IS_PARTIAL_TOKEN

    // Copy the remaining part of the source buffer

    memcpy(curr_ptr, prev_ptr, (srcbuf + oldlen) - prev_ptr);

    // Free the source buffer

    ctx->free_fn((void *) srcbuf, ctx->malloc_data);

    // Point the source parameter to the new buffer

    *source = newbuf;

} // replaceVarname

static char *rewritePixelShader(MOJOSHADER_d3d11Shader *vshader,

                                MOJOSHADER_d3d11Shader *pshader)

{

    const MOJOSHADER_parseData *vpd = vshader->parseData;

    const MOJOSHADER_parseData *ppd = pshader->parseData;

    const char *_Output = "_Output" ENDLINE_STR "{" ENDLINE_STR;

    const char *_Input = "_Input" ENDLINE_STR "{" ENDLINE_STR;

    const char *vsrc = vpd->output;

    const char *psrc = ppd->output;

    const char *a, *b, *vout, *pstart, *vface, *pend;

    size_t substr_len;

    char *pfinal;

    #define MAKE_STRBUF(buf) \

        substr_len = b - a; \

        buf = (const char *) ctx->malloc_fn(substr_len + 1, ctx->malloc_data); \

        memset((void *) buf, '\0', substr_len + 1); \

        memcpy((void *) buf, a, substr_len);

    // Copy the vertex function's output struct into a buffer

    a = strstr(vsrc, _Output) + strlen(_Output);

    b = a;

    while (*(b++) != '}');

    b--;

    MAKE_STRBUF(vout)

    // Split up the pixel shader text...

    // ...everything up to the input contents...

    a = psrc;

    b = strstr(psrc, _Input) + strlen(_Input);

    MAKE_STRBUF(pstart)

    // ...everything after the input contents.

    a = b;

    while (*(a++) != '}');

    a--;

    while (*(b++) != '\0');

    MAKE_STRBUF(pend)

    // Find matching semantics

    int i, j;

    int vfaceidx = -1;

    const char *pvarname, *vvarname;

    for (i = 0; i < ppd->attribute_count; i++)

    {

        for (j = 0; j < vpd->output_count; j++)

        {

            if (ppd->attributes[i].usage == vpd->outputs[j].usage &&

                ppd->attributes[i].index == vpd->outputs[j].index)

            {

                pvarname = ppd->attributes[i].name;

                vvarname = vpd->outputs[j].name;

                if (strcmp(pvarname, vvarname) != 0)

                    replaceVarname(pvarname, vvarname, &pend);

            } // if

            else if (strcmp(ppd->attributes[i].name, "vPos") == 0 &&

                     vpd->outputs[j].usage == MOJOSHADER_USAGE_POSITION &&

                     vpd->outputs[j].index == 0)

            {

                pvarname = ppd->attributes[i].name;

                vvarname = vpd->outputs[j].name;

                if (strcmp(pvarname, vvarname) != 0)

                    replaceVarname(pvarname, vvarname, &pend);

            } // else if

        } // for

        if (strcmp(ppd->attributes[i].name, "vFace") == 0)

            vfaceidx = i;

    } // for

    // Special handling for VFACE

    vface = (vfaceidx != -1) ? "\tbool m_vFace : SV_IsFrontFace;\n" : "";

    // Concatenate the shader pieces together

    substr_len = strlen(pstart) + strlen(vout) + strlen(vface) + strlen(pend);

    pfinal = (char *) ctx->malloc_fn(substr_len + 1, ctx->malloc_data);

    memset((void *) pfinal, '\0', substr_len + 1);

    memcpy(pfinal, pstart, strlen(pstart));

    memcpy(pfinal + strlen(pstart), vout, strlen(vout));

    memcpy(pfinal + strlen(pstart) + strlen(vout), vface, strlen(vface));

    memcpy(pfinal + strlen(pstart) + strlen(vout) + strlen(vface), pend, strlen(pend));

    // Free the temporary buffers

    ctx->free_fn((void *) vout, ctx->malloc_data);

    ctx->free_fn((void *) pstart, ctx->malloc_data);

    ctx->free_fn((void *) pend, ctx->malloc_data);

    #undef MAKE_STRBUF

    return pfinal;

} // spliceVertexShaderInput

static ID3D11PixelShader *compilePixelShader(MOJOSHADER_d3d11Shader *vshader,

                                             MOJOSHADER_d3d11Shader *pshader)

{

    ID3D11PixelShader *retval = NULL;

    const char *source;

    ID3DBlob *blob;

    HRESULT result;

    int needs_free;

    if (pshader->parseData->attribute_count > 0)

    {

        source = rewritePixelShader(vshader, pshader);

        needs_free = 1;

    } // if

    else

    {

        source = pshader->parseData->output;

        needs_free = 0;

    } // else

    result = ctx->D3DCompileFunc(source, strlen(source),

                                 pshader->parseData->mainfn, NULL, NULL,

                                 pshader->parseData->mainfn, "ps_4_0",

                                 D3D_SKIP_OPT, 0, &blob, &blob);

    if (result < 0)

    {

        set_error((const char *) ID3D10Blob_GetBufferPointer(blob));

        ctx->free_fn((void *) source, ctx->malloc_data);

        return NULL;

    } // if

    ID3D11Device_CreatePixelShader(ctx->device,

                                   ID3D10Blob_GetBufferPointer(blob),

                                   ID3D10Blob_GetBufferSize(blob),

                                   NULL, &retval);

    ID3D10Blob_Release(blob);

    if (needs_free)

        ctx->free_fn((void *) source, ctx->malloc_data);

    return retval;

} // compilePixelShader

/* Public API */

int MOJOSHADER_d3d11CreateContext(void *device, void *deviceContext,