/*

 * This file implements a data structure representing a sequence of

 * instructions, which is used by different parts of the compilation

 * pipeline.

 */

#include "Python.h"

#include "pycore_c_array.h" // _Py_CArray_EnsureCapacity

#include "pycore_compile.h" // _PyInstruction

#include "pycore_opcode_utils.h"

#include "pycore_opcode_metadata.h" // OPCODE_HAS_ARG, etc

typedef _PyInstruction instruction;

typedef _PyInstructionSequence instr_sequence;

typedef _Py_SourceLocation location;

#define INITIAL_INSTR_SEQUENCE_SIZE 100

#define INITIAL_INSTR_SEQUENCE_LABELS_MAP_SIZE 10

#include "clinic/instruction_sequence.c.h"

#include <stdbool.h>

#undef SUCCESS

#undef ERROR

#define SUCCESS 0

#define ERROR -1

#define RETURN_IF_ERROR(X)  \

    if ((X) == -1) {        \

        return ERROR;       \

    }

static int

instr_sequence_next_inst(instr_sequence *seq) {

    assert(seq->s_instrs != NULL || seq->s_used == 0);

    _Py_c_array_t array = {

        .array = (void*)seq->s_instrs,

        .allocated_entries = seq->s_allocated,

        .item_size = sizeof(instruction),

        .initial_num_entries = INITIAL_INSTR_SEQUENCE_SIZE,

    };

    RETURN_IF_ERROR(_Py_CArray_EnsureCapacity(&array, seq->s_used + 1));

    seq->s_instrs = array.array;

    seq->s_allocated = array.allocated_entries;

    assert(seq->s_allocated >= 0);

    assert(seq->s_used < seq->s_allocated);

    return seq->s_used++;

}

_PyJumpTargetLabel

_PyInstructionSequence_NewLabel(instr_sequence *seq)

{

    _PyJumpTargetLabel lbl = {++seq->s_next_free_label};

    return lbl;

}

int

_PyInstructionSequence_UseLabel(instr_sequence *seq, int lbl)

{

    int old_size = seq->s_labelmap_size;

    _Py_c_array_t array = {

        .array = (void*)seq->s_labelmap,

        .allocated_entries = seq->s_labelmap_size,

        .item_size = sizeof(int),

        .initial_num_entries = INITIAL_INSTR_SEQUENCE_LABELS_MAP_SIZE,

    };

    RETURN_IF_ERROR(_Py_CArray_EnsureCapacity(&array, lbl));

    seq->s_labelmap = array.array;

    seq->s_labelmap_size = array.allocated_entries;

    for(int i = old_size; i < seq->s_labelmap_size; i++) {

        seq->s_labelmap[i] = -111;  /* something weird, for debugging */

    }

    seq->s_labelmap[lbl] = seq->s_used; /* label refers to the next instruction */

    return SUCCESS;

}

int

_PyInstructionSequence_ApplyLabelMap(instr_sequence *instrs)

{

    if (instrs->s_labelmap == NULL) {

        /* Already applied - nothing to do */

        return SUCCESS;

    }

    /* Replace labels by offsets in the code */

    for (int i=0; i < instrs->s_used; i++) {

        instruction *instr = &instrs->s_instrs[i];

        if (HAS_TARGET(instr->i_opcode)) {

            assert(instr->i_oparg < instrs->s_labelmap_size);

            instr->i_oparg = instrs->s_labelmap[instr->i_oparg];

        }

        _PyExceptHandlerInfo *hi = &instr->i_except_handler_info;

        if (hi->h_label >= 0) {

            assert(hi->h_label < instrs->s_labelmap_size);

            hi->h_label = instrs->s_labelmap[hi->h_label];

        }

    }

    /* Clear label map so it's never used again */

    PyMem_Free(instrs->s_labelmap);

    instrs->s_labelmap = NULL;

    instrs->s_labelmap_size = 0;

    return SUCCESS;

}

#define MAX_OPCODE 511

int

_PyInstructionSequence_Addop(instr_sequence *seq, int opcode, int oparg,

                             location loc)

