/*
 * Copyright (C) 2012 David Anderson
 *
 * This file is part of JITCraft.
 *
 * JITCraft is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.
 * 
 * Foobar is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * JITCraft. If not, see http://www.gnu.org/licenses/.
 */
#include "BytecodeCompiler.h"
#include "Scopes.h"
#include "Modules.h"
#include "Types.h"
#include "Array.h"
#include "Zone.h"
#include "Heap-inl.h"

using namespace jc;

typedef BytecodeEmitter::Label Label;

#define __ emitter_.

BytecodeCompiler::BytecodeCompiler(ParseTree *tree, Module *module)
  : tree_(tree),
    scope_(NULL),
    loop_(NULL),
    module_(module)
{
}

static inline TokenKind
ToCompare(Expression *expr)
{
    if (expr->isBinaryExpression()) {
        BinaryExpression *b = expr->toBinaryExpression();
        if (b->token() >= TOK_EQUALS && b->token() <= TOK_GE)
            return b->token();
    }
    return TOK_EOF;
}

static inline Opcode
CompareTokenToJump(TokenKind kind)
{
    switch (kind) {
      case TOK_EQUALS:
        return OP_JEQ;
      case TOK_NOTEQUALS:
        return OP_JNE;
      case TOK_GT:
        return OP_JG;
      case TOK_GE:
        return OP_JGE;
      case TOK_LT:
        return OP_JL;
      case TOK_LE:
        return OP_JLE;
    }
    assert(false);
    return OP_NOP;
}

static inline Opcode
TokenToOpcode(TokenKind kind)
{
    switch (kind) {
      case TOK_PLUS:
        return OP_ADD;
      case TOK_MINUS:
        return OP_SUB;
      case TOK_STAR:
        return OP_MUL;
      case TOK_SLASH:
        return OP_DIV;
      case TOK_SHR:
        return OP_SHR;
      case TOK_USHR:
        return OP_USHR;
      case TOK_SHL:
        return OP_SHL;
      case TOK_BITAND:
        return OP_BITAND;
      case TOK_BITOR:
        return OP_BITOR;
      case TOK_BITXOR:
        return OP_BITXOR;
    }
    assert(false);
    return OP_NOP;
}

void
BytecodeCompiler::recordLifoSlot(Scope *scope)
{
    if (!scope)
        return;

    if (Variable *var = scope->lifoSlot()) {
        locals_->set(ZONE(), var->slot(), var->type());
        __ savelifo(var->slot());
    }
}

void
BytecodeCompiler::unwindTo(Scope *scope)
{
    // Note: we only need to unwind to the outermost scope's LIFO slot.
    Variable *var = NULL;
    for (Scope *current = scope_; current != scope; current = current->enclosing()) {
        if (Variable *slot = current->lifoSlot())
            var = slot;
    }

    if (var)
        __ unwindlifo(var->slot());
}

void
BytecodeCompiler::visit(NameProxy *name)
{
    switch (name->binding()->kind()) {
      case BoundName::VARIABLE:
      {
        Variable *var = name->binding()->toVariable();

        if (var->location() == Variable::STACK) {
            __ getlocal(var->slot());
            assert(!name->type()->isReference());
        } else {
            assert(var->location() == Variable::PARAMETER);
            __ getarg(var->slot());

            // If this argument is actually a reference, dereference it now.
            if (var->type()->isReference()) {
                if (var->type()->contained()->isFloat()) {
                    __ getref_f();
                } else {
                    assert(var->type()->contained()->isInt32());
                    __ getref_i();
                }
            }
        }
        break;
      }

      default:
      {
        NamedConstant *constant = name->binding()->toNamedConstant();
        __ int_(constant->value().toInt());
        break;
      }
    }
}

void
BytecodeCompiler::visit(ExpressionStatement *node)
{
    // Expressions leave their value on the top of the stack, so we make sure
    // to balance the stack if they're used as a statement.
    node->expression()->accept(this);
    __ pop();
}

