// The MIT License (MIT)
//
// Copyright (C) 2026 Scott Duensing
//
// 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.

// compact.c -- Release build bytecode compaction
//
// Walks the module's bytecode, removes OP_LINE instructions (3 bytes
// each), and rewrites all code-address references so control flow
// still lands on the correct instructions.
//
// A module that uses ON ERROR keeps its statement boundaries: every
// OP_LINE is replaced by a 1-byte OP_STMT instead of being dropped, so
// the VM's error dispatcher can still truncate the eval stack to the
// failing statement and RESUME / RESUME NEXT land on statement starts.
//
// Address references:
//   - BasProcEntryT::codeAddr (absolute)
//   - BasFormVarInfoT::initCodeAddr (absolute, negative = no init)
//   - OP_CALL operand (absolute uint16)
//   - OP_JMP / OP_JMP_TRUE / OP_JMP_FALSE operand (relative int16)
//   - OP_FOR_NEXT loopTop operand (relative int16)
//   - OP_ON_ERROR handler operand (relative int16; 0 = disable, not remapped)
//   - GOSUB return address (emitted as OP_PUSH_INT32 followed by OP_JMP,
//     where the pushed value equals the PC immediately after the JMP)
//
// Safety: if any opcode cannot be sized, any jump overflows int16, or
// the walk doesn't reach codeLen exactly, the module is left untouched.

#include "compact.h"
#include "opcodes.h"

#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define BAS_OPERAND_U16_MAX 0xFFFF  // max absolute address encodable in a uint16 operand

// Instruction sizes of the two statement-boundary encodings.
#define BAS_LINE_INST_SIZE  (1 + BAS_OPERAND_U16)   // OP_LINE [uint16 lineNum]
#define BAS_STMT_INST_SIZE  1                       // OP_STMT


// Function prototypes (alphabetical)
int32_t basCompactBytecode(BasModuleT *mod);
static int32_t *buildRemap(const uint8_t *code, int32_t codeLen, bool keepStmt, int32_t *outNewLen);
static int16_t readI16LE(const uint8_t *p);
static uint16_t readU16LE(const uint8_t *p);
static bool remapAbsU16(uint8_t *newCode, int32_t newOpPos, int32_t operandOffset, uint16_t oldAddr, const int32_t *remap, int32_t codeLen, int32_t newCodeLen);
static bool remapRelI16(uint8_t *newCode, int32_t newOpPos, int32_t operandOffset, int16_t oldOffset, int32_t oldPcAfter, int32_t newPcAfter, const int32_t *remap, int32_t codeLen, int32_t newCodeLen, bool allowZero);
static void writeI16LE(uint8_t *p, int16_t v);
static void writeI32LE(uint8_t *p, int32_t v);
static void writeU16LE(uint8_t *p, uint16_t v);

