zshrs 0.11.42

//! zsh source formatter — syntax-aware reindenter.
//!
//! zshrs extension (no C counterpart: zsh ships no formatter, and
//! shfmt explicitly does not support the zsh dialect). Exposed as:
//!   * `zshrs --fmt [-w] [-i N] [-t] [FILE…]` (bins/zshrs.rs)
//!   * LSP `textDocument/formatting` (src/extensions/lsp.rs), which
//!     the IntelliJ plugin's `LspFormattingSupport` drives via
//!     Reformat Code.
//!
//! Design contract — CONSERVATIVE BY CONSTRUCTION:
//!   * Only leading whitespace (indentation) and trailing whitespace
//!     are ever rewritten. Inner spacing, quoting, line breaks and
//!     token text are preserved byte-for-byte.
//!   * Heredoc bodies (`<<TAG` … `TAG`) pass through verbatim,
//!     including their terminators.
//!   * Idempotent: `format(format(x)) == format(x)`.
//!
//! The line scanner tracks just enough lexical state to find the
//! STRUCTURAL tokens that drive indentation: quotes (`'…'`, `"…"`,
//! `$'…'`), backslash escapes, comments, `${…}` parameter braces
//! (which must NOT count as block braces), command/subshell parens,
//! `[[ … ]]`, and the reserved-word pairs if/fi, do/done (for, while,
//! until, select, repeat), case/esac with pattern-arm handling, and
//! `{ … }` grouping.

/// Formatting options, mirroring the LSP `FormattingOptions` shape.
#[derive(Debug, Clone, Copy)]
pub struct FmtOptions {
    /// Spaces per indent level (ignored for emission when `use_tabs`).
    pub indent_width: usize,
    /// Emit one tab per level instead of spaces.
    pub use_tabs: bool,
}

impl Default for FmtOptions {
    fn default() -> Self {
        FmtOptions {
            indent_width: 4,
            use_tabs: false,
        }
    }
}

