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

// 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 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 {
        return expression_;
    }
    void accept(AstVisitor *visitor) {
        visitor->visit(this);
    }
};

class ExpressionStatement : public Statement
{
    Expression *expression_;

  public:
    ExpressionStatement(Expression *expression)
      : Statement(expression->pos()),
        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(const SourcePosition &pos, Scope *scope, StatementList *statements)
      : Statement(pos),
        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);
    }
};

typedef PoolList<Parameter *> ParameterList;

class FunctionStatement : public Statement
{
    Scope *scope_;
    NamedFunction *fun_;
    Statement *body_;
    ParameterList *parameters_;
    bool explicitReturn_;

  public:
    FunctionStatement(const SourcePosition &pos, Scope *scope, NamedFunction *fun, Statement *body,
                      ParameterList *parameters, bool explicitReturn)
      : Statement(pos),
        scope_(scope),
        fun_(fun),
        body_(body),
        parameters_(parameters),
        explicitReturn_(explicitReturn)
    {
    }

    Kind kind() const {
        return kFunctionStatement;
    }
    Scope *scope() const {
        return scope_;
    }
    NamedFunction *fun() const {
        return fun_;
    }
    Statement *body() const {
        return body_;
    }
    ParameterList *parameters() const {
        return parameters_;
    }
    void accept(AstVisitor *visitor) {
        visitor->visit(this);
    }
    bool hasExplicitReturn() const {
        return explicitReturn_;
    }
};

class IfStatement : public Statement
{
    Expression *condition_;
    Statement *ifTrue_;
    Statement *ifFalse_;

  public:
    IfStatement(const SourcePosition &pos, Expression *condition, Statement *ifTrue)
      : Statement(pos),
        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
{
    ScopedRoot<String> name_;

  public:
    struct Entry {
        NamedConstant *name;
        Expression *expr;

        Entry(NamedConstant *name, Expression *expr)
          : name(name),
            expr(expr)
        {
        }
    };

    typedef PoolList<Entry> EntryList;

  public:
    EnumStatement(const SourcePosition &pos, Handle<String> name, EntryList *entries)
      : Statement(pos),
        name_(name),
        entries_(entries)
    {
    }

    Kind kind() const {
        return kEnumStatement;
    }
    String *name() const {
        return name_;
    }
    EntryList *entries() const {
        return entries_;
    }
    void accept(AstVisitor *visitor) {
        visitor->visit(this);
    }

  private:
    EntryList *entries_;
};

class IncDecExpression : public Expression
{
    TokenKind token_;
    Expression *expression_;
    bool postfix_;

  public:
    IncDecExpression(const SourcePosition &pos, TokenKind token, Expression *expression, bool postfix)
      : Expression(pos),
        token_(token),
        expression_(expression),
        postfix_(postfix)
    {
    }

    Kind kind() const {
        return kIncDecExpression;
    }
    TokenKind token() const {
        return token_;
    }
    Expression *expression() const {
        return expression_;
    }
    bool postfix() const {
        return postfix_;
    }
    void accept(AstVisitor *visitor) {
        visitor->visit(this);
    }
};

class Scope;

typedef PoolList<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_;
    }
};

}

#endif // _include_jitcraft_ast_h_