void
BytecodeCompiler::visit(Assignment *node)
{
    // First compute the left-hand side.
    if (node->lvalue()->isNameProxy()) {
        NameProxy *proxy = node->lvalue()->toNameProxy();
        Variable *var = proxy->binding()->toVariable();

        // Make sure to bind the left-hand side first.
        if (var->type()->isReference() || var->type()->isArray())
            __ get(var);

        node->expression()->accept(this);
    } else if (node->lvalue()->isIndexExpression()) {
        IndexExpression *index = node->lvalue()->toIndexExpression();

        // Note: Eventually, watch out for a whole here. An expression like:
        //
        //   (a[4] = (5, a = []))
        //
        // Could shrink array, so we need to either bind a[i] to a Reference,
        // or keep the bounds check near the set. This is just a semantic
        // decision.
        index->left()->accept(this);
        index->right()->accept(this);
        node->expression()->accept(this);
    }

    // Emit the store.
    if (node->expression()->type()->isArray()) {
        // Unlike references, which can only be computed from a NamePRoxy,
        // arrays can be computed from multiple lvalue types.
        __ arraycopy();
    } else if (node->lvalue()->isNameProxy()) {
        NameProxy *proxy = node->lvalue()->toNameProxy();
        Variable *var = proxy->binding()->toVariable();

        if (var->type()->isReference()) {
            if (var->type()->contained()->isFloat()) {
                __ setref_f();
            } else {
                assert(var->type()->contained()->isInt32());
                __ setref_i();
            }
        } else {
            __ set(var);
        }
    } else {
        IndexExpression *index = node->lvalue()->toIndexExpression();
        __ bounds();
        __ setarray(index->left()->type()->contained()->primitive());
    }
}

void
BytecodeCompiler::visitForALU(Type *needed, Expression *node)
{
    node->accept(this);
    if (needed->isFloat() && !node->type()->isFloat())
        __ cvt_i2f();
}

void
BytecodeCompiler::visit(BinaryExpression *node)
{
    visitForALU(node->type(), node->left());
    visitForALU(node->type(), node->right());

    Opcode op = TokenToOpcode(node->token());
    if (node->type()->isFloat())
        op = (Opcode)(int(op) + (OP_ADD_F - OP_ADD));
    __ binary(op);
}

void
BytecodeCompiler::visit(IndexExpression *node)
{
    node->left()->accept(this);
    node->right()->accept(this);
    __ bounds();

    if (node->type()->isArray())
        __ getarray_o();
    else
        __ getarray(node->type()->primitive());
}

void
BytecodeCompiler::visit(ReturnStatement *node)
{
    if (node->expression())
        node->expression()->accept(this);
    else
        __ int_(0);

    __ return_();
}

void
BytecodeCompiler::visitForTest(Expression *expression, Label *label, bool onTrue)
{
    TokenKind token = ToCompare(expression);
    if (token == TOK_EOF) {
        expression->accept(this);
        if (onTrue)
            __ j(OP_JT, label);
        else
            __ j(OP_JF, label);
        return;
    }

    Opcode op = CompareTokenToJump(token);
    if (!onTrue)
        op = Invert(op);

    assert(expression->isBinaryExpression());
    BinaryExpression *binary = expression->toBinaryExpression();
    binary->left()->accept(this);
    binary->right()->accept(this);
    __ j(op, label);
}

void
BytecodeCompiler::visit(ForStatement *node)
{
    AutoEnterScope forScope(this, node->scope());

    recordLifoSlot(node->scope());

    Label join, update;
    LoopScope loop(scope_, &join, &update, &loop_);

    if (node->declaration())
        node->declaration()->accept(this);

    Label header, test;
    if (node->condition())
        __ jump(&test);

    __ bind(&header);
    node->body()->accept(this);

    __ bind(&update);
    if (node->update())
        node->update()->accept(this);

    if (node->condition()) {
        __ bind(&test);
        visitForTest(node->condition(), &header, true);
    } else {
        __ jump(&header);
    }

    __ bind(&join);

    if (node->scope())
        unwindTo(node->scope()->enclosing());
    assert(__ stackDepth() == 0);
}

