/* vim: set ts=4 sw=4 tw=99 et:
 *
 * Copyright (C) 2012 David Anderson
 *
 * This file is part of SourcePawn.
 *
 * SourcePawn 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.
 * 
 * SourcePawn 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
 * SourcePawn. If not, see http://www.gnu.org/licenses/.
 */
#include <assert.h>
#include "../Zone.h"
#include "../Strings.h"
#include "../NativeTable.h"
#include "../TypeManager.h"
#include "../Modules.h"
#include "../Code.h"
#include "SemanticAnalysis.h"
#include "../Heap-inl.h"

using namespace ke;

#define __ emitter_.

typedef BytecodeEmitter::Label Label;

class AutoEnterScope
{
  public:
    AutoEnterScope(Scope **prevp, Scope *scope)
      : prevp_(prevp),
        prev_(*prevp),
        scope_(scope)
    {
        if (!scope_)
            return;
        prev_ = *prevp_;
        *prevp_ = scope;
        assert(prev_ == scope->enclosing());
    }
    ~AutoEnterScope()
    {
        if (scope_)
            *prevp_ = prev_;
    }
  private:
    Scope **prevp_;
    Scope *prev_;
    Scope *scope_;
};

SemanticAnalysis::SemanticAnalysis(Zone *zone, CompileContext &cc, TranslationUnit *unit)
  : zone_(zone),
    pool_(zone->pool()),
    cc_(cc),
    unit_(unit),
    hir_(NULL),
    emitter_(pool_),
    module_(NULL),
    loop_(NULL)
{
}

SemanticAnalysis::SemanticAnalysis(SemanticAnalysis *parent, Handle<FunctionType> fun)
  : zone_(parent->zone_),
    pool_(parent->pool_),
    cc_(parent->cc_),
    unit_(parent->unit_),
    hir_(NULL),
    scope_(parent->scope_),
    emitter_(pool_),
    module_(parent->module_),
    fun_(fun),
    loop_(NULL)
{

}

void
SemanticAnalysis::analyze()
{
    scope_ = unit_->globalScope();

    module_ = Module::New(zone_);

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

    Local<Struct> globals(zone_, unit_->globalScope()->build(zone_));
    if (!globals)
        return;
    module_->setGlobals(globals);

    Local<Code> code(zone_, __ finish());
    if (!code)
        return;
    code->setModule(module_);
    module_->setCode(code);

#ifndef NDEBUG
//    Disassemble(code, stdout);
#endif

    unit_->module_ = module_;
}

void
SemanticAnalysis::visitFunctionStatement(FunctionStatement *stmt)
{
    FunctionSymbol *sym = stmt->sym();
    Local<FunctionType> type(zone_, sym->type());

    assert(scope_->kind() == Scope::GLOBAL);

    if (stmt->token() == TOK_NATIVE) {
        // Define the native now.
        size_t index;
        if (!zone_->nativeTable()->add(stmt->name(), &index))
            return;

        Local<Native> native(zone_, Native::New(zone_, type, index));
        if (!native)
            return;

        __ object(native);
        __ setglobal(sym);
        __ pop();
    }

    // If the function is a forward or native, we don't have to do anything.
    if (stmt->token() == TOK_NATIVE || stmt->token() == TOK_FORWARD)
        return;

    SemanticAnalysis child(this, type);
    Local<Function> funobj(zone_, child.compile(stmt));
    if (!funobj)
        return;

    __ object(funobj);
    __ setglobal(sym);
    if (stmt->token() == TOK_PUBLIC)
        __ public_();
    else
        __ pop();
}

Function *
SemanticAnalysis::compile(FunctionStatement *stmt)
{
    __ setFirstLine(stmt->pos().line);

    AutoEnterScope argScope(&scope_, stmt->argScope());
    AutoEnterScope varScope(&scope_, stmt->varScope());

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

    Local<Code> code(zone_, __ finish());
    if (!code)
        return NULL;
    code->setModule(module_);

    Local<Function> fun(zone_, Function::New(zone_, code, fun_, stmt->name(), stmt->pos()));
    if (!fun)
        return NULL;

#ifndef NDEBUG
//    Disassemble(code, stdout);
#endif

    return fun;
}

