/*

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

 */

#ifndef _include_jitcraft_ast_h_

#define _include_jitcraft_ast_h_

#include "PoolAllocator.h"

#include "Token.h"

#include "NameBinding.h"

#include "Vector.h"

namespace ke {

#define ASTKINDS(_)             \

    _(VariableDeclaration)      \

    _(ForStatement)             \

    _(ReturnStatement)          \

    _(IntegerLiteral)           \

    _(FloatLiteral)             \

    _(StringLiteral)            \

    _(BinaryExpression)         \

    _(BlockStatement)           \

    _(Assignment)               \

    _(NameProxy)                \

    _(ExpressionStatement)      \

    _(FunctionStatement)        \

    _(CallExpression)           \

    _(IfStatement)              \

    _(IndexExpression)          \

    _(EnumStatement)            \

    _(WhileStatement)           \

    _(BreakStatement)           \

    _(ContinueStatement)        \

    _(IncDecExpression)         \

    _(UnaryExpression)          \

    _(TernaryExpression)

// Forward declarations.

#define _(name) class name;

ASTKINDS(_)

#undef _

class AstVisitor;

// Interface for AST nodes.

class AstNode : public PoolObject

{

    SourcePosition pos_;

  public:

    enum Kind {

#       define _(name) k##name,

        ASTKINDS(_)

#       undef _

        AstKind_Invalid

    };

    AstNode(const SourcePosition &pos)

      : pos_(pos)

    {

    }

    virtual Kind kind() const = 0;

#define _(name)     bool is##name() { return kind() == k##name; }                   \

                    name *to##name() { assert(is##name()); return (name *)this; }

    ASTKINDS(_)

#undef _

    virtual void accept(AstVisitor *visitor) = 0;

    const SourcePosition &pos() const {

        return pos_;

    }

};

class AstVisitor

{

  public:

#define _(name) virtual void visit(name *node) = 0;

    ASTKINDS(_)

#undef _

};

class NameLocation

{

  public:

    enum Kind {

        LOCAL

    };

    static NameLocation Local(VariableDeclaration *decl) {

        NameLocation loc;

        loc.kind_ = LOCAL;

        loc.u_.decl = decl;

        return loc;

    }

    Kind kind() const {

        return kind_;

    }

    VariableDeclaration *local() const {

        assert(kind() == LOCAL);

        return u_.decl;

    }

  private:

    Kind kind_;

    union {

        VariableDeclaration *decl;

    } u_;

};

class Statement : public AstNode

{

  public:

    Statement(const SourcePosition &pos)

      : AstNode(pos)

    {

    }

};

class Expression : public AstNode

{

    ScopedRoot<Type> type_;

  public:

    Expression(const SourcePosition &pos)

      : AstNode(pos),

        type_(NULL)

    {

    }

    void setType(Type *type) {

        type_ = type;

    }

    Type *type() const {

        return type_;

    }

};

typedef PoolList<Statement *> StatementList;

typedef PoolList<Expression *> ExpressionList;

class VariableDeclaration : public Statement

{

    Variable *variable_;

    Expression *initialization_;

    ExpressionList *postDimensions_;

    ScopedRoot<String> tag_;

  public:

    VariableDeclaration(const SourcePosition &pos, Variable *var, Handle<String> tag,

                        ExpressionList *postDimensions, Expression *initialization)

      : Statement(pos),

        variable_(var),

        initialization_(initialization),

        postDimensions_(postDimensions),

        tag_(tag)

    {

    }

    String *tag() const {

        return tag_;

    }

    Kind kind() const {

        return kVariableDeclaration;

    }

    Variable *variable() const {

        return variable_;

    }

    Expression *initialization() const {

        return initialization_;

    }

