/* 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 "../Strings.h"
#include "../Zone.h"
#include "../TypeManager.h"
#include "NameBinding.h"
#include "Scopes.h"
#include "../Spaces-inl.h"
#include "../Heap-inl.h"

using namespace ke;

// Create symbol tables and symbols for every scope and declaration. For type
// declarations and methods, we also pre-allocate a Type object, so that
// circular dependencies between types will work.
class NamePopulator : public AstVisitor
{
  private:
    // This class is analagous to a Scope, however, it allows us to create
    // scopes lazily. If a scope is requested, and the current environment
    // has no scope, then a scope is created. If no scope is ever requested,
    // a scope does not need to be created.
    class SymbolEnv
    {
      public:
        SymbolEnv(SymbolEnv **prevp, Scope *scope = NULL)
          : envp_(prevp),
            prev_(*prevp),
            scope_(scope)
        {
            *envp_ = this;
        }
        ~SymbolEnv() {
            *envp_ = prev_;
        }

        Scope *scope() {
            return scope_;
        }
        void setScope(Scope *scope) {
            assert(!scope_);
            scope_ = scope;
        }
        SymbolEnv *prev() const {
            return prev_;
        }

      private:
        SymbolEnv **envp_;
        SymbolEnv *prev_;
        Scope *scope_;
    };

    enum ReturnStatus {
        NoReturn,
        VoidReturn,
        ValueReturn
    };

    class FunctionScope
    {
      public:
        FunctionScope(FunctionScope **prevp)
          : prevp_(prevp),
            prev_(*prevp),
            status_(NoReturn)
        {
            *prevp_ = this;
        }
        ~FunctionScope() {
            *prevp_ = prev_;
        }
        ReturnStatus returns() const {
            return status_;
        }
        void returns(ReturnStatus status) {
            if (status > status_)
                status_ = status;
        }
      private:
        FunctionScope **prevp_;
        FunctionScope *prev_;
        ReturnStatus status_;
    };

  public:
    NamePopulator(Zone *zone, CompileContext &cc, TranslationUnit *unit)
      : zone_(zone),
        pool_(zone->pool()),
        cc_(cc),
        unit_(unit),
        env_(NULL),
        fun_(NULL)
    {
    }

    bool analyze() {
        ImportScope *importScope = ImportScope::New(pool_);
        if (!importScope)
            return false;

        // At the import level, we declare system types.
        if (!declareSystemTypes(importScope))
            return false;

        GlobalScope *globalScope = GlobalScope::New(pool_, importScope);
        if (!globalScope)
            return false;

        unit_->setImportScope(importScope);
        unit_->setGlobalScope(globalScope);

        SymbolEnv env(&env_, globalScope);

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

        return true;
    }

    bool declareSystemType(ImportScope *scope, const char *name, Handle<Type> type) {
        Local<String> tag(zone_, zone_->makeSymbol(name));
        if (!tag)
            return false;

        TypeSymbol *sym = new (pool_) TypeSymbol(scope, tag, type);
        return scope->addSymbol(sym);
    }

    bool declareSystemType(ImportScope *scope, const char *name, PrimitiveType prim) {
        Local<Type> primitive(zone_, zone_->types()->getPrimitive(prim));
        assert(primitive);
        return declareSystemType(scope, name, primitive);
    }

    bool declareSystemTypes(ImportScope *scope) {
        if (!declareSystemType(scope, "float", PrimitiveType_Float))
            return false;
        if (!declareSystemType(scope, "_", PrimitiveType_Int32))
            return false;
        if (!declareSystemType(scope, "int", PrimitiveType_Int32))
            return false;
        if (!declareSystemType(scope, "bool", PrimitiveType_Bool))
            return false;
        if (!declareSystemType(scope, "char", PrimitiveType_Char))
            return false;

        Local<Type> voidType(zone_, zone_->types()->getVoid());
        if (!declareSystemType(scope, "void", voidType))
            return false;

        return true;
    }

    void visitVariableDeclaration(VariableDeclaration *node) {
        Scope *scope = env_->scope();
        if (!scope || scope->kind() != Scope::GLOBAL) {
            if ((scope = getOrCreateBlockScope()) == NULL)
                return;
        }

        if (scope->kind() == Scope::GLOBAL) {
            // The symbol must be added now, since it could be an export.
            VariableSymbol *sym = new (pool_) VariableSymbol(scope, node->name(), node->pos());
            node->setSymbol(sym);

            if (!registerName(scope, sym))
                return;
        } else {
            // We ensure the scope will exist at the name binding phase, but
            // we don't actually register the variable symbol yet. Otherwise,
            // when binding the initializer, we could encounter:
            //    new x = x;
            //
            // And we don't necessarily want this to bind.
        }

        node->initialization()->accept(this);
    }
    void visitEnumStatement(EnumStatement *node) {
        assert(false);
    }
    void visitFunctionStatement(FunctionStatement *node) {
        // We create function types ahead of time, not for any particular
        // reason. It makes things consistent with FunctionTypeStatement,
        // which actually does produce a type.
        Local<FunctionType> type(zone_, FunctionType::New(zone_, node->token()));
        if (!type)
            return;

        assert(env_->scope());

        // Function statements are always named, so add the name to the outer scope.
        FunctionSymbol *sym = new (pool_) FunctionSymbol(env_->scope(), node->name(), node->pos(), type);
        if (!registerName(env_->scope(), sym))
            return;

        if (node->token() == TOK_PUBLIC ||
            node->token() == TOK_FORWARD ||
            node->token() == TOK_NATIVE)
        {
            // We're in the global scope, for sure.
            assert(env_->scope()->kind() == Scope::GLOBAL);
            GlobalScope *globals = (GlobalScope *)env_->scope();
            if (!globals->addPublic(sym))
                return;
        }

        node->setSymbol(sym);

        // Scope for the function's arguments.
        SymbolEnv argEnv(&env_);
        if (!registerArguments(node->signature()))
            return;

        // Scope for function body.
        SymbolEnv localEnv(&env_);
        FunctionScope fun(&fun_);
        if (node->body())
            node->body()->accept(this);

        // Function statements have quirky semantics around their return
        // type. For reasonable compatibility with SP1, we use the following
        // heuristics for when no explicit type is declared.
        if (!node->signature().returnType()) {
            Local<Type> type(zone_);
            if (node->token() == TOK_FORWARD || node->token() == TOK_NATIVE) {
                // If the function type is a native or forward, we always
                // assume it returns an int32.
                type = zone_->types()->getPrimitive(PrimitiveType_Int32);
            } else if (fun.returns() != ValueReturn) {
                // This function never returns a value, so mark it as
                // returning void.
                type = zone_->types()->getVoid();
            }
            // No need to NULL check here, we'll catch OOM later.
            node->sym()->type()->setReturnType(type);
        }

        node->setScopes(argEnv.scope(), localEnv.scope());
    }
    void visitStructureStatement(StructureStatement *node) {
        assert(false);
    }
    void visiFunctionTypeStatementt(FunctionTypeStatement *node) {
        assert(false);
    }

    void visitNameProxy(NameProxy *name) {
    }
    void visitAssignment(Assignment *node) {
        node->lvalue()->accept(this);
        node->expression()->accept(this);
    }
    void visitBinaryExpression(BinaryExpression *node) {
        node->left()->accept(this);
        node->right()->accept(this);
    }
    void visitReturnStatement(ReturnStatement *node) {
        if (node->expression())
            fun_->returns(ValueReturn);
        else
            fun_->returns(VoidReturn);
    }
    void visitForStatement(ForStatement *node) {
        SymbolEnv env(&env_);

        if (node->initialization())
            node->initialization()->accept(this);
        if (node->condition())
            node->condition()->accept(this);
        if (node->update())
            node->update()->accept(this);
        node->body()->accept(this);

        node->setScope(env.scope());
    }
    void visitBlockStatement(BlockStatement *node) {
        SymbolEnv env(&env_);

        for (size_t i = 0; i < node->statements()->length(); i++)
            node->statements()->at(i)->accept(this);

        node->setScope(env.scope());
    }
    void visitIntegerLiteral(IntegerLiteral *node) {
    }
    void visitExpressionStatement(ExpressionStatement *node) {
        node->expression()->accept(this);
    }
    void visitCallExpression(CallExpression *node) {
        node->callee()->accept(this);
        for (size_t i = 0; i < node->arguments()->length(); i++)
            node->arguments()->at(i)->accept(this);
    }
    void visitIfStatement(IfStatement *node) {
        node->ifTrue()->accept(this);
        if (node->ifFalse())
            node->ifFalse()->accept(this);
    }
    void visitIndexExpression(IndexExpression *node) {
        assert(false);
    }
    void visitFloatLiteral(FloatLiteral *node) {
    }
    void visitWhileStatement(WhileStatement *node) {
        node->condition()->accept(this);
        node->body()->accept(this);
    }
    void visitBreakStatement(BreakStatement *node) {
    }
    void visitContinueStatement(ContinueStatement *node) {
    }
    void visitStringLiteral(StringLiteral *node) {
    }
    void visitIncDecExpression(IncDecExpression *node) {
        assert(false);
    }
    void visitUnaryExpression(UnaryExpression *node) {
        assert(false);
    }
    void visitTernaryExpression(TernaryExpression *node) {
        assert(false);
    }
    void visitBooleanLiteral(BooleanLiteral *node) {
    }
    void visitFieldExpression(FieldExpression *node) {
    }
    void visitSwitchStatement(SwitchStatement *node) {
        if (node->defaultCase())
            node->defaultCase()->accept(this);

        for (size_t i = 0; i < node->cases()->length(); i++)
            node->cases()->at(i)->statement()->accept(this);
    }
    void visitArrayLiteral(ArrayLiteral *node) {
    }
    void visitImportStatement(ImportStatement *node) {
    }

  private:
    Scope *getOrCreateBlockScope() {
        if (env_->scope()) {
            assert(env_->scope()->kind() == Scope::BLOCK);
            return env_->scope();
        }
        
        LocalScope *scope = LocalScope::New(pool_);
        if (!scope)
            return NULL;

        env_->setScope(scope);
        return scope;
    }

    bool registerName(Scope *scope, Symbol *sym) {
        if (!scope)
            return false;

        if (Symbol *other = scope->localLookup(sym->name())) {
            // Report, but allow errors to continue.
            cc_.reportError(sym->pos(), Message_RedeclaredName,
                            sym->name()->chars(),
                            other->pos().line,
                            other->pos().col);
            return true;
        }

        return scope->addSymbol(sym);
    }

    bool registerArguments(const FunctionSignature &sig) {
        if (!sig.parameters()->length())
            return true;

        ArgumentScope *argScope = ArgumentScope::New(pool_);
        if (!argScope)
            return false;
        env_->setScope(argScope);

        for (size_t i = 0; i < sig.parameters()->length(); i++) {
            Parameter *param = sig.parameters()->at(i);
            VariableSymbol *sym = new (pool_) VariableSymbol(argScope, param->name(), param->pos());
            if (!registerName(env_->scope(), sym))
                return false;
            param->setSymbol(sym);
        }

        return true;
    }

  private:
    Zone *zone_;
    PoolAllocator &pool_;
    CompileContext &cc_;
    TranslationUnit *unit_;
    SymbolEnv *env_;
    FunctionScope *fun_;
};

bool
ke::PopulateNamesAndTypes(Zone *zone, CompileContext &cc, TranslationUnit *unit)
{
    NamePopulator populator(zone, cc, unit);
    return populator.analyze();
}

class NameBinder : public AstVisitor
{
  public:
    class AutoLinkScope
    {
      public:
        AutoLinkScope(AutoLinkScope **prevp, Scope *scope)
          : prevp_(prevp),
            prev_(*prevp),
            scope_(scope)
        {
            if (scope_) {
                *prevp_ = this;
                if (prev_)
                    scope_->setParent(prev_->scope());
            }
        }
        ~AutoLinkScope() {
            if (scope_) {
                assert(*prevp_ == this);
                *prevp_ = prev_;
            }
        }
        Scope *scope() const {
            return scope_;
        }

      private:
        AutoLinkScope **prevp_;
        AutoLinkScope *prev_;
        Scope *scope_;
    };

  public:
    NameBinder(Zone *zone, CompileContext &cc, TranslationUnit *unit)
      : zone_(zone),
        pool_(zone->pool()),
        cc_(cc),
        unit_(unit),
        link_(NULL)
    {
    }

    bool bind() {
        // Start with the import scope, and collect all imported names.
        AutoLinkScope importScope(&link_, unit_->importScope());

        for (size_t i = 0; i < unit_->tree()->imports()->length(); i++)
            unit_->tree()->imports()->at(i)->accept(this);

        // Next, enter the global scope and walk all statements.
        AutoLinkScope globalScope(&link_, unit_->globalScope());
        for (size_t i = 0; i < unit_->tree()->statements()->length(); i++) {
            Statement *stmt = unit_->tree()->statements()->at(i);
            stmt->accept(this);
        }

        return true;
    }

    void visitVariableDeclaration(VariableDeclaration *node) {
        // Bind the initializer before registering the declaration, so that we
        // can error on bogus initializers (new x = x).
        node->initialization()->accept(this);

        Local<Type> type(zone_, bindType(node->type()));
        if (!type)
            return;

        // If the node already has a symbol (meaning it was a global), then
        // we don't have to do anything more.
        if (node->sym()) {
            assert(node->sym()->scope()->kind() == Scope::GLOBAL);
            return;
        }

        VariableSymbol *sym = new (pool_) VariableSymbol(link_->scope(), node->name(), node->pos(), type);
        node->setSymbol(sym);
       
        if (Symbol *other = link_->scope()->localLookup(sym->name())) {
            // Report, but allow errors to continue.
            cc_.reportError(sym->pos(), Message_RedeclaredName,
                            sym->name()->chars(),
                            other->pos().line,
                            other->pos().col);
            return;
        }
        
        if (link_->scope()->addSymbol(sym))
            return;
    }
    void visitEnumStatement(EnumStatement *node) {
        assert(false);
    }
    void visitFunctionStatement(FunctionStatement *node) {
        AutoLinkScope argScope(&link_, node->argScope());
        AutoLinkScope varScope(&link_, node->varScope());

        Local<FunctionType> type(zone_, node->sym()->type());
        if (!fillFunctionType(type, node->signature()))
            return;

        if (node->body())
            node->body()->accept(this);
    }
    void visitStructureStatement(StructureStatement *node) {
        assert(false);
    }
    void visitFunctionTypeStatement(FunctionTypeStatement *node) {
        assert(false);
    }

    void visitNameProxy(NameProxy *proxy) {
        Scope *scope = link_->scope();
        Symbol *sym = scope->lookup(proxy->name());
        if (!sym)
            cc_.reportError(proxy->pos(), Message_IdentifierNotFound, proxy->name()->chars());
        proxy->bind(sym);
    }
    void visitAssignment(Assignment *node) {
        node->lvalue()->accept(this);
        node->expression()->accept(this);
    }
    void visitBinaryExpression(BinaryExpression *node) {
        node->left()->accept(this);
        node->right()->accept(this);
    }
    void visitReturnStatement(ReturnStatement *node) {
        node->expression()->accept(this);
    }
    void visitForStatement(ForStatement *node) {
        AutoLinkScope scope(&link_, node->scope());

        if (node->initialization())
            node->initialization()->accept(this);
        if (node->condition())
            node->condition()->accept(this);
        if (node->update())
            node->update()->accept(this);
        node->body()->accept(this);
    }
    void visitBlockStatement(BlockStatement *node) {
        AutoLinkScope scope(&link_, node->scope());
        for (size_t i = 0; i < node->statements()->length(); i++)
            node->statements()->at(i)->accept(this);
    }
    void visitIntegerLiteral(IntegerLiteral *node) {
    }
    void visitExpressionStatement(ExpressionStatement *node) {
        node->expression()->accept(this);
    }
    void visitCallExpression(CallExpression *node) {
        node->callee()->accept(this);
        for (size_t i = 0; i < node->arguments()->length(); i++)
            node->arguments()->at(i)->accept(this);
    }
    void visitIfStatement(IfStatement *node) {
        node->condition()->accept(this);
        node->ifTrue()->accept(this);
        if (node->ifFalse())
            node->ifFalse()->accept(this);
    }
    void visitIndexExpression(IndexExpression *node) {
        node->left()->accept(this);
        node->right()->accept(this);
    }
    void visitFloatLiteral(FloatLiteral *node) {
    }
    void visitWhileStatement(WhileStatement *node) {
        node->condition()->accept(this);
        node->body()->accept(this);
    }
    void visitBreakStatement(BreakStatement *node) {
    }
    void visitContinueStatement(ContinueStatement *node) {
    }
    void visitStringLiteral(StringLiteral *node) {
    }
    void visitIncDecExpression(IncDecExpression *node) {
        node->expression()->accept(this);
    }
    void visitUnaryExpression(UnaryExpression *node) {
        if (node->tag())
            node->tag()->accept(this);
        node->expression()->accept(this);
    }
    void visitTernaryExpression(TernaryExpression *node) {
        node->condition()->accept(this);
        node->left()->accept(this);
        node->right()->accept(this);
    }
    void visitBooleanLiteral(BooleanLiteral *node) {
    }
    void visitFieldExpression(FieldExpression *node) {
        node->left()->accept(this);
    }
    void visitSwitchStatement(SwitchStatement *node) {
        node->expression()->accept(this);
        node->defaultCase()->accept(this);
        for (size_t i = 0; i < node->cases()->length(); i++) {
            // We don't test case expressions because they are literals.
            Case *c = node->cases()->at(i);
            c->statement()->accept(this);
        }
    }
    void visitArrayLiteral(ArrayLiteral *node) {
        for (size_t i = 0; i < node->expressions()->length(); i++)
            node->expressions()->at(i)->accept(this);
    }
    void visitImportStatement(ImportStatement *node) {
        if (node->suffix() == TOK_STAR) {
            // Import every name.
            assert(false);
        } else {
            GlobalScope *other = node->source()->globalScope();
            Symbol *sym = other->findPublic(node->name());
            if (!sym) {
                cc_.reportError(node->pos(), Message_CannotResolveImport, node->name()->chars());
                return;
            }
            link_->scope()->addSymbol(sym);
        }
    }

  private:
    bool fillFunctionType(Handle<FunctionType> fun, const FunctionSignature &sig) {
        Local<Type> type(zone_);
        if (!fun->returnType()) {
            if ((type = bindType(sig.returnType())) == NULL)
                return false;
            fun->setReturnType(type);
        }

        Local<FixedArray> params(zone_, FixedArray::New(zone_, sig.parameters()->length(), Heap::Tenure_Old));
        if (!params)
            return false;

        for (size_t i = 0; i < sig.parameters()->length(); i++) {
            Parameter *param = sig.parameters()->at(i);
            if ((type = bindType(param->type())) == NULL)
                return false;
            params->set(zone_, i, type);
        }

        fun->setParameterTypes(params);

        return true;
    }

    Type *bindType(Expression *expr) {
        if (!expr) {
            // If no type was declared, we return int32.
            return zone_->types()->getPrimitive(PrimitiveType_Int32);
        }
        
        expr->accept(this);

        NameProxy *proxy = expr->asNameProxy();
        assert(proxy);

        Symbol *sym = proxy->sym();
        if (!sym)
            return NULL;

        if (!sym->isType()) {
            cc_.reportError(expr->pos(), Message_IdentifierIsNotAType, sym->name()->chars());
            return NULL;
        }

        return sym->type();
    }

  private:
    Zone *zone_;
    PoolAllocator &pool_;
    CompileContext &cc_;
    TranslationUnit *unit_;
    AutoLinkScope *link_;
};

bool
ke::BindNamesAndTypes(Zone *zone, CompileContext &cc, TranslationUnit *unit)
{
    NameBinder binder(zone, cc, unit);
    return binder.bind();
}