import "@std/mem"
import "@std/arraylist"

import "@source/diagnostics"
import "@source/lexer"
import "@source/ast"

@hide:file
SyntaxError :: enum { unexpected_token }

@hide:file
ParseError :: alias SyntaxError | mem.AllocError | diagnostics.CapacityError

@hide:file
BindingPowerScalar :: distinct u32

@hide:file
BindingPower :: struct {
    left BindingPowerScalar
    right BindingPowerScalar
}

@hide:file
State :: struct {
    builder ast.Builder
    diagnostics @mut diagnostics.Store
    tokens []lexer.Token
    allocator mem.Allocator
    next_token usize
    paren_depth usize = 0
}

#! borrows `tokens` and `source` while parsing; the returned Tree borrows `source` until ast.deinit.
#! `allocator` supplies returned node storage and temporary lists released before return.
#! the caller owns the Tree and must release it with ast.deinit.
parse proc(
    allocator mem.Allocator,
    tokens []lexer.Token,
    source []u8,
    diagnostic_store @mut diagnostics.Store,
) ast.Tree ! (mem.AllocError | diagnostics.CapacityError) {
    state := State{
        builder = ast.init_builder(allocator, source),
        diagnostics = diagnostic_store,
        tokens = tokens,
        allocator = allocator,
        next_token = 0,
    }
    defer ast.deinit_builder(&state.builder)
    root :: try parse_package(&state)
    return ast.finish(&state.builder, root)
}

@hide:file
parse_package proc(state @mut State) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    declarations arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&declarations)

    skip_newlines(state)
    while !next_is(state, .eof) {
        start :: state.next_token
        declaration :: parse_decl(state) catch |err| {
            match err {
                .unexpected_token: {
                    recovered :: try recover_decl(state, start)
                    try arraylist.append(&declarations, recovered)
                    skip_newlines(state)
                    continue
                }
                else: return err
            }
        }
        try arraylist.append(&declarations, declaration)
        skip_newlines(state)
    }

    return try ast.add_package(
        &state.builder,
        ast.Span{ start = 0, end = u32(state.builder.source.len) },
        declarations.items,
    )
}

@hide:file
parse_decl proc(state @mut State) ast.NodeId ! ParseError {
    if (next_is(state, .ident) and next_n_is(state, 1, .proc)) return try parse_func_decl(state)
    return try parse_statement_decl(state)
}

@hide:file
parse_func_decl proc(state @mut State) ast.NodeId ! ParseError {
    name :: try consume(state, .ident)
    _ = try consume(state, .proc)
    _ = try consume(state, .open_paren)

    params arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&params)
    names arraylist.ArrayList(lexer.Token) := arraylist.init(state.allocator)
    defer arraylist.deinit(&names)

    while !next_is(state, .close_paren) {
        arraylist.clear(&names)
        while true {
            token :: try consume(state, .ident)
            try arraylist.append(&names, token)
            if (next_is(state, .ident) or next_is(state, .open_bracket)) break
            _ = try consume(state, .comma)
        }

        type_id :: try parse_type(state)
        type_span :: ast.builder_span(&state.builder, type_id)
        for (names.items) |parameter_name| {
            param_id :: try ast.add_param(
                &state.builder,
                ast.Span{ start = parameter_name.span.start, end = type_span.end },
                parameter_name.span,
                parameter_name.data.str_id,
                type_id,
            )
            try arraylist.append(&params, param_id)
        }
        if next_is(state, .comma) _ = try consume(state, .comma)
    }

    _ = try consume(state, .close_paren)
    return_type :: try parse_type(state)
    body :: try parse_block(state)
    body_span :: ast.builder_span(&state.builder, body)
    return try ast.add_function_decl(
        &state.builder,
        ast.Span{ start = name.span.start, end = body_span.end },
        name.span,
        name.data.str_id,
        return_type,
        params.items,
        body,
    )
}

