#include "moar.h"

#ifdef _WIN32
#define snprintf _snprintf
#endif

static void instrument_graph_with_breakpoints(MVMThreadContext *tc, MVMSpeshGraph *g) {
    MVMSpeshBB *bb = g->entry->linear_next;
    MVMuint16 array_slot = 0;

    MVMint32 last_line_number = -2;
    MVMint32 last_filename = -1;

    char *filename_buf = NULL;

    while (bb) {
        MVMSpeshIns *ins = bb->first_ins;
        MVMSpeshIns *breakpoint_ins;

        MVMBytecodeAnnotation *bbba = MVM_bytecode_resolve_annotation(tc, &g->sf->body, bb->initial_pc);
        MVMint64 line_number = -1;
        MVMint64 filename_string_index = -1;

        MVMuint32 file_bp_idx;

        if (bbba) {
            line_number = bbba->line_number;
            filename_string_index = bbba->filename_string_heap_index;
            MVM_free(bbba);
        } else {
            line_number = -1;
            bb = bb->linear_next;
            continue;
        }

        /* skip PHI instructions, to make sure PHI only occur uninterrupted after start-of-bb */
        while (ins && ins->info->opcode == MVM_SSA_PHI) {
            ins = ins->next;
        }
        if (!ins) ins = bb->last_ins;

        /* Jumplists require the target BB to start in the goto op.
         * We must not break this, or we cause the interpreter to derail */
        if (bb->last_ins->info->opcode == MVM_OP_jumplist) {
            MVMint16 to_skip = bb->num_succ;
            for (; to_skip > 0; to_skip--) {
                bb = bb->linear_next;
            }
            continue;
        }

        if (line_number >= 0) {
            breakpoint_ins = MVM_spesh_alloc(tc, g, sizeof(MVMSpeshIns));
            breakpoint_ins->info        = MVM_op_get_op(MVM_OP_breakpoint);
            breakpoint_ins->operands    = MVM_spesh_alloc(tc, g, 2 * sizeof(MVMSpeshOperand));

            if (last_filename != filename_string_index) {
                if (filename_buf)
                    MVM_free(filename_buf);
                filename_buf = MVM_string_utf8_encode_C_string(tc, MVM_cu_string(tc, g->sf->body.cu, filename_string_index));
            }

            MVM_debugserver_register_line(tc, filename_buf, strlen(filename_buf), line_number, &file_bp_idx);

            breakpoint_ins->operands[0].lit_i32 = file_bp_idx;
            breakpoint_ins->operands[1].lit_i32 = line_number;

            last_filename = filename_string_index;

            MVM_spesh_manipulate_insert_ins(tc, bb, ins->prev, breakpoint_ins);
        }

        /* Now go through instructions to see if any are annotated with a
         * precise filename/lineno as well. */
        while (ins) {
            MVMSpeshAnn *ann = ins->annotations;

            while (ann) {
                if (ann->type == MVM_SPESH_ANN_LINENO) {
                    /* We are very likely to have one instruction here that has
                     * the same annotation as the bb itself. We skip that one.*/
                    if (ann->data.lineno.line_number == line_number && ann->data.lineno.filename_string_index == filename_string_index) {
                        break;
                    }

                    line_number = ann->data.lineno.line_number;
                    filename_string_index = ann->data.lineno.filename_string_index;

                    breakpoint_ins = MVM_spesh_alloc(tc, g, sizeof(MVMSpeshIns));
                    breakpoint_ins->info        = MVM_op_get_op(MVM_OP_breakpoint);
                    breakpoint_ins->operands    = MVM_spesh_alloc(tc, g, 2 * sizeof(MVMSpeshOperand));

                    if (last_filename != filename_string_index) {
                        if (filename_buf)
                            MVM_free(filename_buf);
                        filename_buf = MVM_string_utf8_encode_C_string(tc, MVM_cu_string(tc, g->sf->body.cu, filename_string_index));
                    }

                    MVM_debugserver_register_line(tc, filename_buf, strlen(filename_buf), line_number, &file_bp_idx);

                    breakpoint_ins->operands[0].lit_i32 = file_bp_idx;
                    breakpoint_ins->operands[1].lit_i32 = ann->data.lineno.line_number;

                    /* XXX insert breakpoint op here, too, maybe? */

                    break;
                }

                ann = ann->next;
            }

            ins = ins->next;
        }

        bb = bb->linear_next;
    }

    if (filename_buf)
        MVM_free(filename_buf);
}

