#include "boo/graphicsdev/nx_compiler.hpp"

/*
 * Copyright © 2008, 2009 Intel Corporation
 * Boo Modifications © 2018 Jack Andersen
 *
 * 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 (including the next
 * paragraph) 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.
 */

/** @file nx_compiler.cpp
 *
 * Based on standalone.cpp in compiler/glsl. This file provides a means to
 * compile and link GLSL sources directly into driver programs for the
 * Nouveau GM107 chipset configuration.
 */

#include "ast.h"
#include "glsl_parser_extras.h"
#include "ir_optimization.h"
#include "program.h"
#include "loop_analysis.h"
#include "string_to_uint_map.h"
#include "util/set.h"
#include "linker.h"
#include "ir_builder_print_visitor.h"
#include "builtin_functions.h"
#include "opt_add_neg_to_sub.h"

#include "main/shaderobj.h"
#include "st_program.h"

extern "C" {
#include "nouveau_winsys.h"
#include "nouveau_screen.h"
#include "nvc0/nvc0_program.h"
}

_GLAPI_EXPORT __thread void * _glapi_tls_Context;
_GLAPI_EXPORT __thread struct _glapi_table * _glapi_tls_Dispatch;

int
_glapi_add_dispatch( const char * const * function_names,
                     const char * parameter_signature )
{
    return 0;
}

void
_glapi_destroy_multithread(void)
{
}

void
_glapi_check_multithread(void)
{
}

void
_glapi_set_context(void *context)
{
    _glapi_tls_Context = context;
}

void *
_glapi_get_context()
{
    return _glapi_tls_Context;
}

void
_glapi_set_dispatch(struct _glapi_table *dispatch)
{
    _glapi_tls_Dispatch = dispatch;
}

struct _glapi_table *
_glapi_get_dispatch()
{
    return _glapi_tls_Dispatch;
}

GLuint
_glapi_get_dispatch_table_size(void)
{
    /*
     * The dispatch table size (number of entries) is the size of the
     * _glapi_table struct plus the number of dynamic entries we can add.
     * The extra slots can be filled in by DRI drivers that register new
     * extension functions.
     */
    return 0;
}

class dead_variable_visitor : public ir_hierarchical_visitor {
public:
    dead_variable_visitor()
    {
        variables = _mesa_set_create(NULL,
                                     _mesa_hash_pointer,
                                     _mesa_key_pointer_equal);
    }

    virtual ~dead_variable_visitor()
    {
        _mesa_set_destroy(variables, NULL);
    }

    virtual ir_visitor_status visit(ir_variable *ir)
    {
        /* If the variable is auto or temp, add it to the set of variables that
         * are candidates for removal.
         */
        if (ir->data.mode != ir_var_auto && ir->data.mode != ir_var_temporary)
            return visit_continue;

        _mesa_set_add(variables, ir);

        return visit_continue;
    }

    virtual ir_visitor_status visit(ir_dereference_variable *ir)
    {
        struct set_entry *entry = _mesa_set_search(variables, ir->var);

        /* If a variable is dereferenced at all, remove it from the set of
         * variables that are candidates for removal.
         */
        if (entry != NULL)
            _mesa_set_remove(variables, entry);

        return visit_continue;
    }

    void remove_dead_variables()
    {
        struct set_entry *entry;

        set_foreach(variables, entry) {
            ir_variable *ir = (ir_variable *) entry->key;

            assert(ir->ir_type == ir_type_variable);
            ir->remove();
        }
    }

private:
    set *variables;
};

void nx_compiler::compile_shader(struct gl_context *ctx, struct gl_shader *shader)
{
    struct _mesa_glsl_parse_state *state =
        new(shader) _mesa_glsl_parse_state(ctx, shader->Stage, shader);

    _mesa_glsl_compile_shader(ctx, shader, m_options.dump_ast,
                              m_options.dump_hir, true);

    /* Print out the resulting IR */
    if (!state->error && m_options.dump_lir) {
        _mesa_print_ir(stdout, shader->ir, state);
    }
}