/// One open block on the indent stack.
#[derive(Debug, Clone, Copy, PartialEq)]
enum Block {
    /// `if` … `fi`. `open` flips true once `then` is seen.
    If { open: bool },
    /// for/while/until/select/repeat … `done`. `open` on `do`.
    Loop { open: bool },
    /// `{` … `}` command grouping / function body.
    Brace,
    /// `(` … `)` subshell / array literal / `$(…)` / `((…))` halves.
    Paren,
    /// `[[ … ]]` conditional (multi-line conditions).
    DCond,
    /// `case` … `esac`. `sub` tracks the arm cycle.
    Case { sub: CaseSub },
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum CaseSub {
    /// Between `in` / `;;` and the next `pattern)` — expecting an arm.
    Pattern,
    /// Inside an arm body (after `pattern)`), until `;;` / `;&` / `;|`.
    Body,
}

impl Block {
    /// Whether this block currently indents the lines inside it.
    fn contributes(&self) -> usize {
        match self {
            Block::If { open } | Block::Loop { open } => usize::from(*open),
            Block::Brace | Block::Paren | Block::DCond => 1,
            // case body lines sit two deeper than `case` itself
            // (one for the arm label, one for the body); arm labels
            // sit one deeper. The Pattern state contributes the arm
            // level; Body adds the second level.
            Block::Case { sub } => match sub {
                CaseSub::Pattern => 1,
                CaseSub::Body => 2,
            },
        }
    }
}

/// A heredoc awaiting its terminator.
#[derive(Debug, Clone)]
struct Heredoc {
    tag: String,
    /// `<<-` — terminator may be preceded by tabs.
    strip_tabs: bool,
}

/// Format zsh source: reindent by block structure, strip trailing
/// whitespace, ensure exactly one trailing newline. See module doc
/// for the conservation guarantees.
pub fn format_source(src: &str, opts: &FmtOptions) -> String {
    let mut out = String::with_capacity(src.len() + 64);
    let mut stack: Vec<Block> = Vec::new();
    let mut pending_heredocs: Vec<Heredoc> = Vec::new();
    let mut active_heredoc: Option<Heredoc> = None;
    // Previous line ended with an unquoted `\` — indent one extra.
    let mut continuation = false;

    for line in src.split('\n') {
        // ── Heredoc body passthrough ────────────────────────────────
        if let Some(h) = &active_heredoc {
            let term_candidate = if h.strip_tabs {
                line.trim_start_matches('\t')
            } else {
                line
            };
            out.push_str(line);
            out.push('\n');
            if term_candidate == h.tag {
                active_heredoc = pending_heredocs.pop_front_or_none();
            }
            continue;
        }

        let body = line.trim_start();
        if body.is_empty() {
            // Blank line: no indent whitespace emitted.
            out.push('\n');
            continue;
        }

        // ── Scan the line for structural tokens ─────────────────────
        let scan = scan_line(body, &mut stack, &mut pending_heredocs);

        // ── Compute this line's indent ──────────────────────────────
        // `leading_closers` blocks were popped by tokens at the very
        // start of the line (fi/done/esac/}/)/else/…) — the line
        // itself prints at the dedented level, which is exactly the
        // post-pop stack depth plus any re-opened mid-line blocks
        // counted in `depth_before_line`.
        let mut depth: usize = scan.indent_basis;
        if continuation {
            depth += 1;
        }
        let indent = if opts.use_tabs {
            "\t".repeat(depth)
        } else {
            " ".repeat(depth * opts.indent_width)
        };
        out.push_str(&indent);
        out.push_str(body.trim_end());
        out.push('\n');

        continuation = scan.ends_with_continuation;
        if !pending_heredocs.is_empty() && active_heredoc.is_none() {
            active_heredoc = pending_heredocs.pop_front_or_none();
        }
    }

    // Exactly one trailing newline.
    while out.ends_with("\n\n") {
        out.pop();
    }
    if !out.ends_with('\n') {
        out.push('\n');
    }
    // src.split('\n') yields a trailing "" for newline-terminated
    // input which the blank-line arm turned into one extra \n — the
    // dedup loop above already collapsed it.
    out
}

/// Small helper: Vec-as-FIFO for the heredoc queue.
trait PopFront<T> {
    fn pop_front_or_none(&mut self) -> Option<T>;
}
impl<T> PopFront<T> for Vec<T> {
    fn pop_front_or_none(&mut self) -> Option<T> {
        if self.is_empty() {
            None
        } else {
            Some(self.remove(0))
        }
    }
}

struct LineScan {
    /// Indent level (in block levels) this line should print at.
    indent_basis: usize,
    /// Line ends with an unquoted backslash.
    ends_with_continuation: bool,
}

/// Current stack indent = sum of every block's contribution.
fn stack_indent(stack: &[Block]) -> usize {
    stack.iter().map(|b| b.contributes()).sum()
}

/// Scan one (whitespace-trimmed) line: update `stack` with every
/// structural token, queue heredocs, and report the indent level the
/// line itself should print at.
fn scan_line(
    body: &str,
    stack: &mut Vec<Block>,
    heredocs: &mut Vec<Heredoc>,
) -> LineScan {
    let b = body.as_bytes();
    let n = b.len();
    let mut i = 0usize;
    // Tokens seen so far on this line (false = still at line start —
    // leading closers dedent the line itself).
    let mut seen_token = false;
    // The indent the line prints at. Starts at the current stack
    // level and is lowered by leading closers as they pop.
    let mut indent_basis = stack_indent(stack);
    let mut ends_with_continuation = false;

    // Word accumulator for reserved-word recognition.
    macro_rules! at_line_start {
        () => {
            !seen_token
        };
    }

    while i < n {
        let c = b[i];
        match c {
            b' ' | b'\t' => {
                i += 1;
            }
            b'\\' => {
                // Escape: skip the next char. A `\` as the LAST char
                // is a line continuation.
                if i + 1 >= n {
                    ends_with_continuation = true;
                }
                i += 2;
                seen_token = true;
            }
            b'\'' => {
                // Single-quoted: skip to closing '.
                i += 1;
                while i < n && b[i] != b'\'' {
                    i += 1;
                }
                i += 1;
                seen_token = true;
            }
            b'"' => {
                // Double-quoted: skip to closing unescaped ". `$(…)`
                // inside DQ can span structure, but multi-line DQ
                // strings pass through verbatim anyway (the quote
                // state is line-local by design: a string that spans
                // lines leaves the remainder lines untouched at their
                // original indent only if they parse as body text —
                // acceptable for a conservative formatter).
                i += 1;
                while i < n {
                    if b[i] == b'\\' {
                        i += 2;
                        continue;
                    }
                    if b[i] == b'"' {
                        break;
                    }
                    i += 1;
                }
                i += 1;
                seen_token = true;
            }
            b'`' => {
                i += 1;
                while i < n && b[i] != b'`' {
                    if b[i] == b'\\' {
                        i += 1;
                    }
                    i += 1;
                }
                i += 1;
                seen_token = true;
            }
            b'#' => {
                // Comment for the rest of the line (we are always at
                // a word boundary here: the word arm below consumes
                // `#` inside words like `${#var}` via the `$` arm or
                // the word scanner).
                break;
            }
            b'$' => {
                seen_token = true;
                if i + 1 < n && b[i + 1] == b'\'' {
                    // $'…' ANSI-C string.
                    i += 2;
                    while i < n && b[i] != b'\'' {
                        if b[i] == b'\\' {
                            i += 1;
                        }
                        i += 1;
                    }
                    i += 1;
                } else if i + 1 < n && b[i + 1] == b'{' {
                    // ${…} parameter brace — depth-track to its `}`
                    // so it never counts as a block brace. Nested
                    // quotes inside are skipped coarsely.
                    i += 2;
                    let mut depth = 1usize;
                    while i < n && depth > 0 {
                        match b[i] {
                            b'{' => depth += 1,
                            b'}' => depth -= 1,
                            b'\\' => i += 1,
                            _ => {}
                        }
                        i += 1;
                    }
                } else {
                    // `$(`/`$((` fall through to the paren arms below
                    // on the next iteration; bare `$var` is a word.
                    i += 1;
                }
            }
            b'{' => {
                // Block-open brace. zsh brace expansion `{a,b}` is
                // glued to a word (no preceding whitespace); the
                // word scanner consumes those inline, so a bare `{`
                // here is command grouping.
                stack.push(Block::Brace);
                seen_token = true;
                i += 1;
            }
            b'}' => {
                let popped = matches!(stack.last(), Some(Block::Brace));
                if popped {
                    stack.pop();
                    if at_line_start!() {
                        indent_basis = indent_basis.saturating_sub(1);
                    }
                }
                seen_token = true;
                i += 1;
            }
            b'(' => {
                // In a case arm-pattern position, `(pat)` parens are
                // part of the pattern: push nothing, the matching `)`
                // closes the pattern.
                if let Some(Block::Case { sub: CaseSub::Pattern }) = stack.last() {
                    // Optional leading `(` of an arm pattern: skip.
                    i += 1;
                    seen_token = true;
                    continue;
                }
                stack.push(Block::Paren);
                seen_token = true;
                i += 1;
            }
            b')' => {
                match stack.last_mut() {
                    Some(Block::Case { sub }) if *sub == CaseSub::Pattern => {
                        // `pattern)` — the arm opens its body.
                        *sub = CaseSub::Body;
                        // The arm-label line prints at the Pattern
                        // contribution level; recompute nothing —
                        // label was already being printed at the
                        // pattern level via indent_basis (computed
                        // before this token executed only if the
                        // label began the line; for mid-line `)` the
                        // next lines pick up Body level naturally).
                    }
                    Some(Block::Paren) => {
                        stack.pop();
                        if at_line_start!() {
                            indent_basis = indent_basis.saturating_sub(1);
                        }
                    }
                    _ => {}
                }
                seen_token = true;
                i += 1;
            }
            b';' => {
                // `;;` / `;&` / `;|` terminate a case arm body.
                let two = b.get(i + 1).copied();
                if matches!(two, Some(b';') | Some(b'&') | Some(b'|')) {
                    if let Some(Block::Case { sub }) = stack.last_mut() {
                        if *sub == CaseSub::Body {
                            *sub = CaseSub::Pattern;
                            if at_line_start!() {
                                // `;;` on its own line prints at the
                                // BODY level (one deeper than the arm
                                // label) — matches the dominant zsh
                                // style. indent_basis already holds
                                // the pre-pop (body) level; keep it.
                            }
                        }
                    }
                    i += 2;
                } else {
                    i += 1;
                }
                seen_token = true;
            }
            b'<' => {
                // Heredoc `<<TAG` / `<<-TAG` (but not herestring `<<<`).
                if i + 1 < n && b[i + 1] == b'<' {
                    if i + 2 < n && b[i + 2] == b'<' {
                        i += 3; // <<< herestring
                    } else {
                        i += 2;
                        let strip_tabs = i < n && b[i] == b'-';
                        if strip_tabs {
                            i += 1;
                        }
                        while i < n && (b[i] == b' ' || b[i] == b'\t') {
                            i += 1;
                        }
                        // Delimiter word, possibly quoted.
                        let mut tag = String::new();
                        let mut quote: Option<u8> = None;
                        while i < n {
                            let ch = b[i];
                            match quote {
                                Some(q) if ch == q => {
                                    quote = None;
                                    i += 1;
                                }
                                Some(_) => {
                                    tag.push(ch as char);
                                    i += 1;
                                }
                                None => match ch {
                                    b'\'' | b'"' => {
                                        quote = Some(ch);
                                        i += 1;
                                    }
                                    b'\\' => {
                                        i += 1;
                                        if i < n {
                                            tag.push(b[i] as char);
                                            i += 1;
                                        }
                                    }
                                    b' ' | b'\t' | b';' | b'&' | b'|' | b'<' | b'>'
                                    | b'(' | b')' => break,
                                    _ => {
                                        tag.push(ch as char);
                                        i += 1;
                                    }
                                },
                            }
                        }
                        if !tag.is_empty() {
                            heredocs.push(Heredoc { tag, strip_tabs });
                        }
                    }
                } else {
                    i += 1;
                }
                seen_token = true;
            }
            b'>' => {
                // Redirection operator — structurally inert here, but
                // it MUST have its own arm: `>` is in the word
                // scanner's break set, so falling into the word arm
                // produced a zero-length word and the scan loop never
                // advanced (infinite loop on `(( x > 1 ))`).
                i += 1;
                seen_token = true;
            }
            b'[' => {
                // `[[ … ]]` — push only the DOUBLE form; single `[`
                // is the test command (word).
                if i + 1 < n && b[i + 1] == b'[' && is_word_boundary(b, i, 2) {
                    stack.push(Block::DCond);
                    i += 2;
                } else {
                    i += 1;
                }
                seen_token = true;
            }
            b']' => {
                if i + 1 < n
                    && b[i + 1] == b']'
                    && matches!(stack.last(), Some(Block::DCond))
                {
                    stack.pop();
                    if at_line_start!() {
                        indent_basis = indent_basis.saturating_sub(1);
                    }
                    i += 2;
                } else {
                    i += 1;
                }
                seen_token = true;
            }
            _ => {
                // Word: consume [^ \t;()<>{}'"`\\#]+ and check the
                // reserved words that drive indentation. `#` inside a
                // word (e.g. `a#b`, extendedglob) must not start a
                // comment — include it in the word.
                let start = i;
                while i < n {
                    match b[i] {
                        b' ' | b'\t' | b';' | b'(' | b')' | b'<' | b'>' | b'\''
                        | b'"' | b'`' | b'\\' => break,
                        b'{' | b'}' => {
                            // Brace glued inside a word = brace
                            // expansion / ${…} tail — consume it as
                            // word text (depth-track pairs).
                            i += 1;
                        }
                        _ => i += 1,
                    }
                }
                // Defensive forward-progress guarantee: if a byte is
                // in the word break-set but lacks its own outer arm,
                // a zero-length word would loop forever. Consume one
                // byte and move on instead.
                if i == start {
                    i += 1;
                    seen_token = true;
                    continue;
                }
                let word = &body[start..i];
                let leading = at_line_start!();
                seen_token = true;
                match word {
                    "if" => stack.push(Block::If { open: false }),
                    "then" => {
                        if let Some(Block::If { open }) = stack.last_mut() {
                            if leading && *open {
                                // bare `then` line after a multi-line
                                // condition prints at the `if` level.
                                indent_basis = indent_basis.saturating_sub(1);
                            }
                            *open = true;
                        }
                    }
                    "elif" | "else" => {
                        if leading {
                            if let Some(Block::If { open: true }) = stack.last() {
                                indent_basis = indent_basis.saturating_sub(1);
                            }
                        }
                        // elif re-arms the then-cycle; body depth is
                        // unchanged (If stays open).
                    }
                    "fi" => {
                        if matches!(stack.last(), Some(Block::If { .. })) {
                            let was_open =
                                matches!(stack.last(), Some(Block::If { open: true }));
                            stack.pop();
                            if leading && was_open {
                                indent_basis = indent_basis.saturating_sub(1);
                            }
                        }
                    }
                    "for" | "while" | "until" | "select" | "repeat" => {
                        // `while` also appears as `do … done` driver in
                        // `repeat N; do`; all push a Loop.
                        stack.push(Block::Loop { open: false });
                    }
                    "do" => {
                        if let Some(Block::Loop { open }) = stack.last_mut() {
                            if leading && *open {
                                indent_basis = indent_basis.saturating_sub(1);
                            }
                            *open = true;
                        }
                    }
                    "done" => {
                        if matches!(stack.last(), Some(Block::Loop { .. })) {
                            let was_open =
                                matches!(stack.last(), Some(Block::Loop { open: true }));
                            stack.pop();
                            if leading && was_open {
                                indent_basis = indent_basis.saturating_sub(1);
                            }
                        }
                    }
                    "case" => stack.push(Block::Case {
                        sub: CaseSub::Pattern,
                    }),
                    "esac" => {
                        if let Some(Block::Case { sub }) = stack.last() {
                            let contrib = Block::Case { sub: *sub }.contributes();
                            stack.pop();
                            if leading {
                                indent_basis =
                                    indent_basis.saturating_sub(contrib);
                            }
                        }
                    }
                    _ => {}
                }
            }
        }
    }

    LineScan {
        indent_basis,
        ends_with_continuation,
    }
}

/// True when the token starting at `i` with byte length `len` is
/// delimited by whitespace / line boundaries on both sides.
fn is_word_boundary(b: &[u8], i: usize, len: usize) -> bool {
    let before_ok = i == 0 || matches!(b[i - 1], b' ' | b'\t' | b'(' | b'!' | b'{');
    let after = i + len;
    let after_ok = after >= b.len() || matches!(b[after], b' ' | b'\t');
    before_ok && after_ok
}

#[cfg(test)]
mod tests {
    use super::*;

