#define AST_NEW(T, ikind, ipos, iflags)                     \
    Ast_##T *result = allocate_struct(pctx->perm, Ast_##T); \
    result->flags = iflags;                                 \
    result->kind = AST_##ikind;                             \
    result->parent_scope = pctx->currently_parsed_scope;    \
    result->pos = ipos;                                     \
    result->di = ++pctx->unique_ids

#define ast_new(T, kind, pos, flags) (T *)_ast_new(sizeof(T), kind, pos, flags)
CORE_Static Ast *
_ast_new(size_t size, Ast_Kind kind, Token *pos, Ast_Flag flags = 0) {
    Ast *result = (Ast *)allocate_size(pctx->perm, size);
    result->flags = flags;
    result->kind = kind;
    result->parent_scope = pctx->currently_parsed_scope;
    result->pos = pos;
    result->di = ++pctx->unique_ids;
    return result;
}

CORE_Static void
set_flag_typespec(Ast_Expr *expr) {
    if (!expr) return;
    set_flag(expr->flags, AST_TYPESPEC);

    switch (expr->kind) {
        case AST_UNARY: {
            Ast_Unary *n = (Ast_Unary *)expr;
            set_flag_typespec(n->expr);
        } break;
        case AST_ARRAY: {
            Ast_Array *n = (Ast_Array *)expr;
            set_flag_typespec(n->base);
            set_flag_typespec(n->expr);
        } break;

        case AST_CALL: {
            Ast_Call *n = (Ast_Call *)expr;
            For(n->exprs) set_flag_typespec(it);
            set_flag_typespec(n->name);
        } break;
        case AST_CALL_ITEM: {
            Ast_Call_Item *n = (Ast_Call_Item *)expr;
            set_flag_typespec(n->item);
            set_flag_typespec(n->index);
        } break;

        case AST_IDENT:
            break;
            invalid_default_case;
    }
}

CORE_Static void
propagate_polymorphic(Ast *to, Ast *from) {
    if (is_flag_set(from->flags, AST_IDENT_POLYMORPH)) set_flag(to->flags, AST_IDENT_POLYMORPH);
    if (is_flag_set(from->flags, AST_TYPE_POLYMORPH)) set_flag(to->flags, AST_TYPE_POLYMORPH);
}

CORE_Static void
propagate_polymorphic(Ast *to, Token *from) {
    if (from->kind == TK_Polymorph) set_flag(to->flags, AST_IDENT_POLYMORPH);
}

CORE_Static Ast_Atom *
ast_str(Token *pos, Intern_String string) {
    AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
    result->type = pctx->untyped_string;
    result->intern_val = string;
    return result;
}

CORE_Static Ast_Atom *
ast_ident(Token *pos, Intern_String string) {
    AST_NEW(Atom, IDENT, pos, AST_EXPR | AST_ATOM);
    result->intern_val = string;
    return result;
}

CORE_Static Ast_Atom *
ast_bool(Token *pos, B32 bool_val) {
    AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
    result->bool_val = bool_val;
    result->type = pctx->untyped_bool;
    return result;
}

CORE_Static Ast_Atom *
ast_float(Token *pos, F64 value) {
    AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
    result->type = pctx->untyped_float;
    result->f64_val = value;
    return result;
}

CORE_Static Ast_Atom *
ast_int(Token *pos, BigInt val) {
    AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
    result->type = pctx->untyped_int;
    result->big_int_val = bigint_copy(pctx->perm, &val);
    return result;
}

CORE_Static Ast_Atom *
ast_int(Token *pos, U64 value) {
    return ast_int(pos, bigint_u64(value));
}

CORE_Static Ast_Expr *
ast_expr_binary(Ast_Expr *left, Ast_Expr *right, Token *op) {
    AST_NEW(Binary, BINARY, op, AST_EXPR);
    result->op = op->kind;
    result->left = left;
    result->right = right;
    return result;
}

CORE_Static Ast_Call *
ast_call(Token *pos, Ast_Expr *name, Array<Ast_Call_Item *> exprs) {
    // name here specifies also typespec for compound expressions !
    AST_NEW(Call, CALL, pos, AST_EXPR);
    result->name = name;
    result->exprs = exprs.tight_copy(pctx->perm);
    return result;
}