int32_t basCompactBytecode(BasModuleT *mod) {
    if (!mod || !mod->code || mod->codeLen <= 0) {
        return 0;
    }

    const uint8_t *oldCode    = mod->code;
    int32_t        oldCodeLen = mod->codeLen;

    // Count OP_LINE occurrences and note whether the module traps errors.
    // If there are no OP_LINEs, nothing to do.
    int32_t lineCount = 0;
    bool    keepStmt  = false;
    {
        int32_t pc = 0;

        while (pc < oldCodeLen) {
            uint8_t op      = oldCode[pc];
            int32_t operand = basOpcodeOperandSize(op);

            if (operand < 0 || pc + 1 + operand > oldCodeLen) {
                return 0;   // unknown opcode -- skip compaction
            }

            if (op == OP_LINE) {
                lineCount++;
            }

            if (op == OP_ON_ERROR || op == OP_RESUME || op == OP_RESUME_NEXT) {
                keepStmt = true;
            }

            pc += 1 + operand;
        }

        if (pc != oldCodeLen) {
            return 0;
        }

        if (lineCount == 0) {
            return 0;
        }
    }

    int32_t  newCodeLen = 0;
    int32_t *remap      = buildRemap(oldCode, oldCodeLen, keepStmt, &newCodeLen);

    if (!remap) {
        return 0;
    }

    uint8_t *newCode = (uint8_t *)malloc(newCodeLen > 0 ? newCodeLen : 1);

    if (!newCode) {
        free(remap);
        return 0;
    }

    // Copy bytes (dropping or shrinking OP_LINE) and rewrite address operands.
    bool    ok    = true;
    int32_t oldPc = 0;

    while (oldPc < oldCodeLen && ok) {
        uint8_t op       = oldCode[oldPc];
        int32_t operand  = basOpcodeOperandSize(op);
        int32_t instSize = 1 + operand;
        int32_t newPc    = remap[oldPc];

        if (op == OP_LINE) {
            if (keepStmt) {
                newCode[newPc] = OP_STMT;
            }

            oldPc += instSize;
            continue;
        }

        // Copy the instruction verbatim first; we'll overwrite operands that
        // need remapping below.
        memcpy(newCode + newPc, oldCode + oldPc, instSize);

        switch (op) {
            case OP_CALL: {
                uint16_t oldAddr = readU16LE(oldCode + oldPc + 1);

                if (!remapAbsU16(newCode, newPc, 1, oldAddr, remap, oldCodeLen, newCodeLen)) {
                    ok = false;
                }

                break;
            }

            case OP_JMP:
            case OP_JMP_TRUE:
            case OP_JMP_FALSE: {
                int16_t oldOff = readI16LE(oldCode + oldPc + 1);

                if (!remapRelI16(newCode, newPc, 1, oldOff,
                                 oldPc + 3, newPc + 3,
                                 remap, oldCodeLen, newCodeLen, false)) {
                    ok = false;
                }

                break;
            }

            case OP_FOR_NEXT: {
                int16_t oldOff = readI16LE(oldCode + oldPc + 4);

                if (!remapRelI16(newCode, newPc, 4, oldOff,
                                 oldPc + 6, newPc + 6,
                                 remap, oldCodeLen, newCodeLen, false)) {
                    ok = false;
                }

                break;
            }

            case OP_FOR_INIT: {
                // [uint16 varIdx] [uint8 scope] [int16 skipOffset]
                int16_t oldOff = readI16LE(oldCode + oldPc + 4);

                if (!remapRelI16(newCode, newPc, 4, oldOff,
                                 oldPc + 6, newPc + 6,
                                 remap, oldCodeLen, newCodeLen, false)) {
                    ok = false;
                }

                break;
            }

            case OP_ON_ERROR: {
                int16_t oldOff = readI16LE(oldCode + oldPc + 1);

                if (!remapRelI16(newCode, newPc, 1, oldOff,
                                 oldPc + 3, newPc + 3,
                                 remap, oldCodeLen, newCodeLen, true)) {
                    ok = false;
                }

                break;
            }

            case OP_PUSH_INT32: {
                // Detect GOSUB return-address push and remap the absolute address.
                if (basIsGosubPush(oldCode, oldCodeLen, oldPc)) {
                    int32_t oldAddr = basReadI32LE(oldCode + oldPc + 1);

                    if (oldAddr < 0 || oldAddr > oldCodeLen) {
                        ok = false;
                        break;
                    }

                    int32_t newAddr = remap[oldAddr];

                    if (newAddr < 0 || newAddr > newCodeLen) {
                        ok = false;
                        break;
                    }

                    writeI32LE(newCode + newPc + 1, newAddr);
                }

                break;
            }

            default:
                break;
        }

        oldPc += instSize;
    }

    if (!ok) {
        free(newCode);
        free(remap);
        return 0;
    }

    // Rewrite proc entry points
    for (int32_t i = 0; i < mod->procCount; i++) {
        int32_t oldAddr = mod->procs[i].codeAddr;

        if (oldAddr < 0 || oldAddr > oldCodeLen) {
            free(newCode);
            free(remap);
            return 0;
        }

        mod->procs[i].codeAddr = remap[oldAddr];
    }

    // Rewrite form-var init code addresses. Negative means "no init code".
    for (int32_t i = 0; i < mod->formVarInfoCount; i++) {
        int32_t oldAddr = mod->formVarInfo[i].initCodeAddr;

        if (oldAddr < 0) {
            continue;
        }

        if (oldAddr > oldCodeLen) {
            free(newCode);
            free(remap);
            return 0;
        }

        int32_t oldLen = mod->formVarInfo[i].initCodeLen;
        int32_t oldEnd = oldAddr + oldLen;

        if (oldEnd > oldCodeLen) {
            oldEnd = oldCodeLen;
        }

        int32_t newAddr = remap[oldAddr];
        int32_t newEnd  = remap[oldEnd];

        mod->formVarInfo[i].initCodeAddr = newAddr;
        mod->formVarInfo[i].initCodeLen  = newEnd - newAddr;
    }

    // Rewrite entry point
    if (mod->entryPoint >= 0 && mod->entryPoint <= oldCodeLen) {
        mod->entryPoint = remap[mod->entryPoint];
    }

    // Swap in the new code
    free(mod->code);
    mod->code    = newCode;
    mod->codeLen = newCodeLen;

    int32_t removed = oldCodeLen - newCodeLen;

    free(remap);
    return removed;
}