static void instrument_graph(MVMThreadContext *tc, MVMSpeshGraph *g) {
    MVMSpeshBB *bb = g->entry->linear_next;
    MVMuint16 array_slot = 0;

    MVMint32 last_line_number = -2;
    MVMint32 last_filename = -1;

    MVMuint16 allocd_slots  = g->num_bbs * 2;
    char *line_report_store = MVM_calloc(allocd_slots, sizeof(char));

    /* Since we don't know the right size for the line report store
     * up front, we will have to realloc it along the way. After that
     * we havee to fix up the arguments to the coverage log instructions */
    MVMuint32 fixup_alloc = g->num_bbs * 2;
    MVMuint32 fixup_elems = 0;
    MVMuint32 fixup_idx; /* for iterating over the fixup array */
    MVMSpeshIns **to_fixup = MVM_malloc(fixup_alloc * sizeof(MVMSpeshIns*));

    while (bb) {
        MVMSpeshIns *ins = bb->first_ins;
        MVMSpeshIns *log_ins;

        MVMBytecodeAnnotation *bbba = MVM_bytecode_resolve_annotation(tc, &g->sf->body, bb->initial_pc);
        MVMint64 line_number;
        MVMint64 filename_string_index;
        if (bbba) {
            line_number = bbba->line_number;
            filename_string_index = bbba->filename_string_heap_index;
            MVM_free(bbba);
        } else {
            line_number = -1;
            bb = bb->linear_next;
            continue;
        }

        /* skip PHI instructions, to make sure PHI only occur uninterrupted after start-of-bb */
        while (ins && ins->info->opcode == MVM_SSA_PHI) {
            ins = ins->next;
        }
        if (!ins) ins = bb->last_ins;

        /* Jumplists require the target BB to start in the goto op.
         * We must not break this, or we cause the interpreter to derail */
        if (bb->last_ins->info->opcode == MVM_OP_jumplist) {
            MVMint16 to_skip = bb->num_succ;
            for (; to_skip > 0; to_skip--) {
                bb = bb->linear_next;
            }
            continue;
        }

        log_ins = MVM_spesh_alloc(tc, g, sizeof(MVMSpeshIns));
        log_ins->info        = MVM_op_get_op(MVM_OP_coverage_log);
        log_ins->operands    = MVM_spesh_alloc(tc, g, 4 * sizeof(MVMSpeshOperand));

        log_ins->operands[0].lit_str_idx = filename_string_index;
        log_ins->operands[1].lit_i32 = line_number;

        if (last_line_number == line_number && last_filename == filename_string_index) {
            /* Consecutive BBs with the same line number and filename should
             * share one "already reported" slot. */
            log_ins->operands[2].lit_i32 = array_slot;
        } else {
            log_ins->operands[2].lit_i32 = array_slot++;
            last_line_number = line_number;
            last_filename = filename_string_index;

            if (array_slot == allocd_slots) {
                allocd_slots *= 2;
                line_report_store = MVM_realloc(line_report_store, sizeof(char) * allocd_slots);
            }
        }

        to_fixup[fixup_elems++] = log_ins;
        if (fixup_elems == fixup_alloc) {
            fixup_alloc *= 2;
            to_fixup = MVM_realloc(to_fixup, sizeof(MVMSpeshIns*) * fixup_alloc);
        }
        MVM_spesh_manipulate_insert_ins(tc, bb, ins, log_ins);

        /* Now go through instructions to see if any are annotated with a
         * precise filename/lineno as well. */
        while (ins) {
            MVMSpeshAnn *ann = ins->annotations;

            while (ann) {
                if (ann->type == MVM_SPESH_ANN_LINENO) {
                    /* We are very likely to have one instruction here that has
                     * the same annotation as the bb itself. We skip that one.*/
                    if (ann->data.lineno.line_number == line_number && ann->data.lineno.filename_string_index == filename_string_index) {
                        break;
                    }

                    log_ins = MVM_spesh_alloc(tc, g, sizeof(MVMSpeshIns));
                    log_ins->info        = MVM_op_get_op(MVM_OP_coverage_log);
                    log_ins->operands    = MVM_spesh_alloc(tc, g, 4 * sizeof(MVMSpeshOperand));

                    log_ins->operands[0].lit_str_idx = ann->data.lineno.filename_string_index;
                    log_ins->operands[1].lit_i32 = ann->data.lineno.line_number;
                    log_ins->operands[2].lit_i32 = array_slot++;

                    if (array_slot == allocd_slots) {
                        allocd_slots *= 2;
                        line_report_store = MVM_realloc(line_report_store, sizeof(char) * allocd_slots);
                    }

                    to_fixup[fixup_elems++] = log_ins;
                    if (fixup_elems == fixup_alloc) {
                        fixup_alloc *= 2;
                        to_fixup = MVM_realloc(to_fixup, sizeof(MVMSpeshIns*) * fixup_alloc);
                    }
                    break;
                }

                ann = ann->next;
            }

            ins = ins->next;
        }

        bb = bb->linear_next;
    }

    line_report_store = MVM_realloc(line_report_store, sizeof(char) * (array_slot + 1));

    for (fixup_idx = 0; fixup_idx < fixup_elems; fixup_idx++) {
        MVMSpeshIns *ins = to_fixup[fixup_idx];

        ins->operands[3].lit_i64 = (MVMint64)line_report_store;
    }

    if (array_slot == 0) {
        MVM_free(line_report_store);
    }
}

