#ifndef _ast_h_
#define _ast_h_

#include "pre.h"
#include "datatype.h"
#include "symbol.h"
#include "location.h"

#define ast_node_is_atom(nk) \
	(nk == AST_IDENT || nk == AST_NUMBER || nk == AST_STRLIT || nk == AST_PROCCALL)
#define ast_node_is_unary(nk) \
	(nk == AST_UNARY || ast_node_is_atom(nk))
#define ast_node_is_expr(nk) \
	(nk == AST_BINEXPR || ast_node_is_unary(nk))

enum AstType
{
	AST_INVALID, /* For use as a placeholder until the actual type is decided */
	AST_NUMBER, /* number */
	AST_IDENT, /* ident */
	AST_STRLIT, /* strlit */
	AST_PROCDEF, /* proc */
	AST_PROCCALL, /* call */
	AST_PROCCALL_ARGS, /* */
	AST_VARDECL, /* var */
	AST_VARASSIGN, /* varassgn */
	AST_IF, /* ifse */
	AST_RETURN, /* ret */
	AST_BREAK,
	AST_LOOP, /* loop */
	AST_STMTS, /* stmts */
	AST_EXPRS, /* exprs */
	AST_BINEXPR, /* bin */
	AST_UNARY, /* unary */
	AST_PRAGMA, /* pragma */
	AST_DISCARD,
};

typedef struct Ast Ast;

typedef struct {
	Str op;
	Ast *left, *right;
	DataType *type; /* filled in by sema */
} AstBinop;

typedef struct {
	Str op;
	Ast *atom;
	DataType *type; /* filled in by sema */
} AstUnary;

typedef struct {
	Str ident;
	Str dtype;
	/* Symbol kind for this parameter, `SymVar` would represent a mutable
	 * parameter and `SymLet` a immutable one. */
	enum SymbolKind kind;
	Location ident_loc, dtype_loc;
} AstIdentTypePair;

typedef struct {
	Str name;
	bool ispublic;
	Ast *body;
	Vec(AstIdentTypePair) args;
	Ast *rettype;

	DataType *type;
} AstProc;

typedef struct {
	Str name;
	Ast *args;
} AstProcCall;

typedef struct {
	Str name;
	/* Data type, nil if no type was explicitly stated, meaning that
	 * type deduction must be made from the expression, also implying that
	 * if this field is nil, `expr` MUSN'T be nil. */
	Ast *datatype;
	Ast *expr; /* if the declaration assigns a value */
	enum SymbolKind kind; /* whether is a let, var or const... */
	DataType *type; /* filled in by sema */
} AstVarDecl;

typedef struct {
	Str name;
	Ast *expr;
} AstVarAssign;

typedef struct {
	u64 n;
	DataType *type; /* filled in by the sema */
} AstNumber;

typedef struct {
	Ast *cond;
	Ast *body;
} AstElif;

typedef struct {
	Ast *cond;
	Ast *true_body;
	Ast *false_body;
	Vec(AstElif) elifs;
} AstIf;

/* Abstract representation of a loop, providing a pre and post condition.
 * `while` loops are modelled as a loop with a precondition only.
 * For infinite loops both `precond` and `postcond` are nil. */
typedef struct {
	Ast *precond, *postcond, *body;
} AstLoop;

typedef struct {
	/* Attributes for now can only be identifiers */
	Vec(Str) attrs;
	Ast *node; /* The decorated node */
} AstPragma;

typedef struct {
	Ast *expr;
} AstDiscard;

struct Ast {
	enum AstType type;
	union {
		AstBinop bin; /* binary expression */
		AstUnary unary; /* unary operator */
		AstNumber number; /* number (this is an atom) */
		Str ident; /* identifier (this is an atom too) */
		AstProc proc; /* procedure definition */
		AstProcCall call; /* procedure call */
		AstVarDecl var; /* variable declaration */
		AstVarAssign varassgn;
		Ast *ret; /* return statement, this points to its expression (if any) */
		AstIf ifse; /* if statement/expression */
		AstLoop loop;
		Vec(Ast *) stmts;
		Vec(Ast *) exprs;
		Str strlit; /* String literal */
		AstPragma pragma;
		AstDiscard discard;
	};
	Location loc; /* location in the source code of this node */
};
_Static_assert(sizeof(Ast) <= 512, "AST node got too bloated");

#endif