CORE_Static Ast_Call_Item *
ast_call_item(Token *pos, Ast_Atom *name, Ast_Expr *index, Ast_Expr *item) {
    AST_NEW(Call_Item, CALL_ITEM, pos, AST_EXPR);
    result->name = name;
    result->item = item;
    result->index = index;
    return result;
}

CORE_Static Ast_Expr *
ast_expr_unary(Token *pos, Token_Kind op, Ast_Expr *expr) {
    AST_NEW(Unary, UNARY, pos, AST_EXPR);
    result->flags = AST_EXPR;
    result->expr = expr;
    result->op = op;
    return result;
}

CORE_Static Ast_Expr *
ast_expr_index(Token *pos, Ast_Expr *expr, Ast_Expr *index) {
    AST_NEW(Index, INDEX, pos, AST_EXPR);
    result->flags = AST_EXPR;
    result->expr = expr;
    result->index = index;
    return result;
}

CORE_Static Ast_Lambda *
ast_lambda(Token *pos, Array<Ast_Decl *> params, Array<Ast_Expr *> ret, Ast_Scope *scope) {
    AST_NEW(Lambda, LAMBDA_EXPR, pos, AST_EXPR);
    result->flags = AST_EXPR;
    result->args = params.tight_copy(pctx->perm);
    result->ret = ret.tight_copy(pctx->perm);
    result->scope = scope;
    return result;
}

CORE_Static Ast_If *
ast_if(Token *pos, Array<Ast_If_Node *> ifs) {
    AST_NEW(If, IF, pos, AST_STMT);
    result->ifs = ifs.tight_copy(pctx->perm);
    return result;
}

CORE_Static Ast_For *
ast_for(Token *pos, Ast_Expr *init, Ast_Expr *cond, Ast_Expr *iter, Ast_Scope *scope) {
    AST_NEW(For, FOR, pos, AST_STMT);
    result->init = init;
    result->cond = cond;
    result->iter = iter;
    result->scope = scope;
    return result;
}

CORE_Static Ast_Pass *
ast_pass(Token *pos) {
    AST_NEW(Pass, PASS, pos, AST_STMT);
    return result;
}

CORE_Static Ast_Break *
ast_break(Token *pos) {
    AST_NEW(Break, BREAK, pos, AST_STMT);
    return result;
}

CORE_Static Ast_Return *
ast_return(Token *pos, Array<Ast_Expr *> expr) {
    AST_NEW(Return, RETURN, pos, AST_STMT);
    if (expr.len) {
        For(expr) assert(is_flag_set(it->flags, AST_EXPR));
        result->expr = expr.tight_copy(pctx->perm);
    }
    return result;
}

CORE_Static Ast_If_Node *
ast_if_node(Token *pos, Ast_Expr *init, Ast_Expr *expr, Ast_Scope *scope) {
    AST_NEW(If_Node, IF_NODE, pos, AST_STMT);
    result->scope = scope;
    result->expr = expr;
    result->init = (Ast_Binary *)init;
    if (result->init) {
        assert(init->kind == AST_VAR);
    }
    return result;
}

CORE_Static Ast_Array *
ast_array(Token *pos, Ast_Expr *expr) {
    AST_NEW(Array, ARRAY, pos, AST_EXPR);
    result->expr = expr;
    return result;
}

CORE_Static Ast_Scope *
begin_decl_scope(Allocator *scratch, Token *pos) {
    AST_NEW(Scope, SCOPE, pos, AST_DECL);
    result->file = pctx->currently_parsed_file;
    result->module = pctx->currently_parsed_file->module;
    result->scope_id = pctx->scope_ids++;
    result->debug_name = pos->string;
    assert(result->file);
    pctx->currently_parsed_scope = result;
    return result;
}

CORE_Static void
finalize_decl_scope(Ast_Scope *scope) {
    pctx->currently_parsed_scope = scope->parent_scope;
}

CORE_Static Ast_Scope *
begin_stmt_scope(Allocator *scratch, Token *pos) {
    AST_NEW(Scope, SCOPE, pos, AST_STMT);
    result->stmts = {scratch};
    result->file = pctx->currently_parsed_file;
    result->module = pctx->currently_parsed_file->module;
    result->scope_id = pctx->scope_ids++;
    result->debug_name = pos->string;
    assert(result->file);
    pctx->currently_parsed_scope = result;
    return result;
}

