/**

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

 */

#if (defined(__APPLE__) && defined(__MACH__))

#define PLATFORM_APPLE 1

#include "TargetConditionals.h"

#define OBJC_OLD_DISPATCH_PROTOTYPES 1

#include <objc/message.h>

#define objc_msgSend_STR ((void* (*)(void*, void*, const char*))objc_msgSend)

#define objc_msgSend_PTR ((void* (*)(void*, void*, void*))objc_msgSend)

#define objc_msgSend_INT_PTR ((void* (*)(void*, void*, int, void*))objc_msgSend)

#define objc_msgSend_PTR_PTR_PTR ((void* (*)(void*, void*, void*, void*, void*))objc_msgSend)

#endif /* (defined(__APPLE__) && defined(__MACH__)) */

#define __MOJOSHADER_INTERNAL__ 1

#include "mojoshader_internal.h"

typedef struct MOJOSHADER_mtlUniformBuffer MOJOSHADER_mtlUniformBuffer;

typedef struct MOJOSHADER_mtlShader

{

    const MOJOSHADER_parseData *parseData;

    MOJOSHADER_mtlUniformBuffer *ubo;

    void *library; // MTLLibrary*

    int numInternalBuffers;

} MOJOSHADER_mtlShader;

// 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

// profile-specific implementations...

#if SUPPORT_PROFILE_METAL && PLATFORM_APPLE

#ifdef MOJOSHADER_EFFECT_SUPPORT

/* Structs */

typedef struct MOJOSHADER_mtlEffect

{

    MOJOSHADER_effect *effect;

    unsigned int num_shaders;

    MOJOSHADER_mtlShader *shaders;

    unsigned int *shader_indices;

    unsigned int num_preshaders;

    unsigned int *preshader_indices;

    MOJOSHADER_mtlShader *current_vert;

    MOJOSHADER_mtlShader *current_frag;

    MOJOSHADER_effectShader *current_vert_raw;

    MOJOSHADER_effectShader *current_frag_raw;

    MOJOSHADER_mtlShader *prev_vert;

    MOJOSHADER_mtlShader *prev_frag;

    void *library; // MTLLibrary*

} MOJOSHADER_mtlEffect;

typedef struct MOJOSHADER_mtlUniformBuffer

{

    void *device; // MTLDevice*

    int bufferSize;

    int numInternalBuffers;

    void **internalBuffers; // MTLBuffer*

    int internalBufferSize;

    int internalOffset;

    int currentFrame;

    int alreadyWritten;

} MOJOSHADER_mtlUniformBuffer;

typedef struct MOJOSHADER_mtlShaderState

{

    MOJOSHADER_mtlShader *vertexShader;

    MOJOSHADER_mtlShader *fragmentShader;

    void *vertexUniformBuffer; // MTLBuffer*

    void *fragmentUniformBuffer; // MTLBuffer*

    int vertexUniformOffset;

    int fragmentUniformOffset;

} MOJOSHADER_mtlShaderState;

/* Objective-C selector references */

static void *classNSString = NULL;

static void *selAlloc = NULL;

static void *selInitWithUTF8String = NULL;

static void *selUTF8String = NULL;

static void *selLength = NULL;

static void *selContents = NULL;

static void *selNewBufferWithLength = NULL;

static void *selRelease = NULL;

static void *selNewLibraryWithSource = NULL;

static void *selLocalizedDescription = NULL;

static void *selNewFunctionWithName = NULL;

static void *selRetain = NULL;

/* Helper functions */

static void initSelectors(void)

{

    classNSString = (void*) objc_getClass("NSString");

    selAlloc = sel_registerName("alloc");

    selInitWithUTF8String = sel_registerName("initWithUTF8String:");

    selUTF8String = sel_registerName("UTF8String");

    selLength = sel_registerName("length");

    selContents = sel_registerName("contents");

    selNewBufferWithLength = sel_registerName("newBufferWithLength:options:");

    selRelease = sel_registerName("release");

    selNewLibraryWithSource = sel_registerName("newLibraryWithSource:options:error:");

    selLocalizedDescription = sel_registerName("localizedDescription");

    selNewFunctionWithName = sel_registerName("newFunctionWithName:");

    selRetain = sel_registerName("retain");

} // initSelectors