nx_compiler::nx_compiler()
{
    m_options.glsl_version = 330;
    m_options.do_link = true;
}

nx_compiler::~nx_compiler()
{
    if (m_ownsCtx)
    {
        _mesa_glsl_release_types();
        _mesa_glsl_release_builtin_functions();
        if (m_st)
            st_destroy_context(m_st);
        if (m_screen)
            m_screen->destroy(m_screen);
    }
}

bool nx_compiler::initialize(struct pipe_screen *screen, struct st_context *st,
                             const struct standalone_options *o)
{
    m_screen = screen;
    m_st = st;
    if (o)
        memcpy(&m_options, o, sizeof(*o));
    return true;
}

bool nx_compiler::initialize(const struct standalone_options* o)
{
    m_ownsCtx = true;
    bool glsl_es;

    if (o)
        memcpy(&m_options, o, sizeof(*o));

    switch (m_options.glsl_version) {
    case 100:
    case 300:
        glsl_es = true;
        break;
    case 110:
    case 120:
    case 130:
    case 140:
    case 150:
    case 330:
    case 400:
    case 410:
    case 420:
    case 430:
    case 440:
    case 450:
    case 460:
        glsl_es = false;
        break;
    default:
        fprintf(stderr, "Unrecognized GLSL version `%d'\n", m_options.glsl_version);
        return false;
    }

    gl_api use_api;
    if (glsl_es) {
        use_api = API_OPENGLES2;
    } else {
        use_api = m_options.glsl_version > 130 ? API_OPENGL_CORE : API_OPENGL_COMPAT;
    }

    struct nouveau_screen *(*init)(struct nouveau_device *);

    struct nouveau_drm *fakedrm = (struct nouveau_drm *)malloc(sizeof(struct nouveau_drm));
    if (!fakedrm)
        return false;
    memset(fakedrm, 0, sizeof(*fakedrm));
    nouveau_device *ndev;
    if (nouveau_device_new(&fakedrm->client, 0, nullptr, 0, &ndev))
        return false;

    switch (ndev->chipset & ~0xf) {
#if 0
    case 0x30:
    case 0x40:
    case 0x60:
        init = nv30_screen_create;
        break;
    case 0x50:
    case 0x80:
    case 0x90:
    case 0xa0:
        init = nv50_screen_create;
        break;
#endif
    default:
    case 0xc0:
    case 0xd0:
    case 0xe0:
    case 0xf0:
    case 0x100:
    case 0x110:
    case 0x120:
    case 0x130:
        init = nvc0_screen_create;
        break;
    }

    struct nouveau_screen *screen = init(ndev);
    if (!screen)
        return false;
    screen->refcount = 1;
    struct pipe_context *p_ctx = screen->base.context_create(&screen->base, nullptr, 0);
    if (!p_ctx)
    {
        screen->base.destroy(&screen->base);
        return false;
    }

    st_config_options opts = {};
    struct st_context *st = st_create_context(use_api, p_ctx, nullptr, nullptr, &opts, false);
    if (!st)
    {
        screen->base.destroy(&screen->base);
        return false;
    }

    return initialize(&screen->base, st);
}

nx_shader_stage_object::nx_shader_stage_object(const nx_shader_stage_object& other)
: m_parent(other.m_parent)
{
    if (!other.m_shader || !m_parent)
        return;
    struct gl_context *ctx = m_parent->m_st->ctx;
    _mesa_reference_shader(ctx, &m_shader, other.m_shader);
}

nx_shader_stage_object& nx_shader_stage_object::operator=(const nx_shader_stage_object& other)
{
    m_parent = other.m_parent;
    if (!other.m_shader || !m_parent)
        return *this;
    struct gl_context *ctx = m_parent->m_st->ctx;
    _mesa_reference_shader(ctx, &m_shader, other.m_shader);
    return *this;
}

