/*

  codegen_common.c -- Common functions for the code generators

  Copyright © 2022-2024 Samuel Lidén Borell <samuel@kodafritt.se>

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

*/

#include "codegen_common.h"
#include "../outformat/outformat_common.h"
#include <assert.h>
#include <stdio.h>
#include <string.h>

void cg_ir_op_iter_init(struct IrOpIter *it, struct EbbIter *ebb_iter)
{
    const struct FuncIR *func = ebb_iter->func;
    unsigned num = ebb_iter->ebb->first_op;
    it->end_num = ebb_iter->ebb->end_op;
    it->current_num = num;
    it->func = func;
    if (num != it->end_num) {
        it->op = &func->opchunk_by_id[num/OPS_PER_CHUNK]->ops[num & (OPS_PER_CHUNK-1)];
    }
}

int cg_ir_op_iter_valid(struct IrOpIter *it)
{
    return it->current_num < it->end_num;
}

int cg_ir_op_iter_is_last(struct IrOpIter *it)
{
    return it->current_num == it->end_num-1;
}

void cg_ir_op_iter_next(struct IrOpIter *it)
{
    unsigned num = ++it->current_num;
    if ((num & (OPS_PER_CHUNK-1)) != 0) {
        it->op++;
    } else if (it->current_num < it->end_num) {
        it->op = &it->func->opchunk_by_id[num/OPS_PER_CHUNK]->ops[0];
    }
}

const struct Op *cg_lookahead(struct IrOpIter *it)
{
    unsigned num = it->current_num + 1;
    if ((num & (OPS_PER_CHUNK-1)) != 0) {
        return &it->op[1];
    } else if (it->current_num < it->end_num) {
        return &it->func->opchunk_by_id[num/OPS_PER_CHUNK]->ops[0];
    } else {
        return NULL;
    }
}

/*const struct Op *cg_lookahead_is(struct IrOpIter *it, unsigned optype)
{
    const struct Op *op = cg_lookahead(it);
    if (op && op->op == optype) return op;
    else return NULL;
}*/

/**
 * Allocates a new chunk for machine code.
 *
 * \param   Compilation context
 * \return  Pointer to the start of the new chunk.
 */
unsigned char *cg_morespace(struct CgCtx *cg)
{
    /* XXX if temporary per-function arenas are added, then this should allocate in "persistent" arena */
    struct CgCodeChunk *chunk = allocp(cg->csbe, sizeof(struct CgCodeChunk));
    if (!chunk) return NULL;
    chunk->next = NULL;
    if (cg->code.last_chunk) {
        unsigned size = CG_CHUNK_SIZE - (cg->code.end - cg->code.ptr);
        cg->code.last_chunk->next = chunk;
        cg->code.last_chunk->size = size;
        cg->code.size += size;
    } else {
        cg->code.chunks = chunk;
    }
    cg->code.last_chunk = chunk;
    cg->code.ptr = &chunk->code[0];
    cg->code.end = &chunk->code[CG_CHUNK_SIZE];
    return cg->code.ptr;
}

void cgcode_reset(struct CgCode *code)
{
    code->chunks = NULL;
    code->last_chunk = NULL;
    code->ptr = NULL;
    code->end = NULL;
    code->size = 0;
}

void cg_lastchunk(struct CgCtx *cg)
{
    if (cg->code.last_chunk != NULL) {
        unsigned size = CG_CHUNK_SIZE - (cg->code.end - cg->code.ptr);
        cg->code.last_chunk->size = size;
        cg->code.size += size;
    }
}

uint32 cg_get_position(const struct CgCtx *cg)
{
    if (cg->code.last_chunk != NULL) {
        unsigned size = CG_CHUNK_SIZE - (cg->code.end - cg->code.ptr);
        return cg->code.size + size;
    } else {
        return 0;
    }
}

void cg_align_stacksize(struct CgCtx *cg, unsigned add, unsigned stackalign)
{
    /* FIXME detect overflows */
    uint32 size = cg->func_stacksize + add;
    cg->func_stacksize = UINT_ALIGN(size, stackalign);
}