CORE_Static void
finalize_stmt_scope(Ast_Scope *scope) {
    scope->stmts = scope->stmts.tight_copy(pctx->perm);
    pctx->currently_parsed_scope = scope->parent_scope;
}

CORE_Static Ast_Decl *
ast_struct(Token *pos, Ast_Scope *scope, Ast_Kind kind, Array<Ast_Decl *> polymorph_parameters) {
    assert(kind == AST_STRUCT || kind == AST_UNION);
    AST_NEW(Decl, STRUCT, pos, AST_DECL | AST_AGGREGATE);
    result->kind = kind;
    result->scope = scope;
    if (polymorph_parameters.len) {
        result->polymorph_parameters = polymorph_parameters;
        set_flag(result->flags, AST_POLYMORPH);
        result->polymorphs.allocator = pctx->heap;
    }
    return result;
}

CORE_Static Ast_Decl *
ast_enum(Token *pos, Ast_Expr *typespec, Ast_Scope *scope) {
    AST_NEW(Decl, ENUM, pos, AST_DECL);
    result->scope = scope;
    result->typespec = typespec;
    set_flag_typespec(typespec);
    return result;
}

CORE_Static Ast_Decl *
ast_var(Token *pos, Ast_Expr *typespec, Intern_String name, Ast_Expr *expr) {
    AST_NEW(Decl, VAR, pos, AST_DECL);
    result->name = name;
    result->typespec = typespec;
    set_flag_typespec(typespec);
    result->expr = expr;
    return result;
}

CORE_Static Ast_Decl *
ast_const(Token *pos, Intern_String name, Value value) {
    AST_NEW(Decl, CONST, pos, AST_DECL);
    result->value = value;
    result->name = name;
    return result;
}

CORE_Static Ast_Decl *
ast_const(Token *pos, Intern_String name, Ast_Expr *expr) {
    AST_NEW(Decl, CONST, pos, AST_DECL);
    result->expr = expr;
    result->name = name;
    return result;
}

CORE_Static Ast_Decl *
ast_type(Token *pos, Intern_String name, Ast_Type *type) {
    AST_NEW(Decl, TYPE, pos, AST_DECL);
    result->type = pctx->type_type;
    result->type_val = type;
    result->name = name;
    return result;
}

CORE_Static Ast_Scope *
ast_decl_scope(Token *pos, Allocator *allocator, Ast_File *file) {
    AST_NEW(Scope, SCOPE, pos, AST_DECL);
    result->file = file;

    result->scope_id = pctx->scope_ids++;
    assert(result->file);
    return result;
}

CORE_Static Ast_Decl *
ast_namespace(Token *pos, Ast_Scope *module, Intern_String name) {
    AST_NEW(Decl, NAMESPACE, pos, AST_DECL);
    result->scope = module;
    result->name = name;
    return result;
}

CORE_Static Ast_Builtin *
ast_runtime_assert(Token *pos, Ast_Expr *expr, Intern_String message) {
    AST_NEW(Builtin, RUNTIME_ASSERT, pos, AST_EXPR);
    result->expr = expr;
    result->assert_message = message;
    return result;
}

CORE_Static Ast_Builtin *
ast_constant_assert(Token *pos, Ast_Expr *expr, Intern_String message) {
    AST_NEW(Builtin, CONSTANT_ASSERT, pos, AST_EXPR);
    result->expr = expr;
    result->assert_message = message;
    return result;
}

CORE_Static Ast_Builtin *
ast_sizeof(Token *pos, Ast_Expr *expr) {
    AST_NEW(Builtin, SIZE_OF, pos, AST_EXPR);
    result->expr = expr;
    return result;
}

CORE_Static Ast_Builtin *
ast_len(Token *pos, Ast_Expr *expr) {
    AST_NEW(Builtin, LENGTH_OF, pos, AST_EXPR);
    result->expr = expr;
    return result;
}

CORE_Static Ast_Builtin *
ast_alignof(Token *pos, Ast_Expr *expr) {
    AST_NEW(Builtin, ALIGN_OF, pos, AST_EXPR);
    result->expr = expr;
    return result;
}