    fn fmt(s: &str) -> String {
        format_source(s, &FmtOptions::default())
    }

    /// Indentation by block structure: if/for/case nesting.
    #[test]
    fn nested_blocks_reindent() {
        let src = "if true; then\nfor x in a b; do\nprint $x\ndone\nfi\n";
        let want = "if true; then\n    for x in a b; do\n        print $x\n    done\nfi\n";
        assert_eq!(fmt(src), want);
    }

    /// case arms at +1, bodies at +2, `;;` at body level, `esac` at
    /// case level. Arm-pattern `)` must not pop a paren block.
    #[test]
    fn case_arms_and_bodies() {
        let src = "case $x in\na)\nprint a\n;;\n(b|c)\nprint bc\n;;\nesac\n";
        let want = "case $x in\n    a)\n        print a\n        ;;\n    (b|c)\n        print bc\n        ;;\nesac\n";
        assert_eq!(fmt(src), want);
    }

    /// `${…}` braces never open a block; `{a,b}` brace expansion glued
    /// to a word never opens a block.
    #[test]
    fn param_and_expansion_braces_ignored() {
        let src = "print ${name:-x} file{1,2}\nprint after\n";
        assert_eq!(fmt(src), src);
    }

    /// Function bodies via `{ … }` indent; closing `}` dedents.
    #[test]
    fn function_braces() {
        let src = "f() {\nprint hi\n}\n";
        let want = "f() {\n    print hi\n}\n";
        assert_eq!(fmt(src), want);
    }