@hide:file
parse_block proc(state @mut State) ast.NodeId ! ParseError {
    open :: try consume(state, .open_curly)
    statements arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&statements)

    skip_newlines(state)
    while !next_is(state, .close_curly) and !next_is(state, .eof) {
        start :: state.next_token
        statement :: parse_statement(state, false) catch |err| {
            match err {
                .unexpected_token: {
                    recovered :: try recover_statement(state, start)
                    try arraylist.append(&statements, recovered)
                    skip_newlines(state)
                    continue
                }
                else: return err
            }
        }
        try arraylist.append(&statements, statement)
        skip_newlines(state)
    }

    close :: try consume(state, .close_curly)
    return try ast.add_block(
        &state.builder,
        ast.Span{ start = open.span.start, end = close.span.end },
        statements.items,
    )
}

@hide:file
parse_statement proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    if (
        next_is(state, .ident)
        and (
            next_n_is(state, 1, .ident)
            or (next_n_is(state, 1, .open_bracket) and array_starts_type(state))
            or next_n_is(state, 1, .double_colon)
            or next_n_is(state, 1, .colon_equal)
        )
    ) {
        return try parse_statement_decl(state)
    }
    match current(state).kind {
        .if: return try parse_if(state)
        .while: return try parse_while(state, allow_else)
        .for: return try parse_for(state, allow_else)
        .return: return try parse_return(state, allow_else)
        .break, .continue: return try parse_loop_control(state, allow_else)
        .ident, .int, .float, .string, .true, .false, .minus, .bang, .open_paren, .open_bracket: {
            expression :: try parse_expr(state)
            if next_is(state, .equal) return try parse_assignment(state, expression, allow_else)
            return try ast.add_expression_stmt(&state.builder, expression)
        }
        else: return .unexpected_token
    }
}

@hide:file
parse_return proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    keyword :: try consume(state, .return)
    if (
        next_is(state, .newline) or next_is(state, .close_curly) or next_is(state, .eof)
        or (allow_else and next_is(state, .else))
    ) {
        return try ast.add_return_stmt(&state.builder, keyword.span, null)
    }
    value :: try parse_expr(state)
    if (
        !next_is(state, .newline) and !next_is(state, .close_curly) and !next_is(state, .eof)
        and !(allow_else and next_is(state, .else))
    ) {
        return .unexpected_token
    }
    # include consumed grouping punctuation, which has no separate expression node.
    end :: state.tokens[state.next_token - 1].span.end
    return try ast.add_return_stmt(
        &state.builder,
        ast.Span{ start = keyword.span.start, end = end },
        value,
    )
}

@hide:file
parse_loop_control proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    keyword :: current(state)
    state.next_token += 1
    if (
        !next_is(state, .newline)
        and !next_is(state, .close_curly)
        and !next_is(state, .eof)
        and !(allow_else and next_is(state, .else))
    ) return .unexpected_token

    if (keyword.kind == .break)
        return try ast.add_break_stmt(&state.builder, keyword.span)

    return try ast.add_continue_stmt(&state.builder, keyword.span)
}

@hide:file
parse_assignment proc(state @mut State, target ast.NodeId, allow_else bool) ast.NodeId ! ParseError {
    _ = try consume(state, .equal)
    value :: try parse_expr(state)
    if (
        !next_is(state, .newline)
        and !next_is(state, .close_curly)
        and !next_is(state, .eof)
        and !(allow_else and next_is(state, .else))
    ) {
        return .unexpected_token
    }
    end :: state.tokens[state.next_token - 1].span.end
    return try ast.add_assignment_stmt(
        &state.builder,
        ast.Span{ start = ast.builder_span(&state.builder, target).start, end = end },
        target,
        value,
    )
}