CORE_Static Ast_Var_Unpack *
ast_var_unpack(Token *pos, Array<Ast_Decl *> vars, Ast_Expr *expr) {
    AST_NEW(Var_Unpack, VAR_UNPACK, pos, AST_STMT);
    result->vars = vars.tight_copy(pctx->perm);
    result->expr = expr;
    return result;
}

//
// Value
//

CORE_Static Value
value_bool(B32 v) {
    Value value;
    value.bool_val = v;
    value.type = pctx->untyped_bool;
    return value;
}

CORE_Static Value
value_int(BigInt b) {
    Value value;
    value.big_int_val = b;
    value.type = pctx->untyped_int;
    return value;
}

CORE_Static Value
value_int(S64 s64) {
    Value value;
    value.type = pctx->untyped_int;
    bigint_init_signed(&value.big_int_val, s64);
    return value;
}

CORE_Static Value
value_float(F64 b) {
    Value value;
    value.f64_val = b;
    value.type = pctx->untyped_float;
    return value;
}

CORE_Static Value
value_float(BigInt a) {
    Value value;
    value.f64_val = bigint_as_float(&a);
    value.type = pctx->untyped_float;
    return value;
}

CORE_Static B32
is_ident(Ast *ast) {
    B32 result = ast->kind == AST_IDENT;
    return result;
}

CORE_Static B32
is_binary(Ast *ast) {
    B32 result = ast->kind == AST_BINARY;
    return result;
}

CORE_Static B32
is_atom(Ast *ast) {
    B32 result = is_flag_set(ast->flags, AST_ATOM);
    return result;
}

//
// Iterator
//

const unsigned AST_NODE_END = 128;

struct Ast_Iter {
    Array<Ast *> stack;
    Ast *ast;
    Ast_Kind kind;
    bool skip_children;
    uint32_t visit_id;
    uint32_t di;
};

uint32_t Ast_Iter_VisitIDGen;
Ast_Iter iterate_depth_first(Allocator *a, Ast *ast) {
    assert(ast);
    Ast_Iter result = {};
    result.stack = {a};
    result.ast = ast;
    result.kind = ast->kind;
    result.visit_id = ++Ast_Iter_VisitIDGen;
    return result;
}

