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

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

 }