#include "bytecode.h"

// Implementation of the Slang bytecode (BC)
// (most notably including conversion from IR to BC)

#include "compiler.h"
#include "ir.h"
#include "ir-insts.h"
#include "lower-to-ir.h"

namespace Slang
{
struct SharedBytecodeGenerationContext;

// Representation of a `BCPtr<T>` during actual bytecode generation.
// This representation is to deal with the fact that the actual
// storage for the bytecode data might get reallocated during emission
// so that we need to be careful and not work with raw `BCPtr<T>`.
template<typename T>
struct BytecodeGenerationPtr
{
    SharedBytecodeGenerationContext*    sharedContext;
    UInt                                offset;

    BytecodeGenerationPtr()
        : sharedContext(nullptr)
        , offset(0)
    {}


    BytecodeGenerationPtr(
        SharedBytecodeGenerationContext*    sharedContext,
        UInt                                offset)
        : sharedContext(sharedContext)
        , offset(offset)
    {}

    BytecodeGenerationPtr(
        BytecodeGenerationPtr<T> const& ptr)
        : sharedContext(ptr.sharedContext)
        , offset(ptr.offset)
    {}

    template<typename U>
    BytecodeGenerationPtr(
        BytecodeGenerationPtr<U> const& ptr,
        typename EnableIf<IsConvertible<T*, U*>::Value, void>::type * = 0)
        : sharedContext(ptr.sharedContext)
        , offset(ptr.offset)
    {}

    operator BCPtr<T>()
    {
        return BCPtr<T>(getPtr());
    }

    T* operator->()
    {
        return getPtr();
    }

    T& operator*()
    {
        return *getPtr();
    }

    T& operator[](UInt index)
    {
        return getPtr()[index];
    }

    BytecodeGenerationPtr<T> operator+(Int index)
    {
        return BytecodeGenerationPtr<T>(
            sharedContext,
            offset + index*sizeof(T));
    }

    T* getPtr() const;
};

#if 0
template<typename T>
void BCPtr<T>::operator=(BytecodeGenerationPtr<T> const& ptr)
{
    fprintf(stderr, "0x%p: operator=BGP 0x%p\n", this, ptr.getPtr());
    *this = ptr.getPtr();
}
#endif

struct SharedBytecodeGenerationContext
{
    // The final generated bytecode stream
    List<uint8_t>   bytecode;

    // Map from a global symbol to its global ID
    Dictionary<IRInst*, Int> mapGlobalSymbolToGLobalID;
};

struct BytecodeGenerationContext
{
    SharedBytecodeGenerationContext*    shared;

    // The function that is in scope for this context
    IRFunc* currentIRFunc;

    // Counter for global symbols that have been assigned
    // so that they can be used by this function
    List<BCConst>  remappedGlobalSymbols;

    // Map an instruction to its ID for use local
    // to the current context
    Dictionary<IRInst*, Int> mapInstToLocalID;