    ExpressionList *postDimensions() const {

        return postDimensions_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class Parameter : public PoolObject

{

    Variable *variable_;

    ScopedRoot<String> tag_;

    bool reference_;

    bool const_;

    ExpressionList *dimensions_;

  public:

    Parameter(Variable *var, Handle<String> tag, bool reference, bool isConst, ExpressionList *dimensions)

      : variable_(var),

        tag_(tag),

        reference_(reference),

        const_(isConst),

        dimensions_(dimensions)

    {

    }

    String *tag() const {

        return tag_;

    }

    Variable *variable() const {

        return variable_;

    }

    bool reference() const {

        return reference_;

    }

    ExpressionList *dimensions() const {

        return dimensions_;

    }

    bool isConst() const {

        return const_;

    }

};

class NameProxy : public Expression

{

    ScopedRoot<String> name_;

    BoundName *binding_;

  public:

    NameProxy(const SourcePosition &pos, Handle<String> name, BoundName *binding)

      : Expression(pos),

        name_(name),

        binding_(binding)

    {

    }

    Kind kind() const {

        return kNameProxy;

    }

    String *name() const {

        return name_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

    BoundName *binding() const {

        return binding_;

    }

    void bind(BoundName *name) {

        assert(!binding());

        binding_ = name;

    }

};

class IntegerLiteral : public Expression

{

    int value_;

  public:

    IntegerLiteral(const SourcePosition &pos, int value)

      : Expression(pos),

        value_(value)

    {

    }

    Kind kind() const {

        return kIntegerLiteral;

    }

    int value() const {

        return value_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class FloatLiteral : public Expression

{

    double value_;

  public:

    FloatLiteral(const SourcePosition &pos, double value)

      : Expression(pos),

        value_(value)

    {

    }

    Kind kind() const {

        return kFloatLiteral;

    }

    double value() const {

        return value_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class StringLiteral : public Expression

{

    ScopedRoot<ByteArray> string_;

  public:

    StringLiteral(const SourcePosition &pos, ByteArray *array)

      : Expression(pos),

        string_(array)

    {

    }

    Kind kind() const {

        return kStringLiteral;

    }

    ByteArray *string() const {

        return string_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class UnaryExpression : public Expression

{

    Expression *expression_;

    TokenKind token_;

  public:

    UnaryExpression(const SourcePosition &pos, TokenKind token, Expression *expr)

      : Expression(pos),

        expression_(expr),

        token_(token)

    {

    }

    Kind kind() const {

        return kUnaryExpression;

    }

    Expression *expression() const {

        return expression_;

    }

    TokenKind token() const {

        return token_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class BinaryExpression : public Expression

{

    Expression *left_;

    Expression *right_;

    TokenKind token_;

  public:

    BinaryExpression(const SourcePosition &pos, TokenKind token, Expression *left, Expression *right)

      : Expression(pos),

        left_(left),

        right_(right),

        token_(token)

    {

    }

    Kind kind() const {

        return kBinaryExpression;

    }

    Expression *left() const {

        return left_;

    }

    Expression *right() const {

        return right_;

    }

    TokenKind token() const {

        return token_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class TernaryExpression : public Expression

{

    Expression *condition_;

    Expression *left_;

    Expression *right_;

  public:

    TernaryExpression(const SourcePosition &pos, Expression *condition, Expression *left, Expression *right)

      : Expression(pos),

        condition_(condition),

        left_(left),

        right_(right)

    {

    }

    Kind kind() const {

        return kTernaryExpression;

    }

    Expression *condition() const {

        return condition_;

    }

    Expression *left() const {

        return left_;

    }

    Expression *right() const {

        return right_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class IndexExpression : public Expression

{

    Expression *left_;

    Expression *right_;

  public:

    IndexExpression(const SourcePosition &pos, Expression *left, Expression *right)

      : Expression(pos),

        left_(left),

        right_(right)

    {

    }

    Kind kind() const {

        return kIndexExpression;

    }

    Expression *left() const {

        return left_;

    }

    Expression *right() const {

        return right_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class CallExpression : public Expression

{

    Expression *callee_;

    ExpressionList *arguments_;

  public:

    CallExpression(const SourcePosition &pos, Expression *callee, ExpressionList *arguments)

      : Expression(pos),

        callee_(callee),

        arguments_(arguments)

    {

    }

    Kind kind() const {

        return kCallExpression;

    }

    Expression *callee() const {

        return callee_;

    }

    ExpressionList *arguments() const {

        return arguments_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class ForStatement : public Statement

{

    Scope *scope_;

    Statement *initialization_;

    Expression *condition_;

    Statement *update_;

    Statement *body_;

  public:

    ForStatement(const SourcePosition &pos, Scope *scope, Statement *initialization,

                 Expression *condition, Statement *update, Statement *body)

      : Statement(pos),

        scope_(scope),

        initialization_(initialization),

        condition_(condition),

        update_(update),

        body_(body)

    {

    }

    Kind kind() const {

        return kForStatement;

    }

    Statement *initialization() const {

        return initialization_;

    }

    Scope *scope() const {

        return scope_;

    }

    Expression *condition() const {

        return condition_;

    }

    Statement *update() const {

        return update_;

    }

    Statement *body() const {

        return body_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class WhileStatement : public Statement

{

    TokenKind token_;

    Expression *condition_;

    Statement *body_;

  public:

    WhileStatement(const SourcePosition &pos, TokenKind kind, Expression *condition, Statement *body)

      : Statement(pos),

        token_(kind),

        condition_(condition),

        body_(body)

    {

    }

    Kind kind() const {

        return kWhileStatement;

    }

    TokenKind token() const {

        return token_;

    }

    Expression *condition() const {

        return condition_;

    }

    Statement *body() const {

        return body_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class ReturnStatement : public Statement

{

    Expression *expression_;

  public:

    ReturnStatement(const SourcePosition &pos, Expression *expression)

      : Statement(pos),

        expression_(expression)

    {

    }

    Kind kind() const {

        return kReturnStatement;

    }

    Expression *expression() const {

        return expression_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class BreakStatement : public Statement

{

  public:

    BreakStatement(const SourcePosition &pos)

      : Statement(pos)

    {

    }

    Kind kind() const {

        return kBreakStatement;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class ContinueStatement : public Statement

{

  public:

    ContinueStatement(const SourcePosition &pos)

      : Statement(pos)

    {

    }

    Kind kind() const {

        return kContinueStatement;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class Assignment : public Expression

{

    TokenKind token_;

    Expression *lvalue_;

    Expression *expression_;

  public:

    Assignment(const SourcePosition &pos, TokenKind token, Expression *left, Expression *right)

      : Expression(pos),

        token_(token),

        lvalue_(left),

        expression_(right)

    {

    }

    Kind kind() const {

        return kAssignment;

    }

    TokenKind token() const {

        return token_;

    }

    Expression *lvalue() const {

        return lvalue_;

    }

    Expression *expression() const {