static void *cstr_to_nsstr(const char *str)

{

    return objc_msgSend_STR(

        objc_msgSend(classNSString, selAlloc),

        selInitWithUTF8String,

        str

    );

} // cstr_to_nsstr

static const char *nsstr_to_cstr(void *str)

{

    return (char *) objc_msgSend(str, selUTF8String);

} // nssstr_to_cstr

/* Linked list */

typedef struct LLNODE {

    MOJOSHADER_mtlUniformBuffer *data;

    struct LLNODE *next;

} LLNODE;

static LLNODE *LL_append_node(LLNODE **baseNode,

                              MOJOSHADER_malloc m,

                              void *d)

{

    LLNODE *prev = NULL;

    LLNODE *node = *baseNode;

    /* Append a node to the linked list. */

    while (node != NULL)

    {

        prev = node;

        node = node->next;

    } // while

    node = m(sizeof(LLNODE), d);

    node->next = NULL;

    /* Connect the old to the new. */

    if (prev != NULL)

        prev->next = node;

    /* Special case for the first node. */

    if (*baseNode == NULL)

        *baseNode = node;

    return node;

} // LL_append_node

static void LL_remove_node(LLNODE **baseNode,

                           MOJOSHADER_mtlUniformBuffer *data,

                           MOJOSHADER_free f,

                           void *d)

{

    LLNODE *prev = NULL;

    LLNODE *node = *baseNode;

    /* Search for node with matching data pointer. */

    while (node != NULL && node->data != data)

    {

        prev = node;

        node = node->next;

    } // while

    if (node == NULL)

    {

        /* This should never happen. */

        assert(0);

    } // if

    /* Clear data pointer. The data must be freed separately. */

    node->data = NULL;

    /* Connect the old to the new. */

    if (prev != NULL)

        prev->next = node->next;

    /* Special cases where the first node is removed. */

    if (prev == NULL)

        *baseNode = (node->next != NULL) ? node->next : NULL;

    /* Free the node! */

    f(node, d);

} // LL_remove_node

/* Internal register utilities */

// 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

// 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];

static inline void copy_parameter_data(MOJOSHADER_effectParam *params,

                                       unsigned int *param_loc,

                                       MOJOSHADER_symbol *symbols,

                                       unsigned int symbol_count,

                                       float *regf, int *regi, uint8 *regb)

{

    int i, j, r, c;

    i = 0;

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

    {

        const MOJOSHADER_symbol *sym = &symbols[i];

        const MOJOSHADER_effectValue *param = &params[param_loc[i]].value;

        // float/int registers are vec4, so they have 4 elements each

        const uint32 start = sym->register_index << 2;

        if (param->type.parameter_type == MOJOSHADER_SYMTYPE_FLOAT)

            memcpy(regf + start, param->valuesF, sym->register_count << 4);

        else if (sym->register_set == MOJOSHADER_SYMREGSET_FLOAT4)

        {

            // Structs are a whole different world...

            if (param->type.parameter_class == MOJOSHADER_SYMCLASS_STRUCT)

                memcpy(regf + start, param->valuesF, sym->register_count << 4);

            else

            {

                // Sometimes int/bool parameters get thrown into float registers...

                j = 0;

                do

                {

                    c = 0;

                    do

                    {

                        regf[start + (j << 2) + c] = (float) param->valuesI[(j << 2) + c];

                    } while (++c < param->type.columns);

                } while (++j < sym->register_count);

            } // else

        } // else if

        else if (sym->register_set == MOJOSHADER_SYMREGSET_INT4)

            memcpy(regi + start, param->valuesI, sym->register_count << 4);

        else if (sym->register_set == MOJOSHADER_SYMREGSET_BOOL)

        {

            j = 0;

            r = 0;

            do

            {

                c = 0;

                do

                {

                    // regb is not a vec4, enjoy that 'start' bitshift! -flibit

                    regb[(start >> 2) + r + c] = param->valuesI[(j << 2) + c];

                    c++;

                } while (c < param->type.columns && ((r + c) < sym->register_count));

                r += c;

                j++;

            } while (r < sym->register_count);

        } // else if

    } // for

} // copy_parameter_data