@hide:file
parse_if proc(state @mut State) ast.NodeId ! ParseError {
    keyword :: try consume(state, .if)
    condition :: try parse_guard(state)
    braced :: next_is(state, .open_curly)
    then_block :: try parse_arm(state, true)
    # only a braced arm permits its else on a later line.
    after_then :: state.next_token
    if (braced) skip_newlines(state)
    alternative ?ast.NodeId := null
    if next_is(state, .else) {
        state.next_token += 1
        alternative = try parse_arm(state, true)
    } else {
        state.next_token = after_then
    }
    end :: if (alternative) |block|
        ast.builder_span(&state.builder, block).end
    else
        ast.builder_span(&state.builder, then_block).end
    return try ast.add_if_stmt(
        &state.builder,
        ast.Span{ start = keyword.span.start, end = end },
        condition,
        then_block,
        alternative,
    )
}

@hide:file
parse_while proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    keyword :: try consume(state, .while)
    condition :: try parse_guard(state)
    body :: try parse_arm(state, allow_else)
    return try ast.add_while_stmt(
        &state.builder,
        ast.Span{ start = keyword.span.start, end = ast.builder_span(&state.builder, body).end },
        condition,
        body,
    )
}

@hide:file
CaptureList :: struct {
    item ?ast.NodeId = null
    index ?ast.NodeId = null
    pointer bool = false
    indexed bool = false
}

@hide:file
parse_capture_name proc(state @mut State) ?ast.NodeId ! ParseError {
    name :: try consume(state, .ident)
    capture ?ast.NodeId := null
    if (
        name.span.end != name.span.start + 1
        or state.builder.source[usize(name.span.start)] != '_'
    ) capture = try ast.add_identifier(&state.builder, name.span, name.data.str_id)
    return capture
}

@hide:file
parse_captures proc(state @mut State) CaptureList ! ParseError {
    _ = try consume(state, .pipe)
    captures := CaptureList{}
    if next_is(state, .at) {
        captures.pointer = true
        state.next_token += 1
    }
    captures.item = try parse_capture_name(state)
    if next_is(state, .comma) {
        captures.indexed = true
        state.next_token += 1
        captures.index = try parse_capture_name(state)
    }
    _ = try consume(state, .pipe)
    return captures
}

@hide:file
parse_for proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    keyword :: try consume(state, .for)
    skip_newlines(state)
    iterable_start :: state.next_token
    iterable := try parse_expr(state)
    grouped :: parenthesized_guard(state, iterable_start)
    skip_newlines(state)
    captures :: parse_captures(state) catch |err| {
        match err {
            .unexpected_token: {}
            else: return err
        }
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        iterable = try ast.add_invalid_expr(
            &state.builder, ast.builder_span(&state.builder, iterable), diagnostic,
        )
        while (
            !next_is(state, .open_curly) and !next_is(state, .close_curly)
            and !next_is(state, .newline) and !next_is(state, .eof)
            and !next_is(state, .else)
        ) state.next_token += 1
        yield CaptureList{}
    }
    after_captures :: state.next_token
    skip_newlines(state)
    if (
        !next_is(state, .open_curly) and !grouped
        and ast.builder_expr_kind(&state.builder, iterable) != .call
        and ast.builder_expr_kind(&state.builder, iterable) != .invalid
    ) {
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        iterable = try ast.add_invalid_expr(
            &state.builder, ast.builder_span(&state.builder, iterable), diagnostic,
        )
    }
    # a malformed capture list must not swallow the next line's sibling.
    if (
        ast.builder_expr_kind(&state.builder, iterable) == .invalid
        and !next_is(state, .open_curly)
    ) state.next_token = after_captures
    body :: try parse_arm(state, allow_else)
    return try ast.add_for_stmt(
        &state.builder,
        ast.Span{ start = keyword.span.start, end = ast.builder_span(&state.builder, body).end },
        iterable, captures.item, captures.index, captures.pointer, captures.indexed, body,
    )
}