int cg_add_addrslot(struct CgCtx *cg, struct CgAddrSlot **aslist,
                    unsigned addrtype)
{
    struct CgAddrSlot *as = allocp(cg->csbe, sizeof(struct CgAddrSlot));
    ZCHK(as);
    assert(addrtype);
    as->next = *aslist;
    as->chunk = cg->code.last_chunk;
    assert(!cg->code.ptr || cg->code.ptr < cg->code.end);
    as->offset_in_chunk = cg->code.ptr ?
            CG_CHUNK_SIZE - (cg->code.end - cg->code.ptr) : 0;
    assert(as->offset_in_chunk < CG_CHUNK_SIZE);
    as->location = cg_get_position(cg);
    as->addrtype = addrtype;
    *aslist = as;
    return 1;
}

void cg_edit_bytes_le(struct CgCtx *cg, struct CgAddrSlot *as, int offset,
                      uint64 data, uint32 length)
{
    struct CgCodeChunk *chunk = as->chunk;
    unsigned char *code;
    if (!chunk) {
        assert(as->location == 0);
        assert(as->offset_in_chunk == 0);
        assert(cg->code.chunks);
        chunk = cg->code.chunks;
    }
    while (UNLIKELY(as->offset_in_chunk + offset >= CG_CHUNK_SIZE)) {
        offset -= CG_CHUNK_SIZE;
        chunk = chunk->next;
    }
    code = &chunk->code[as->offset_in_chunk + offset];
    while (length--) {
        if (code == cg->code.end) {
            chunk = chunk->next;
            code = chunk->code;
        }
        *(code++) |= data & 0xFFU;
        data >>= 8;
    }
}

/**
 * Adds a relocation.
 *
 * \param cg            Compilation context
 * \param reloc_kind    Type of relocation
 * \param target_sym_id Target function or global data item.
 * \param addend        Addend.
 * \return  1 if successful, 0 on out-of-memory error.
 */
int cg_add_reloc(struct CgCtx *cg, enum CgRelocKind reloc_kind,
                 unsigned target_sym_id, int addend)
{
    struct CgRelocEntry *re = allocp(cg->csbe, sizeof(struct CgRelocEntry));
    ZCHK(re);
    assert(target_sym_id <= CG_SPECIALFUNC(0xF));
    re->next = NULL;
    re->location = cg_get_position(cg);
    re->kind = reloc_kind;
    if (CG_IS_DATA_RELOC(re)) {
        assert(target_sym_id < cg->csbe->max_datadef_count);
    } else if (CG_IS_CODE_RELOC(re)) {
        assert(target_sym_id < cg->csbe->max_funcdef_count ||
                CG_IS_SPECIALFUNC(target_sym_id));
    }
    re->target_id = target_sym_id;
    re->addend = addend;
    *cg->reloc_nextptr = re;
    cg->reloc_nextptr = &re->next;
    return 1;
}

struct SymInfo *cg_helper_func_start(struct CgCtx *cg, const char *symname)
{
    struct CsbeFuncdef *dummy_funcdef;
    struct SymInfo *sym;

    dummy_funcdef = allocp(cg->csbe, sizeof(struct CsbeFuncdef));
    if (UNLIKELY(!dummy_funcdef)) return NULL;
    dummy_funcdef->flags = CSBEFD_INTERN_USED|CSBEFD_INTERN_EXPORTED;
    sym = allocp(cg->csbe, sizeof(struct SymInfo));
    if (UNLIKELY(!sym)) return NULL;
    sym->funcdef = dummy_funcdef;
    sym->datadef = NULL;
    sym->offset = 0;
    sym->size = 0;
    sym->mode = SYMMODE_LOCAL; /* TODO could be OBJGLOBAL if placed into a separate section */
    sym->name = symname;
    sym->namelen = strlen(symname);
    sym->strtab_offs = 0;
    sym->next = NULL;
    *cg->syminfo_nextptr = sym;
    cg->syminfo_nextptr = &sym->next;
    sym->offset = cg_get_position(cg);
    return sym;
}

