//---------------------------------------------------------------------------//
//  1100101 |_ mov #0, r4         __                                         //
//     11   |_ <0xb380 %5c4>     / _|_ _____ ___                             //
//     0110 |_ 3.50 -> 3.60     |  _\ \ / _ (_-<                             //
//          |_ base# + offset   |_| /_\_\___/__/                             //
//---------------------------------------------------------------------------//

#include <fxos/function.h>
#include <fxos/util/format.h>
#include <fxos/util/log.h>

namespace FxOS {

//=== Function ===//

Function::Function(Binary &binary, u32 address):
    BinaryObject(binary, BinaryObject::Function, address, 0)
{
    /* Size is not determined at first. */

    /* Default unambiguous name */
    setName(format("fun.%08x", address));
}

/* Add a basic block to the function. The entry block must be added first. */
BasicBlock &Function::addBasicBlock(BasicBlock &&bb)
{
    m_blocks.push_back(bb);
    return m_blocks.back();
}

/* Update the function's BinaryObject size by finding the last address covered
   by any instruction in the function. */
void Function::updateFunctionSize()
{
    u32 max_address = this->address();

    for(BasicBlock &bb: *this) {
        if(bb.instructionCount() == 0)
            continue;
        Instruction &insn = bb.instructionAtIndex(bb.instructionCount() - 1);
        max_address = std::max(max_address, insn.address() + insn.size());
    }

    this->setSize(max_address - this->address());
}

/* The first step in building function CFGs is delimiting the blocks. Starting
   from the entry point, we generate "superblocks" by reading instructions
   linearly until we find a terminator.

   In general, a superblock will be split into multiple basic blocks, with a
   cut at every target of a jump inside the superblock. We record these as we
   explore, and generate basic blocks at the end. */
struct Superblock
{
    /* Addresses of all instructions in the superblock. */
    std::vector<u32> addresses;
    /* Addresses of all basic block leaders in the superblock. */
    std::set<u32> leaders;

    /* Whether the superblock ends with a dynamic jump */
    bool mustDynamicJump = false;
    /* Whether the superblock may end with a jump to a static target */
    bool mayStaticJump = false;
    /* Whether the superblock may end by a fallthrough */
    bool mayFallthrough = false;
    /* Whether the superblock ends with a return */
    bool mustReturn = false;