@hide:file
parse_guard proc(state @mut State) ast.NodeId ! ParseError {
    condition_start :: state.next_token
    condition := parse_expr(state) catch |err| {
        match err {
            .unexpected_token: {}
            else: return err
        }
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        yield try ast.add_invalid_expr(&state.builder, current(state).span, diagnostic)
    }
    if (ast.builder_expr_kind(&state.builder, condition) == .invalid) {
        # discard a malformed guard's tail, but leave its body and enclosing block intact.
        while (
            !next_is(state, .open_curly) and !next_is(state, .close_curly)
            and !next_is(state, .newline) and !next_is(state, .eof)
            and !next_is(state, .else) and !next_is(state, .return)
            and !next_is(state, .if) and !next_is(state, .while)
            and !next_is(state, .for)
            and !next_is(state, .break) and !next_is(state, .continue)
        ) {
            if (next_is(state, .ident) and next_n_is(state, 1, .equal)) break
            state.next_token += 1
        }
    }
    if (
        !next_is(state, .open_curly)
        and !parenthesized_guard(state, condition_start)
        and ast.builder_expr_kind(&state.builder, condition) != .call
        and ast.builder_expr_kind(&state.builder, condition) != .invalid
    ) {
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        condition = try ast.add_invalid_expr(
            &state.builder,
            ast.builder_span(&state.builder, condition),
            diagnostic,
        )
    }
    return condition
}

@hide:file
parenthesized_guard proc(state @State, start usize) bool {
    if (state.tokens[start].kind != .open_paren) return false
    depth usize := 0
    for start..state.next_token |index| {
        kind :: state.tokens[index].kind
        if (kind == .open_paren) depth += 1
        if (kind == .close_paren) {
            depth -= 1
            if (depth == 0) return index + 1 == state.next_token
        }
    }
    return false
}

@hide:file
parse_arm proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    if next_is(state, .open_curly) return try parse_block(state)
    start :: state.next_token
    statement :: parse_statement(state, allow_else) catch |err| {
        match err {
            .unexpected_token: {}
            else: return err
        }
        diagnostic :: try unexpected_diagnostic(state)
        span :: synchronize(state, start, true, true, false)
        yield try ast.add_invalid_stmt(&state.builder, span, diagnostic)
    }
    statements [1]ast.NodeId :: [statement]
    return try ast.add_block(
        &state.builder,
        ast.builder_span(&state.builder, statement),
        statements[..],
    )
}


@hide:file
parse_statement_decl proc(state @mut State) ast.NodeId ! ParseError {
    name :: try consume(state, .ident)
    type_id ?ast.NodeId :: if (next_is(state, .ident) or next_is(state, .open_bracket))
        try parse_type(state)
    else
        null

    assignment :: try consume_either(state, &[.double_colon, .colon_equal])
    initializer :: try parse_expr(state)
    end :: ast.builder_span(&state.builder, initializer).end
    mutability :: if (assignment.kind == .double_colon) .constant else .variable
    return try ast.add_statement_decl(
        &state.builder,
        ast.Span{ start = name.span.start, end = end },
        name.span,
        name.data.str_id,
        type_id,
        mutability,
        initializer,
    )
}

@hide:file
parse_type proc(state @mut State) ast.NodeId ! ParseError {
    if next_is(state, .open_bracket) {
        open :: try consume(state, .open_bracket)
        count ?ast.NodeId := null
        if (
            next_is(state, .ident)
            and current(state).span.end == current(state).span.start + 1
            and state.builder.source[usize(current(state).span.start)] == '_'
        ) {
            state.next_token += 1
        } else {
            count = try parse_expr(state)
        }
        _ = try consume(state, .close_bracket)
        mutable :: next_is(state, .mut)
        if (mutable) state.next_token += 1
        element :: try parse_type(state)
        return try ast.add_array_type(
            &state.builder,
            ast.Span{ start = open.span.start, end = ast.builder_span(&state.builder, element).end },
            count, element, mutable,
        )
    }
    token :: try consume(state, .ident)
    return try ast.add_identifier(&state.builder, token.span, token.data.str_id)
}