void nx_shader_stage_object::reset()
{
    if (!m_shader || !m_parent)
        return;
    struct gl_context *ctx = m_parent->m_st->ctx;
    _mesa_reference_shader(ctx, &m_shader, nullptr);
}

nx_shader_stage_object::operator bool() const
{
    if (!m_shader)
        return false;
    return m_shader->CompileStatus;
}

nx_shader_stage nx_shader_stage_object::stage() const
{
    return nx_shader_stage(m_shader->Stage);
}

const char* nx_shader_stage_object::info_log() const
{
    if (!m_shader)
        return nullptr;
    return m_shader->InfoLog;
}

nx_linked_shader::nx_linked_shader(const nx_linked_shader& other)
: m_parent(other.m_parent)
{
    if (!other.m_program || !m_parent)
        return;
    struct gl_context *ctx = m_parent->m_st->ctx;
    _mesa_reference_shader_program(ctx, &m_program, other.m_program);
}

nx_linked_shader& nx_linked_shader::operator=(const nx_linked_shader& other)
{
    m_parent = other.m_parent;
    if (!other.m_program || !m_parent)
        return *this;
    struct gl_context *ctx = m_parent->m_st->ctx;
    _mesa_reference_shader_program(ctx, &m_program, other.m_program);
    return *this;
}

void nx_linked_shader::reset()
{
    if (!m_program || !m_parent)
        return;
    struct gl_context *ctx = m_parent->m_st->ctx;
    _mesa_reference_shader_program(ctx, &m_program, nullptr);
}

nx_shader_stage_object nx_compiler::compile(nx_shader_stage type, const char *source)
{
    struct gl_context *ctx = m_st->ctx;

    nx_shader_stage_object ret(*this);
    ret.m_shader = rzalloc(nullptr, gl_shader);
    assert(ret.m_shader != NULL);
    ret.m_shader->RefCount = 1;

    ret.m_shader->Stage = gl_shader_stage(type);
    ret.m_shader->Source = source;

    compile_shader(ctx, ret.m_shader);

    /* Mesa doesn't actually own the source, so take it away here */
    ret.m_shader->Source = nullptr;

    return ret;
}

