intellijai/src/main/kotlin/dev/hgh/jai/parser/JaiParserUtil.kt

package dev.hgh.jai.parser

import com.intellij.lang.PsiBuilder
import com.intellij.lang.parser.GeneratedParserUtilBase
import com.intellij.psi.TokenType
import com.intellij.psi.tree.IElementType
import dev.hgh.jai.lexer.JaiTokenTypes
import dev.hgh.jai.lexer.JaiTokenTypes as T

/**
 * Hand-written helpers referenced from `Jai.bnf` as external rules.
 *
 * These exist because a few Jai constructs are not expressible as pure BNF over our
 * token set:
 *
 *  - `#ident` is a single DIRECTIVE token, so "is this `#if`?" is a *text* test.
 *  - `if x == { case ...; }` reuses `==` as a switch introducer, so the expression
 *    parser must refuse `==` when a `{` follows.
 *  - `(` starts both a parenthesised expression and a procedure header.
 *  - A declaration has no introducer keyword, so `foo :: ...` vs `foo(...)` needs
 *    lookahead.
 *
 * Everything here is pure lookahead except the `dir*` family, which consume exactly
 * one DIRECTIVE token.
 */
object JaiParserUtil : GeneratedParserUtilBase() {
    /**
     * Directives whose operand is a brace-delimited block, e.g. `#asm { ... }`.
     * Tried before [DIRECTIVES_WITH_OPERAND] so `#run { ... }` takes the block form
     * while `#run foo()` falls through to the expression form.
     */
    private val DIRECTIVES_WITH_BLOCK =
        setOf(
            "asm",
            "modify",
            "run",
            "insert",
            "code",
            "run_and_insert",
            "compile_time",
            "bytes",
            "no_reset",
        )

    /**
     * Directives that take a single expression operand, e.g. `#import "Basic"`.
     * The operand is optional: `#foreign` appears both bare and as `#foreign lib`.
     * Anything not listed here is parsed as a bare directive so that, for example,
     * `#scope_file` cannot swallow the declaration that follows it.
     */
    private val DIRECTIVES_WITH_OPERAND =
        setOf(
            "import",
            "load",
            "run",
            "assert",
            "insert",
            "code",
            "type",
            "char",
            "library",
            "system_library",
            "foreign_library",
            "foreign_system_library",
            "align",
            "placeholder",
            "placeholders",
            "bake",
            "bake_arguments",
            "bake_constants",
            "deprecated",
            "add_context",
            "module_parameters",
            "place",
            "overlay",
            "procedure_of_call",
            "body_text",
            "version",
            "poke_name",
            "dump",
            "foreign",
            "define",
            "include",
            "undef",
            "ifdef",
            "pragma",
            "elsewhere",
            "expand",
            "specified",
            "location",
            "this",
        )

    // ---------------------------------------------------------------- directives

    /**
     * Consumes the DIRECTIVE token if its text matches exactly.
     *
     * [text] includes the `#`: Grammar-Kit rewrites a `<<dir "if">>` argument into a
     * reference to the `if` *token* because `'if'` is a declared token text, so the
     * grammar has to say `<<dir "#if">>`.
     */
    @JvmStatic
    fun dir(
        b: PsiBuilder,
        level: Int,
        text: String,
    ): Boolean {
        if (b.tokenType !== T.DIRECTIVE) return false
        if (b.tokenText != text) return false
        b.advanceLexer()
        return true
    }

    /** True (without consuming) if the current token is the directive [text] (with `#`). */
    @JvmStatic
    fun dirAhead(
        b: PsiBuilder,
        level: Int,
        text: String,
    ): Boolean = b.tokenType === T.DIRECTIVE && b.tokenText == text

    @JvmStatic
    fun dirWithBlock(
        b: PsiBuilder,
        level: Int,
    ): Boolean = consumeDirectiveIn(b, DIRECTIVES_WITH_BLOCK)

    @JvmStatic
    fun dirWithOperand(
        b: PsiBuilder,
        level: Int,
    ): Boolean = consumeDirectiveIn(b, DIRECTIVES_WITH_OPERAND)

    /** Any DIRECTIVE token; used for the bare `#complete` / `#scope_file` form. */
    @JvmStatic
    fun dirPlain(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        if (b.tokenType !== T.DIRECTIVE) return false
        b.advanceLexer()
        return true
    }

    /**
     * A directive that takes neither an operand nor a block, so it stands alone as a
     * statement without a `;` — `#scope_file`, `#no_padding`, `#compiler`.
     */
    @JvmStatic
    fun dirBare(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        if (b.tokenType !== T.DIRECTIVE) return false
        val text = b.tokenText ?: return false
        val name = text.removePrefix("#")
        if (name in DIRECTIVES_WITH_OPERAND || name in DIRECTIVES_WITH_BLOCK) return false
        b.advanceLexer()
        return true
    }