    /* If mayStaticJump is set, target address */
    u32 staticTarget = 0xffffffff;
    /* If mayFallthrough is set, fallthrough address */
    u32 fallthroughTarget = 0xffffffff;
};

// TODO: Unclear what the exit status of the superblock is in case of error
static Superblock exploreSuperblock(Function &function, u32 entry)
{
    Superblock sb;
    sb.leaders.insert(entry);

    VirtualSpace &vspace = function.parentBinary().vspace();
    bool inDelaySlot = false;
    bool terminatorFound = false;
    u32 pc = entry;

    while(!terminatorFound || inDelaySlot) {
        sb.addresses.push_back(pc);

        /* Read the next instruction from memory */
        // TODO: Handle 32-bit DSP instructions
        if(!vspace.covers(pc, 2)) {
            FxOS_log(ERR, "superblock %08x exits vspace at %08x", entry, pc);
            break;
        }
        u32 opcodeBits = vspace.read_u16(pc);

        Instruction ins(function, pc, opcodeBits);
        AsmInstruction opcode = ins.opcode();

        if(inDelaySlot && !opcode.isValidDelaySlot()) {
            FxOS_log(ERR, "superblock %08x has invalid delay slot at %08x",
                entry, pc);
            break;
        }

        /* Set exit properties when finding the terminator */
        if(opcode.isBlockTerminator()) {
            sb.mustDynamicJump = opcode.isDynamicJump();
            sb.mayStaticJump = opcode.isAnyStaticJump();
            sb.mayFallthrough = opcode.isConditionalJump();
            sb.mustReturn = opcode.isReturn();

            if(sb.mayStaticJump)
                sb.staticTarget = opcode.getPCRelativeTarget(pc);
        }

        terminatorFound = terminatorFound || opcode.isBlockTerminator();
        inDelaySlot = !inDelaySlot && opcode.hasDelaySlot();
        pc += 2;
    }

    if(sb.mayFallthrough)
        sb.fallthroughTarget = pc;
    return sb;
}

/* Cut a superblock in the list and returns true if one contains provided
   address, otherwise returns false. */
static bool cutSuperblockAt(std::vector<Superblock> &blocks, u32 address)
{
    for(auto &b: blocks) {
        auto const &a = b.addresses;
        if(std::find(a.begin(), a.end(), address) != a.end()) {
            b.leaders.insert(address);
            return true;
        }
    }
    return false;
}

void Function::exploreFunctionAt(u32 functionAddress)
{
    assert(!(functionAddress & 1) && "function starts at unaligned address");

    std::vector<Superblock> blocks;

    std::queue<u32> queue;
    queue.push(functionAddress);

    while(!queue.empty()) {
        u32 entry = queue.front();
        queue.pop();

        /* If this address was found by another superblock that was explored
           while [entry] was in the queue, perform the cut now */
        if(cutSuperblockAt(blocks, entry))
            continue;

        Superblock sb = exploreSuperblock(*this, entry);

        /* Process static jump targets and fallthrough targets to queue new
           superblocks or cut existing ones */
        if(sb.mayFallthrough) {
            if(!cutSuperblockAt(blocks, sb.fallthroughTarget))
                queue.push(sb.fallthroughTarget);
        }
        if(sb.mayStaticJump) {
            if(!cutSuperblockAt(blocks, sb.staticTarget))
                queue.push(sb.staticTarget);
        }

        blocks.push_back(std::move(sb));
    }

    /* Cut superblocks. The loop on b.leaders schedules the construction of new
       BasicBlock objects but the iteration is really the multi-part do loop
       using the iterator on b.addresses. */
    for(auto &b: blocks) {
        auto it = b.addresses.begin();

        for(u32 _: b.leaders) {
            (void)_;
            BasicBlock bb0(*this, *it, *it == functionAddress);
            BasicBlock &bb = addBasicBlock(std::move(bb0));

            do {
                // TODO: Support 32-bit instructions
                u32 opcode = parentBinary().vspace().read_u16(*it);
                Instruction ins(*this, *it, opcode);
                bb.addInstruction(std::move(ins));
                it++;
            }
            while(it != b.addresses.end() && !b.leaders.count(*it));

            bb.finalizeBlock();
        }
    }

    // TODO: Set successors and predecessors
}

//=== BasicBlock ===//

BasicBlock::BasicBlock(Function &function, u32 address, bool isEntryBlock):
    m_function {function}, m_address {address}, m_flags {0}
{
    if(isEntryBlock)
        m_flags |= Flags::IsEntryBlock;
}

uint BasicBlock::blockIndex() const
{
    for(uint i = 0; i < m_function.blockCount(); i++) {
        BasicBlock &bb = m_function.basicBlockByIndex(i);
        if(&bb == this)
            return i;
    }
    assert(false && "blockIndex: block not in its own parent");
}

bool BasicBlock::mayFallthrough() const
{
    Instruction const *ins = terminatorInstruction();
    return !ins || ins->opcode().isConditionalJump();
}

bool BasicBlock::hasStaticTarget() const
{
    Instruction const *ins = terminatorInstruction();
    return ins && ins->opcode().isAnyStaticJump();
}

u32 BasicBlock::staticTarget() const
{
    Instruction const *ins = terminatorInstruction();
    if(!ins || !ins->opcode().isAnyStaticJump())
        return 0xffffffff;
    return ins->opcode().getPCRelativeTarget(ins->address());
}
bool BasicBlock::hasDynamicTarget() const
{
    Instruction const *ins = terminatorInstruction();
    return ins && ins->opcode().isDynamicJump();
}

void BasicBlock::addInstruction(Instruction &&insn)
{
    insn.setBlockContext(this->blockIndex(), m_instructions.size());
    m_instructions.push_back(std::move(insn));
}

void BasicBlock::finalizeBlock()
{
    /* Ensure a bunch of invariants. */

    /* Instruction must be sequential. */
    u32 pc = this->address();
    for(Instruction &insn: *this) {
        assert(insn.address() == pc && "non-sequential instructions in bb");
        pc += insn.size();
    }

    /* The block must have no more than one terminator. */
    Instruction *term = nullptr;
    for(Instruction &insn: *this) {
        bool isReturn = insn.opcode().isBlockTerminator();
        assert(!(term && isReturn) && "bb with multiple terminators");
    }

    /* The block must have a delay slot iff the terminator has one. */
    bool hasDelaySlot = false;
    if(term) {
        hasDelaySlot = term->opcode().hasDelaySlot();
        assert(
            term->indexInBlock() == this->instructionCount() - hasDelaySlot - 1
            && "incorrectly placed bb terminator");
    }

    /* Set structural flags. */

    if(hasDelaySlot)
        m_flags |= Flags::HasDelaySlot;
    if(!term)
        m_flags |= Flags::NoTerminator;
    if(term && term->opcode().isReturn())
        m_flags |= Flags::IsTerminator;

    if(hasDelaySlot) {
        Instruction *DSI = delaySlotInstruction();
        DSI->setFlags(DSI->flags() | Instruction::Flags::InDelaySlot);
    }
}

//=== Instruction ===//

Instruction::Instruction(Function &function, u32 address, u32 opcode):
    m_function {function}, m_address {address}, m_opcode {opcode}
{
    /* Start with no flags; they will be set as needed */
    m_flags = 0;
}

void Instruction::setBlockContext(uint blockIndex, uint insnIndex)
{
    m_blockIndex = blockIndex;
    m_insnIndex = insnIndex;
}

} /* namespace FxOS */