nx_linked_shader nx_compiler::link(unsigned num_stages, const nx_shader_stage_object **stages, std::string* infoLog)
{
    nx_linked_shader ret(*this);
    int status = EXIT_SUCCESS;
    struct gl_context *ctx = m_st->ctx;

    struct gl_shader_program *whole_program;

    whole_program = rzalloc (NULL, struct gl_shader_program);
    assert(whole_program != NULL);
    whole_program->Type = GL_SHADER_PROGRAM_MESA;
    whole_program->data = rzalloc(whole_program, struct gl_shader_program_data);
    assert(whole_program->data != NULL);
    whole_program->data->RefCount = 1;
    whole_program->data->InfoLog = ralloc_strdup(whole_program->data, "");
    ret.m_program = whole_program;

    whole_program->Shaders = (struct gl_shader **)calloc(num_stages, sizeof(struct gl_shader *));
    assert(whole_program->Shaders != NULL);

    for (unsigned i = 0; i < num_stages; i++) {
        whole_program->Shaders[whole_program->NumShaders] = stages[i]->m_shader;
        stages[i]->m_shader->RefCount++;
        whole_program->NumShaders++;

        if (!stages[i]->m_shader->CompileStatus) {
            status = EXIT_FAILURE;
            break;
        }
    }

    if (status == EXIT_SUCCESS) {
        _mesa_clear_shader_program_data(ctx, whole_program);

        if (m_options.do_link)  {
            link_shaders(ctx, whole_program);
            for (int i = 0; i < MESA_SHADER_STAGES; ++i) {
                if (whole_program->_LinkedShaders[i])
                    whole_program->_LinkedShaders[i]->Program->Target = _mesa_shader_stage_to_program(i);
            }
        } else {
            const gl_shader_stage stage = whole_program->Shaders[0]->Stage;

            whole_program->data->LinkStatus = LINKING_SUCCESS;
            whole_program->_LinkedShaders[stage] =
                link_intrastage_shaders(whole_program /* mem_ctx */,
                                        ctx,
                                        whole_program,
                                        whole_program->Shaders,
                                        1,
                                        true);
            whole_program->_LinkedShaders[stage]->Program->Target = _mesa_shader_stage_to_program(stage);

            /* Par-linking can fail, for example, if there are undefined external
             * references.
             */
            if (whole_program->_LinkedShaders[stage] != NULL) {
                assert(whole_program->data->LinkStatus);

                struct gl_shader_compiler_options *const compiler_options =
                    &ctx->Const.ShaderCompilerOptions[stage];

                exec_list *const ir =
                    whole_program->_LinkedShaders[stage]->ir;

                bool progress;
                do {
                    progress = do_function_inlining(ir);

                    progress = do_common_optimization(ir,
                                                      false,
                                                      false,
                                                      compiler_options,
                                                      true)
                               && progress;
                } while(progress);
            }
        }

        status = (whole_program->data->LinkStatus) ? EXIT_SUCCESS : EXIT_FAILURE;

        if (infoLog)
            *infoLog = whole_program->data->InfoLog;

        if (status == EXIT_SUCCESS) {
            for (unsigned i = 0; i < MESA_SHADER_STAGES; i++) {
                struct gl_linked_shader *shader = whole_program->_LinkedShaders[i];

                if (!shader)
                    continue;

                add_neg_to_sub_visitor v;
                visit_list_elements(&v, shader->ir);

                dead_variable_visitor dv;
                visit_list_elements(&dv, shader->ir);
                dv.remove_dead_variables();
            }

            if (m_options.dump_builder) {
                for (unsigned i = 0; i < MESA_SHADER_STAGES; i++) {
                    struct gl_linked_shader *shader = whole_program->_LinkedShaders[i];

                    if (!shader)
                        continue;

                    _mesa_print_builder_for_ir(stdout, shader->ir);
                }
            }

            ctx->_Shader = &ctx->Shader;
            st_link_shader(ctx, whole_program);
            ctx->_Shader = nullptr;
            return ret;
        }
    }

    return nx_linked_shader(*this);
}

static void
SizeProgramBuffer(const nvc0_program *prog, size_t &sz)
{
    sz += 140;
    sz += prog->code_size;
}

template<class T> static void
OutputField(T f, uint8_t *&ptr)
{
    memcpy(ptr, &f, sizeof(f));
    ptr += sizeof(f);
}

static void
BuildProgramBuffer(const nvc0_program *prog, uint8_t *&ptr)
{
    OutputField(prog->type, ptr);
    OutputField(prog->translated, ptr);
    OutputField(prog->need_tls, ptr);
    OutputField(prog->num_gprs, ptr);
    OutputField<uint32_t>(prog->code_base, ptr);
    OutputField<uint32_t>(prog->code_size, ptr);
    OutputField<uint32_t>(prog->parm_size, ptr);
    for (const auto& h : prog->hdr)
        OutputField(h, ptr);
    for (const auto& h : prog->flags)
        OutputField(h, ptr);
    OutputField(prog->vp.clip_mode, ptr);
    OutputField(prog->vp.clip_enable, ptr);
    OutputField(prog->vp.cull_enable, ptr);
    OutputField(prog->vp.num_ucps, ptr);
    OutputField(prog->vp.edgeflag, ptr);
    OutputField(prog->vp.need_vertex_id, ptr);
    OutputField(prog->vp.need_draw_parameters, ptr);
    OutputField(prog->fp.early_z, ptr);
    OutputField(prog->fp.colors, ptr);
    OutputField(prog->fp.color_interp[0], ptr);
    OutputField(prog->fp.color_interp[1], ptr);
    OutputField(prog->fp.sample_mask_in, ptr);
    OutputField(prog->fp.force_persample_interp, ptr);
    OutputField(prog->fp.flatshade, ptr);
    OutputField(prog->fp.reads_framebuffer, ptr);
    OutputField(prog->fp.post_depth_coverage, ptr);
    OutputField(prog->tp.tess_mode, ptr);
    OutputField(prog->tp.input_patch_size, ptr);
    OutputField(prog->cp.lmem_size, ptr);
    OutputField(prog->cp.smem_size, ptr);
    OutputField(prog->num_barriers, ptr);
    memcpy(ptr, prog->code, prog->code_size);
    ptr += prog->code_size;
}

