/* 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 "ConstantEvaluator.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 (node->dims()) {

        }

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

            node->sym()->setType(type);

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