/*

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

 {

 }

 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;

     }

     assert(false);

     return OP_NOP;

 }

 void

 BytecodeCompiler::leaveScope(Scope *scope)

 {

     if (!scope)

         return;

     if (Variable *var = scope->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);

         __ st_sta_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())

         __ ld_sta_p();

     else

         __ ld_sta_i();

 }

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

     if (node->declaration())

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

     Label header, test;

     if (node->condition())

         __ jump(&test);

     __ bind(&header);

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

     if (node->update())

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

     if (node->condition()) {

         __ bind(&test);

         visitForTest(node->condition(), &header, true);

     } else {

         __ jump(&header);

     }

     leaveScope(node->scope());

     assert(__ stackDepth() == 0);

 }

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

 }