    /// Heredoc bodies and terminators pass through verbatim — no
    /// reindent, no trailing-space strip inside.
    #[test]
    fn heredoc_verbatim() {
        let src = "if true; then\ncat <<EOF\n  raw  spaces  \n\tand tabs\nEOF\nfi\n";
        let want = "if true; then\n    cat <<EOF\n  raw  spaces  \n\tand tabs\nEOF\nfi\n";
        assert_eq!(fmt(src), want);
    }

    /// `<<-TAG` terminator with leading tabs is recognized; body kept.
    #[test]
    fn heredoc_dash_tab_terminator() {
        let src = "cat <<-EOF\n\tbody\n\tEOF\nprint after\n";
        let want = "cat <<-EOF\n\tbody\n\tEOF\nprint after\n";
        assert_eq!(fmt(src), want);
    }

    /// Keywords inside quotes/comments are inert.
    #[test]
    fn quoted_keywords_inert() {
        let src = "print 'if then fi'\nprint \"do done\" # case in\nprint x\n";
        assert_eq!(fmt(src), src);
    }

    /// Line continuations indent one extra level.
    #[test]
    fn continuation_indents() {
        let src = "print one \\\ntwo\nprint three\n";
        let want = "print one \\\n    two\nprint three\n";
        assert_eq!(fmt(src), want);
    }