    private fun consumeDirectiveIn(
        b: PsiBuilder,
        names: Set<String>,
    ): Boolean {
        if (b.tokenType !== T.DIRECTIVE) return false
        val text = b.tokenText ?: return false
        if (text.length < 2 || text.substring(1) !in names) return false
        b.advanceLexer()
        return true
    }

    /**
     * A `,flag` suffix on a directive: `#import,file`, `#library,system,link_always`,
     * `#run,stallable`.
     *
     * Adjacency is required. Without it, in
     * `(callback: (*GUID) #c_call, lpContext: *void)` the flag would swallow
     * `, lpContext` and the parameter list would fall apart. Every flag in the corpus
     * is written tight against the directive; a separator comma always has a space.
     */
    @JvmStatic
    fun directiveFlag(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        if (b.tokenType !== T.COMMA) return false
        if (b.rawLookup(-1) === TokenType.WHITE_SPACE) return false
        if (b.rawLookup(1) !== T.IDENT) return false
        b.advanceLexer()
        b.advanceLexer()
        return true
    }

    /** `name =` inside an argument list, where `name` may be a keyword (`remove=`). */
    @JvmStatic
    fun isNamedArgumentAhead(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        val t = b.tokenType ?: return false
        if (t !== T.IDENT && t !in JaiTokenTypes.KEYWORDS) return false
        return b.lookAhead(1) === T.EQ
    }

    // ------------------------------------------------------------------ operators

    /**
     * An expression that stops short of `=`.
     *
     * Needed wherever a *type* is followed by `= value`: in `a : float = 37.0` the
     * type slot must yield `float`, not the assignment `float = 37.0`. Grammar-Kit
     * only ever calls the expression root at the lowest priority from BNF, so this
     * re-enters it one level up.
     */
    @JvmStatic
    fun exprNoAssign(
        b: PsiBuilder,
        level: Int,
    ): Boolean = JaiParser.expr(b, level, ASSIGNMENT_PRIORITY)

    /**
     * Consumes `==` unless a `{` follows it, because `if x == { case ...; }` is the
     * switch form and the `==` belongs to the statement, not to an equality expression.
     */
    @JvmStatic
    fun eqEqNotSwitch(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        if (b.tokenType !== T.EQ_EQ) return false
        if (b.lookAhead(1) === T.LBRACE) return false
        b.advanceLexer()
        return true
    }

    /**
     * True if the previous token already closes a construct, so no `;` is needed:
     * `}` ends a block, struct or enum body, a here-string ends on its own terminator
     * line (`BODY :: #string DONE … DONE`), and a `;` was already taken by a nested
     * statement (`code :: #code a := 1;`).
     */
    @JvmStatic
    fun prevEndsConstruct(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        var i = -1
        while (true) {
            val t: IElementType = b.rawLookup(i) ?: return false
            if (t !== TokenType.WHITE_SPACE && t !== T.LINE_COMMENT && t !== T.BLOCK_COMMENT) {
                return t === T.RBRACE || t === T.HERE_STRING || t === T.SEMICOLON
            }
            i--
        }
    }

    // ------------------------------------------------------------------ lookahead

    /**
     * `(`…`)` followed by `->`, a directive, or a `{` body, where the parameter list
     * is either empty or contains a `:`/`$` at depth 1. Distinguishes a procedure
     * header from `if (x) { }`, whose parenthesised condition looks similar.
     */
    @JvmStatic
    fun isProcHeaderAhead(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        if (b.tokenType !== T.LPAREN) return false
        var depth = 0
        var i = 0
        var sawParamMarker = false
        var empty = true
        while (true) {
            val t = b.lookAhead(i) ?: return false
            when (t) {
                T.LPAREN, T.LBRACKET, T.BEGIN_ARRAY_LITERAL, T.LBRACE, T.BEGIN_STRUCT_LITERAL -> {
                    depth++
                }

                T.RPAREN -> {
                    depth--
                    if (depth == 0) {
                        // A `name:` or `$T` inside the parentheses can only be a
                        // parameter list, so the header stands on its own: a bare
                        // proc *type* like `proc: (info: *Info, data: T)` has no
                        // arrow, directive or body after it.
                        if (sawParamMarker) return true
                        val next = b.lookAhead(i + 1)
                        return when (next) {
                            T.RIGHT_ARROW, T.DIRECTIVE -> true
                            T.LBRACE -> empty
                            else -> false
                        }
                    }
                }

                T.RBRACKET, T.RBRACE -> {
                    depth--
                }

                else -> {
                    if (depth == 1) {
                        empty = false
                        if (t === T.COLON || t === T.DOLLAR || t === T.DOUBLE_DOLLAR ||
                            t === T.COLON_COLON || t === T.COLON_EQ
                        ) {
                            sawParamMarker = true
                        }
                    }
                }
            }
            i++
            if (i > MAX_LOOKAHEAD) return false
        }
    }

    /** `(` that starts a parenthesised *return list*, i.e. not a procedure type. */
    @JvmStatic
    fun isReturnGroupAhead(
        b: PsiBuilder,
        level: Int,
    ): Boolean = b.tokenType === T.LPAREN && !isProcHeaderAhead(b, level)