{

    assert(0 <= opcode && opcode <= MAX_OPCODE);

    assert(IS_WITHIN_OPCODE_RANGE(opcode));

    assert(OPCODE_HAS_ARG(opcode) || HAS_TARGET(opcode) || oparg == 0);

    assert(0 <= oparg && oparg < (1 << 30));

    int idx = instr_sequence_next_inst(seq);

    RETURN_IF_ERROR(idx);

    instruction *ci = &seq->s_instrs[idx];

    ci->i_opcode = opcode;

    ci->i_oparg = oparg;

    ci->i_loc = loc;

    return SUCCESS;

}

int

_PyInstructionSequence_InsertInstruction(instr_sequence *seq, int pos,

                                         int opcode, int oparg, location loc)

{

    assert(pos >= 0 && pos <= seq->s_used);

    int last_idx = instr_sequence_next_inst(seq);

    RETURN_IF_ERROR(last_idx);

    for (int i=last_idx-1; i >= pos; i--) {

        seq->s_instrs[i+1] = seq->s_instrs[i];

    }

    instruction *ci = &seq->s_instrs[pos];

    ci->i_opcode = opcode;

    ci->i_oparg = oparg;

    ci->i_loc = loc;

    /* fix the labels map */

    for(int lbl=0; lbl < seq->s_labelmap_size; lbl++) {

        if (seq->s_labelmap[lbl] >= pos) {

            seq->s_labelmap[lbl]++;

        }

    }

    return SUCCESS;

}

int

_PyInstructionSequence_SetAnnotationsCode(instr_sequence *seq,

                                          instr_sequence *annotations)

{

    assert(seq->s_annotations_code == NULL);

    seq->s_annotations_code = annotations;

    return SUCCESS;

}

int

_PyInstructionSequence_AddNested(instr_sequence *seq, instr_sequence *nested)

{

    if (seq->s_nested == NULL) {

        seq->s_nested = PyList_New(0);

        if (seq->s_nested == NULL) {

            return ERROR;

        }

    }

    if (PyList_Append(seq->s_nested, (PyObject*)nested) < 0) {

        return ERROR;

    }

    return SUCCESS;

}

void

PyInstructionSequence_Fini(instr_sequence *seq) {

    Py_XDECREF(seq->s_nested);

    PyMem_Free(seq->s_labelmap);

    seq->s_labelmap = NULL;

    PyMem_Free(seq->s_instrs);

    seq->s_instrs = NULL;

    if (seq->s_annotations_code != NULL) {

        PyInstructionSequence_Fini(seq->s_annotations_code);

        Py_CLEAR(seq->s_annotations_code);

    }

}

/*[clinic input]

class InstructionSequenceType "_PyInstructionSequence *" "&_PyInstructionSequence_Type"

[clinic start generated code]*/

/*[clinic end generated code: output=da39a3ee5e6b4b0d input=589963e07480390f]*/

static _PyInstructionSequence*

inst_seq_create(void)

{

    _PyInstructionSequence *seq;

    seq = PyObject_GC_New(_PyInstructionSequence, &_PyInstructionSequence_Type);

    if (seq == NULL) {

        return NULL;

    }

    seq->s_instrs = NULL;

    seq->s_allocated = 0;

    seq->s_used = 0;

    seq->s_next_free_label = 0;

    seq->s_labelmap = NULL;

    seq->s_labelmap_size = 0;

    seq->s_nested = NULL;

    seq->s_annotations_code = NULL;

    PyObject_GC_Track(seq);

    return seq;

}

PyObject*

_PyInstructionSequence_New(void)

{

    _PyInstructionSequence *seq = inst_seq_create();

    if (seq == NULL) {

        return NULL;

    }

    return (PyObject*)seq;

}

/*[clinic input]

@classmethod

InstructionSequenceType.__new__ as inst_seq_new

Create a new InstructionSequence object.

[clinic start generated code]*/

static PyObject *

inst_seq_new_impl(PyTypeObject *type)

