/* vim: set ts=4 sw=4 tw=99 et:
 *
 * 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 "SemanticAnalysis.h"
#include "Types.h"
#include "Zone.h"
#include "TypeManager.h"
#include "Scopes.h"
#include "Heap-inl.h"

using namespace jc;

static inline bool
IsPrimitiveType(Type *type, PrimitiveType primitive)
{
    if (!type)
        return false;

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

    return primitive == type->primitive();
}

static inline bool
EvaluateBinaryOperation(TokenKind kind, const BoxedPrimitive &left, const BoxedPrimitive &right,
                        BoxedPrimitive *out)
{
    switch (kind) {
      case TOK_PLUS:
        *out = BoxedPrimitive::Int(left.toInt() + right.toInt());
        return true;

      case TOK_SHL:
        *out = BoxedPrimitive::Int(left.toInt() << right.toInt());
        return true;

      default:
	return false;
    }
}

static bool
EvaluateForConstant(Expression *expression, BoxedPrimitive *out)
{
    switch (expression->kind()) {
      case AstNode::kBinaryExpression:
      {
        BoxedPrimitive left;
        BoxedPrimitive right;
        BinaryExpression *b = expression->toBinaryExpression();
        if (!EvaluateForConstant(b->left(), &left) ||
            !EvaluateForConstant(b->right(), &right))
        {
            return false;
        }
        return EvaluateBinaryOperation(b->token(), left, right, out);
      }

      case AstNode::kNameProxy:
      {
        NameProxy *proxy = expression->toNameProxy();
        if (proxy->binding()->kind() != BoundName::CONSTANT)
            return false;

        NamedConstant *c = proxy->binding()->toNamedConstant();
        if (!IsPrimitiveType(c->type(), PrimitiveType_Int32))
            return false;

        *out = c->value();
        return true;
      }

      case AstNode::kIntegerLiteral:
        *out = BoxedPrimitive::Int(expression->toIntegerLiteral()->value());
        return true;

      default:
        return false;
    }
}

SemanticAnalysis::SemanticAnalysis(Zone *zone, CompileContext &cc, ParseTree *tree)
  : zone_(zone),
    cc_(cc),
    tree_(tree),
    loopDepth_(0),
    scope_(NULL)
{
}

bool
SemanticAnalysis::analyze()
{
    if (cc_.failures())
        return false;

    for (size_t i = 0; i < tree_->statements()->length(); i++) {
        Statement *statement = tree_->statements()->at(i);
        if (statement->isFunctionStatement()) {
            FunctionStatement *funstmt = statement->toFunctionStatement();
            if (!preanalyze(funstmt))
                return false;
        }
    }

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

    return !cc_.failures();
}

Type *
SemanticAnalysis::visitAndApplyType(Expression *expression)
{
    expression->accept(this);
    return expression->type();
}

Type *
SemanticAnalysis::visitForLoad(Expression *expression)
{
    // :TODO: use Local
    Type *type = visitAndApplyType(expression);
    if (!type)
        return NULL;

    if (type->isReference()) {
        expression->setType(type->contained());
    } else if (expression->isIndexExpression()) {
        if (type->isArray())
            expression->setType(type->contained());
    }

    return expression->type();
}

Type *
SemanticAnalysis::visitForRef(Expression *expression)
{
    expression->accept(this);
    return expression->type();
}

static inline const char *
GetTypeName(Type *type)
{
    if (type->isArray())
        return "array";
    if (type->isFunction())
        return "function";
    if (type->isReference())
        type = type->contained();
    if (type->tag())
        return type->tag();
    return "int";
}

void
SemanticAnalysis::unaryTypeError(AstNode *node, TokenKind token, Type *type)
{
    const char *op = TokenNames[token];
    const char *name = GetTypeName(type);

    cc_.reportError(node->pos(), Message_UnaryTypeError, op, name);
}

void
SemanticAnalysis::binaryTypeError(AstNode *node, TokenKind token, Type *left, Type *right)
{
    const char *op = TokenNames[token];
    const char *leftName = GetTypeName(left);
    const char *rightName = GetTypeName(right);

    cc_.reportError(node->pos(), Message_BinaryTypeError, op, leftName, rightName);
}

Type *
SemanticAnalysis::typeCheckBinary(AstNode *node, TokenKind token, Type *left, Type *right)
{
    if ((token >= TOK_PLUS && token <= TOK_SLASH) ||
        (token >= TOK_ASSIGN_ADD && token <= TOK_ASSIGN_DIV))
    {
        // If either side is a float, the other side must be a float or int.
        if (left->isFloat() || right->isFloat()) {
            Type *floatType = left->isFloat() ? left : right;
            Type *other = left->isFloat() ? right : left;
            if (other->isFloat() || (other->isInt32() && !other->tag()))
                return floatType;
        }
    }

    if ((!left->isInt32() || left->tag()) ||
        (!right->isInt32() || right->tag()))
    {
        binaryTypeError(node, token, left, right);
        return NULL;
    }

    return left;
}

bool
SemanticAnalysis::canTest(AstNode *node, Type *type)
{
    if (type->isArray()) {
        cc_.reportError(node->pos(), Message_ArrayIsNotBoolean);
        return false;
    }
    return true;
}

Type *
SemanticAnalysis::visitForLValue(Expression *expression)
{
    if (expression->isIndexExpression()) {
        Type *type = visitAndApplyType(expression);
        if (!type)
            return NULL;

        return type->contained();
    }

    if (!expression->isNameProxy()) {
        cc_.reportError(expression->pos(), Message_ExpectedLValue);
        return NULL;
    }

    NameProxy *proxy = expression->toNameProxy();
    if (!proxy->binding()->kind() == BoundName::VARIABLE) {
        cc_.reportError(expression->pos(), Message_ExpectedLValue);
        return NULL;
    }

    Variable *var = proxy->binding()->toVariable();
    if (!var->type())
        return NULL;

    if  (var->type()->isReference())
        return var->type()->contained();
    return var->type();
}

void
SemanticAnalysis::visit(NameProxy *name)
{
    name->setType(name->binding()->type());
}

void
SemanticAnalysis::visit(Assignment *node)
{
    Type *exprType = visitForLoad(node->expression());
    if (!exprType)
        return ;

    // :TODO: use Local
    Type *type = visitForLValue(node->lvalue());
    if (!type)
        return;

    if (node->token() != TOK_ASSIGN) {
        if (!typeCheckBinary(node, node->token(), type, exprType))
            return;
    }

    if (!typeCheckAssignment(node, exprType, type, true))
        return;

    node->setType(type);
}

void
SemanticAnalysis::visit(IncDecExpression *node)
{
    // :TODO: use Local
    Type *type = visitForLValue(node->expression());
    if (!type)
        return;

    assert(!type->isReference());

    if ((!type->isInt32() || (type->isInt32() && type->tag())) && !type->isFloat()) {
        unaryTypeError(node, node->token(), type);
        return;
    }

    node->setType(type);
}

void
SemanticAnalysis::visit(BinaryExpression *node)
{
    // :TODO: use Local
    Type *left = visitForLoad(node->left());
    Type *right = visitForLoad(node->right());
    if (!left || !right)
        return;

    Type *result = typeCheckBinary(node, node->token(), left, right);
    if (!result)
        return;

    node->setType(result);
}

void
SemanticAnalysis::visit(ReturnStatement *node)
{
    Type *type = visitForLoad(node->expression());
    if (!type)
        return;

    if (!typeCheckAssignment(node, type, function_->returnType(), true))
        return;
}

void
SemanticAnalysis::visit(ForStatement *node)
{
    loopDepth_++;
    if (node->scope())
        scope_ = node->scope();

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

    if (node->condition()) {
        Type *type = visitForLoad(node->condition());
        if (!type || !canTest(node->condition(), type)) {
            loopDepth_--;
            return;
        }
    }

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

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

    if (node->scope())
        scope_ = scope_->enclosing();
    loopDepth_--;
}

void 
SemanticAnalysis::visit(WhileStatement *node)
{
    loopDepth_++;

    Type *type = visitForLoad(node->condition());
    if (!type || !canTest(node->condition(), type)) {
        loopDepth_--;
        return;
    }

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

void
SemanticAnalysis::visit(BlockStatement *node)
{
    if (node->scope())
        scope_ = node->scope();
    for (size_t i = 0; i < node->statements()->length(); i++) {
        Statement *statement = node->statements()->at(i);
        statement->accept(this);
    }
    if (node->scope())
        scope_ = scope_->enclosing();
}

void
SemanticAnalysis::visit(VariableDeclaration *node)
{
    Local<Type> type(zone_, zone_->types()->getPrimitiveForTag(node->tag()));
    if (!type)
        return;

    if (node->postDimensions()) {
        Local<Type> init(zone_);
        if (node->initialization()) {
            if ((init = visitAndApplyType(node->initialization())) == NULL)
                return;
        }

        if (node->postDimensions()->size() > Type::MAX_UNSAFE_ARRAY_DEPTH) {
            cc_.reportError(node->pos(), Message_MaximumArrayDepth);
            return;
        }

        // If we have an initializer, make sure it's the same type.
        if (init && node->postDimensions()->size() != init->levels()) {
            cc_.reportError(node->pos(), Message_ArrayInitializerDepth);
            return;
        }

        Local<Type> check(zone_, init);
        for (size_t i = node->postDimensions()->size() - 1; i < node->postDimensions()->size(); i--) {
            Expression *expr = node->postDimensions()->at(i);
            int size = -1;

            if (expr) {
                Local<Type> exprType(zone_, visitForLoad(expr));
                if (!exprType)
                    return;

                if (!IsPrimitiveType(exprType, PrimitiveType_Int32)) {
                    cc_.reportError(expr->pos(), Message_ArrayDimensionNotInteger);
                    return;
                }

                BoxedPrimitive box;
                if (EvaluateForConstant(expr, &box) && box.type() == PrimitiveType_Int32) {
                    size = box.toInt();
                    if (size <= 0) {
                        cc_.reportError(expr->pos(), Message_BadArraySize);
                        return;
                    }
                    if (size >= INT_MAX / sizeof(int)) {
                        cc_.reportError(expr->pos(), Message_ArraySizeTooBig);
                        return;
                    }

                    if (check) {
                        if (check->elements() == -1) {
                            cc_.reportError(expr->pos(), Message_ArraySizeCannotBeDetermined);
                            return;
                        }

                        if (check->isCharArray()) {
                            if (check->elements() > size) {
                                cc_.reportError(node->pos(), Message_ArrayAssignmentTooBig);
                                return;
                            }
                        } else if (check->elements() != size) {
                            cc_.reportError(node->pos(), Message_ArrayDimensionSizeMismatch, unsigned(i));
                            return;
                        }
                    }
                } else {
                    if (init) {
                        cc_.reportError(node->pos(), Message_DynamicArrayHasInitializer);
                        return;
                    }

                    // here until we get dynamic array support.
                    cc_.reportError(expr->pos(), Message_ArrayDimensionNotInteger);
                    return;
                }
            } else {
                // Infer the size from the initializer.
                if (!init) {
                    cc_.reportError(node->pos(), Message_ArraySizeNotSpecified);
                    return;
                }
                if (init->elements() == -1) {
                    cc_.reportError(node->pos(), Message_ArraySizeCannotBeDetermined);
                    return;
                }
                size = init->elements();
            }

            type = zone_->types()->unsafeArray(type, size);
            if (!type)
                return;

            if (check)
                check = check->contained();
        }

        // Right now, any array is allowed on the LIFO heap. In case the
        // runtime actually decides to use the LIFO heap, we must reserve
        // a slot so we can unwind the heap when the current scope leaves.
        scope_->setUsesLifoHeap();
    } else {
        if (node->initialization()) {
            Type *right = visitForLoad(node->initialization());
            if (!right)
                return;
            if (!typeCheckAssignment(node, type, right, true))
                return;
        }
        if (type->primitive() == PrimitiveType_Char) {
            cc_.reportError(node->pos(), Message_CannotDeclareStringCell);
            return;
        }
    }

    node->variable()->setType(type);
}

void
SemanticAnalysis::visit(IntegerLiteral *node)
{
    node->setType(zone_->types()->getPrimitive(PrimitiveType_Int32));
}

void
SemanticAnalysis::visit(ExpressionStatement *node)
{
    node->expression()->accept(this);
}

bool
SemanticAnalysis::preanalyze(FunctionStatement *node)
{
    Local<FixedArray> args(zone_, FixedArray::New(zone_, node->parameters()->size(), Heap::Tenure_Old));
    if (!args)
        return false;

    // Apply types to arguments.
    for (size_t i = 0; i < node->parameters()->size(); i++) {
        Parameter *parameter = node->parameters()->at(i);
        
        Local<Type> type(zone_, zone_->types()->getPrimitiveForTag(parameter->tag()));
        if (!type)
            return false;

        if (type->primitive() == PrimitiveType_Char) {
            cc_.reportError(node->pos(), Message_CannotDeclareStringCell);
            return false;
        }

        if (parameter->reference()) {
            type = zone_->types()->getReference(type);
            if (!type)
                return false;
        }

        args->set(zone_, i, type);
        parameter->variable()->setType(type);
    }

    Local<Type> returnType(zone_, zone_->types()->getPrimitive(PrimitiveType_Int32));
    if (!returnType)
        return false;

    Local<Type> funtype(zone_, zone_->types()->newFunction(returnType, args));
    if (!funtype)
        return false;

    node->fun()->setType(funtype);
    return true;
}

void
SemanticAnalysis::visit(FunctionStatement *node)
{
    function_ = node->fun()->type();
    scope_ = node->scope();
    node->body()->accept(this);
    scope_ = NULL;
    function_ = NULL;
}

bool
SemanticAnalysis::typeCheckArrayAssignment(AstNode *node, Type *given, Type *actual, bool forAssign)
{
    if (!given->isArray()) {
        cc_.reportError(node->pos(), Message_ExpectedArray);
        return false;
    }
    if (given->levels() != actual->levels()) {
        cc_.reportError(node->pos(), Message_ArrayDimensionsMismatch);
        return false;
    }

    unsigned level = 0;
    while (given->isArray()) {
        if (actual->elements() != -1) {
            if (forAssign) {
                if (given->elements() == -1 || actual->elements() == -1) {
                    cc_.reportError(node->pos(), Message_ArraySizeNotKnown);
                    return false;
                }

                if (actual->isCharArray()) {
                    // String assignment can have a truncated |given|.
                    if (given->elements() > actual->elements()) {
                        cc_.reportError(node->pos(), Message_ArrayAssignmentTooBig);
                        return false;
                    }
                } else {
                    if (given->elements() != actual->elements()) {
                        cc_.reportError(node->pos(), Message_ArrayDimensionSizeMismatch, level);
                        return false;
                    }
                }
            } else {
                if (actual->elements() != -1 && given->elements() != actual->elements()) {
                    cc_.reportError(node->pos(), Message_ArrayDimensionSizeMismatch, level);
                    return false;
                }
            }
        }
        given = given->contained();
        actual = actual->contained();
    }

    if (given->primitive() != actual->primitive()) {
        const char *givenTag = given->tag() ? given->tag() : "_";
        const char *actualTag = given->tag() ? given->tag() : "_";
        cc_.reportError(node->pos(), Message_TagMismatch, actualTag, givenTag);
        return false;
    }

    return true;
}

bool
SemanticAnalysis::typeCheckAssignment(AstNode *node, Type *given, Type *actual, bool forAssign)
{
    if (actual->isArray())
        return typeCheckArrayAssignment(node, given, actual, forAssign);

    if (given->isArray()) {
        cc_.reportError(node->pos(), Message_CannotAssignArrayToScalar);
        return false;
    }

    if (!given->isPrimitive() || !actual->primitive()) {
        cc_.reportError(node->pos(), Message_UnsupportedType);
        return false;
    }
    
    // The given type may not have been dereferenced.
    if (given->isReference())
        given = given->contained();
    if (actual->isReference())
        actual = given->contained();

    if (given->primitive() != actual->primitive()) {
        const char *givenTag = given->tag() ? given->tag() : "_";
        const char *actualTag = given->tag() ? given->tag() : "_";
        cc_.reportError(node->pos(), Message_TagMismatch, actualTag, givenTag);
        return false;
    }

    return true;
}

bool
SemanticAnalysis::typeCheckParameter(AstNode *node, Type *given, Type *actual)
{
    return typeCheckAssignment(node, given, actual, false);
}

void
SemanticAnalysis::visit(CallExpression *node)
{
    Type *type = visitForLoad(node->callee());
    if (!type)
        return;

    if (!type->isFunction()) {
        cc_.reportError(node->pos(), Message_CalleeNotFunction);
        return;
    }
    if (node->arguments()->size() != type->parameters()->length()) {
        cc_.reportError(node->pos(), Message_ArgumentCountMismatch);
        return;
    }

    for (unsigned i = 0; i < node->arguments()->size(); i++) {
        // Note: we don't compute references unless the callee needs a
        // reference type.
        Expression *expression = node->arguments()->at(i);
        Type *given = visitAndApplyType(expression);
        if (!given)
            return;

        // If it's an IndexExpression, we special case the ability to compute
        // a reference.
        if (expression->isIndexExpression())
            given = given->contained();

        Type *actual = type->parameterAt(i);
        if (!typeCheckParameter(expression, given, actual))
            return;

        // If the actual is a reference, and the given is not, we'll need to
        // compute a reference, which may be allocated on the stack.
        if (actual->isReference() && !given->isReference())
            scope_->setUsesLifoHeap();
    }
}

void
SemanticAnalysis::visit(IfStatement *node)
{
    Type *type = visitForLoad(node->condition());
    if (!type)
        return;

    if (!canTest(node, type))
        return;

    node->ifTrue()->accept(this);
    if (node->ifFalse())
        node->ifFalse()->accept(this);
}

void
SemanticAnalysis::visit(IndexExpression *node)
{
    // :TODO: use locals
    Type *left = visitForRef(node->left());
    Type *right = visitForLoad(node->right());

    if (!left || !right)
        return;

    if (!right->isInt32()) {
        cc_.reportError(node->right()->pos(), Message_IndexMustBeInteger);
        return;
    }
    if (!left->isArray()) {
        cc_.reportError(node->left()->pos(), Message_IndexBaseMustBeArray);
        return;
    }
    node->setType(left);
}

void
SemanticAnalysis::visit(EnumStatement *node)
{
    // Forbid the blessed types.
    Name *tag = node->name();
    if (!strcmp(tag, "Float") ||
        !strcmp(tag, "String") ||
        !strcmp(tag, "any") ||
        !strcmp(tag, "bool"))
    {
        cc_.reportError(node->pos(), Message_InvalidEnumName);
        return;
    }

    Type *type = zone_->types()->getPrimitiveForTag(tag);
    if (!type)
        return;

    int value = 0;
    for (size_t i = 0; i < node->entries()->size(); i++) {
        EnumStatement::Entry &entry = node->entries()->at(i);

        BoxedPrimitive box = BoxedPrimitive::Int(value);
        if (entry.expr) {
            entry.expr->accept(this);
            if (!EvaluateForConstant(entry.expr, &box)) {
                cc_.reportError(entry.expr->pos(), Message_EnumFieldNotConstant);
                return;
            }

            assert(box.type() == PrimitiveType_Int32);
            value = box.toInt();
        }

        entry.name->set(type, box);
        value++;
    }
}

void
SemanticAnalysis::visit(StringLiteral *node)
{
    Type *type = zone_->types()->getPrimitive(PrimitiveType_Char);
    if (!type)
        return;

    type = Type::NewArray(zone_, type, node->string()->length() + 1, true);
    if (!type)
        return;

    node->setType(type);
}

void
SemanticAnalysis::visit(FloatLiteral *node)
{
    Type *type = zone_->types()->getPrimitive(PrimitiveType_Float);
    if (!type)
        return;

    node->setType(type);
}

void
SemanticAnalysis::visit(BreakStatement *node)
{
    if (loopDepth_ == 0)
        cc_.reportError(node->pos(), Message_BreakOutsideLoop);
}

void
SemanticAnalysis::visit(ContinueStatement *node)
{
    if (loopDepth_ == 0)
        cc_.reportError(node->pos(), Message_ContinueOutsideLoop);
}