    /**
     * `[#directive [(…)] | using | `]* IDENT (',' IDENT)* (':' | '::' | ':=')` — a
     * declaration. Jai declarations have no introducer keyword, so this is what tells
     * `foo :: () {}` apart from `foo();`.
     *
     * The prefix skip covers `#as using base: Document;`,
     * `#overlay (unknown.vtable) using vtable: *Vtable;`,
     * `#as using,except(vtable) iunknown: IUnknown;` and the macro-exported
     * `` `it := entry.value; ``.
     */
    @JvmStatic
    fun isDeclarationAhead(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        var i = 0
        var inPrefix = false
        loop@ while (true) {
            when (b.lookAhead(i)) {
                T.DIRECTIVE, T.KW_USING -> {
                    i++
                    inPrefix = true
                }

                T.LPAREN, T.LBRACKET, T.BEGIN_ARRAY_LITERAL, T.BEGIN_STRUCT_LITERAL -> {
                    i = skipBalancedGroup(b, i)
                    if (i < 0) return false
                }

                // Only a `,flag` belonging to the prefix, as in `using,except(x)`.
                // A comma between declared names is handled by the loop below.
                T.COMMA -> {
                    if (!inPrefix || b.lookAhead(i + 1) !== T.IDENT) break@loop
                    i += 2
                    inPrefix = false
                }

                else -> {
                    break@loop
                }
            }
        }
        while (true) {
            if (b.lookAhead(i) === T.BACKTICK) i++
            if (b.lookAhead(i) !== T.IDENT) return false
            i++
            when (b.lookAhead(i)) {
                T.COLON, T.COLON_COLON, T.COLON_EQ -> return true
                T.COMMA -> i++
                else -> return false
            }
        }
    }

    /**
     * Index just past the closer matching the bracket at [start], or -1 if unbalanced.
     * Covers `(…)`, `[…]`, `.[…]` and `.{…}`, which all appear in declaration
     * prefixes: `#overlay (x) using …`, `using,except .["x", "y"] …`.
     */
    private fun skipBalancedGroup(
        b: PsiBuilder,
        start: Int,
    ): Int {
        var depth = 0
        var i = start
        while (i - start <= MAX_LOOKAHEAD) {
            when (b.lookAhead(i) ?: return -1) {
                T.LPAREN, T.LBRACKET, T.BEGIN_ARRAY_LITERAL, T.LBRACE, T.BEGIN_STRUCT_LITERAL -> {
                    depth++
                }

                T.RPAREN, T.RBRACKET, T.RBRACE -> {
                    depth--
                    if (depth == 0) return i + 1
                }

                else -> {}
            }
            i++
        }
        return -1
    }

    /**
     * Consumes a `{ … }` block as opaque tokens.
     *
     * Used for `#asm`, whose body is x86-64 assembly (`movd source:, byte;`) and not
     * Jai at all. Giving it structure is a separate job; tiling it keeps the
     * surrounding file parseable.
     */
    @JvmStatic
    fun opaqueBlock(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        if (b.tokenType !== T.LBRACE) return false
        var depth = 0
        while (!b.eof()) {
            val t = b.tokenType
            if (t === T.LBRACE) {
                depth++
            } else if (t === T.RBRACE) {
                depth--
                if (depth == 0) {
                    b.advanceLexer()
                    return true
                }
            }
            b.advanceLexer()
        }
        return false
    }

    /** `[$|$$] IDENT (',' [$|$$] IDENT)* (':' | '::' | ':=')` — a named parameter. */
    @JvmStatic
    fun isParamNameAhead(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        var i = 0
        while (true) {
            var t = b.lookAhead(i)
            if (t === T.DOLLAR || t === T.DOUBLE_DOLLAR) {
                i++
                t = b.lookAhead(i)
            }
            if (t !== T.IDENT) return false
            i++
            when (b.lookAhead(i)) {
                T.COLON, T.COLON_COLON, T.COLON_EQ -> return true
                T.COMMA -> i++
                else -> return false
            }
        }
    }

    /** `IDENT (',' IDENT)* ':'` — the named-variable prefix of a `for` header. */
    @JvmStatic
    fun isForNamesAhead(
        b: PsiBuilder,
        level: Int,
    ): Boolean {
        var i = 0
        while (true) {
            if (b.lookAhead(i) === T.BACKTICK) i++
            if (b.lookAhead(i) !== T.IDENT) return false
            i++
            when (b.lookAhead(i)) {
                T.COLON -> return true
                T.COMMA -> i++
                else -> return false
            }
        }
    }

    /** `IDENT ':'` — a loop label. */
    @JvmStatic
    fun isLabelAhead(
        b: PsiBuilder,
        level: Int,
    ): Boolean = b.tokenType === T.IDENT && b.lookAhead(1) === T.COLON

    private const val MAX_LOOKAHEAD = 4096

    /** Priority of `assignExpr` in the generated expression parser; see `Jai.bnf`. */
    private const val ASSIGNMENT_PRIORITY = 0
}