/*

 * 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 "StringPool.h"

#include "NameBinding.h"

namespace jc {

#define ASTKINDS(_)             \

    _(VariableDeclaration)      \

    _(ForStatement)             \

    _(ReturnStatement)          \

    _(IntegerLiteral)           \

    _(FloatLiteral)             \

    _(BinaryExpression)         \

    _(BlockStatement)           \

    _(Assignment)               \

    _(NameProxy)                \

    _(ExpressionStatement)      \

    _(FunctionStatement)        \

    _(CallExpression)           \

    _(IfStatement)              \

    _(IndexExpression)          \

    _(EnumStatement)            \

    _(WhileStatement)

// Forward declarations.

#define _(name) class name;

ASTKINDS(_)

#undef _

class AstVisitor;

// Interface for AST nodes.

class AstNode : public PoolObject

{

  public:

    enum Kind {

#       define _(name) k##name,

        ASTKINDS(_)

#       undef _

        AstKind_Invalid

    };

    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;

};

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

{

};

class Expression : public AstNode

{

    Type *type_;

  public:

    Expression()

      : type_(NULL)

    {

    }

    void setType(Type *type) {

        type_ = type;

    }

    Type *type() const {

        return type_;

    }

};

class Declaration : public Statement

{

};

typedef PoolList<Statement *> StatementList;

typedef std::vector<Expression *> ExpressionList;

class VariableDeclaration : public Declaration

{

    Variable *variable_;

    Expression *initialization_;

    ExpressionList *postDimensions_;

    Name *tag_;

  public:

    VariableDeclaration(Variable *var, Name *tag, ExpressionList *postDimensions, Expression *initialization)

      : variable_(var),

        postDimensions_(postDimensions),

        initialization_(initialization),

        tag_(tag)

    {

    }

    Name *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 NameProxy : public Expression

{

    Name *name_;

    BoundName *binding_;

  public:

    NameProxy(Name *name, BoundName *binding)

      : name_(name),

        binding_(binding)

    {

    }

    Kind kind() const {

        return kNameProxy;

    }

    Name *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(int value)

      : 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(double value)

      : value_(value)

    {

    }

    Kind kind() const {

        return kFloatLiteral;

    }

    double value() const {

        return value_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class BinaryExpression : public Expression

{

    Expression *left_;

    Expression *right_;

    TokenKind token_;

  public:

    BinaryExpression(TokenKind token, Expression *left, Expression *right)

      : token_(token),

        left_(left),

        right_(right)

    {

    }

    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 IndexExpression : public Expression

{

    Expression *left_;

    Expression *right_;

  public:

    IndexExpression(Expression *left, Expression *right)

      : 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(Expression *callee, ExpressionList *arguments)

      : 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_;

    VariableDeclaration *declaration_;

    Expression *condition_;

    Statement *update_;

    Statement *body_;

  public:

    ForStatement(Scope *scope, VariableDeclaration *declaration,

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

      : scope_(scope),

        declaration_(declaration),

        condition_(condition),

        update_(update),

        body_(body)

    {

    }

    Kind kind() const {

        return kForStatement;

    }

    VariableDeclaration *declaration() const {

        return declaration_;

    }

    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(TokenKind kind, Expression *condition, Statement *body)

      : 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(Expression *expression)

      : expression_(expression)

    {

    }

    Kind kind() const {

        return kReturnStatement;

    }

    Expression *expression() const {

        return expression_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class Assignment : public Expression

{

    Expression *lvalue_;

    Expression *expression_;

  public:

    Assignment(Expression *left, Expression *right)

      : lvalue_(left),

        expression_(right)

    {

    }

    Kind kind() const {

        return kAssignment;

    }

    Expression *lvalue() const {

        return lvalue_;

    }

    Expression *expression() const {

        return expression_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class ExpressionStatement : public Statement

{

    Expression *expression_;

  public:

    ExpressionStatement(Expression *expression)

      : expression_(expression)

    {

    }

    Kind kind() const {

        return kExpressionStatement;

    }

    Expression *expression() const {

        return expression_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class BlockStatement : public Statement

{

    Scope *scope_;

    StatementList *statements_;

  public:

    BlockStatement(Scope *scope, StatementList *statements)

      : scope_(scope),

        statements_(statements)

    {

    }

    Kind kind() const {

        return kBlockStatement;

    }

    Scope *scope() const {

        return scope_;

    }

    StatementList *statements() const {

        return statements_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class FunctionStatement : public Statement

{

    Scope *scope_;

    NamedFunction *fun_;

    Statement *body_;

  public:

    FunctionStatement(Scope *scope, NamedFunction *fun, Statement *body)

      : scope_(scope),

        fun_(fun),

        body_(body)

    {

    }

    Kind kind() const {

        return kFunctionStatement;

    }

    Scope *scope() const {

        return scope_;

    }

    NamedFunction *fun() const {

        return fun_;

    }

    Statement *body() const {

        return body_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class IfStatement : public Statement

{

    Expression *condition_;

    Statement *ifTrue_;

    Statement *ifFalse_;

  public:

    IfStatement(Expression *condition, Statement *ifTrue)

      : condition_(condition),

        ifTrue_(ifTrue),

        ifFalse_(NULL)

    {

    }

    Kind kind() const {

        return kIfStatement;

    }

    Expression *condition() const {

        return condition_;

    }

    Statement *ifTrue() const {

        return ifTrue_;

    }

    Statement *ifFalse() const {

        return ifFalse_;

    }

    void setIfFalse(Statement *ifFalse) {

        ifFalse_ = ifFalse;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

};

class EnumStatement : public Statement

{

    Name *name_;

  public:

    struct Entry {

        NamedConstant *name;

        Expression *expr;

        Entry(NamedConstant *name, Expression *expr)

          : name(name),

            expr(expr)

        {

        }

    };

    typedef std::vector<Entry> EntryList;

  public:

    EnumStatement(Name *name, EntryList *entries)

      : name_(name),

        entries_(entries)

    {

    }

    Kind kind() const {

        return kEnumStatement;

    }

    Name *name() const {

        return name_;

    }

    EntryList *entries() const {

        return entries_;

    }

    void accept(AstVisitor *visitor) {

        visitor->visit(this);

    }

  private:

    EntryList *entries_;

};

class Scope;

typedef std::vector<NameProxy *> NameProxyList;

class ParseTree : public PoolObject

{

    StatementList *statements_;

    Scope *scope_;

    NameProxyList *unbound_;

  public:

    ParseTree(Scope *scope, StatementList *statements, NameProxyList *unbound)

      : statements_(statements),

        scope_(scope),

        unbound_(unbound)

    {

    }

    void dump(FILE *fp);

    StatementList *statements() const {

        return statements_;

    }

    Scope *scope() const {

        return scope_;

    }

    NameProxyList *unbound() const {

        return unbound_;

    }

};

}