void 
SemanticAnalysis::visitBlockStatement(BlockStatement *node)
{
    AutoEnterScope scope(&scope_, node->scope());

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

void
SemanticAnalysis::visitExpressionStatement(ExpressionStatement *node)
{
    // While building HIR, we enter new pool and root scopes, since HIR
    // creation is purely temporary.
    PoolScope poolScope(zone_->pool());
    RootScope rootScope(zone_);

    HIR *hir = eval(node->expression());
    if (!hir)
        return;

    discharge(hir);
    __ pop();
}

void
SemanticAnalysis::visitIfStatement(IfStatement *node)
{
    Label join;

    if (!node->ifFalse()) {
        Label success;

        visitForTest(node->condition(), &success, &join, &success);
        __ bind(&success);
        node->ifTrue()->accept(this);
    } else {
        Label success, next;

        visitForTest(node->condition(), &success, &next, &success);
        __ bind(&success);
        node->ifTrue()->accept(this);
        __ jump(&join);
        __ bind(&next);
        node->ifFalse()->accept(this);
    }

    __ bind(&join);
}

void
SemanticAnalysis::visitForStatement(ForStatement *node)
{
    AutoEnterScope scope(&scope_, node->scope());
    PoolScope poolScope(zone_->pool());
    RootScope rootScope(zone_);

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

    if (node->initialization())
        node->initialization()->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, &join, &join);
    } else {
        __ jump(&header);
    }

    __ bind(&join);

    assert(__ stackDepth() == 0);
}

void
SemanticAnalysis::visitWhileStatement(WhileStatement *node)
{
    PoolScope poolScope(zone_->pool());
    RootScope rootScope(zone_);

    Label header, join;

    __ bind(&header);

    LoopScope loop(scope_, &join, &header, &loop_);
    if (node->token() == TOK_WHILE) {
        Label body;
        visitForTest(node->condition(), &body, &join, &body);
        node->body()->accept(this);
        __ jump(&header);
    } else {
        node->body()->accept(this);
        visitForTest(node->condition(), &header, &join, &join);
    }

    __ bind(&join);

    assert(__ stackDepth() == 0);
}

static inline Heap::Tenure
InferTenure(Symbol *sym)
{
    if (sym->scope()->kind() == Scope::GLOBAL)
        return Heap::Tenure_Old;
    // :TODO: need to use LIFO!
    return Heap::Tenure_Default;
}

void
SemanticAnalysis::visitVariableDeclaration(VariableDeclaration *node)
{
    PoolScope poolScope(zone_->pool());
    RootScope rootScope(zone_);

    if (node->initialization()) {
        HIR *hir = NULL;

        if ((hir = eval(node->initialization())) == NULL)
            return;
        if ((hir = coerce(hir, node->sym()->type(), Coerce_Assign)) == NULL)
            return;

        discharge(hir);
    } else {
        if (node->sym()->type()->isInt32() ||
            node->sym()->type()->isEnum() ||
            node->sym()->type()->isChar() ||
            node->sym()->type()->isBool())
        {
            __ int_(0);
        } else if (node->sym()->type()->isFloat()) {
            __ float_(0);
        } else if (node->sym()->type()->isArray()) {
            Local<ArrayType> type(zone_, ArrayType::cast(node->sym()->type()));
            Local<ArrayMap> map(zone_, ArrayMap::Attach(zone_, type));
            if (!map)
                return;

            if (type->isFixedLength()) {
                __ newarray(OP_NEWFIXED, map, InferTenure(node->sym()));
            } else if (node->dims()) {
                assert(false);
                __ newarray(OP_NEWSIZED, map, InferTenure(node->sym()));
            } else {
                __ newarray(OP_NEWEMPTY, map, InferTenure(node->sym()));
            }
        }
    }

    if (!__ allocate(node->sym()))
        return;

    set(node->sym());
    __ pop();
}