    /// Multi-line `$( … )` and array literals indent inside parens.
    #[test]
    fn multiline_cmdsubst_and_array() {
        let src = "files=(\na\nb\n)\nprint $files\n";
        let want = "files=(\n    a\n    b\n)\nprint $files\n";
        assert_eq!(fmt(src), want);
    }

    /// else / elif dedent to the `if` level.
    #[test]
    fn else_elif_level() {
        let src = "if a; then\nb\nelif c; then\nd\nelse\ne\nfi\n";
        let want = "if a; then\n    b\nelif c; then\n    d\nelse\n    e\nfi\n";
        assert_eq!(fmt(src), want);
    }

    /// Trailing whitespace stripped; exactly one final newline.
    #[test]
    fn trailing_ws_and_final_newline() {
        assert_eq!(fmt("print x   \n\n\n"), "print x\n");
        assert_eq!(fmt("print x"), "print x\n");
    }

    /// Idempotence on a representative composite.
    #[test]
    fn idempotent() {
        let src = "f() {\nif x; then\ncase $1 in\na) y ;;\nesac\nfi\n}\ncat <<EOF\nkeep\nEOF\n";
        let once = fmt(src);
        assert_eq!(fmt(&once), once);
    }

    /// Tabs mode emits one tab per level.
    #[test]
    fn tabs_mode() {
        let opts = FmtOptions {
            indent_width: 4,
            use_tabs: true,
        };
        let got = format_source("if x; then\ny\nfi\n", &opts);
        assert_eq!(got, "if x; then\n\ty\nfi\n");
    }

    /// extendedglob `#` inside a word doesn't start a comment and
    /// `(( … ))` arithmetic stays balanced.
    #[test]
    fn hash_in_word_and_arith() {
        let src = "print ${#arr}\nif (( x > 1 )); then\ny\nfi\n";
        let want = "print ${#arr}\nif (( x > 1 )); then\n    y\nfi\n";
        assert_eq!(fmt(src), want);
    }

    /// Multi-line [[ … ]] conditions indent their continuation.
    #[test]
    fn multiline_dcond() {
        let src = "if [[ -n $a &&\n-n $b ]]; then\nx\nfi\n";
        let want = "if [[ -n $a &&\n    -n $b ]]; then\n    x\nfi\n";
        assert_eq!(fmt(src), want);
    }
}