//
//Copyright (C) 2014 LunarG, Inc.
//
//All rights reserved.
//
//Redistribution and use in source and binary forms, with or without
//modification, are permitted provided that the following conditions
//are met:
//
//    Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//
//    Redistributions in binary form must reproduce the above
//    copyright notice, this list of conditions and the following
//    disclaimer in the documentation and/or other materials provided
//    with the distribution.
//
//    Neither the name of 3Dlabs Inc. Ltd. nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
//THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
//"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
//LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTAstreamITY AND FITNESS
//FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
//COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
//INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
//BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
//LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
//CAUSED AND ON ANY THEORY OF LIAstreamITY, WHETHER IN CONTRACT, STRICT
//LIAstreamITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
//ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
//POSSIstreamITY OF SUCH DAMAGE.

//
// Author: John Kessenich, LunarG
//

//
// Disassembler for SPIR-V.
//

#include <unordered_map>
#include <iomanip>
#include <stack>
#include <sstream>

#include "GLSL450Lib.h"
extern const char* GlslStd450DebugNames[GLSL_STD_450::Count];

#include "disassemble.h"
#include "doc.h"

namespace spv {

void Kill(std::ostream& out, const char* message)
{
    out << std::endl << "Disassembly failed: " << message << std::endl;
    exit(1);
}

// Container class for a single instance of a SPIR-V stream, with methods for disassembly.
class SpirvStream {
public:
    SpirvStream(std::ostream& out, const std::vector<unsigned int>& stream) : out(out), stream(stream), word(0), nextNestedControl(0) { }
    virtual ~SpirvStream() { }

    void validate();
    void processInstructions();

protected:
    OpCode getOpCode(int id) const { return idInstruction[id] ? (OpCode)(stream[idInstruction[id]] & OpCodeMask) : OpNop; }

    // Output methods
    void outputIndent();
    void formatId(Id id, std::stringstream&);
    void outputResultId(Id id);
    void outputTypeId(Id id);
    void outputId(Id id);
    void disassembleImmediates(int numOperands);
    void disassembleIds(int numOperands);
    void disassembleString();
    void disassembleInstruction(Id resultId, Id typeId, OpCode opCode, int numOperands);

    // Data
    std::ostream& out;                       // where to write the disassembly
    const std::vector<unsigned int>& stream; // the actual word stream
    int size;                                // the size of the word stream
    int word;                                // the next word of the stream to read

    // map each <id> to the instruction that created it
    Id bound;
    std::vector<unsigned int> idInstruction;  // the word offset into the stream where the instruction for result [id] starts; 0 if not yet seen (forward reference or function parameter)

    std::vector<std::string> idDescriptor;    // the best text string known for explaining the <id>

    // schema
    unsigned int schema;