/*[clinic end generated code: output=98881de92c8876f6 input=b393150146849c74]*/

{

    return (PyObject*)inst_seq_create();

}

/*[clinic input]

InstructionSequenceType.use_label

  label: int

Place label at current location.

[clinic start generated code]*/

static PyObject *

InstructionSequenceType_use_label_impl(_PyInstructionSequence *self,

                                       int label)

/*[clinic end generated code: output=4c06bbacb2854755 input=da55f49bb91841f3]*/

{

    if (_PyInstructionSequence_UseLabel(self, label) < 0) {

        return NULL;

    }

    Py_RETURN_NONE;

}

/*[clinic input]

InstructionSequenceType.addop

  opcode: int

  oparg: int

  lineno: int

  col_offset: int

  end_lineno: int

  end_col_offset: int

Append an instruction.

[clinic start generated code]*/

static PyObject *

InstructionSequenceType_addop_impl(_PyInstructionSequence *self, int opcode,

                                   int oparg, int lineno, int col_offset,

                                   int end_lineno, int end_col_offset)

/*[clinic end generated code: output=af0cc22c048dfbf3 input=012762ac88198713]*/

{

    _Py_SourceLocation loc = {lineno, col_offset, end_lineno, end_col_offset};

    if (_PyInstructionSequence_Addop(self, opcode, oparg, loc) < 0) {

        return NULL;

    }

    Py_RETURN_NONE;

}

/*[clinic input]

InstructionSequenceType.new_label -> int

Return a new label.

[clinic start generated code]*/

static int

InstructionSequenceType_new_label_impl(_PyInstructionSequence *self)

/*[clinic end generated code: output=dcb0589e4f5bf4bd input=c66040b9897bc327]*/

{

    _PyJumpTargetLabel lbl = _PyInstructionSequence_NewLabel(self);

    return lbl.id;

}

/*[clinic input]

InstructionSequenceType.add_nested

  nested: object

Add a nested sequence.

[clinic start generated code]*/

static PyObject *

InstructionSequenceType_add_nested_impl(_PyInstructionSequence *self,

                                        PyObject *nested)

/*[clinic end generated code: output=14540fad459f7971 input=f2c482568b3b3c0f]*/

{

    if (!_PyInstructionSequence_Check(nested)) {

        PyErr_Format(PyExc_TypeError,

                     "expected an instruction sequence, not %T",

                     Py_TYPE(nested));

        return NULL;

    }

    if (_PyInstructionSequence_AddNested(self, (_PyInstructionSequence*)nested) < 0) {

        return NULL;

    }

    Py_RETURN_NONE;

}

/*[clinic input]

InstructionSequenceType.get_nested

Add a nested sequence.

[clinic start generated code]*/

static PyObject *

InstructionSequenceType_get_nested_impl(_PyInstructionSequence *self)

/*[clinic end generated code: output=f415112c292630cb input=e429e474c57b95b4]*/

{

    if (self->s_nested == NULL) {

        return PyList_New(0);

    }

    return Py_NewRef(self->s_nested);

}

/*[clinic input]

InstructionSequenceType.get_instructions

Return the instructions as a list of tuples or labels.

[clinic start generated code]*/

static PyObject *

InstructionSequenceType_get_instructions_impl(_PyInstructionSequence *self)

/*[clinic end generated code: output=23f4f3f894c301b3 input=fbadb5dadb611291]*/

{

    if (_PyInstructionSequence_ApplyLabelMap(self) < 0) {

        return NULL;

    }

    PyObject *instructions = PyList_New(0);

    if (instructions == NULL) {

        return NULL;

    }

    for (int i = 0; i < self->s_used; i++) {

        instruction *instr = &self->s_instrs[i];

        location loc = instr->i_loc;

        PyObject *inst_tuple;

        if (OPCODE_HAS_ARG(instr->i_opcode)) {

            inst_tuple = Py_BuildValue(

                "(iiiiii)", instr->i_opcode, instr->i_oparg,

                loc.lineno, loc.end_lineno,

                loc.col_offset, loc.end_col_offset);

        }

        else {

            inst_tuple = Py_BuildValue(

                "(iOiiii)", instr->i_opcode, Py_None,

                loc.lineno, loc.end_lineno,

                loc.col_offset, loc.end_col_offset);

        }

        if (inst_tuple == NULL) {

            goto error;

        }

        int res = PyList_Append(instructions, inst_tuple);

        Py_DECREF(inst_tuple);

        if (res != 0) {

            goto error;

        }

    }

    return instructions;

error:

    Py_XDECREF(instructions);

    return NULL;

}