    // Map an instruction to  the ID for its auxiliary
    // symbol data
    Dictionary<IRInst*, UInt> mapInstToNestedID;
};

template<typename T>
T* BytecodeGenerationPtr<T>::getPtr() const
{
    if(!sharedContext) return nullptr;
    return (T*)(sharedContext->bytecode.Buffer() + offset);
}


BCPtr<void>::RawVal allocateRaw(
    BytecodeGenerationContext*  context,
    size_t                      size,
    size_t                      alignment)
{
    size_t currentOffset = context->shared->bytecode.Count();

    size_t beginOffset = (currentOffset + (alignment-1)) & ~(alignment-1);

    size_t endOffset = beginOffset + size;

    for(size_t ii = currentOffset; ii < endOffset; ++ii)
        context->shared->bytecode.Add(0);

    return beginOffset;
}

template<typename T>
BytecodeGenerationPtr<T> allocate(
    BytecodeGenerationContext*  context)
{
    return BytecodeGenerationPtr<T>(context->shared, allocateRaw(context, sizeof(T), alignof(T)));
}

template<typename T>
BytecodeGenerationPtr<T> allocateArray(
    BytecodeGenerationContext*  context,
    UInt                        count)
{
    return BytecodeGenerationPtr<T>(context->shared, allocateRaw(context, count * sizeof(T), alignof(T)));
}

template<typename T>
BytecodeGenerationPtr<T> getPtr(
    BytecodeGenerationContext*  context)
{
    return BytecodeGenerationPtr<T>(context->shared, context->shared->bytecode.Count());
}


void encodeUInt8(
    BytecodeGenerationContext*  context,
    uint8_t                     value)
{
    context->shared->bytecode.Add(value);
}

void encodeUInt(
    BytecodeGenerationContext*  context,
    UInt                        value)
{
    if( value < 128 )
    {
        encodeUInt8(context, value);
        return;
    }

    uint8_t bytes[16];
    UInt count = 0;

    for(;;)
    {
        UInt index = count++;
        bytes[index] = value & 0x7F;
        value = value >> 7;
        if (!value)
            break;

        bytes[index] |= 0x80;
    }

    UInt index = count;
    while (index--)
    {
        encodeUInt8(context, bytes[index]);
    }
}

void encodeSInt(
    BytecodeGenerationContext*  context,
    Int                         value)
{
    UInt uValue;
    if( value < 0 )
    {
        uValue = (~UInt(value) << 1) | 1;
    }
    else
    {
        uValue = UInt(value) << 1;
    }

    encodeUInt(context, uValue);
}

Int getLocalID(
    BytecodeGenerationContext*  context,
    IRInst*                     inst)
{
    Int localID = 0;
    if( context->mapInstToLocalID.TryGetValue(inst, localID) )
    {
        return localID;
    }

    Int globalID = 0;
    if( context->shared->mapGlobalSymbolToGLobalID.TryGetValue(inst, globalID) )
    {
        BCConst bcConst;
        bcConst.globalID = globalID;

        UInt remappedSymbolIndex = context->remappedGlobalSymbols.Count();
        context->remappedGlobalSymbols.Add(bcConst);

        localID = ~remappedSymbolIndex;
        context->mapInstToLocalID.Add(inst, localID);
        return localID;
    }

    SLANG_UNEXPECTED("no ID for inst");
    return -9999;
}

void encodeOperand(
    BytecodeGenerationContext*  context,
    IRInst*                     operand)
{
    auto id = getLocalID(context, operand);
    encodeSInt(context, id);
}

bool opHasResult(IRInst* inst)
{
    auto type = inst->getType();
    if( !type || type->op != kIROp_VoidType )
    {
        return true;
    }
    return false;
}

void generateBytecodeForInst(
    BytecodeGenerationContext*  context,
    IRInst*                     inst)
{
    // We are generating bytecode for a local instruction
    // inside a function or similar context.
    switch( inst->op )
    {
    default:
        {
            // As a default case, we will assume that bytecode ops
            // and the IR's internal opcodes are the same, and then
            // encode the necessary extra info:
            //

            auto argCount = inst->getArgCount();
            encodeUInt(context, inst->op);
            encodeUInt(context, argCount);
            for( UInt aa = 0; aa < argCount; ++aa )
            {
                encodeOperand(context, inst->getArg(aa));
            }

            auto type = inst->getType();
            if( type && type->op == kIROp_VoidType )
            {
                // This instructions has no type, so don't emit a destination
            }
            else
            {
                // The instruction can be encoded
                // as its own operand for the destination.
                encodeOperand(context, inst);
            }
        }
        break;

    case kIROp_ReturnVoid:
        // Trivial encoding here
        encodeUInt(context, inst->op);
        break;

    case kIROp_IntLit:
        {
            auto ii = (IRConstant*) inst;
            encodeUInt(context, ii->op);
            encodeOperand(context, ii->getType());

            // TODO: probably want distinct encodings
            // for signed vs. unsigned here.
            encodeUInt(context, UInt(ii->u.intVal));

            // destination:
            encodeOperand(context, inst);
        }
        break;

    case kIROp_FloatLit:
        {
            auto ii = (IRConstant*) inst;
            encodeUInt(context, ii->op);
            encodeOperand(context, ii->getType());

            static const UInt size = sizeof(IRFloatingPointValue);
            unsigned char buffer[size];
            memcpy(buffer, &ii->u.floatVal, sizeof(buffer));

            for(UInt ii = 0; ii < size; ++ii)
            {
                encodeUInt8(context, buffer[ii]);
            }

            // destination:
            encodeOperand(context, inst);
        }
        break;

    case kIROp_boolConst:
        {
            auto ii = (IRConstant*) inst;
            encodeUInt(context, ii->op);
            encodeUInt(context, ii->u.intVal ? 1 : 0);

            // destination:
            encodeOperand(context, inst);
        }
        break;

    case kIROp_Func:
        {
            encodeUInt(context, inst->op);

            // We just want to encode the ID for the function
            // symbol data, and then do the rest on the decode side
            UInt nestedID = 0;
            context->mapInstToNestedID.TryGetValue(inst, nestedID);
            encodeUInt(context, nestedID);

            // destination:
            encodeOperand(context, inst);
        }
        break;

    case kIROp_Store:
        {
            encodeUInt(context, inst->op);

            // We need to encode the type being stored, to make
            // our lives easier.
            encodeOperand(context, inst->getArg(2)->getType());
            encodeOperand(context, inst->getArg(1));
            encodeOperand(context, inst->getArg(2));
        }
        break;

    case kIROp_Load:
        {
            encodeUInt(context, inst->op);
            encodeOperand(context, inst->getType());
            encodeOperand(context, inst->getArg(1));
            encodeOperand(context, inst);
        }
        break;
    }
}

Int getIDForGlobalSymbol(
    BytecodeGenerationContext*  context,
    IRInst*                     inst)
{
    Int globalID;
    if(context->shared->mapGlobalSymbolToGLobalID.TryGetValue(inst, globalID))
        return globalID;

    SLANG_UNEXPECTED("no such ID");
}

uint32_t getTypeForGlobalSymbol(
    BytecodeGenerationContext*  context,
    IRInst*                     inst)
{
    auto type = inst->getType();
    if(!type)
        return 0;

    return getIDForGlobalSymbol(context, type);
}

BytecodeGenerationPtr<char> allocateString(
    BytecodeGenerationContext*  context,
    char const*                 data,
    UInt                        size)
{
    BytecodeGenerationPtr<char> ptr = allocateArray<char>(context, size + 1);
    memcpy(ptr.getPtr(), data, size);
    return ptr;
}

BytecodeGenerationPtr<char> allocateString(
    BytecodeGenerationContext*  context,
    String const&               str)
{
    return allocateString(context,
        str.Buffer(),
        str.Length());
}

BytecodeGenerationPtr<char> allocateString(
    BytecodeGenerationContext*  context,
    Name*                       name)
{
    return allocateString(context, name->text);
}

BytecodeGenerationPtr<char> tryGenerateNameForSymbol(
    BytecodeGenerationContext*      context,
    IRInst*                         inst)
{
    // TODO: this is gross, and the IR should probably have
    // a more direct means of querying a name for a symbol.
    if (auto highLevelDeclDecoration = inst->findDecoration<IRHighLevelDeclDecoration>())
    {
        auto decl = highLevelDeclDecoration->decl;
        if (auto reflectionNameMod = decl->FindModifier<ParameterBlockReflectionName>())
        {
            return allocateString(context, reflectionNameMod->name);
        }
        else if(auto name = decl->nameAndLoc.name)
        {
            return allocateString(context, name);
        }
    }

    return BytecodeGenerationPtr<char>();
}

BytecodeGenerationPtr<BCSymbol> generateBytecodeSymbolForInst(
    BytecodeGenerationContext*  context,
    IRInst*                     inst)
{
    switch( inst->op )
    {
    case kIROp_Func:
        {
            auto irFunc = (IRFunc*) inst;
            BytecodeGenerationPtr<BCFunc> bcFunc = allocate<BCFunc>(context);

            bcFunc->op = inst->op;
            bcFunc->typeGlobalID = getTypeForGlobalSymbol(context, inst);

            BytecodeGenerationContext   subContextStorage;
            BytecodeGenerationContext*  subContext = &subContextStorage;
            subContext->shared = context->shared;
            subContext->currentIRFunc = irFunc;

            // First we need to enumerate our basic blocks, so that they
            // can reference one another (basic blocks can forward reference
            // blocks that haven't been seen yet).
            //
            // Note: we allow the IDs of blocks to overlap with ordinary
            // "register" numbers, because there is no case where an operand
            // could be either a block or an ordinary register.
            //
            UInt blockCounter = 0;
            for( auto bb = irFunc->getFirstBlock(); bb; bb = bb->getNextBlock() )
            {
                Int blockID = Int(blockCounter++);
                subContext->mapInstToLocalID.Add(bb, blockID);
            }
            UInt blockCount = blockCounter;

            // Allocate the array of block objects to be stored in the
            // bytecode file.
            auto bcBlocks = allocateArray<BCBlock>(context, blockCount);
            bcFunc->blockCount = blockCount;
            bcFunc->blocks = bcBlocks;

            // Now loop through the blocks again, and allocate the storage
            // for any parameters, variables, or registers used inside
            // each block.
            //
            // We'll count in a first pass, and then fill things in
            // using a second pass
            Int regCounter = 0;
            blockCounter = 0;
            for( auto bb = irFunc->getFirstBlock(); bb; bb = bb->getNextBlock() )
            {
                UInt blockID = blockCounter++;
                UInt paramCount = 0;
                for( auto ii = bb->firstChild; ii; ii = ii->nextInst )
                {
                    switch( ii->op )
                    {
                    default:
                        // Default behavior: if an op has a result,
                        // then it needs a register to store it.
                        if(opHasResult(ii))
                        {
                            regCounter++;
                        }
                        break;

                    case kIROp_Param:
                        // A parameter always uses a register.
                        regCounter++;
                        //
                        // We also want to keep a count of the parameters themselves.
                        paramCount++;
                        break;


                    case kIROp_Var:
                        // A `var` (`alloca`) node needs two registers:
                        // one to hold the actual storage, and another
                        // to hold the pointer.
                        regCounter += 2;
                        break;
                    }
                }

                bcBlocks[blockID].paramCount = paramCount;
            }

            // Okay, we've counted how many registers we need for each block,
            // and now we can allocate the contiguous array we will use.
            UInt regCount = regCounter;
            auto bcRegs = allocateArray<BCReg>(context, regCount);

            bcFunc->regCount = regCount;
            bcFunc->regs = bcRegs;

            // Now we will loop over things again to fill in the information
            // on all these registers.

            regCounter = 0;
            blockCounter = 0;
            for( auto bb = irFunc->getFirstBlock(); bb; bb = bb->getNextBlock() )
            {
                UInt blockID = blockCounter++;

                // Loop over just the parameters first, to ensure they
                // are always the first N registers of a block.
                //
                // This means the parameters of the function itself
                // are always the first N registers in the overall list.
                //
                bcBlocks[blockID].params = bcRegs + regCounter;
                for( auto ii = bb->firstChild; ii; ii = ii->nextInst )
                {
                    if(ii->op != kIROp_Param)
                        continue;

                    Int localID = regCounter++;
                    subContext->mapInstToLocalID.Add(ii, localID);

                    bcRegs[localID].op = ii->op;
                    bcRegs[localID].name = tryGenerateNameForSymbol(context, ii);
                    bcRegs[localID].previousVarIndexPlusOne = localID;
                    bcRegs[localID].typeGlobalID = getTypeForGlobalSymbol(context, ii);
                }

                // Now loop over the non-parameter instructions and
                // allocate actual register locations to them.
                for( auto ii = bb->firstChild; ii; ii = ii->nextInst )
                {
                    switch(ii->op)
                    {
                    case kIROp_Param:
                        // Already handled.
                        break;

                    default:
                        // For an ordinary instruction with a result,
                        // allocate it here.
                        if( opHasResult(ii) )
                        {
                            Int localID = regCounter++;
                            subContext->mapInstToLocalID.Add(ii, localID);

                            bcRegs[localID].op = ii->op;
                            bcRegs[localID].name = tryGenerateNameForSymbol(context, ii);
                            bcRegs[localID].previousVarIndexPlusOne = localID;
                            bcRegs[localID].typeGlobalID = getTypeForGlobalSymbol(context, ii);
                        }
                        break;

                    case kIROp_Var:
                        {
                            // As handled in the earlier loop, we are
                            // allocating *two* locations for each `var`
                            // instruction. The first of these will be
                            // the actual pointer value, while the second
                            // will be the storage for the variable value.
                            Int localID = regCounter;
                            regCounter += 2;

                            subContext->mapInstToLocalID.Add(ii, localID);
                            bcRegs[localID].op = ii->op;
                            bcRegs[localID].name = tryGenerateNameForSymbol(context, ii);
                            bcRegs[localID].previousVarIndexPlusOne = localID;
                            bcRegs[localID].typeGlobalID = getTypeForGlobalSymbol(context, ii);

                            bcRegs[localID+1].op = ii->op;
                            bcRegs[localID+1].previousVarIndexPlusOne = localID+1;
                            bcRegs[localID+1].typeGlobalID = getIDForGlobalSymbol(context,
                                ((IRPtrType*) ii->getType())->getValueType());
                        }
                        break;
                    }
                }
            }

            // Now that we've allocated our blocks and our registers
            // we can go through the actual process of emitting instructions. Hooray!
            blockCounter = 0;
            for( auto bb = irFunc->getFirstBlock(); bb; bb = bb->getNextBlock() )
            {
                UInt blockID = blockCounter++;
                bcBlocks[blockID].code = getPtr<BCOp>(context);


                for( auto ii = bb->firstChild; ii; ii = ii->nextInst )
                {
                    // What we do with each instruction depends a bit on the
                    // kind of instruction it is.
                    switch( ii->op )
                    {
                    default:
                        // For most instructions we just emit their bytecode
                        // ops directly.
                        generateBytecodeForInst(subContext, ii);
                        break;

                    case kIROp_Param:
                        // Don't actually emit code for these, because
                        // there isn't really anything to *execute*
                        //
                        // Note that we *do* allow the `var` nodes
                        // to be executed, just because they need
                        // to set up a register with the pointer value.
                        break;
                    }
                }
            }

            // Finally, after emitting all the instructions we can
            // build a table of global symbols taht need to be
            // imported into the current function as constants.
            UInt constCount = subContext->remappedGlobalSymbols.Count();
            auto bcConsts = allocateArray<BCConst>(context, constCount);

            bcFunc->constCount = constCount;
            bcFunc->consts = bcConsts;

            for( UInt cc = 0; cc < constCount; ++cc )
            {
                bcConsts[cc] = subContext->remappedGlobalSymbols[cc];
            }

            return bcFunc;
        }
        break;

    default:
        // Most instructions don't need a custom representation.
        return BytecodeGenerationPtr<BCSymbol>();
    }
}

BytecodeGenerationPtr<BCModule> generateBytecodeForModule(
    BytecodeGenerationContext*  context,
    IRModule*                   irModule)
{
    // The module will get encoded much like a function,
    // and then that function will be "invoked" to load
    // the module.
    //
    auto bcModule = allocate<BCModule>(context);
    bcModule->op = irModule->op;
    bcModule->typeGlobalID = 0;

    // The logical function that the module representats
    // will only have a single block, so we can allocate it now.
    //
    auto bcBlock = allocate<BCBlock>(context);
    bcBlock->paramCount = 0;
    bcBlock->params = BytecodeGenerationPtr<BCReg>();

    bcModule->blockCount = 1;
    bcModule->blocks = bcBlock;

    // Because the module is the top-most level, there is
    // no need for it to have "constants" that represent
    // values imported from the next outer scope.
    //
    bcModule->constCount = 0;
    bcModule->consts = BytecodeGenerationPtr<BCConst>();

    // We need to compute how many "registers" to allocate
    // for the module, where the registers represent the
    // values being computed at the global scope.
    UInt regCount = 0;
    for( auto inst = irModule->firstChild; inst; inst = inst->nextInst )
    {
        if(!opHasResult(inst))
            continue;

        Int globalID = Int(regCount++);

        context->shared->mapGlobalSymbolToGLobalID.Add(inst, globalID);

        // In the global scope, global IDs are also the local IDs
        context->mapInstToLocalID.Add(inst, globalID);
    }

    auto bcRegs = allocateArray<BCReg>(context, regCount);

    bcModule->regCount = regCount;
    bcModule->regs = bcRegs;

    // Now lets walk through and initialize all those bytecode
    // register representations so that we can use them.
    UInt regCounter= 0;
    for( auto inst = irModule->firstChild; inst; inst = inst->nextInst )
    {
        if(!opHasResult(inst))
            continue;

        UInt regIndex = *context->mapInstToLocalID.TryGetValue(inst);

        BytecodeGenerationPtr<char> name = tryGenerateNameForSymbol(context, inst);

        bcRegs[regIndex].op = inst->op;
        bcRegs[regIndex].name = name;
        bcRegs[regIndex].typeGlobalID = getTypeForGlobalSymbol(context, inst);
        bcRegs[regIndex].previousVarIndexPlusOne = regIndex;
    }

    // Some instructions represent "nested" symbols that will need
    // custom handling, and we will represent those here.
    List<BytecodeGenerationPtr<BCSymbol>> nestedSymbols;
    for( auto inst = irModule->firstChild; inst; inst = inst->nextInst )
    {
        UInt regIndex = *context->mapInstToLocalID.TryGetValue(inst);

        auto bcSymbol = generateBytecodeSymbolForInst(context, inst);
        if (!bcSymbol.getPtr())
            continue;

        UInt nestedSymbolID = nestedSymbols.Count();
        nestedSymbols.Add(bcSymbol);

        context->mapInstToNestedID.Add(inst, nestedSymbolID);

        bcSymbol->name = bcRegs[regIndex].name;
    }

    auto nestedSymbolCount = nestedSymbols.Count();
    auto bcNestedSymbols = allocateArray<BCPtr<BCSymbol>>(context, nestedSymbolCount);

    bcModule->nestedSymbolCount = nestedSymbolCount;
    bcModule->nestedSymbols = bcNestedSymbols;
    for (UInt ii = 0; ii < nestedSymbolCount; ++ii)
    {
        bcNestedSymbols[ii] = nestedSymbols[ii];
    }


    // Finally, we can go through and emit the actual code for
    // the initialization step.
    bcBlock->code = getPtr<BCOp>(context);
    for( auto inst = irModule->firstChild; inst; inst = inst->nextInst )
    {
        // Generate bytecode for global-scope inst
        generateBytecodeForInst(context, inst);
    }
    // Need to encode a terminator here, just to keep the encoding valid
    encodeUInt(context, kIROp_ReturnVoid);

#if 0

    // Now we can go through and generate the bytecode object
    // that will represent each of these global symbols

    List<BytecodeGenerationPtr<BCSymbol>> globalSymbols;

    for( auto inst = irModule->firstChild; inst; inst = inst->nextInst )
    {
        // Generate bytecode for global-scope inst
        auto globalSymbol = generateBytecodeForGlobalSymbol(context, inst);
        globalSymbols.Add(globalSymbol);
    }
#endif

    return bcModule;
}

void generateBytecodeStream(
    BytecodeGenerationContext*  context,
    IRModule*                   irModule)
{
    // Header must be the very first thing in the bytecode stream
    BytecodeGenerationPtr<BCHeader> header = allocate<BCHeader>(context);

    memcpy(header->magic, "slang\0bc", sizeof(header->magic));
    header->version = 0;

    // HACK: ensure that a NULL pointer in an operand field can
    // be encoded.
    context->shared->mapGlobalSymbolToGLobalID.Add(nullptr, -1);

    header->module = generateBytecodeForModule(context, irModule);
}

List<uint8_t> emitSlangIRForEntryPoint(
    EntryPointRequest*  entryPoint)
{
    auto compileRequest = entryPoint->compileRequest;
    auto irModule = lowerEntryPointToIR(
        entryPoint,
        compileRequest->layout.Ptr(),
        // TODO: we need to pick the target more carefully here
        CodeGenTarget::HLSL);

#if 0
    String irAsm = getSlangIRAssembly(irModule);
    fprintf(stderr, "%s\n", irAsm.Buffer());
#endif

    // Now we need to encode that IR into a binary format
    // for transmission/serialization/etc.

    SharedBytecodeGenerationContext sharedContext;

    BytecodeGenerationContext context;
    context.shared = &sharedContext;

    generateBytecodeStream(&context, irModule);

    return sharedContext.bytecode;
}

} // namespace Slang