/* Uniform buffer utilities */

static inline int next_highest_alignment(int n)

{

    #if TARGET_OS_IOS || TARGET_OS_TV || TARGET_OS_SIMULATOR

    int align = 16;

    #else

    int align = 256;

    #endif

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

} // next_highest_alignment

static int UBO_buffer_length(void *buffer)

{

    return (int) objc_msgSend(buffer, selLength);

} // UBO_buffer_length

static void *UBO_buffer_contents(void *buffer)

{

    return (void *) objc_msgSend(buffer, selContents);

} // UBO_buffer_contents

static void *UBO_create_backing_buffer(MOJOSHADER_mtlUniformBuffer *ubo, int f)

{

    void *oldBuffer = ubo->internalBuffers[f];

    void *newBuffer = objc_msgSend_INT_PTR(

        ubo->device,

        selNewBufferWithLength,

        ubo->internalBufferSize,

        NULL

    );

    if (oldBuffer != NULL)

    {

        // Copy over data from old buffer

        memcpy(

            UBO_buffer_contents(newBuffer),

            UBO_buffer_contents(oldBuffer),

            UBO_buffer_length(oldBuffer)

        );

        // Free the old buffer

        objc_msgSend(oldBuffer, selRelease);

    } //if

    return newBuffer;

} // UBO_create_backing_buffer

static void UBO_predraw(MOJOSHADER_mtlUniformBuffer *ubo)

{

    if (!ubo->alreadyWritten)

    {

        ubo->alreadyWritten = 1;

        return;

    } // if

    ubo->internalOffset += ubo->bufferSize;

    int buflen = UBO_buffer_length(ubo->internalBuffers[ubo->currentFrame]);

    if (ubo->internalOffset >= buflen)

    {

        // Double capacity when we're out of room

        if (ubo->internalOffset >= ubo->internalBufferSize)

            ubo->internalBufferSize *= 2;

        ubo->internalBuffers[ubo->currentFrame] =

            UBO_create_backing_buffer(ubo, ubo->currentFrame);

    } //if

} // UBO_predraw

static void UBO_end_frame(MOJOSHADER_mtlUniformBuffer *ubo)

{

    ubo->internalOffset = 0;

    ubo->currentFrame = (ubo->currentFrame + 1) % ubo->numInternalBuffers;

    ubo->alreadyWritten = 0;

} // UBO_end_frame

LLNODE *ubos = NULL; /* global linked list of all active UBOs */

static MOJOSHADER_mtlUniformBuffer *create_ubo(MOJOSHADER_mtlShader *shader,

                                               void *mtlDevice,

                                               MOJOSHADER_malloc m,

                                               void *d)

