/*

 * 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"

using namespace jc;

typedef BytecodeEmitter::Label Label;

#define __ emitter_.

BytecodeCompiler::BytecodeCompiler(ParseTree *tree)

  : tree_(tree),

    scope_(NULL),

    loop_(NULL)

{

}

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::leaveScope(Scope *scope)

{

    if (!scope)

        return;

    if (Variable *var = scope->firstUnsafeArray())

        __ unwind(var->slot());

}

void

BytecodeCompiler::unwindTo(Scope *scope)

{

    for (Scope *current = scope_; current != scope; current = current->enclosing()) {

        if (Variable *var = current->firstUnsafeArray())

            __ unwind(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());

        } else {

            assert(var->location() == Variable::PARAMETER);

            __ getarg(var->slot());

        }

        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)

{

    if (node->lvalue()->isNameProxy()) {

        NameProxy *proxy = node->lvalue()->toNameProxy();

        Variable *var = proxy->binding()->toVariable();

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

        if (var->location() == Variable::STACK) {

            __ setlocal(var->slot());

        } else {

            assert(var->location() == Variable::PARAMETER);

            __ setarg(var->slot());

        }

    } else if (node->lvalue()->isIndexExpression()) {

        IndexExpression *index = node->lvalue()->toIndexExpression();

        index->left()->accept(this);

        index->right()->accept(this);

        __ bounds_sta();

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

        if (index->left()->type()->contained()->isFloat()) {

            __ setarray_f();

        } else {

            assert(index->left()->type()->contained()->isInt());

            __ setarray_i();

        }

    }

}

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_sta();

    if (node->type()->isArray()) {

        __ getarray_p();

    } else if (node->type()->isFloat()) {

        __ getarray_f();

    } else {

        __ getarray_i();

        assert(node->type()->isInt());

    }

}

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());

    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);

    leaveScope(node->scope());

    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());

    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)

        leaveScope(node->scope());

}

void

BytecodeCompiler::visit(IntegerLiteral *node)

{

    __ int_(node->value());

}

void

BytecodeCompiler::visit(FloatLiteral *node)

{

    __ float_((float)node->value());

}

void

BytecodeCompiler::visit(VariableDeclaration *node)

{

    Variable *var = node->variable();

    assert(var->location() == Variable::STACK);

    Type *type = var->type();

    if (type->isArray()) {

        bool isConstant = true;

        for (Type *child = type; child->isArray(); child = child->contained()) {

            if (child->elements() == Type::UNKNOWN_SIZE) {

                isConstant = false;

                break;

            }

        }

        if (!isConstant) {

            for (size_t i = 0; i < node->postDimensions()->size(); i++) {

                Expression *expr = node->postDimensions()->at(i);

                expr->accept(this);

            }

        }

        __ stackarray(type);

    } else {

        if (node->initialization()) {

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

        } else {

            if (!node->postDimensions())

                __ int_(0);

        }

    }

    __ setlocal(var->slot());

    __ pop();

    locals_[var->slot()] = var->type();

}

void

BytecodeCompiler::visit(CallExpression *node)

{

    NameProxy *fun = node->callee()->toNameProxy();

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

        Expression *arg = node->arguments()->at(i);

        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_);

    Code *code = child.compileFunction(module_, node);

    assert(!module_->getFunction(node->fun()->slot()));

    module_->setFunction(node->fun()->slot(), code);

}

void

BytecodeCompiler::visit(EnumStatement *node)

{

    // :TODO: eventually, we'll export these.

}

Code *

BytecodeCompiler::compileFunction(Module *module, FunctionStatement *node)

{

    AutoEnterScope functionScope(this, node->scope());

    locals_ = new Type *[node->scope()->maxContainedLocals()];

    memset(locals_, 0, sizeof(Type *) * node->scope()->maxContainedLocals());

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

    __ stop();

    if (__ oom())

        return false;

    bytecode *start = (bytecode *)malloc(__ length() * sizeof(bytecode));

    memcpy(start, __ buffer(), __ length());

    unsigned nargs = node->scope()->numArguments();

    Type **args = new Type *[nargs];

    for (size_t i = 0; i < node->scope()->numArguments(); i++)

        args[i] = node->scope()->getArgument(i)->type();

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

        assert(locals_[i]);

    Code *code = new Code(module,

                          start,

                          __ length(),

                          args,

                          nargs,

                          locals_,

                          node->scope()->maxContainedLocals(),

                          __ maxStackDepth());

    return code;

}

Module *

BytecodeCompiler::compile()

{

    Code **functions = new Code *[tree_->scope()->numFunctions()];

    memset(functions, 0, sizeof(Code *) * tree_->scope()->numFunctions());

    module_ = new Module(functions, tree_->scope()->numFunctions());

    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(functions[i]);

#endif

    return module_;

}