void
BytecodeCompiler::visit(WhileStatement *node)
{
    Label header, join;

    __ bind(&header);

    LoopScope loop(scope_, &join, &header, &loop_);

    if (node->token() == TOK_WHILE) {
        visitForTest(node->condition(), &join, false);
        node->body()->accept(this);
        __ jump(&header);
    } else {
        node->body()->accept(this);
        visitForTest(node->condition(), &header, true);
    }

    __ bind(&join);
}

void
BytecodeCompiler::visit(BreakStatement *node)
{
    unwindTo(loop_->scope());
    __ jump(loop_->break_());
}

void
BytecodeCompiler::visit(ContinueStatement *node)
{
    unwindTo(loop_->scope());
    __ jump(loop_->continue_());
}

void
BytecodeCompiler::visit(BlockStatement *node)
{
    AutoEnterScope blockScope(this, node->scope());

    recordLifoSlot(node->scope());

    for (size_t i = 0; i < node->statements()->length(); i++) {
        Statement *statement = node->statements()->at(i);
        statement->accept(this);
        assert(__ stackDepth() == 0);
    }

    // We don't bother emitting an "unwind" instruction for function scopes,
    // since this is taken care of by the calling convention.
    if (node->scope() && node->scope()->enclosing()->kind() != Scope::FUNCTION)
        unwindTo(node->scope()->enclosing());
}

void
BytecodeCompiler::visit(IntegerLiteral *node)
{
    __ int_(node->value());
}

void
BytecodeCompiler::visit(FloatLiteral *node)
{
    __ float_((float)node->value());
}

static inline ArrayMap *
BuildArrayMap(Zone *zone, Type *type)
{
    // :TODO: use local
    ArrayMap *contained = NULL;
    if (type->levels() > 1) {
        contained = BuildArrayMap(zone, type->contained());
        if (!contained)
            return NULL;
    }

    return ArrayMap::New(zone, type, contained);
}

void
BytecodeCompiler::visit(StringLiteral *node)
{
    // :TODO: use Local
    ArrayMap *map = BuildArrayMap(ZONE(), node->type());
    if (!map)
        return;

    Array *array = Array::New(ZONE(), map, node->string()->length(), Heap::Tenure_Old);
    if (!array)
        return;

    assert(array->nelements() >= node->string()->length());
    memcpy(array->chars(), node->string()->raw(), node->string()->length() * sizeof(char));

    __ object(array);
}

void
BytecodeCompiler::visit(VariableDeclaration *node)
{
    Variable *var = node->variable();
    assert(var->location() == Variable::STACK);
    
    Type *type = var->type();
    if (type->isArray()) {
        // :TODO: use Local
        ArrayMap *map = BuildArrayMap(ZONE(), type);
        __ newarray_lifo(map);

        if (node->initialization()) {
            node->initialization()->accept(this);
            __ arraycopy();
        }
    } else {
        if (node->initialization()) {
            node->initialization()->accept(this);
        } else {
            if (!node->postDimensions())
                __ int_(0);
        }
    }
    __ setlocal(var->slot());
    __ pop();

    locals_->set(ZONE(), var->slot(), var->type());
}