    // stack of structured-merge points
    std::stack<Id> nestedControl;
    Id nextNestedControl;         // need a slight delay for when we are nested
};

void SpirvStream::validate()
{
    size = (int)stream.size();
    if (size < 4)
        Kill(out, "stream is too short");

    // Magic number
    if (stream[word++] != MagicNumber) {
        out << "Bad magic number";
        return;
    }

    // Version
    out << "// Module Version " << stream[word++] << std::endl;

    // Generator's magic number
    out << "// Generated by (magic number): " << std::setbase(16) << stream[word++] << std::setbase(10) << std::endl;

    // Result <id> bound
    bound = stream[word++];
    idInstruction.resize(bound);
    idDescriptor.resize(bound);
    out << "// Id's are bound by " << bound << std::endl;
    out << std::endl;

    // Reserved schema, must be 0 for now
    schema = stream[word++];
    if (schema != 0)
        Kill(out, "bad schema, must be 0");
}

// Loop over all the instructions, in order, processing each.
// Boiler plate for each is handled here directly, the rest is dispatched.
void SpirvStream::processInstructions()
{
    // Instructions
    while (word < size) {
        int instructionStart = word;

        // Instruction wordCount and opcode
        unsigned int firstWord = stream[word];
        unsigned wordCount = firstWord >> WordCountShift;
        OpCode opCode = (OpCode)(firstWord & OpCodeMask);
        int nextInst = word + wordCount;
        ++word;

        // Presence of full instruction
        if (nextInst > size)
            Kill(out, "stream instruction terminated too early");

        // Base for computing number of operands; will be updated as more is learned
        unsigned numOperands = wordCount - 1;

        // Type <id>
        Id typeId = 0;
        if (InstructionDesc[opCode].hasType()) {
            typeId = stream[word++];
            --numOperands;
        }

        // Result <id>
        Id resultId = 0;
        if (InstructionDesc[opCode].hasResult()) {
            resultId = stream[word++];
            --numOperands;

            // save instruction for future reference
            idInstruction[resultId] = instructionStart;
        }

        outputResultId(resultId);
        outputTypeId(typeId);
        outputIndent();

        // Hand off the OpCode and all its operands
        disassembleInstruction(resultId, typeId, opCode, numOperands);
        if (word != nextInst) {
            out << " ERROR, incorrect number of operands consumed.  At " << word << " instead of " << nextInst << " instruction start was " << instructionStart;
            word = nextInst;
        }
        out << std::endl;
    }
}

void SpirvStream::outputIndent()
{
    for (int i = 0; i < (int)nestedControl.size(); ++i)
        out << "  ";
}

void SpirvStream::formatId(Id id, std::stringstream& idStream)
{
    if (id >= bound)
        Kill(out, "Bad <id>");

    if (id != 0) {
        idStream << id;
        if (idDescriptor[id].size() > 0)
            idStream << "(" << idDescriptor[id] << ")";
    }
}

void SpirvStream::outputResultId(Id id)
{
    const int width = 16;
    std::stringstream idStream;
    formatId(id, idStream);
    out << std::setw(width) << std::right << idStream.str();
    if (id != 0)
        out << ":";
    else
        out << " ";

    if (nestedControl.size() && id == nestedControl.top())
        nestedControl.pop();
}

void SpirvStream::outputTypeId(Id id)
{
    const int width = 12;
    std::stringstream idStream;
    formatId(id, idStream);
    out << std::setw(width) << std::right << idStream.str() << " ";
}

void SpirvStream::outputId(Id id)
{
    if (id >= bound)
        Kill(out, "Bad <id>");

    out << id;
    if (idDescriptor[id].size() > 0)
        out << "(" << idDescriptor[id] << ")";
}

void SpirvStream::disassembleImmediates(int numOperands)
{
    for (int i = 0; i < numOperands; ++i) {
        out << stream[word++];
        if (i < numOperands - 1)
            out << " ";
    }
}

void SpirvStream::disassembleIds(int numOperands)
{
    for (int i = 0; i < numOperands; ++i) {
        outputId(stream[word++]);
        if (i < numOperands - 1)
            out << " ";
    }
}

void SpirvStream::disassembleString()
{
    out << " \"";

    char* wordString;
    bool done = false;
    do {
        unsigned int content = stream[word];
        wordString = (char*)&content;
        for (int charCount = 0; charCount < 4; ++charCount) {
            if (*wordString == 0) {
                done = true;
                break;
            }
            out << *(wordString++);
        }
        ++word;
    } while (! done);

    out << "\"";
}

void SpirvStream::disassembleInstruction(Id resultId, Id typeId, OpCode opCode, int numOperands)
{
    // Process the opcode

    if (opCode < 0 || opCode >= OpCount)
        Kill(out, "Bad opcode");
    else
        out << InstructionDesc[opCode].opName + 2; // Skip the "Op"

    if (opCode == OpLoopMerge || opCode == OpSelectionMerge)
        nextNestedControl = stream[word];
    else if (opCode == OpBranchConditional || opCode == OpSwitch) {
        if (nextNestedControl) {
            nestedControl.push(nextNestedControl);
            nextNestedControl = 0;
        }
    } else if (opCode == OpExtInstImport)
        idDescriptor[resultId] = (char*)(&stream[word]);
    else {
        if (idDescriptor[resultId].size() == 0) {
            switch (opCode) {
            case OpTypeInt:
                idDescriptor[resultId] = "int";
                break;
            case OpTypeFloat:
                idDescriptor[resultId] = "float";
                break;
            case OpTypeBool:
                idDescriptor[resultId] = "bool";
                break;
            case OpTypeStruct:
                idDescriptor[resultId] = "struct";
                break;
            case OpTypePointer:
                idDescriptor[resultId] = "ptr";
                break;
            case OpTypeVector:
                if (idDescriptor[stream[word]].size() > 0)
                    idDescriptor[resultId].append(idDescriptor[stream[word]].begin(), idDescriptor[stream[word]].begin() + 1);
                idDescriptor[resultId].append("vec");
                switch (stream[word + 1]) {
                case 2:   idDescriptor[resultId].append("2");   break;
                case 3:   idDescriptor[resultId].append("3");   break;
                case 4:   idDescriptor[resultId].append("4");   break;
                case 8:   idDescriptor[resultId].append("8");   break;
                case 16:  idDescriptor[resultId].append("16");  break;
                case 32:  idDescriptor[resultId].append("32");  break;
                default: break;
                }
                break;
            default:
                break;
            }
        }
    }

    // Process the operands.  Note, a new context-dependent set could be
    // swapped in mid-traversal.

    // Handle textures specially, so can put out helpful strings.
    if (opCode == OpTypeSampler) {
        disassembleIds(1);
        out << " " << DimensionString((Dimensionality)stream[word++]);
        switch (stream[word++]) {
        case 0: out << " texture";        break;
        case 1: out << " image";          break;
        case 2: out << " filter+texture"; break;
        }
        out << (stream[word++] != 0 ? " array" : "");
        out << (stream[word++] != 0 ? " depth" : "");
        out << (stream[word++] != 0 ? " multi-sampled" : "");
        return;
    }

    // Handle all the parameterized operands
    for (int op = 0; op < InstructionDesc[opCode].operands.getNum(); ++op) {
        out << " ";
        switch (InstructionDesc[opCode].operands.getClass(op)) {
        case OperandId:
            disassembleIds(1);
            // Get names for printing "(XXX)" for readability, *after* this id
            if (opCode == OpName)
                idDescriptor[stream[word - 1]] = (char*)(&stream[word]);
            break;
        case OperandOptionalId:
        case OperandVariableIds:
            disassembleIds(numOperands);
            return;
        case OperandVariableLiterals:
            if (opCode == OpDecorate && stream[word - 1] == DecBuiltIn) {
                out << BuiltInString(stream[word++]);
                --numOperands;
                ++op;
            }
            disassembleImmediates(numOperands);
            return;
        case OperandVariableLiteralId:
            while (numOperands > 0) {
                out << std::endl;
                outputResultId(NoResult);
                outputTypeId(NoType);
                outputIndent();
                out << "     case ";
                disassembleImmediates(1);
                out << ": ";
                disassembleIds(1);
                numOperands -= 2;
            }
            return;
        case OperandLiteralNumber:
            disassembleImmediates(1);
            if (opCode == OpExtInst) {
                unsigned entrypoint = stream[word - 1];
                if (entrypoint < GLSL_STD_450::Count)
                    out << "(" << GlslStd450DebugNames[entrypoint] << ")";
            }
            break;
        case OperandLiteralString:
            disassembleString();
            return;
        case OperandSource:
            out << SourceString((SourceLanguage)stream[word++]);
            break;
        case OperandExecutionModel:
            out << ExecutionModelString((ExecutionModel)stream[word++]);
            break;
        case OperandAddressing:
            out << AddressingString((AddressingModel)stream[word++]);
            break;
        case OperandMemory:
            out << MemoryString((MemoryModel)stream[word++]);
            break;
        case OperandExecutionMode:
            out << ExecutionModeString((ExecutionMode)stream[word++]);
            break;
        case OperandStorage:
            out << StorageClassString((StorageClass)stream[word++]);
            break;
        case OperandDimensionality:
            out << DimensionString((Dimensionality)stream[word++]);
            break;
        case OperandDecoration:
            out << DecorationString((Decoration)stream[word++]);
            break;
        case OperandBuiltIn:
            out << BuiltInString((BuiltIn)stream[word++]);
            break;
        case OperandSelect:
            out << SelectControlString((SelectControl)stream[word++]);
            break;
        case OperandLoop:
            out << LoopControlString((LoopControl)stream[word++]);
            break;
        case OperandFunction:
        {
            unsigned int control = stream[word++];
            if (control == 0)
                out << FunctionControlString(control);
            else {
                for (int m = 0; m < FunctionControlCount; ++m) {
                    if (control & (1 << m))
                        out << FunctionControlString(m);
                }
            }
            break;
        }
        case OperandMemorySemantics:
            for (int shift = 0; shift < MemorySemanticsCount; ++shift) {
                unsigned lit = (stream[word] & (1 << shift));
                if (lit)
                    out << MemorySemanticsString(lit) << " ";
            }
            word++;
            break;
        case OperandMemoryAccess:
            out << MemoryAccessString(stream[word++]);
            break;
        case OperandExecutionScope:
            out << ExecutionScopeString(stream[word++]);
            break;
        case OperandGroupOperation:
            out << GroupOperationString(stream[word++]);
            break;
        case OperandKernelEnqueueFlags:
            out << KernelEnqueueFlagsString(stream[word++]);
            break;
        case OperandKernelProfilingInfo:
            out << KernelProfilingInfoString(stream[word++]);
            break;
        default:
            break;
        }
        --numOperands;
    }

    return;
}

void Disassemble(std::ostream& out, const std::vector<unsigned int>& stream)
{
    SpirvStream SpirvStream(out, stream);
    SpirvStream.validate();
    SpirvStream.processInstructions();
}

}; // end namespace spv