std::pair<std::shared_ptr<uint8_t[]>, size_t>
nx_compiler::offline_link(unsigned num_stages, const nx_shader_stage_object **stages, std::string *infoLog)
{
    std::pair<std::shared_ptr<uint8_t[]>, size_t> ret = {};
    auto whole_program = link(num_stages, stages, infoLog);
    if (!whole_program)
        return ret;

    for (unsigned i = 0; i < MESA_SHADER_STAGES; i++) {
        struct gl_linked_shader *shader = whole_program.m_program->_LinkedShaders[i];
        if (!shader)
            continue;
        struct gl_program *prog = shader->Program;

        switch (prog->Target) {
        case GL_VERTEX_PROGRAM_ARB: {
            struct st_vertex_program *p = (struct st_vertex_program *)prog;
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            SizeProgramBuffer(dp, ret.second);
            break;
        }
        case GL_TESS_CONTROL_PROGRAM_NV: {
            struct st_common_program *p = st_common_program(prog);
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            SizeProgramBuffer(dp, ret.second);
            break;
        }
        case GL_TESS_EVALUATION_PROGRAM_NV: {
            struct st_common_program *p = st_common_program(prog);
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            SizeProgramBuffer(dp, ret.second);
            break;
        }
        case GL_GEOMETRY_PROGRAM_NV: {
            struct st_common_program *p = st_common_program(prog);
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            SizeProgramBuffer(dp, ret.second);
            break;
        }
        case GL_FRAGMENT_PROGRAM_ARB: {
            struct st_fragment_program *p = (struct st_fragment_program *)prog;
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            SizeProgramBuffer(dp, ret.second);
            break;
        }
        default:
            assert(0);
        }
    }

    ret.first.reset(new uint8_t[ret.second]);
    uint8_t *pbuf = ret.first.get();

    for (unsigned i = 0; i < MESA_SHADER_STAGES; i++) {
        struct gl_linked_shader *shader = whole_program.m_program->_LinkedShaders[i];
        if (!shader)
            continue;
        struct gl_program *prog = shader->Program;

        switch (prog->Target) {
        case GL_VERTEX_PROGRAM_ARB: {
            struct st_vertex_program *p = (struct st_vertex_program *)prog;
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            BuildProgramBuffer(dp, pbuf);
            break;
        }
        case GL_TESS_CONTROL_PROGRAM_NV: {
            struct st_common_program *p = st_common_program(prog);
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            BuildProgramBuffer(dp, pbuf);
            break;
        }
        case GL_TESS_EVALUATION_PROGRAM_NV: {
            struct st_common_program *p = st_common_program(prog);
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            BuildProgramBuffer(dp, pbuf);
            break;
        }
        case GL_GEOMETRY_PROGRAM_NV: {
            struct st_common_program *p = st_common_program(prog);
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            BuildProgramBuffer(dp, pbuf);
            break;
        }
        case GL_FRAGMENT_PROGRAM_ARB: {
            struct st_fragment_program *p = (struct st_fragment_program *)prog;
            nvc0_program *dp = (nvc0_program *)p->variants->driver_shader;
            BuildProgramBuffer(dp, pbuf);
            break;
        }
        default:
            assert(0);
        }
    }

    return ret;
}