void
SemanticAnalysis::visitReturnStatement(ReturnStatement *node)
{
    if (!node->expression()) {
        if (!fun_->returnType()->isVoid()) {
            cc_.reportError(node->pos(), Message_UsedVoidReturn);
            return;
        }

        __ returnvoid();
        return;
    }

    PoolScope poolScope(zone_->pool());
    RootScope rootScope(zone_);

    HIR *hir = eval(node->expression());
    if (!hir)
        return;

    Local<Type> returnType(zone_, fun_->returnType());
    if ((hir = coerce(hir, returnType, Coerce_Return)) == NULL)
        return;

    discharge(hir);
    __ return_();
}

void
SemanticAnalysis::visitBinaryExpression(BinaryExpression *node)
{
    HIR *left = eval(node->left());
    HIR *right = eval(node->right());
    if (!left || !right)
        return;

    hir_ = binary(node->token(), left, right);
}

bool
SemanticAnalysis::checkArgumentCount(Handle<FunctionType> type, unsigned actual)
{
    if (type->parameters()->length() == actual)
        return true;

    if (actual > type->parameters()->length()) {
        if (!type->isNative() || !type->isNativeVariadic())
            return false;
        return true;
    }

    if (!type->defaults())
        return false;

    // Missing arguments. See if there's a default argument for each missing
    // actual.
    for (unsigned i = actual; i < type->parameters()->length(); i++) {
        if (!type->defaults()->at(i))
            return false;
    }

    return true;
}

void
SemanticAnalysis::visitNameProxy(NameProxy *proxy)
{
    Symbol *sym = proxy->sym();
    VariableSymbol *var = sym->asVariable();

    Scope *in = sym->scope();
    switch (in->kind()) {
      case Scope::GLOBAL:
        hir_ = new (pool_) HGlobal(sym);
        return;

      default:
        assert(in->kind() == Scope::BLOCK);
        assert(var->storage() == VariableSymbol::Local);
        hir_ = new (pool_) HLocal(var);
        return;
    }
}

void
SemanticAnalysis::visitIntegerLiteral(IntegerLiteral *lit)
{
    Local<Type> type(zone_, zone_->types()->getPrimitive(PrimitiveType_Int32));
    hir_ = new (pool_) HInteger(type, lit->value());
}

void
SemanticAnalysis::visitFloatLiteral(FloatLiteral *node)
{
    Local<Type> type(zone_, zone_->types()->getPrimitive(PrimitiveType_Float));
    hir_ = new (pool_) HFloat(type, (float)node->value());
}

void
SemanticAnalysis::visitCallExpression(CallExpression *node)
{
    HIR *callee = eval(node->callee());
    if (!callee)
        return;

    if (!callee->type()->isFunction()) {
        cc_.reportError(node->pos(), Message_CalleeNotFunction);
        return;
    }

    Local<FunctionType> fun(zone_, FunctionType::cast(callee->type()));
    
    if (!checkArgumentCount(fun, node->arguments()->length())) {
        cc_.reportError(node->pos(), Message_ArgumentCountMismatch);
        return;
    }

    if (fun->isForward()) {
        cc_.reportError(node->pos(), Message_ForwardNotImplemented, fun->name()->chars());
        return;
    }

    HIRList *args = new (pool_) HIRList;
    for (unsigned i = 0; i < node->arguments()->length(); i++) {
        Expression *expression = node->arguments()->at(i);
        HIR *arg = eval(expression);
        if (!arg)
            return;

        Handle<Type> given = arg->type();
        
        bool mustComputeReference;
        if (i >= fun->parameters()->length()) {
            // Special handling for natives. Arguments past the formal must
            // always have references computed. However, arrays are already
            // references for natives.
            assert(fun->isNative() && fun->isNativeVariadic());
            mustComputeReference = !given->isReference() &&
                                   !given->isArray();
        } else {
            Local<Type> actual(zone_, fun->parameterAt(i));

            if (expression->isIndexExpression()) {
                // Temporarily disabled.
                assert(!actual->isArray());
                if (actual->isReference()) {
                    cc_.reportError(expression->pos(), Message_CannotComputeIndexRef);
                    return;
                }
            }

            // When we start overloading, coercion will have to happen later.
            arg = coerce(arg, actual, Coerce_Arg);
            if (!arg)
                return;

            mustComputeReference = actual->isReference() && !given->isReference();
        }

        // :TODO: handle mustComputeReference
        assert(!mustComputeReference);

        if (!args->append(arg))
            return;
    }

    hir_ = new (pool_) HCall(fun->returnType(), callee, args);
}