@hide:file
parse_expr proc(state @mut State) ast.NodeId ! ParseError {
    return try parse_expr_bp(state, BindingPowerScalar(0))
}

@hide:file
parse_expr_bp proc(state @mut State, min_bp BindingPowerScalar) ast.NodeId ! ParseError {
    skip_operand_newlines(state)
    token := current(state)
    expression_start :: state.next_token
    left := if (prefix_binding_power(token.kind)) |right_bp| {
        state.next_token += 1
        operand :: try parse_expr_bp(state, right_bp)
        operand_span :: ast.builder_span(&state.builder, operand)
        operator ast.UnaryOperator :: if (token.kind == .bang) .logical_not else .negate
        yield try ast.add_unary(
            &state.builder,
            ast.Span{ start = token.span.start, end = operand_span.end },
            operator,
            operand,
        )
    } else {
        yield try parse_primary(state)
    }
    left_grouped := parenthesized_guard(state, expression_start)

    while true {
        skip_continuation_newlines(state)
        token = current(state)
        if (token.kind == .open_paren) {
            left = try parse_call(state, left, state.tokens[expression_start].span.start)
            left_grouped = false
            continue
        }
        if (token.kind == .open_bracket) {
            left = try parse_index(state, left, state.tokens[expression_start].span.start)
            left_grouped = false
            continue
        }
        if (token.kind == .range or token.kind == .range_inclusive) {
            if (min_bp > BindingPowerScalar(0)) break
            left = try validate_range_bound(state, left, left_grouped)
            state.next_token += 1
            right_start :: state.next_token
            right := try parse_expr_bp(state, BindingPowerScalar(1))
            right = try validate_range_bound(state, right, parenthesized_guard(state, right_start))
            left = try ast.add_range(
                &state.builder,
                ast.Span{
                    start = state.tokens[expression_start].span.start,
                    end = state.tokens[state.next_token - 1].span.end,
                },
                left, right, token.kind == .range_inclusive,
            )
            left_grouped = false
            continue
        }
        if (infix_binding_power(token.kind)) |power| {
            if (power.left < min_bp) break
            state.next_token += 1
            right :: try parse_expr_bp(state, power.right)
            left_span :: ast.builder_span(&state.builder, left)
            right_span :: ast.builder_span(&state.builder, right)
            left = try ast.add_binary(
                &state.builder,
                ast.Span{ start = left_span.start, end = right_span.end },
                binary_operator(token.kind),
                left,
                right,
            )
            left_grouped = false
            continue
        }
        break
    }
    return left
}

@hide:file
validate_range_bound proc(
    state @mut State, bound ast.NodeId, grouped bool,
) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    if (grouped) return bound
    match ast.builder_expr_kind(&state.builder, bound) {
        .unary, .binary, .range: {
            span :: ast.builder_span(&state.builder, bound)
            diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, span)
            return try ast.add_invalid_expr(&state.builder, span, diagnostic)
        }
        else: return bound
    }
}

@hide:file
parse_call proc(state @mut State, callee ast.NodeId, start u32) ast.NodeId ! ParseError {
    _ = try consume(state, .open_paren)
    state.paren_depth += 1
    defer state.paren_depth -= 1
    skip_operand_newlines(state)
    arguments arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&arguments)

    while (!is_call_boundary(current(state).kind)) {
        argument :: try parse_expr(state)
        try arraylist.append(&arguments, argument)
        if (!next_is(state, .comma)) break
        state.next_token += 1
        skip_operand_newlines(state)
    }
    close :: try consume(state, .close_paren)
    return try ast.add_call(
        &state.builder,
        ast.Span{
            start = start,
            end = close.span.end,
        },
        callee,
        arguments.items,
    )
}

@hide:file
parse_index proc(state @mut State, base ast.NodeId, start u32) ast.NodeId ! ParseError {
    _ = try consume(state, .open_bracket)
    index :: try parse_expr(state)
    if !next_is(state, .close_bracket) {
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        return try ast.add_invalid_expr(
            &state.builder,
            ast.Span{
                start = start,
                end = ast.builder_span(&state.builder, index).end,
            },
            diagnostic,
        )
    }
    close :: try consume(state, .close_bracket)
    return try ast.add_index(
        &state.builder,
        ast.Span{ start = start, end = close.span.end },
        base, index,
    )
}