/* Adds instrumented version of the unspecialized bytecode. */
static void add_instrumentation(MVMThreadContext *tc, MVMStaticFrame *sf, MVMuint8 want_coverage) {
    MVMSpeshCode  *sc;
    MVMStaticFrameInstrumentation *ins;
    MVMSpeshGraph *sg = MVM_spesh_graph_create(tc, sf, 1, 0);
    if (want_coverage)
        instrument_graph(tc, sg);
    else
        instrument_graph_with_breakpoints(tc, sg);
    sc = MVM_spesh_codegen(tc, sg);
    ins = MVM_calloc(1, sizeof(MVMStaticFrameInstrumentation));
    ins->instrumented_bytecode        = sc->bytecode;
    ins->instrumented_handlers        = sc->handlers;
    ins->instrumented_bytecode_size   = sc->bytecode_size;
    ins->uninstrumented_bytecode      = sf->body.bytecode;
    ins->uninstrumented_handlers      = sf->body.handlers;
    ins->uninstrumented_bytecode_size = sf->body.bytecode_size;
    sf->body.instrumentation = ins;
    MVM_spesh_graph_destroy(tc, sg);
    MVM_free(sc);
}


/* Instruments code with per-line logging of code coverage */
static void line_numbers_instrument(MVMThreadContext *tc, MVMStaticFrame *sf, MVMuint8 want_coverage) {
    if (!sf->body.instrumentation || sf->body.bytecode != sf->body.instrumentation->instrumented_bytecode) {
        /* Handle main, non-specialized, bytecode. */
        if (!sf->body.instrumentation)
            add_instrumentation(tc, sf, want_coverage);
        sf->body.bytecode      = sf->body.instrumentation->instrumented_bytecode;
        sf->body.handlers      = sf->body.instrumentation->instrumented_handlers;
        sf->body.bytecode_size = sf->body.instrumentation->instrumented_bytecode_size;

        /* Throw away any argument guard so we'll never resolve prior
         * specializations again. */
        MVM_spesh_arg_guard_discard(tc, sf);
    }
}

/* Instruments code with per-line logging of code coverage */
void MVM_line_coverage_instrument(MVMThreadContext *tc, MVMStaticFrame *sf) {
    line_numbers_instrument(tc, sf, 1);
}

/* Instruments code with a breakpoint check instruction af every line number change */
void MVM_breakpoint_instrument(MVMThreadContext *tc, MVMStaticFrame *sf) {
    line_numbers_instrument(tc, sf, 0);
}

void MVM_line_coverage_report(MVMThreadContext *tc, MVMString *filename, MVMuint32 line_number, MVMuint16 cache_slot, char *cache) {
    if (tc->instance->coverage_control == 2 || (!tc->instance->coverage_control && cache[cache_slot] == 0)) {
        char *encoded_filename;
        char composed_line[256];
        size_t length;

        cache[cache_slot] = 1;

        encoded_filename = MVM_string_utf8_encode_C_string(tc, filename);
        if ((length = snprintf(composed_line, 255, "HIT  %s  %d\n", encoded_filename, line_number)) > 0) {
            fputs(composed_line, tc->instance->coverage_log_fh);
        }
        MVM_free(encoded_filename);
    }
}