void cg_helper_func_end(struct CgCtx *cg, struct SymInfo *sym)
{
    sym->size = cg_get_position(cg) - sym->offset;
}

void cg_internal_error(struct CgCtx *cg, const char *sourcefile, int line)
{
    cg->csbe->error = CSBEE_INTERNAL_ERROR;
    fprintf(stderr, "internal codegen error at %s:%d\n", sourcefile, line);
}

int cg_stack_alloc_vars(struct CgCtx *cg, const struct FuncIR *f)
{
    struct CgVar *cv;
    const struct Var *v;
    unsigned remaining;
    unsigned reg_size, tempregs, savedregs, paramregs;

    assert(cg->arch != NULL);

    /* TODO allocate stack space for outbound stack params.
            (this works like a union type, so only space/alignment for
             the largest set of params need to be taken into account)
            this must be allocated just above the stack pointer. */

    /* TODO consider doing reg-alloc per arch? that way,
            it can take reg-sets into account (e.g. float vs int).
            and this code could perhaps be shared for all arch'es. */
    if (f->num_vars) {
        cv = allocpa(cg->csbe, sizeof(struct CgVar), f->num_vars);
        ZCHK(cv);
        cg->vars = cv;
    }

    /* TODO determine the location of incoming stack-allocated paramters */
    /* TODO overlay all single-EBB vars (CSBEVD_INTERN_SINGLE_EBB) */
    /* XXX this does not support arch'es where the stack grows up */
    v = f->vars;
    remaining = f->num_vars;
    reg_size = cg->arch->reg_size;
    tempregs = cg->arch->tempreg_count;
    savedregs = cg->arch->savedreg_count;
    paramregs = cg->arch->paramreg_count;
    while (remaining--) {
        unsigned size, align;
        if ((v->flags & (CSBEVD_INTERN_USED|CSBEVD_INTERN_FUNCPARAM)) != 0) {
            int alloc;
            ZCHK(cg_get_type_size(cg, &v->type, &size, &align));
            cv->size = size;
            cv->align = align;
            assert(v->lane);
            /* TODO there is some duplication with this and assign_reg() in the codegen */
            if ((v->flags & CSBEVD_INTERN_HAS_ADDRESS) != 0) {
                /* A pointer is created. Must copied to stack (TODO) */
                /* TODO ...except for passed-on-stack function parameters */
                alloc = 1;
            } else if ((v->flags & CSBEVD_INTERN_FUNCPARAM) != 0 &&
                       v->lane - CSBE_INTERN_VARLANE_FIRSTPARAM < paramregs) {
                /* Function params that need to go on the stack should
                   already be on the stack. */
                alloc = 0;
            } else if (size > reg_size) {
                /* Too large to fit in a register */
                alloc = 1;
            } else if (v->lane < f->first_callee_saved_varlane &&
                       v->lane > tempregs + savedregs) {
                /* Temporary variable: Completely out of registers */
                alloc = 1;
            } else if (v->lane >= f->first_callee_saved_varlane &&
                       v->lane - MIN(f->first_callee_saved_varlane-1,
                                     tempregs) > savedregs) {
                /* Saved variable: Out of callee-saved registers */
                alloc = 1;
            } else if (v->type.is_struct || v->type.is_array) {
                /* TODO place small structs/array in registers */
                alloc = 1;
            } else {
                /* Scalar with low varlane. It will get a register */
                alloc = 0;
            }

            cv->has_stack_space = alloc;
            if (alloc) {
                cg_align_stacksize(cg, 0, align);
                cv->stackoffs = cg->func_stacksize;
                cg_align_stacksize(cg, size, align);
            }
        }
        cv++;
        v++;
    }

    /* Allocate stack space for saving regs */
    cg->regsave_stackend = cg->func_stacksize;
    {
        /* Callee-saved regs */
        unsigned regs = f->first_addressable_varlane -
                        f->first_callee_saved_varlane;
        unsigned num_single_ebb =
                        f->first_callee_saved_varlane -
                        f->first_non_temporary_varlane;
        if (num_single_ebb > tempregs) {
            /* Single-EBB variables have spilled over
               to the callee-saved regs */
            regs += num_single_ebb - tempregs;
        }
        regs = MIN(regs, savedregs);
        cg->alloced_calleesaved = regs;
        /* Caller saved params */
        if (f->contains_likely_calls) {
            regs += MIN(f->funcdef->num_params, paramregs);
        }
        if (regs) {
            cg_align_stacksize(cg, regs * reg_size, reg_size);
        }
    }
    cg->regsave_stackoffset = cg->func_stacksize;
    return 1;
}