// ============================================================
// Walk bytecode, build remap
// ============================================================
//
// remap[oldPos] = newPos for every byte position in [0, oldCodeLen].
// For OP_LINE bytes: when keepStmt is set the opcode byte maps to the
// 1-byte OP_STMT that replaces it; otherwise (and for the operand bytes)
// remap points at where the NEXT instruction starts in the new code.
// Final entry remap[oldCodeLen] = newCodeLen.
//
// Returns malloc'd array of size (oldCodeLen + 1), or NULL on failure.
static int32_t *buildRemap(const uint8_t *code, int32_t codeLen, bool keepStmt, int32_t *outNewLen) {
    int32_t *remap = (int32_t *)malloc((codeLen + 1) * sizeof(int32_t));

    if (!remap) {
        return NULL;
    }

    int32_t oldPc  = 0;
    int32_t newPc  = 0;

    while (oldPc < codeLen) {
        uint8_t op      = code[oldPc];
        int32_t operand = basOpcodeOperandSize(op);

        if (operand < 0) {
            free(remap);
            return NULL;
        }

        int32_t instSize = 1 + operand;

        if (oldPc + instSize > codeLen) {
            free(remap);
            return NULL;
        }

        if (op == OP_LINE) {
            // Dropped entirely, or shrunk to OP_STMT: the opcode byte maps
            // to the boundary marker (if kept) and the operand bytes to the
            // next instruction.
            remap[oldPc] = newPc;

            if (keepStmt) {
                newPc += BAS_STMT_INST_SIZE;
            }

            for (int32_t i = 1; i < BAS_LINE_INST_SIZE; i++) {
                remap[oldPc + i] = newPc;
            }
        } else {
            for (int32_t i = 0; i < instSize; i++) {
                remap[oldPc + i] = newPc + i;
            }

            newPc += instSize;
        }

        oldPc += instSize;
    }

    if (oldPc != codeLen) {
        free(remap);
        return NULL;
    }

    remap[codeLen] = newPc;
    *outNewLen = newPc;
    return remap;
}


// ============================================================
// Little-endian helpers (bytecode is always LE regardless of host)
// ============================================================
static int16_t readI16LE(const uint8_t *p) {
    return (int16_t)((uint16_t)p[0] | ((uint16_t)p[1] << 8));
}


static uint16_t readU16LE(const uint8_t *p) {
    return (uint16_t)p[0] | ((uint16_t)p[1] << 8);
}


// ============================================================
// Apply remap to a single instruction's operand
// ============================================================
//
// Returns true on success, false if an offset overflows int16 or a
// target doesn't land on a valid instruction in the old code.
static bool remapAbsU16(uint8_t *newCode, int32_t newOpPos, int32_t operandOffset, uint16_t oldAddr, const int32_t *remap, int32_t codeLen, int32_t newCodeLen) {
    // oldAddr is an arbitrary uint16 operand; remap[] only has codeLen+1
    // entries, so a corrupt CALL address would index out of bounds.
    if (oldAddr > codeLen) {
        return false;
    }

    int32_t newAddr = remap[oldAddr];

    if (newAddr < 0 || newAddr > newCodeLen) {
        return false;
    }

    if (newAddr > BAS_OPERAND_U16_MAX) {
        return false;
    }

    writeU16LE(newCode + newOpPos + operandOffset, (uint16_t)newAddr);
    return true;
}


// Rewrite a relative int16 offset.
//   oldOpPos, oldPcAfter: position of opcode and PC after reading offset
//   newOpPos, newPcAfter: same in the new code
//   operandOffset: byte offset from opcode to the int16 operand
//   oldOffset: the offset as stored in the old code
//
// Handles ON_ERROR's special case (offset == 0 means "disable").
// For ON_ERROR the caller passes allowZero=true; the zero is preserved as-is.
static bool remapRelI16(uint8_t *newCode, int32_t newOpPos, int32_t operandOffset, int16_t oldOffset, int32_t oldPcAfter, int32_t newPcAfter, const int32_t *remap, int32_t codeLen, int32_t newCodeLen, bool allowZero) {
    if (allowZero && oldOffset == 0) {
        writeI16LE(newCode + newOpPos + operandOffset, 0);
        return true;
    }

    int32_t oldTarget = oldPcAfter + oldOffset;

    if (oldTarget < 0 || oldTarget > codeLen) {
        return false;
    }

    int32_t newTarget = remap[oldTarget];

    if (newTarget < 0 || newTarget > newCodeLen) {
        return false;
    }

    int32_t newOffset = newTarget - newPcAfter;

    // A nonzero ON ERROR handler offset that collapses to 0 would be misread
    // by the VM as "disable handler" (0 is the disable sentinel).  The
    // oldOffset == 0 disable case was already handled above, so a 0 here means
    // a real handler collapsed -- abort compaction so the original
    // (uncompacted) bytecode with a valid offset is retained.
    if (allowZero && newOffset == 0) {
        return false;
    }

    if (newOffset < INT16_MIN || newOffset > INT16_MAX) {
        return false;
    }

    writeI16LE(newCode + newOpPos + operandOffset, (int16_t)newOffset);
    return true;
}


static void writeI16LE(uint8_t *p, int16_t v) {
    uint16_t u = (uint16_t)v;
    p[0] = (uint8_t)(u & 0xFF);
    p[1] = (uint8_t)((u >> 8) & 0xFF);
}


static void writeI32LE(uint8_t *p, int32_t v) {
    uint32_t u = (uint32_t)v;
    p[0] = (uint8_t)(u & 0xFF);
    p[1] = (uint8_t)((u >> 8) & 0xFF);
    p[2] = (uint8_t)((u >> 16) & 0xFF);
    p[3] = (uint8_t)((u >> 24) & 0xFF);
}


static void writeU16LE(uint8_t *p, uint16_t v) {
    p[0] = (uint8_t)(v & 0xFF);
    p[1] = (uint8_t)((v >> 8) & 0xFF);
}