@hide:file
array_starts_type proc(state @State) bool {
    depth usize := 0
    for (state.next_token)..(state.tokens.len) |i| {
        kind :: state.tokens[i].kind
        if (kind == .open_bracket) depth += 1
        if (kind == .close_bracket) {
            depth -= 1
            if (depth == 0) {
                if (i + 1 >= state.tokens.len) return false
                next :: state.tokens[i + 1].kind
                return next == .ident or next == .mut or next == .open_bracket
            }
        }
        if (kind == .open_curly or kind == .close_curly or kind == .eof) return false
    }
    return false
}

@hide:file
parse_array_literal proc(state @mut State) ast.NodeId ! ParseError {
    open :: try consume(state, .open_bracket)
    elements arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&elements)
    after_open :: state.next_token
    skip_newlines(state)
    if named_binding_starts(state) {
        state.next_token = after_open
        return try invalid_array_literal(state, open.span.start, open.span.end)
    }
    while !next_is(state, .close_bracket) {
        element :: try parse_expr(state)
        try arraylist.append(&elements, element)
        after_element :: state.next_token
        skip_newlines(state)
        if next_is(state, .close_bracket) break
        if !next_is(state, .comma) {
            state.next_token = after_element
            return try invalid_array_literal(state, open.span.start, ast.builder_span(&state.builder, element).end)
        }
        state.next_token += 1
        after_comma :: state.next_token
        skip_newlines(state)
        if named_binding_starts(state) {
            state.next_token = after_comma
            return try invalid_array_literal(state, open.span.start, ast.builder_span(&state.builder, element).end)
        }
    }
    close :: try consume(state, .close_bracket)
    return try ast.add_array_literal(
        &state.builder, ast.Span{ start = open.span.start, end = close.span.end }, elements.items,
    )
}

@hide:file
invalid_array_literal proc(state @mut State, start, end u32) ast.NodeId ! ParseError {
    diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
    return try ast.add_invalid_expr(&state.builder, ast.Span{ start = start, end = end }, diagnostic)
}

@hide:file
named_binding_starts proc(state @State) bool {
    return (
        next_is(state, .ident)
        and (
            next_n_is(state, 1, .double_colon) or next_n_is(state, 1, .colon_equal)
            or next_n_is(state, 1, .ident) or next_n_is(state, 1, .proc)
        )
    )
}

@hide:file
parse_primary proc(state @mut State) ast.NodeId ! ParseError {
    token :: current(state)
    match token.kind {
        .open_bracket: {
            if array_starts_type(state) return try parse_type(state)
            return try parse_array_literal(state)
        }
        .open_paren: {
            state.next_token += 1
            state.paren_depth += 1
            defer state.paren_depth -= 1
            expression :: try parse_expr(state)
            if next_is(state, .close_paren) {
                state.next_token += 1
                return expression
            }
            diagnostic :: try diagnostics.add(
                state.diagnostics,
                .unexpected_token,
                current(state).span,
            )
            return try ast.add_invalid_expr(&state.builder, ast.Span{
                start = token.span.start,
                end = ast.builder_span(&state.builder, expression).end,
            }, diagnostic)
        }
        .int: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .int, token.span)
        }
        .float: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .float, token.span)
        }
        .string: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .string, token.span)
        }
        .true, .false: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .bool, token.span)
        }
        .ident: {
            state.next_token += 1
            return try ast.add_identifier(&state.builder, token.span, token.data.str_id)
        }
        else: {
            diagnostic :: if (token.data.diagnostic != diagnostics.no_diagnostic)
                token.data.diagnostic
            else
                try diagnostics.add(state.diagnostics, .unexpected_token, token.span)

            if !is_expression_recovery_boundary(token.kind) {
                state.next_token += 1
            }

            return try ast.add_invalid_expr(&state.builder, token.span, diagnostic)
        }
    }
}

