formatter.kt

fun blockifyLineCommentRuns(code: String): String {
    // If ktfmt handed us something we can't re-parse, the transform simply doesn't apply.
    val file = try {
        FieldKtParser.parse(code)
    } catch (t: Throwable) {
        return code
    }

    val ownLine = file.collectDescendantsOfType<PsiComment>()
        .filter { it.tokenType == KtTokens.EOL_COMMENT }
        .sortedBy { it.startOffset }
        .mapNotNull { comment ->
            // Fragment markers are execution boundaries. Leave them as line comments,
            // separating prose runs on either side so the splitter still sees them.
            if (isFragmentBreak(comment.text)) return@mapNotNull null
            // A `// {{ name }}` head marker is read as a line comment (KotlinJsHeadProfiles.resolve).
            if (HEAD_MARKER.matches(comment.text)) return@mapNotNull null
            val start = comment.startOffset
            val lineStart = code.lastIndexOf('\n', start - 1) + 1
            val indent = code.substring(lineStart, start)
            // Anything other than leading whitespace makes this a trailing comment on a line of
            // code, which stays as it is.
            if (indent.isNotBlank()) return@mapNotNull null
            LineCommentLine(comment.endOffset, lineStart, indent, comment.text.removePrefix("//"))
        }

    val runs = mutableListOf<MutableList<LineCommentLine>>()
    ownLine.forEach { line ->
        val open = runs.lastOrNull()
        val previous = open?.lastOrNull()
        // Exactly one newline between them: same indentation, no blank line, nothing else.
        val adjacent = previous != null &&
            previous.indent == line.indent &&
            code.substring(previous.end, line.lineStart) == "\n"
        if (open != null && adjacent) open.add(line) else runs.add(mutableListOf(line))
    }

    // Back to front, so the offsets of the runs we haven't rewritten yet stay valid.
    val out = StringBuilder(code)
    runs.asReversed()
        .filter { it.size >= MINIMUM_RUN }
        .filter { run -> run.none { it.body.contains("/*") || it.body.contains("*/") } }
        .forEach { run ->
            val indent = run.first().indent
            val block = buildString {
                append(indent).append("/*\n")
                run.forEach { line ->
                    val text = line.body.removePrefix(" ").trimEnd()
                    // An empty line stays empty rather than becoming an indent's worth of spaces.
                    if (text.isEmpty()) append('\n') else append(indent).append(text).append('\n')
                }
                append(indent).append("*/")
            }
            out.replace(run.first().lineStart, run.last().end, block)
        }
    return out.toString()
}

private val glslProperties = setOf(
    "fragment",
    "vertex",
    "compute",
    "task",
    "mesh",
    "raygen",
    "miss",
    "closesthit",
    "anyhit",
    "intersection",
    "callable"
)

fun formatGlslCode(code: String): String {
    // Keep literal text and comment layout out of spacing/statement rewriting. In particular,
    // shader headers often contain aligned tables, and string contents are not GLSL whitespace.
    var marker = "FIELD_GLSL_PROTECTED_"
    while (code.contains(marker)) marker = "_" + marker
    val protected = mutableListOf<Pair<String, String>>()
    val masked = GLSL_PROTECTED.replace(code) { match ->
        val text = match.value
        val id = marker + protected.size + "_END"
        val token = when {
            text.startsWith("//") -> "//$id"
            text.startsWith("/*") -> "/*$id*/"
            else -> "\"$id\""
        }
        protected += token to text
        token
    }
    val normalized = splitGlslStatements(masked)
    val out = StringBuilder()
    var indent = 0
    var previousBlank = false
    var previousLine = ""

    normalized.lineSequence().forEach { raw ->
        val line = raw.trim()
        if (line.isEmpty()) {
            if (!previousBlank && out.isNotEmpty()) out.appendLine()
            previousBlank = true
            return@forEach
        }
        previousBlank = false
        val structure = line.replace(Regex("""/\*.*?\*/|//.*"""), "").trim()
        if (structure.startsWith("}")) indent = (indent - 1).coerceAtLeast(0)
        val effectiveIndent = if (structure.startsWith("case ") || structure.startsWith("default:")) {
            (indent - 1).coerceAtLeast(0)
        } else {
            indent
        }
        if (effectiveIndent == 0 && previousLine == "};" && looksLikeGlslFunctionDeclaration(line)) {
            out.appendLine()
        }
        repeat(effectiveIndent) { out.append("    ") }
        val cleaned = cleanGlslSpacing(line)
        out.appendLine(cleaned)
        previousLine = cleaned
        if (structure.endsWith("{")) indent += 1
    }

    var formatted = out.toString().trimEnd() + "\n"
    protected.forEach { (token, original) -> formatted = formatted.replace(token, original) }
    return formatted
}

private val GLSL_PROTECTED = Regex("""//[^\r\n]*|/\*[\s\S]*?\*/|"(?:\\.|[^"\\])*"|'(?:\\.|[^'\\])*'""")