{

    int uniformCount = shader->parseData->uniform_count;

    if (uniformCount == 0)

        return NULL;

    // Calculate how big we need to make the buffer

    int buflen = 0;

    for (int i = 0; i < uniformCount; i += 1)

    {

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

        int uniformSize = 16;

        if (shader->parseData->uniforms[i].type == MOJOSHADER_UNIFORM_BOOL)

            uniformSize = 1;

        buflen += (arrayCount ? arrayCount : 1) * uniformSize;

    } // for

    // Make the UBO

    MOJOSHADER_mtlUniformBuffer *ubo = (MOJOSHADER_mtlUniformBuffer *) m(sizeof(MOJOSHADER_mtlUniformBuffer), d);

    ubo->device = mtlDevice;

    ubo->alreadyWritten = 0;

    ubo->bufferSize = next_highest_alignment(buflen);

    ubo->currentFrame = 0;

    ubo->numInternalBuffers = shader->numInternalBuffers;

    ubo->internalBufferSize = ubo->bufferSize * 16; // pre-allocate some extra room!

    ubo->internalBuffers = m(ubo->numInternalBuffers * sizeof(void*), d);

    ubo->internalOffset = 0;

    for (int i = 0; i < ubo->numInternalBuffers; i++)

    {

        ubo->internalBuffers[i] = NULL;

        ubo->internalBuffers[i] = UBO_create_backing_buffer(ubo, i);

    } // for

    /* Add the UBO to the global list so it can be updated. */

    LLNODE *node = LL_append_node(&ubos, m, d);

    node->data = ubo;

    return ubo;

} // create_ubo

static void dealloc_ubo(MOJOSHADER_mtlShader *shader,

                        MOJOSHADER_free f,

                        void* d)

{

    if (shader->ubo == NULL)

        return;

    LL_remove_node(&ubos, shader->ubo, f, d);

    for (int i = 0; i < shader->ubo->numInternalBuffers; i++)

    {

        objc_msgSend(shader->ubo->internalBuffers[i], selRelease);

        shader->ubo->internalBuffers[i] = NULL;

    } // for

    f(shader->ubo->internalBuffers, d);

    f(shader->ubo, d);

} // dealloc_ubo

static void *get_uniform_buffer(MOJOSHADER_mtlShader *shader)

{

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

        return NULL;

    return shader->ubo->internalBuffers[shader->ubo->currentFrame];

} // get_uniform_buffer

static int get_uniform_offset(MOJOSHADER_mtlShader *shader)

{

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

        return 0;

    return shader->ubo->internalOffset;

} // get_uniform_offset

static void update_uniform_buffer(MOJOSHADER_mtlShader *shader)

{

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

        return;

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

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

    {

        regF = vs_reg_file_f;

        regI = vs_reg_file_i;

        regB = vs_reg_file_b;

    } // if

    else

    {

        regF = ps_reg_file_f;

        regI = ps_reg_file_i;

        regB = ps_reg_file_b;

    } // else

    UBO_predraw(shader->ubo);

    void *buf = shader->ubo->internalBuffers[shader->ubo->currentFrame];

    void *contents = UBO_buffer_contents(buf) + shader->ubo->internalOffset;

    int offset = 0;

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

    {

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

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

        int size = arrayCount ? arrayCount : 1;

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

        {

            case MOJOSHADER_UNIFORM_FLOAT:

                memcpy(

                    contents + (offset * 16),

                    &regF[4 * idx],

                    size * 16

                );

                break;

            case MOJOSHADER_UNIFORM_INT:

                // !!! FIXME: Need a test case

                memcpy(

                    contents + (offset * 16),

                    &regI[4 * idx],

                    size * 16

                );

                break;

            case MOJOSHADER_UNIFORM_BOOL:

                // !!! FIXME: Need a test case

                memcpy(

                    contents + offset,

                    &regB[idx],

                    size

                );

                break;

            default:

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

                break;

        } // switch

        offset += size;

    } // for

} // update_uniform_buffer

/* Public API */

MOJOSHADER_mtlEffect *MOJOSHADER_mtlCompileEffect(MOJOSHADER_effect *effect,

                                                  void *mtlDevice,

                                                  int numBackingBuffers)

{

    int i;

    MOJOSHADER_malloc m = effect->malloc;

    MOJOSHADER_free f = effect->free;

    void *d = effect->malloc_data;

    int current_shader = 0;

    int current_preshader = 0;

    int src_len = 0;

    // Make sure the Objective-C selectors have been initialized...

    if (selAlloc == NULL)

        initSelectors();

    MOJOSHADER_mtlEffect *retval = (MOJOSHADER_mtlEffect *) m(sizeof (MOJOSHADER_mtlEffect), d);

    if (retval == NULL)

    {