int cg_get_elemsize(const struct CgCtx *cg, const struct CsbeType *type,
                    unsigned *size_out, unsigned *align_out)
{
    unsigned size;
    unsigned align;

    assert(cg->arch != NULL);
    assert(!type->is_unbound);
    if (type->is_struct) {
        const struct CsbeType *fieldtype = type->elemtypes;
        unsigned remaining = type->num_fields;
        unsigned maxsize = 0;
        int is_union = (type->packmode == CSBEPM_UNION);
        align = 1;
        size = 0;
        while (remaining--) {
            unsigned field_size, field_align;
            ZCHK(cg_get_type_size(cg, fieldtype, &field_size, &field_align));
            size = UINT_ALIGN(size, field_align);
            size += field_size;
            if (field_align > align) { align = field_align; }
            if (is_union) {
                if (field_size > maxsize) { maxsize = field_size; }
                size = 0;
            }
            fieldtype++;
        }
        if (is_union) {
            size = maxsize;
        }
    } else {
        enum CsbeTypeKind tk = type->simple_kind;
        assert(tk >= 0 && tk < NUM_TYPEKINDS);
        size = cg->arch->types_sizes[tk];
        align = cg->arch->types_align[tk];
    }
    assert(align >= 1);
    /* C rounds up the size to the alignment even for non-arrays.
       Do the same here, to ensure compatiblity */
    size = UINT_ALIGN(size, align);
    *size_out = size;
    *align_out = align;
    return 1;
}

int cg_get_type_size(const struct CgCtx *cg, const struct CsbeType *type,
                     unsigned *size_out, unsigned *align_out)
{
    unsigned size;
    unsigned align;

    if (!cg_get_elemsize(cg, type, &size, &align)) return 0;
    if (type->is_array) {
        unsigned total_len;
        if (UNLIKELY(csbe_multiply_arraylengths(cg->csbe, &total_len,
                            size, type->array_length) != CSBEE_OK)) {
            return 0;
        }
        size = total_len;
    }
    *size_out = size;
    *align_out = align;
    return 1;
}

int cg_get_elem_offset(const struct CgCtx *cg, const struct CsbeType *type,
                       unsigned index, unsigned fieldnum,
                       unsigned *offset_out)
{
    unsigned offs = 0;
    /* Array element offset */
    if (type->is_array) {
        if (index != 0) {
            unsigned elem_size, elem_align;
            if (!cg_get_elemsize(cg, type, &elem_size, &elem_align)) return 0;
            if (UNLIKELY(csbe_multiply_arraylengths(cg->csbe, &offs,
                            elem_size, index) != CSBEE_OK)) {
               return 0;
            }
        }
    } else {
        assert(index == 0);
    }
    /* Struct field offset */
    if (type->is_struct && type->packmode != CSBEPM_UNION) {
        const struct CsbeType *fieldtype = type->elemtypes;
        unsigned remaining = fieldnum;
        assert(fieldnum < type->num_fields);
        for (;;) {
            unsigned field_size, field_align;
            ZCHK(cg_get_type_size(cg, fieldtype, &field_size, &field_align));
            offs = UINT_ALIGN(offs, field_align);
            if (remaining-- == 0) break;
            offs += field_size;
            fieldtype++;
        }
    } else {
        assert(fieldnum == 0);
    }
    *offset_out = offs;
    return 1;
}