private fun splitGlslStatements(code: String): String {
    val out = StringBuilder()
    var i = 0
    var parenDepth = 0
    var atLineStart = true
    var inLineComment = false
    var inBlockComment = false
    var inString: Char? = null
    var escaped = false

    // Whether the newline `out` ends with was put there by us, after a `;` or a brace, rather than
    // copied from the source. The source's own newline after it is then the same line break; a
    // further one is a blank line the author wrote, and is kept.
    var inserted = false

    fun appendNewlineIfNeeded() {
        while (out.endsWith(' ') || out.endsWith('\t')) out.setLength(out.length - 1)
        if (!out.endsWith('\n')) {
            out.append('\n')
            inserted = true
        }
        atLineStart = true
    }

    while (i < code.length) {
        val c = code[i]
        val next = code.getOrNull(i + 1)

        if (inLineComment) {
            out.append(c)
            if (c == '\n') {
                inLineComment = false
                atLineStart = true
            }
            i++
            continue
        }

        if (inBlockComment) {
            out.append(c)
            if (c == '*' && next == '/') {
                out.append(next)
                i += 2
                inBlockComment = false
            } else {
                if (c == '\n') atLineStart = true
                i++
            }
            continue
        }

        if (inString != null) {
            out.append(c)
            if (escaped) {
                escaped = false
            } else if (c == '\\') {
                escaped = true
            } else if (c == inString) {
                inString = null
            }
            i++
            continue
        }

        if (atLineStart && c == '#') {
            val end = code.indexOf('\n', i).let { if (it < 0) code.length else it }
            out.append(code.substring(i, end).trimEnd())
            appendNewlineIfNeeded()
            i = if (end < code.length) end + 1 else end
            // The directive's source newline was consumed above. A following newline is
            // an authored blank line, rather than the duplicate of an inserted break.
            inserted = false
            continue
        }

        // A template span ({% property %}, {{ value }}, {# note #}) is Field's, not GLSL's: copy it as
        // it stands rather than reading its braces as a block.
        if (c == '{' && (next == '%' || next == '{' || next == '#')) {
            val closer = when (next) { '%' -> "%}"; '#' -> "#}"; else -> "}}" }
            val end = code.indexOf(closer, i + 2).let { if (it < 0) code.length else it + 2 }
            out.append(code, i, end)
            atLineStart = false
            i = end
            continue
        }

        when {
            c == '/' && next == '/' -> {
                // A semicolon/brace may just have inserted a newline. A comment authored on
                // that same source line still belongs to the preceding statement.
                val sourceLineStart = code.lastIndexOf('\n', i - 1) + 1
                if (inserted && out.endsWith('\n') && code.substring(sourceLineStart, i).isNotBlank()) {
                    out.setLength(out.length - 1)
                    out.append(' ')
                    inserted = false
                    atLineStart = false
                }
                out.append(c).append(next)
                inLineComment = true
                i += 2
            }
            c == '/' && next == '*' -> {
                out.append(c).append(next)
                inBlockComment = true
                i += 2
            }
            c == '"' || c == '\'' -> {
                out.append(c)
                inString = c
                i++
            }
            c == '(' -> {
                parenDepth += 1
                out.append(c)
                atLineStart = false
                i++
            }
            c == ')' -> {
                parenDepth = (parenDepth - 1).coerceAtLeast(0)
                out.append(c)
                atLineStart = false
                i++
            }
            c == '{' -> {
                if (out.isNotEmpty() && !out.endsWith(' ') && !out.endsWith('\n')) out.append(' ')
                out.append(c)
                appendNewlineIfNeeded()
                i++
            }
            c == '}' -> {
                appendNewlineIfNeeded()
                // `} pc;` closes a block that declares an instance; the name stays on the brace's line.
                val instance = BLOCK_INSTANCE.find(code, i + 1)?.takeIf { it.range.first == i + 1 }
                if (next == ';') {
                    out.append("};")
                    i += 2
                } else if (instance != null) {
                    out.append("} ").append(instance.groupValues[1].replace(Regex("""\s+"""), "")).append(';')
                    i = instance.range.last + 1
                } else {
                    out.append(c)
                    i++
                }
                appendNewlineIfNeeded()
            }
            c == ';' && parenDepth == 0 -> {
                out.append(c)
                appendNewlineIfNeeded()
                i++
            }
            c == '\n' -> {
                while (out.endsWith(' ') || out.endsWith('\t')) out.setLength(out.length - 1)
                when {
                    !out.endsWith('\n') -> out.append('\n')
                    inserted -> inserted = false
                    // A blank line; formatGlslCode collapses a run of them to one.
                    else -> out.append('\n')
                }
                atLineStart = true
                i++
            }
            c.isWhitespace() -> {
                if (!out.endsWith(' ') && !out.endsWith('\n')) out.append(' ')
                i++
            }
            else -> {
                out.append(c)
                atLineStart = false
                inserted = false
                i++
            }
        }
    }

    return out.toString()
}

private val BLOCK_INSTANCE = Regex("""[ \t]*([A-Za-z_]\w*(?:\s*\[[^\]\n]*\])?)[ \t]*;""")

private fun cleanGlslSpacing(line: String): String {
    return line
        .replace(Regex("""\s+([,;)\]])"""), "$1")
        .replace(Regex("""([(\[])\s+"""), "$1")
        .replace(Regex("""\s+\{"""), " {")
        .replace(Regex("""\s{2,}"""), " ")
}

private fun looksLikeGlslFunctionDeclaration(line: String): Boolean {
    if (!line.endsWith("{")) return false
    if (!line.contains('(') || !line.contains(')')) return false
    val head = line.substringBefore('(').trim()
    if (head.isEmpty()) return false
    val name = head.substringAfterLast(' ')
    if (name in setOf("if", "for", "while", "switch")) return false
    return head.any { it.isWhitespace() }
}

private fun StringBuilder.endsWith(c: Char): Boolean = isNotEmpty() && this[length - 1] == c