void
SemanticAnalysis::visitAssignment(Assignment *node)
{
    HLValue *lval = lvalue(node->lvalue());
    if (!lval)
        return;

    HIR *hir = eval(node->expression());
    if ((hir = coerce(hir, lval->type(), Coerce_Assign)) == NULL)
        return;

    hir_ = new (pool_) HStore(node->token(), lval, hir);
}

void
SemanticAnalysis::visitBreakStatement(BreakStatement *node)
{
    if (!loop_) {
        cc_.reportError(node->pos(), Message_BreakOutsideLoop);
        return;
    }

    // Note, if we re-introduce LIFO allocs, we need to unwind here.
    __ jump(loop_->break_());
}

void
SemanticAnalysis::visitContinueStatement(ContinueStatement *node)
{
    if (!loop_) {
        cc_.reportError(node->pos(), Message_ContinueOutsideLoop);
        return;
    }

    // Note, if we re-introduce LIFO allocs, we need to unwind here.
    __ jump(loop_->continue_());
}

void
SemanticAnalysis::visitIndexExpression(IndexExpression *node)
{
    HIR *left = eval(node->left());
    if (!left)
        return;

    HIR *right = eval(node->right());
    if (!right)
        return;

    // Hrm... yes, for now, let's coerce here.
    Local<Type> i32(zone_, zone_->types()->getPrimitive(PrimitiveType_Int32));
    if ((right = coerce(right, i32, Coerce_Arg)) == NULL)
        return;
    
    // Do not coerce here, there's no reason to yet.
    if (!left->type()->isArray()) {
        cc_.reportError(node->left()->pos(), Message_IndexBaseMustBeArray);
        return;
    }

    hir_ = new (pool_) HIndex(ArrayType::cast(left->type())->contained(),
                              left,
                              right);
}

void
SemanticAnalysis::visitIncDecExpression(IncDecExpression *node)
{
    HLValue *lval = lvalue(node->expression());
    if (!lval)
        return;

    if (!lval->type()->isInt32() && !lval->type()->isFloat()) {
        assert(false);
        return;
    }

    HIR *rval;
    if (lval->type()->isInt32())
        rval = new (pool_) HInteger(lval->type(), 1);
    else
        rval = new (pool_) HFloat(lval->type(), 1.0);

    if (node->postfix())
        hir_ = new (pool_) HPostIncDec(node->token(), lval, rval);
    else
        hir_ = new (pool_) HStore(node->token(), lval, rval);
}

Type *
SemanticAnalysis::coercionType(TokenKind token, HIR *left, HIR *right)
{
    Handle<Type> l = left->type();
    Handle<Type> r = right->type();

    if (token >= TOK_PLUS && token <= TOK_SLASH) {
        // We allow: [float|int32] + [float|int32]
        if (l->isFloat() && r->isInt32())
            return l;
        if (l->isInt32() && r->isFloat())
            return r;
        if (l->isInt32() && r->isInt32())
            return l;
    }

    assert(l == r);

    return left->type();
}

HIR *
SemanticAnalysis::coerce(HIR *hir, Handle<Type> to, CoercionKind kind)
{
    Handle<Type> from = hir->type();

    if (from->isInt32() && to->isFloat())
        return new (pool_) HConvert(to, OP_CVT_I2F, hir);
    if (from->isInt32() && to->isBool())
        return new (pool_) HConvert(to, OP_CVT_I2B, hir);
    
    assert(from == to);
    
    return hir;
}

void
SemanticAnalysis::set(VariableSymbol *sym)
{
    if (sym->scope()->kind() == Scope::GLOBAL) {
        __ setglobal(sym);
    } else {
        assert(sym->storage() == VariableSymbol::Local);
        __ setlocal(sym);
    }
}