@hide:file
is_expression_boundary proc(kind lexer.TokenKind) bool {
    return match kind {
        .close_paren, .newline, .close_curly, .open_curly, .else, .return, .if, .while, .for, .eof: true
        else: false
    }
}

@hide:file
is_call_boundary proc(kind lexer.TokenKind) bool {
    return (
        kind == .break
        or kind == .continue
        or is_expression_boundary(kind)
    )
}

@hide:file
is_expression_recovery_boundary proc(kind lexer.TokenKind) bool {
    return (
        kind == .comma
        or kind == .close_bracket
        or kind == .and
        or kind == .or
        or kind == .pipe
        or is_expression_boundary(kind)
    )
}

@hide:file
skip_operand_newlines proc(state @mut State) void {
    if (state.paren_depth == 0 or !next_is(state, .newline)) return
    start :: state.next_token
    skip_newlines(state)
    kind :: current(state).kind
    if (kind != .close_paren and is_expression_boundary(kind)) state.next_token = start
    # a missing operand must not consume the declaration or assignment on the next line.
    if (
        named_binding_starts(state)
        or (next_is(state, .ident) and next_n_is(state, 1, .equal))
    ) state.next_token = start
}

@hide:file
skip_continuation_newlines proc(state @mut State) void {
    if (state.paren_depth == 0 or !next_is(state, .newline)) return
    start :: state.next_token
    skip_newlines(state)
    kind :: current(state).kind
    if (
        kind != .close_paren and kind != .comma and kind != .open_paren
        and kind != .close_bracket and kind != .open_bracket
        and kind != .range and kind != .range_inclusive
        and infix_binding_power(kind) == null
    ) state.next_token = start
}

@hide:file
recover_decl proc(
    state @mut State,
    start usize,
) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    diagnostic :: try unexpected_diagnostic(state)
    span :: synchronize(state, start, false, false, true)
    return try ast.add_invalid_decl(&state.builder, span, diagnostic)
}

@hide:file
recover_statement proc(
    state @mut State,
    start usize,
) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    diagnostic :: try unexpected_diagnostic(state)
    span :: synchronize(state, start, true, false, true)
    return try ast.add_invalid_stmt(&state.builder, span, diagnostic)
}

@hide:file
synchronize proc(
    state @mut State, start usize, stop_at_curly, allow_else, consume_newline bool,
) ast.Span {
    first :: state.tokens[start].span
    end := first.end
    curly usize := 0
    paren usize := 0
    square usize := 0
    for (start..state.next_token) |i| {
        kind :: state.tokens[i].kind
        if (kind == .open_paren) paren += 1
        if (kind == .close_paren and paren > 0) paren -= 1
        if (kind == .open_bracket) square += 1
        if (kind == .close_bracket and square > 0) square -= 1
    }
    while !next_is(state, .eof) {
        if (curly == 0) {
            if (stop_at_curly and next_is(state, .close_curly)) break
            if (allow_else and next_is(state, .else)) break
            if (
                next_is(state, .newline)
                and !recovery_closes_delimiters(state, paren, square)
            ) break
        }
        kind :: current(state).kind
        if (kind == .open_curly) curly += 1
        if (kind == .close_curly and curly > 0) curly -= 1
        if (kind == .open_paren) paren += 1
        if (kind == .close_paren and paren > 0) paren -= 1
        if (kind == .open_bracket) square += 1
        if (kind == .close_bracket and square > 0) square -= 1
        end = current(state).span.end
        state.next_token += 1
    }
    if (consume_newline and next_is(state, .newline)) {
        end = current(state).span.end
        state.next_token += 1
    }
    return ast.Span{ start = first.start, end = end }
}