static PyMethodDef inst_seq_methods[] = {

   INSTRUCTIONSEQUENCETYPE_ADDOP_METHODDEF

   INSTRUCTIONSEQUENCETYPE_NEW_LABEL_METHODDEF

   INSTRUCTIONSEQUENCETYPE_USE_LABEL_METHODDEF

   INSTRUCTIONSEQUENCETYPE_ADD_NESTED_METHODDEF

   INSTRUCTIONSEQUENCETYPE_GET_NESTED_METHODDEF

   INSTRUCTIONSEQUENCETYPE_GET_INSTRUCTIONS_METHODDEF

   {NULL, NULL, 0, NULL},

};

static PyMemberDef inst_seq_memberlist[] = {

    {NULL}      /* Sentinel */

};

static PyGetSetDef inst_seq_getsetters[] = {

    {NULL}      /* Sentinel */

};

static void

inst_seq_dealloc(PyObject *op)

{

    _PyInstructionSequence *seq = (_PyInstructionSequence *)op;

    PyObject_GC_UnTrack(seq);

    PyInstructionSequence_Fini(seq);

    PyObject_GC_Del(seq);

}

static int

inst_seq_traverse(PyObject *op, visitproc visit, void *arg)

{

    _PyInstructionSequence *seq = (_PyInstructionSequence *)op;

    Py_VISIT(seq->s_nested);

    Py_VISIT((PyObject *)seq->s_annotations_code);

    return 0;

}

static int

inst_seq_clear(PyObject *op)

{

    _PyInstructionSequence *seq = (_PyInstructionSequence *)op;

    Py_CLEAR(seq->s_nested);

    Py_CLEAR(seq->s_annotations_code);

    return 0;

}

PyTypeObject _PyInstructionSequence_Type = {

    PyVarObject_HEAD_INIT(&PyType_Type, 0)

    "InstructionSequence",

    sizeof(_PyInstructionSequence),

    0,

    inst_seq_dealloc,   /*tp_dealloc*/

    0,                  /*tp_vectorcall_offset*/

    0,                  /*tp_getattr*/

    0,                  /*tp_setattr*/

    0,                  /*tp_as_async*/

    0,                  /*tp_repr*/

    0,                  /*tp_as_number*/

    0,                  /*tp_as_sequence*/

    0,                  /*tp_as_mapping*/

    0,                  /* tp_hash */

    0,                  /* tp_call */

    0,                  /* tp_str */

    PyObject_GenericGetAttr,  /* tp_getattro */

    0,                  /* tp_setattro */

    0,                  /* tp_as_buffer */

    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */

    inst_seq_new__doc__,                    /* tp_doc */

    inst_seq_traverse,                      /* tp_traverse */

    inst_seq_clear,                         /* tp_clear */

    0,                                      /* tp_richcompare */

    0,                                      /* tp_weaklistoffset */

    0,                                      /* tp_iter */

    0,                                      /* tp_iternext */

    inst_seq_methods,                       /* tp_methods */

    inst_seq_memberlist,                    /* tp_members */

    inst_seq_getsetters,                    /* tp_getset */

    0,                                      /* tp_base */

    0,                                      /* tp_dict */

    0,                                      /* tp_descr_get */

    0,                                      /* tp_descr_set */

    0,                                      /* tp_dictoffset */

    0,                                      /* tp_init */

    0,                                      /* tp_alloc */

    inst_seq_new,                           /* tp_new */

};