HIR *
SemanticAnalysis::binary(TokenKind token, HIR *left, HIR *right)
{
    Local<Type> coercion(zone_, coercionType(token, left, right));
    if (!coercion)
        return NULL;
    if ((left = coerce(left, coercion, Coerce_Arg)) == NULL)
        return NULL;
    if ((right = coerce(right, coercion, Coerce_Arg)) == NULL)
        return NULL;
    
    switch (token) {
     case TOK_PLUS:
      case TOK_MINUS:
      case TOK_STAR:
      case TOK_SLASH:
      case TOK_PERCENT:
      case TOK_AMPERSAND:
      case TOK_BITOR:
      case TOK_BITXOR:
      case TOK_SHR:
      case TOK_USHR:
      case TOK_SHL:
        return new (pool_) HBinary(left->type(), token, left, right);

      case TOK_EQUALS:
      case TOK_NOTEQUALS:
      case TOK_LT:
      case TOK_LE:
      case TOK_GT:
      case TOK_GE:
      case TOK_AND:
      case TOK_OR:
      {
        Local<Type> result(zone_, zone_->types()->getPrimitive(PrimitiveType_Bool));
        return new (pool_) HBinary(result, token, left, right);
      }

      default:
        assert(false);
        return NULL;
    }
}

static inline TokenKind
ToCompare(Expression *expr)
{
    if (expr->isBinaryExpression()) {
        BinaryExpression *b = expr->toBinaryExpression();
        if (b->token() >= TOK_EQUALS && b->token() <= TOK_OR)
            return b->token();
    } else if (expr->isUnaryExpression()) {
        UnaryExpression *u = expr->toUnaryExpression();
        if (u->token() == TOK_NOT)
            return TOK_NOT;
    }
    return TOK_EOF;
}

void
SemanticAnalysis::visitForTest(Expression *expr,
                               Label *trueBranch,
                               Label *falseBranch,
                               Label *fallthrough)
{
    Local<Type> boolType(zone_, zone_->types()->getPrimitive(PrimitiveType_Bool));

    TokenKind token = ToCompare(expr);
    if (token == TOK_EOF || token == TOK_NOT) {
        HIR *hir;
        if (token == TOK_EOF)
            hir = eval(expr);
        else
            hir = eval(expr->toUnaryExpression()->expression());
        if (!hir)
            return;
        if ((hir = coerce(hir, boolType, Coerce_Assign)) == NULL)
            return;

        discharge(hir);

        bool jumpOnTrue = (fallthrough == falseBranch) ^ (token == TOK_NOT);
        Opcode op = jumpOnTrue ? OP_JT : OP_JF;
        
        if (fallthrough == falseBranch)
            __ j(op, trueBranch);
        else
            __ j(op, falseBranch);
        return;
    }

    BinaryExpression *bin = expr->toBinaryExpression();
    if (token == TOK_AND || token == TOK_OR) {
        assert(false);
        //visitForLogicTest(expr, trueBranch, falseBranch, fallthrough);
        return;
    }

    HIR *left = eval(bin->left());
    HIR *right = eval(bin->right());
    if (!left || !right)
        return;

    HIR *hir = binary(bin->token(), left, right);
    if (!hir)
        return;
    if ((hir = coerce(hir, boolType, Coerce_Assign)) == NULL)
        return;

    discharge(hir);
    if (fallthrough == trueBranch)
        __ j(OP_JF, falseBranch);
    else
        __ j(OP_JT, trueBranch);
}

HLValue *
SemanticAnalysis::lvalue(Expression *expr)
{
    HIR *hir = eval(expr);
    if (!hir)
        return NULL;
    
    if (!hir->isLValue()) {
        cc_.reportError(expr->pos(), Message_ExpectedLValue);
        return NULL;
    }

    return static_cast<HLValue *>(hir);
}

HIR *
SemanticAnalysis::eval(Expression *expr)
{
    assert(!hir_);
    expr->accept(this);

    HIR *hir = hir_;
    hir_ = NULL;

    return hir;
}

void
SemanticAnalysis::discharge(HIR *hir)
{
    EmitHIR(zone_, emitter_, hir);
}