void next(Ast_Iter *iter) {
    Ast *ast = iter->ast;
    ast->visit_id = iter->visit_id;

    if (iter->skip_children) {
        iter->skip_children = false;
        goto end_of_switch;
    }

    switch (iter->kind) {
        case AST_SCOPE: {
            Ast_Scope *node = (Ast_Scope *)ast;
            iter->stack.add(node);
            For(node->stmts) iter->stack.add(it);
            For(node->decls) iter->stack.add(it);
        } break;

        case AST_MODULE: break; // This happens when we import stuff
        case AST_FILE: invalid_codepath; break;

        case AST_IDENT:
        case AST_VALUE: {
            Ast_Atom *node = (Ast_Atom *)ast;
        } break;

        case AST_INDEX: {
            Ast_Index *node = (Ast_Index *)ast;
            iter->stack.add(node);
            iter->stack.add(node->index);
            iter->stack.add(node->expr);
            assert(node->index);
            assert(node->expr);
        } break;

        case AST_UNARY: {
            Ast_Unary *node = (Ast_Unary *)ast;
            iter->stack.add(node);
            iter->stack.add(node->expr);
            assert(node->expr);
        } break;

        case AST_BINARY: {
            Ast_Binary *node = (Ast_Binary *)ast;
            iter->stack.add(node);
            iter->stack.add(node->right);
            iter->stack.add(node->left);
            assert(node->right);
            assert(node->left);
        } break;

        case AST_CALL_ITEM: {
            Ast_Call_Item *node = (Ast_Call_Item *)ast;
            iter->stack.add(node);
            iter->stack.add(node->item);
            assert(node->item);
            if (node->call_flags & CALL_INDEX) {
                iter->stack.add(node->index);
                assert(node->index);
            }
            else if (node->call_flags & CALL_NAME) {
                iter->stack.add(node->name);
                assert(node->name);
            }
        } break;

        case AST_COMPOUND:
        case AST_CALL: {
            Ast_Call *node = (Ast_Call *)ast;
            iter->stack.add(node);
            For(node->exprs) iter->stack.add(it);
            if (node->name) iter->stack.add(node->name);
        } break;

        case AST_TYPE_OF:
        case AST_LENGTH_OF:
        case AST_ALIGN_OF:
        case AST_SIZE_OF:
        case AST_RUNTIME_ASSERT:
        case AST_CONSTANT_ASSERT: {
            Ast_Builtin *node = (Ast_Builtin *)ast;
            iter->stack.add(node);
            iter->stack.add(node->expr);
            assert(node->expr);
        } break;

        case AST_SWITCH: {
            Ast_Switch *node = (Ast_Switch *)ast;
            iter->stack.add(node);
            if (node->default_scope) iter->stack.add(node->default_scope);
            For(node->cases) iter->stack.add(it);
            iter->stack.add(node->value);
            assert(node->value);
        } break;

        case AST_SWITCH_CASE: {
            Ast_Switch_Case *node = (Ast_Switch_Case *)ast;
            iter->stack.add(node);
            iter->stack.add(node->scope);
            assert(node->scope);
            For(node->labels) iter->stack.add(it);
        } break;

        case AST_VAR_UNPACK: {
            Ast_Var_Unpack *node = (Ast_Var_Unpack *)ast;
            iter->stack.add(node);
            iter->stack.add(node->expr);
            assert(node->expr);
            For(node->vars) iter->stack.add(it);
        } break;

        case AST_PASS:
        case AST_BREAK: {
            Ast *node = (Ast *)ast;
        } break;

        case AST_NAMESPACE:
        case AST_STRUCT:
        case AST_UNION:
        case AST_ENUM:
        case AST_LAMBDA:
        case AST_TYPE: // @cleanup: what is this used for?
        case AST_CONST:
        case AST_VAR: {
            Ast_Decl *node = (Ast_Decl *)ast;

            iter->stack.add(node);
            if (node->scope) iter->stack.add(node->scope);
            if (node->expr) iter->stack.add(node->expr);
            if (node->typespec) iter->stack.add(node->typespec);

            // omitting polymorphs
            // omitting polymorph parameters
        } break;

        case AST_ARRAY: {
            Ast_Array *node = (Ast_Array *)ast;
            iter->stack.add(node);
            if (node->expr) iter->stack.add(node->expr);
        } break;

        case AST_FOR: {
            Ast_For *node = (Ast_For *)ast;
            iter->stack.add(node);
            iter->stack.add(node->scope);
            assert(node->scope);
            if (node->iter) iter->stack.add(node->iter);
            if (node->cond) iter->stack.add(node->cond);
            if (node->init) iter->stack.add(node->init);
        } break;

        case AST_IF: {
            Ast_If *node = (Ast_If *)ast;
            iter->stack.add(node);
            For(node->ifs) iter->stack.add(it);
        } break;

        case AST_IF_NODE: {
            Ast_If_Node *node = (Ast_If_Node *)ast;
            iter->stack.add(node);
            iter->stack.add(node->scope);
            assert(node->scope);
            if (node->expr) iter->stack.add(node->expr);
            if (node->init) iter->stack.add(node->init);
        } break;

        case AST_RETURN: {
            Ast_Return *node = (Ast_Return *)ast;
            iter->stack.add(node);
            For(node->expr) iter->stack.add(it);
        } break;

        case AST_LAMBDA_EXPR: {
            Ast_Lambda *node = (Ast_Lambda *)ast;
            iter->stack.add(node);
            if (node->scope) iter->stack.add(node->scope);
            For(node->ret) iter->stack.add(it);
            For(node->args) iter->stack.add(it);
        } break;

        default: assert(!"Invalid default case");
    }
end_of_switch:

    if (iter->stack.len <= 0) {
        iter->ast = 0;
        iter->kind = AST_NONE;
        iter->stack.dealloc();
        return;
    }

    iter->ast = iter->stack.pop();
    assert(iter->ast != 0);
    iter->di += 1;
    iter->kind = iter->ast->kind;
    if (iter->ast->visit_id == iter->visit_id) {
        iter->kind = (Ast_Kind)((unsigned)iter->kind + AST_NODE_END);
    }
}