@hide:file
recovery_closes_delimiters proc(state @State, paren_depth, square_depth usize) bool {
    if (paren_depth == 0 and square_depth == 0) return false
    paren := paren_depth
    square := square_depth
    for (state.next_token)..(state.tokens.len) |i| {
        kind :: state.tokens[i].kind
        match kind {
            .open_curly, .close_curly, .eof, .if, .else, .while, .for,
            .return, .break, .continue: return false
            .ident: {
                if (i + 1 < state.tokens.len) {
                    next :: state.tokens[i + 1].kind
                    if (next == .double_colon or next == .colon_equal or next == .ident or next == .proc) {
                        return false
                    }
                }
            }
            else: {}
        }
        if (kind == .open_paren) paren += 1
        if (kind == .close_paren and paren > 0) paren -= 1
        if (kind == .open_bracket) square += 1
        if (kind == .close_bracket and square > 0) square -= 1
        if (paren == 0 and square == 0) return true
    }
    return false
}

@hide:file
unexpected_diagnostic proc(state @mut State) diagnostics.DiagnosticId ! (mem.AllocError | diagnostics.CapacityError) {
    token :: current(state)
    if (token.kind == .invalid and token.data.diagnostic != diagnostics.no_diagnostic) {
        return token.data.diagnostic
    }
    return try diagnostics.add(state.diagnostics, .unexpected_token, token.span)
}

@hide:file
skip_newlines proc(state @mut State) void {
    while next_is(state, .newline) state.next_token += 1
}

@hide:file
current proc(state @State) lexer.Token {
    if (state.next_token >= state.tokens.len) return state.tokens[state.tokens.len - 1]
    return state.tokens[state.next_token]
}

@hide:file
next_is proc(state @State, kind lexer.TokenKind) bool {
    return current(state).kind == kind
}

@hide:file
next_n_is proc(state @State, count usize, kind lexer.TokenKind) bool {
    index :: state.next_token + count
    return index < state.tokens.len and state.tokens[index].kind == kind
}

@hide:file
consume proc(state @mut State, kind lexer.TokenKind) lexer.Token ! SyntaxError {
    token :: current(state)
    if (token.kind != kind) return .unexpected_token
    state.next_token += 1
    return token
}

@hide:file
consume_either proc(state @mut State, kinds []lexer.TokenKind) lexer.Token ! SyntaxError {
    token :: current(state)
    for (kinds) |kind| if token.kind == kind {
        state.next_token += 1
        return token
    }
    return .unexpected_token
}

@hide:file
prefix_binding_power proc(kind lexer.TokenKind) ?BindingPowerScalar {
    return match kind {
        .minus, .bang: BindingPowerScalar(50)
        else: null
    }
}

@hide:file
infix_binding_power proc(kind lexer.TokenKind) ?BindingPower {
    return match kind {
        .or: BindingPower{ left = BindingPowerScalar(2), right = BindingPowerScalar(3) }
        .and: BindingPower{ left = BindingPowerScalar(4), right = BindingPowerScalar(5) }
        .equal_equal, .not_equal: BindingPower{ left = BindingPowerScalar(10), right = BindingPowerScalar(11) }
        .less, .less_equal, .greater, .greater_equal: BindingPower{ left = BindingPowerScalar(20), right = BindingPowerScalar(21) }
        .plus, .minus: BindingPower{ left = BindingPowerScalar(30), right = BindingPowerScalar(31) }
        .star, .slash: BindingPower{ left = BindingPowerScalar(40), right = BindingPowerScalar(41) }
        else: null
    }
}

@hide:file
binary_operator proc(kind lexer.TokenKind) ast.BinaryOperator {
    return match kind {
        .and: .logical_and
        .or: .logical_or
        .plus: .add
        .minus: .subtract
        .star: .multiply
        .slash: .divide
        .equal_equal: .equal
        .not_equal: .not_equal
        .less: .less
        .less_equal: .less_equal
        .greater: .greater
        .greater_equal: .greater_equal
        else: unreachable
    }
}