void
BytecodeCompiler::visit(CallExpression *node)
{
    NameProxy *fun = node->callee()->toNameProxy();

    // :TODO: use Local
    Type *funType = node->callee()->type();

    for (size_t i = 0; i < node->arguments()->size(); i++) {
        Expression *arg = node->arguments()->at(i);

        Type *given = arg->type();
        Type *actual = Type::cast(funType->parameters()->at(i));
        if (actual->isReference()) {
            if (given->isReference()) {
                // The input type is already a reference, so we're just passing
                // it on to the callee.
                arg->accept(this);
            } else if (arg->isNameProxy()) {
                // :TODO: test with functions and globals
                // Compute a reference to the name.
                Variable *var = arg->toNameProxy()->binding()->toVariable();
                __ newref(var);
            } else if (arg->isIndexExpression()) {
                IndexExpression *expr = arg->toIndexExpression();
                expr->left()->accept(this);
                expr->right()->accept(this);
                __ bounds();
                __ arrayref_lifo();
            } else {
                // We don't know how to, or can't, build a direct reference to
                // this type. Make an indirect one, by taking the top of the
                // stack and boxing it up.
                arg->accept(this);
                __ slotref_lifo();
            }
        } else {
            arg->accept(this);
        }
    }

    __ callglobal(node->arguments()->size(), 0);
}

void
BytecodeCompiler::visit(IfStatement *node)
{
    if (!node->ifFalse()) {
        Label fail;
        visitForTest(node->condition(), &fail, false);
        node->ifTrue()->accept(this);
        __ bind(&fail);
    } else {
        Label fail, join;
        visitForTest(node->condition(), &fail, false);
        node->ifTrue()->accept(this);
        __ jump(&join);
        __ bind(&fail);
        node->ifFalse()->accept(this);
        __ bind(&join);
    }
}

void
BytecodeCompiler::visit(FunctionStatement *node)
{
    BytecodeCompiler child(tree_, module_);
    Code *code = child.compileFunction(module_, node);
    
    assert(!module_->getFunction(node->fun()->slot()));
    module_->setFunction(ZONE(), node->fun()->slot(), code);
}

void
BytecodeCompiler::visit(EnumStatement *node)
{
    // :TODO: eventually, we'll export these.
}

Code *
BytecodeCompiler::compileFunction(Module *module, FunctionStatement *node)
{
    // Note: No need to record a life slot here, since it's part of the Frame.
    AutoEnterScope functionScope(this, node->scope());

    // :TODO: use locals
    locals_ = FixedArray::New(ZONE(), node->scope()->maxContainedLocals(), Heap::Tenure_Old);

    node->body()->accept(this);

    __ stop();

    if (__ oom())
        return NULL;

    ByteArray *bytes = ByteArray::New(ZONE(), __ length() * sizeof(bytecode), Heap::Tenure_Old);
    if (!bytes)
        return NULL;
    memcpy(bytes->raw(), __ buffer(), __ length());

    unsigned nargs = node->scope()->numArguments();
    FixedArray *args = FixedArray::New(ZONE(), nargs, Heap::Tenure_Old);
    if (!args)
        return NULL;

    for (size_t i = 0; i < node->scope()->numArguments(); i++)
        args->set(ZONE(), i, node->scope()->getArgument(i)->type());

    for (size_t i = 0; i < node->scope()->maxContainedLocals(); i++)
        assert(locals_->at(i));

    FixedArray *objects = FixedArray::New(ZONE(), __ objects()->size(), Heap::Tenure_Old);
    if (!objects)
        return NULL;

    for (size_t i = 0; i < __ objects()->size(); i++)
        objects->set(ZONE(), i, __ objects()->at(i));

    Code *code = Code::New(ZONE(), bytes, locals_, args, objects, __ maxStackDepth(), module_);
    return code;
}

Module *
BytecodeCompiler::compile()
{
    // :TODO: root
    module_ = Module::New(ZONE(), tree_->scope()->numFunctions());
    if (!module_)
        return NULL;

    for (size_t i = 0; i < tree_->statements()->length(); i++) {
        Statement *statement = tree_->statements()->at(i);
        statement->accept(this);
    }

#ifndef NDEBUG
    for (size_t i = 0; i < tree_->scope()->numFunctions(); i++)
        assert(module_->getFunction(i));
#endif

    return module_;
}