/* 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.

         }

         if (node->initialization())

             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) {

         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) {

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

 }

 static const int INFER_ARRAY_SIZE = 0;

 static const int EVAL_ARRAY_SIZE = -1;

 static const int DYNAMIC_ARRAY_SIZE = -2;

 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).

         if (node->initialization())

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

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

         if (!type)

             return;

         if (node->dims() || type->isArray()) {

             int dims[MAX_ARRAY_DEPTH];

             int levels = evaluateDimensions(node->pos(), type, node->dims(), dims);

             if (levels < 0)

                 return;

             // Do some extra inference based on the initializer, if present.

             Expression *init = node->initialization();

             if (init &&

                 init->isArrayLiteral() &&

                 init->toArrayLiteral()->isFixed())

             {

                 Local<Type> contained(zone_, type->isArray()

                                              ? ArrayType::cast(type)->contained()

                                              : type);

                 if (!inferFixedArrayDimensions(init->toArrayLiteral(), contained, dims, levels))

                     return;

             }

             // If we got extra dimensions, we need to build a new type.

             if (node->dims()) {

                 if ((type = buildArrayType(type, dims, node->dims()->length())) == NULL)

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

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

         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:

     int evaluateDimensions(const SourcePosition &pos, Handle<Type> typeArg, ExpressionList *exprs, int dims[MAX_ARRAY_DEPTH]) {

         unsigned level = 0;

         Local<Type> type(zone_, typeArg);

         while (type->isArray()) {

             ArrayType *atype = ArrayType::cast(type);

             dims[level] = atype->isFixedLength()

                           ? atype->fixedLength()

                           : DYNAMIC_ARRAY_SIZE;

             type = atype->contained();

             level++;

         }

         if (!exprs)

             return level;

         if (level + exprs->length() > MAX_ARRAY_DEPTH) {

             cc_.reportError(pos, Message_MaximumArrayDepth);

             return -1;

         }

         bool hadEmptySize = false;

         for (size_t i = 0; i < exprs->length(); i++) {

             Expression *expr = exprs->at(i);

             if (!expr) {

                 dims[i] = INFER_ARRAY_SIZE;

                 hadEmptySize = true;

                 continue;

             }

             BoxedPrimitive box;

             if (EvaluateForConstant(expr, &box) && box.type() == PrimitiveType_Int32) {

                 int size = box.toInt();

                 if (size <= 0) {

                     cc_.reportError(expr->pos(), Message_BadArraySize);

                     return false;

                 }

                 dims[i] = size;

                 continue;

             }

             if (hadEmptySize) {

                 cc_.reportError(expr->pos(), Message_BadDynamicInitializer);

                 return false;

             }

             dims[i] = EVAL_ARRAY_SIZE;

         }

         return level + exprs->length();

     }

     // This function computes missing dimension sizes based on a literal

     // initializer. We don't really handle any size rules here; we just

     // do a really dumb inference.

     bool inferFixedArrayDimensions(ArrayLiteral *lit, Handle<Type> contained, int dims[MAX_ARRAY_DEPTH], int levels) {

         for (int i = 0; i < levels; i++) {

             if (dims[i] == EVAL_ARRAY_SIZE || dims[i] == DYNAMIC_ARRAY_SIZE) {

                 // The following are invalid:

                 //   new x[<expr>] = {...}

                 //   new Type:x = {...}  // where Type is a non-fixed array type

                 cc_.reportError(lit->pos(), Message_CannotInitializeEvalArray);

                 return false;

             }

             if (dims[i] == INFER_ARRAY_SIZE) {

                 // Character arrays are always dynamic when inferred from string literals.

                 if (i == levels - 1 && contained->isChar())

                     dims[i] = DYNAMIC_ARRAY_SIZE;

                 else

                     dims[i] = lit->expressions()->length();

             }

             if (i != levels - 1) {

                 // Fixed array literals are guaranteed to have at least one expression.

                 Expression *next = lit->expressions()->at(0);

                 if (!next->isArrayLiteral()) {

                     // Whelp... this array is just malformed. Not the best error message

                     // but whatever.

                     cc_.reportError(next->pos(), Message_ArraySizeCannotBeDetermined);

                     return false;

                 }

                 lit = next->toArrayLiteral();

             }

         }

         return true;

     }

     ArrayType *buildArrayType(Handle<Type> contained, int dims[MAX_ARRAY_DEPTH], size_t postlevels) {

         Local<Type> type(zone_, contained);

         for (size_t i = postlevels - 1; i < postlevels; i--) {

             int size = dims[i] > 0

                        ? dims[i]

                        : ArrayType::DYNAMIC_ARRAY_SIZE;

             type = zone_->types()->newArray(contained, size);

             if (!type)

                 return NULL;

         }

         return ArrayType::cast(type);

     }