Skip to main content

sql_dialect_fmt_formatter/
doc.rs

1//! A small, self-contained pretty-printing engine in the Wadler → Prettier → rome/biome → ruff
2//! lineage. The grammar-agnostic core is a `Doc` intermediate representation plus a width-aware
3//! [`print`]er; SQL-specific rules live in [`crate::sql`] and only ever build `Doc`s.
4//!
5//! The model has three moving parts:
6//! * **Builders** ([`text`], [`concat`], [`group`], [`indent`], [`line`], [`soft_line`],
7//!   [`hard_line`], [`space`], [`join`]) construct the IR.
8//! * **Groups** are the unit of layout choice: a group is printed *flat* (its soft lines become
9//!   spaces or nothing) when it fits within the remaining width, otherwise it *breaks* (its soft
10//!   lines become newlines + indentation).
11//! * A **hard line** forces every group that contains it to break — the classic `breakParent`
12//!   propagation — so caller-mandated line breaks are never silently collapsed.
13//!
14//! See `docs/research/prior-art.md` for the design rationale and references.
15
16use std::borrow::Cow;
17
18/// The pretty-printer intermediate representation.
19///
20/// Construct values through the builder functions rather than the variants directly; the shape is
21/// public only so tests and future SQL rules can pattern-match if needed.
22#[derive(Clone, Debug, PartialEq, Eq)]
23pub enum Doc {
24    /// Verbatim text. Must not contain newlines — use the line builders for those so width
25    /// tracking and indentation stay correct. The display width is folded in at construction time
26    /// (via [`text`]) so the fit/break measurement never re-scans the string.
27    Text(Cow<'static, str>, usize),
28    /// A verbatim slice of source that may contain newlines, reproduced byte-for-byte (modulo the
29    /// printer's per-line right trim in [`finalize`]). This is the vehicle for verbatim fallback
30    /// regions without smuggling embedded `\n`s through [`Doc::Text`], which would corrupt column
31    /// tracking. Each contained newline re-bases the column to the current indentation.
32    ///
33    /// The cached width is the display width of the slice's last line. Build through
34    /// [`source_code_slice`], which folds the width in.
35    SourceCodeSlice(Cow<'static, str>, usize),
36    /// A line break whose rendering depends on the enclosing group's mode (see [`LineKind`]).
37    Line(LineKind),
38    /// A sequence of documents laid out one after another.
39    Concat(Vec<Doc>),
40    /// Try candidate layouts in order and print the first one that fits on the current line; if no
41    /// candidate fits, print the last one. This is the small Doc-IR escape hatch used by Prettier /
42    /// Biome for constructs with genuinely different layouts rather than just flat-vs-break lines.
43    BestFitting(Vec<Doc>),
44    /// A layout-choice boundary. Printed flat if it fits, otherwise broken. When `expand` is set
45    /// the group always breaks (Prettier's `shouldBreak`) — used for "explode this collection"
46    /// decisions like a magic trailing comma. Unlike a hard line, `expand` does **not** propagate
47    /// to enclosing groups, so an inner collection can explode while its parent stays flat.
48    ///
49    /// `must_break` is `expand || has_forced_break(content)`, computed once when the group is built
50    /// (see [`group`]). Because groups are constructed bottom-up, each group carries its own
51    /// answer and ancestors read it in O(1) instead of re-walking the whole subtree on every
52    /// fit/break decision.
53    Group {
54        content: Box<Doc>,
55        expand: bool,
56        must_break: bool,
57    },
58    /// Picks one of two layouts by the enclosing group's mode: `broken` when that group breaks,
59    /// `flat` when it stays flat (Prettier's `ifBreak`). Built through [`if_group_breaks`].
60    ///
61    /// The broken arm is not consulted by [`has_forced_break`], so an `if_group_breaks` whose broken
62    /// arm contains a hard line does not force its own group to break.
63    IfBreak { broken: Box<Doc>, flat: Box<Doc> },
64    /// Increases the indentation level applied to line breaks inside it.
65    Indent(Box<Doc>),
66    /// Content deferred to just before the next newline (or the document's end). The vehicle for
67    /// trailing line comments, which must not have code emitted after them on the same line.
68    LineSuffix(Box<Doc>),
69    /// A zero-width marker that forces every enclosing group to break, without itself emitting a
70    /// newline. Pairs with [`LineSuffix`] so a trailing `--` comment actually ends its line.
71    BreakParent,
72}
73
74/// How a [`Doc::Line`] renders, as a function of the enclosing group's mode.
75#[derive(Clone, Copy, Debug, PartialEq, Eq)]
76pub enum LineKind {
77    /// Flat: a single space. Broken: a newline + indentation. (Prettier's `line`.)
78    Space,
79    /// Flat: nothing. Broken: a newline + indentation. (Prettier's `softline`.)
80    Soft,
81    /// Always a newline + indentation, and forces every enclosing group to break. (`hardline`.)
82    Hard,
83}
84
85// ---- builders ----
86
87/// Verbatim text (a borrowed `&'static str` or an owned `String`). The display width is measured
88/// once here and cached on the node, so the printer's fit/break measurement never re-scans it.
89pub fn text(s: impl Into<Cow<'static, str>>) -> Doc {
90    let s = s.into();
91    let width = text_width(&s);
92    Doc::Text(s, width)
93}
94
95/// A verbatim slice of source that may span multiple lines. Use this — not [`text`] — for fallback
96/// regions that can contain embedded newlines: it re-bases the column at each newline so the slice
97/// stays aligned under indentation, and caches its last line's display width for fit decisions.
98pub fn source_code_slice(s: impl Into<Cow<'static, str>>) -> Doc {
99    let s = s.into();
100    let width = last_line_width(&s);
101    Doc::SourceCodeSlice(s, width)
102}
103
104/// Lay out the parts one after another.
105pub fn concat(parts: Vec<Doc>) -> Doc {
106    Doc::Concat(parts)
107}
108
109/// A layout-choice group: flat when it fits on the line, broken otherwise.
110pub fn group(inner: Doc) -> Doc {
111    let must_break = has_forced_break(&inner);
112    Doc::Group {
113        content: Box::new(inner),
114        expand: false,
115        must_break,
116    }
117}
118
119/// A group forced to break (its soft lines always become newlines), without forcing enclosing
120/// groups to break. The building block for magic-trailing-comma "keep this exploded".
121pub fn group_expanded(inner: Doc) -> Doc {
122    Doc::Group {
123        content: Box::new(inner),
124        expand: true,
125        must_break: true,
126    }
127}
128
129/// Indent every line break that occurs inside `inner` by one more level.
130pub fn indent(inner: Doc) -> Doc {
131    Doc::Indent(Box::new(inner))
132}
133
134/// A hard-line-delimited, indented block: `inner` is placed on its own indented line(s), framed by a
135/// hard line before and after (Prettier/biome's `block_indent`).
136pub fn block_indent(inner: Doc) -> Doc {
137    concat(vec![indent(concat(vec![hard_line(), inner])), hard_line()])
138}
139
140/// Choose `broken` when the enclosing group breaks and `flat` when it stays flat (Prettier's
141/// `ifBreak`). Neither arm emits a newline by itself; this only selects which content appears.
142pub fn if_group_breaks(broken: Doc, flat: Doc) -> Doc {
143    Doc::IfBreak {
144        broken: Box::new(broken),
145        flat: Box::new(flat),
146    }
147}
148
149/// A space when flat, a newline when broken.
150pub fn line() -> Doc {
151    Doc::Line(LineKind::Space)
152}
153
154/// Nothing when flat, a newline when broken.
155pub fn soft_line() -> Doc {
156    Doc::Line(LineKind::Soft)
157}
158
159/// Always a newline; forces enclosing groups to break.
160pub fn hard_line() -> Doc {
161    Doc::Line(LineKind::Hard)
162}
163
164/// A literal, non-collapsible space.
165pub fn space() -> Doc {
166    Doc::Text(Cow::Borrowed(" "), 1)
167}
168
169/// The empty document (renders to nothing).
170pub fn empty() -> Doc {
171    Doc::Concat(Vec::new())
172}
173
174/// Defer `inner` to just before the next newline (used for trailing line comments).
175pub fn line_suffix(inner: Doc) -> Doc {
176    Doc::LineSuffix(Box::new(inner))
177}
178
179/// Force enclosing groups to break without emitting a newline here.
180pub fn break_parent() -> Doc {
181    Doc::BreakParent
182}
183
184/// Interleave `sep` between `items` (no separator before the first or after the last).
185pub fn join(sep: Doc, items: Vec<Doc>) -> Doc {
186    let mut parts = Vec::with_capacity(items.len().saturating_mul(2));
187    for (i, item) in items.into_iter().enumerate() {
188        if i > 0 {
189            parts.push(sep.clone());
190        }
191        parts.push(item);
192    }
193    Doc::Concat(parts)
194}
195
196/// Try the candidate documents in order, printing the first that fits in the remaining width and
197/// falling back to the last candidate when none fit. Empty candidates render as empty.
198pub fn best_fitting(candidates: Vec<Doc>) -> Doc {
199    Doc::BestFitting(candidates)
200}
201
202// ---- write-style composition layer ----
203
204/// A growable sink of [`Doc`] elements, the target of the [`doc_write!`] macro and of [`Format`]
205/// implementors.
206#[derive(Default)]
207pub struct DocBuffer {
208    parts: Vec<Doc>,
209}
210
211impl DocBuffer {
212    /// A new, empty buffer.
213    pub fn new() -> Self {
214        DocBuffer { parts: Vec::new() }
215    }
216
217    /// Append one already-built [`Doc`] element in order.
218    pub fn write_element(&mut self, doc: Doc) {
219        self.parts.push(doc);
220    }
221
222    /// Append a value by letting it format itself into this buffer.
223    pub fn write_fmt_value<T: Format + ?Sized>(&mut self, value: &T) {
224        value.fmt(self);
225    }
226
227    /// Append a value formatted by an external [`FormatRule`].
228    pub fn write_with_rule<T: ?Sized, R: FormatRule<T>>(&mut self, rule: &R, item: &T) {
229        rule.fmt(item, self);
230    }
231
232    /// Consume the buffer, returning the assembled document.
233    pub fn finish(mut self) -> Doc {
234        if self.parts.len() == 1 {
235            self.parts.pop().expect("len checked == 1")
236        } else {
237            Doc::Concat(self.parts)
238        }
239    }
240}
241
242/// A value that knows how to render itself into a [`DocBuffer`].
243pub trait Format {
244    /// Push this value's document representation onto `buffer`, in source order.
245    fn fmt(&self, buffer: &mut DocBuffer);
246}
247
248/// Formatting logic for a `T` kept outside `T`, matching the biome/ruff rule pattern.
249pub trait FormatRule<T: ?Sized> {
250    /// Push `item`'s document representation onto `buffer`.
251    fn fmt(&self, item: &T, buffer: &mut DocBuffer);
252}
253
254impl Format for Doc {
255    fn fmt(&self, buffer: &mut DocBuffer) {
256        buffer.write_element(self.clone());
257    }
258}
259
260impl<T: Format + ?Sized> Format for &T {
261    fn fmt(&self, buffer: &mut DocBuffer) {
262        (**self).fmt(buffer);
263    }
264}
265
266/// Run a single [`Format`] value to completion, returning its assembled [`Doc`].
267pub fn format_value<T: Format + ?Sized>(value: &T) -> Doc {
268    let mut buffer = DocBuffer::new();
269    buffer.write_fmt_value(value);
270    buffer.finish()
271}
272
273/// Push a comma-bracketed list of [`Format`] elements into a [`DocBuffer`], in order.
274#[macro_export]
275macro_rules! doc_write {
276    ($buffer:expr, [ $( $element:expr ),* $(,)? ]) => {{
277        $( $buffer.write_fmt_value(&$element); )*
278    }};
279}
280
281#[doc(inline)]
282pub use crate::doc_write;
283
284// ---- printing ----
285
286/// Knobs for the printer. Opinionated by design: just a target width and an indent step.
287#[derive(Clone, Copy, Debug)]
288pub struct PrintOptions {
289    /// The column the printer tries to keep lines within.
290    pub line_width: usize,
291    /// Number of spaces added per indentation level.
292    pub indent_width: usize,
293}
294
295impl Default for PrintOptions {
296    fn default() -> Self {
297        PrintOptions {
298            line_width: 100,
299            indent_width: 4,
300        }
301    }
302}
303
304#[derive(Clone, Copy, PartialEq, Eq)]
305enum Mode {
306    Flat,
307    Break,
308}
309
310#[derive(Clone, Copy)]
311struct Cmd<'a> {
312    indent: usize,
313    mode: Mode,
314    doc: &'a Doc,
315}
316
317/// Display width of a string in terminal columns: the sum of each character's column width, so a
318/// line of CJK text is measured by how wide it actually renders, not by its scalar count. This only
319/// feeds the fit/break decision, so refining it never changes which tokens are emitted.
320///
321/// Pure ASCII — the overwhelmingly common case in SQL — has width equal to its byte length, so we
322/// short-circuit on it and only decode/classify code points when a non-ASCII byte is present.
323fn text_width(s: &str) -> usize {
324    if s.is_ascii() {
325        return s.len();
326    }
327    s.chars().map(char_width).sum()
328}
329
330/// Display width of the last line of a possibly multi-line slice: everything after the final
331/// newline, or the whole string when it has none.
332fn last_line_width(s: &str) -> usize {
333    match s.rfind('\n') {
334        Some(nl) => text_width(&s[nl + 1..]),
335        None => text_width(s),
336    }
337}
338
339/// Column width of a single character: `2` for East Asian Wide / Fullwidth code points (per Unicode
340/// Annex #11 — CJK ideographs, kana, Hangul syllables, fullwidth forms, most emoji), `1` otherwise.
341/// Combining marks are not special-cased (rare in SQL), so this is an upper bound there.
342fn char_width(c: char) -> usize {
343    let cp = c as u32;
344    let wide = matches!(cp,
345        0x1100..=0x115F   // Hangul Jamo
346        | 0x2E80..=0x303E // CJK radicals, Kangxi, CJK symbols & punctuation
347        | 0x3041..=0x33FF // Hiragana, Katakana, Bopomofo, CJK compat
348        | 0x3400..=0x4DBF // CJK Unified Ideographs Extension A
349        | 0x4E00..=0x9FFF // CJK Unified Ideographs
350        | 0xA000..=0xA4CF // Yi syllables
351        | 0xAC00..=0xD7A3 // Hangul syllables
352        | 0xF900..=0xFAFF // CJK compatibility ideographs
353        | 0xFE10..=0xFE19 // vertical forms
354        | 0xFE30..=0xFE6F // CJK compatibility & small form variants
355        | 0xFF00..=0xFF60 // fullwidth forms
356        | 0xFFE0..=0xFFE6 // fullwidth signs
357        | 0x1F300..=0x1FAFF // symbols, pictographs & emoji (wide)
358        | 0x20000..=0x3FFFD // supplementary ideographic planes
359    );
360    if wide {
361        2
362    } else {
363        1
364    }
365}
366
367/// Does `doc` contain a hard line anywhere within (propagating through nested groups, exactly like
368/// Prettier's `breakParent`)? Such a document can never be printed flat.
369///
370/// Nested groups are consulted via their cached `must_break` flag rather than re-walked. Because
371/// the builders compute this bottom-up (a group's flag is set when it is constructed, by which time
372/// its children already carry theirs), the whole tree is classified in O(nodes) overall.
373fn has_forced_break(doc: &Doc) -> bool {
374    match doc {
375        Doc::Line(LineKind::Hard) | Doc::BreakParent => true,
376        Doc::Line(_) | Doc::Text(..) | Doc::SourceCodeSlice(..) => false,
377        // A line suffix's own content is deferred and must not force the current line to break.
378        Doc::LineSuffix(_) => false,
379        // `if_group_breaks` resolves against its group's mode. Letting its broken arm force that
380        // group to break would be circular; the flat arm is what can force flat layout impossible.
381        Doc::IfBreak { flat, .. } => has_forced_break(flat),
382        Doc::Concat(parts) => parts.iter().any(has_forced_break),
383        Doc::BestFitting(candidates) => {
384            !candidates.is_empty() && candidates.iter().all(has_forced_break)
385        }
386        Doc::Indent(inner) => has_forced_break(inner),
387        // An exploded group propagates to ancestors: a multiline collection can't sit inline, so
388        // every group containing it must break too (cf. Black's magic trailing comma). The answer
389        // was precomputed when the group was built.
390        Doc::Group { must_break, .. } => *must_break,
391    }
392}
393
394/// Would `next`, followed by the not-yet-processed `rest` of the print stack, fit on the current
395/// line within `remaining` columns? Everything is measured as if flat until the first newline that
396/// is actually taken (a hard line, or a soft line in an already-broken group).
397fn fits(mut remaining: isize, rest: &[Cmd], next: Cmd, opts: &PrintOptions) -> bool {
398    if remaining < 0 {
399        return false;
400    }
401    let mut stack: Vec<Cmd> = vec![next];
402    let mut rest_idx = rest.len();
403    loop {
404        let cmd = match stack.pop() {
405            Some(cmd) => cmd,
406            None => {
407                if rest_idx == 0 {
408                    return true;
409                }
410                rest_idx -= 1;
411                rest[rest_idx]
412            }
413        };
414        match cmd.doc {
415            Doc::Text(_, width) => {
416                remaining -= *width as isize;
417                if remaining < 0 {
418                    return false;
419                }
420            }
421            Doc::SourceCodeSlice(s, _) => {
422                let mut pieces = s.split('\n');
423                let first = pieces.next().unwrap_or_default();
424                remaining -= text_width(first) as isize;
425                if remaining < 0 {
426                    return false;
427                }
428                for piece in pieces {
429                    remaining =
430                        opts.line_width as isize - cmd.indent as isize - text_width(piece) as isize;
431                    if remaining < 0 {
432                        return false;
433                    }
434                }
435            }
436            Doc::IfBreak { broken, flat } => {
437                let chosen = if cmd.mode == Mode::Break {
438                    broken
439                } else {
440                    flat
441                };
442                stack.push(Cmd { doc: chosen, ..cmd });
443            }
444            Doc::Concat(parts) => {
445                for part in parts.iter().rev() {
446                    stack.push(Cmd { doc: part, ..cmd });
447                }
448            }
449            Doc::BestFitting(candidates) => {
450                if let Some(candidate) = candidates.first() {
451                    stack.push(Cmd {
452                        doc: candidate,
453                        ..cmd
454                    });
455                }
456            }
457            Doc::Indent(inner) => stack.push(Cmd {
458                indent: cmd.indent + opts.indent_width,
459                doc: inner,
460                mode: cmd.mode,
461            }),
462            Doc::Group {
463                content,
464                must_break,
465                ..
466            } => {
467                let mode = if *must_break { Mode::Break } else { Mode::Flat };
468                stack.push(Cmd {
469                    indent: cmd.indent,
470                    mode,
471                    doc: content,
472                });
473            }
474            // A line suffix is deferred to the next newline; it does not consume current width.
475            // A break parent is a zero-width marker. Neither affects whether the line fits.
476            Doc::LineSuffix(_) | Doc::BreakParent => {}
477            Doc::Line(kind) => match cmd.mode {
478                // A newline is taken here, so everything up to it fit.
479                Mode::Break => return true,
480                Mode::Flat => match kind {
481                    LineKind::Hard => return true,
482                    LineKind::Soft => {}
483                    LineKind::Space => {
484                        remaining -= 1;
485                        if remaining < 0 {
486                            return false;
487                        }
488                    }
489                },
490            },
491        }
492    }
493}
494
495/// Render `doc` to a string. Trailing whitespace is trimmed from every line and the result ends
496/// with exactly one newline (or is empty), so the output is stable under re-formatting.
497pub fn print(doc: &Doc, opts: &PrintOptions) -> String {
498    let mut out = String::new();
499    let mut col = 0usize;
500    let mut cmds: Vec<Cmd> = vec![Cmd {
501        indent: 0,
502        mode: Mode::Break,
503        doc,
504    }];
505    // Content deferred by `LineSuffix`, flushed in order just before the next newline (or at EOF).
506    let mut line_suffixes: Vec<Cmd> = Vec::new();
507
508    loop {
509        let cmd = match cmds.pop() {
510            Some(cmd) => cmd,
511            None if !line_suffixes.is_empty() => {
512                // Flush remaining suffixes at the document's end (no trailing newline followed).
513                while let Some(suffix) = line_suffixes.pop() {
514                    cmds.push(suffix);
515                }
516                continue;
517            }
518            None => break,
519        };
520        match cmd.doc {
521            Doc::Text(s, width) => {
522                out.push_str(s);
523                col += *width;
524            }
525            Doc::SourceCodeSlice(s, width) => {
526                let mut first = true;
527                for piece in s.split('\n') {
528                    if !first {
529                        out.push('\n');
530                        for _ in 0..cmd.indent {
531                            out.push(' ');
532                        }
533                    }
534                    out.push_str(piece);
535                    first = false;
536                }
537                col = if s.contains('\n') {
538                    cmd.indent + *width
539                } else {
540                    col + *width
541                };
542            }
543            Doc::IfBreak { broken, flat } => {
544                let chosen = if cmd.mode == Mode::Break {
545                    broken
546                } else {
547                    flat
548                };
549                cmds.push(Cmd { doc: chosen, ..cmd });
550            }
551            Doc::Concat(parts) => {
552                for part in parts.iter().rev() {
553                    cmds.push(Cmd { doc: part, ..cmd });
554                }
555            }
556            Doc::BestFitting(candidates) => {
557                if candidates.is_empty() {
558                    continue;
559                }
560                let mut chosen = candidates.last().expect("non-empty candidates");
561                for candidate in candidates {
562                    if fits(
563                        opts.line_width as isize - col as isize,
564                        &cmds,
565                        Cmd {
566                            indent: cmd.indent,
567                            mode: cmd.mode,
568                            doc: candidate,
569                        },
570                        opts,
571                    ) {
572                        chosen = candidate;
573                        break;
574                    }
575                }
576                cmds.push(Cmd { doc: chosen, ..cmd });
577            }
578            Doc::LineSuffix(inner) => line_suffixes.push(Cmd { doc: inner, ..cmd }),
579            Doc::BreakParent => {}
580            Doc::Indent(inner) => cmds.push(Cmd {
581                indent: cmd.indent + opts.indent_width,
582                doc: inner,
583                mode: cmd.mode,
584            }),
585            Doc::Group {
586                content,
587                must_break,
588                ..
589            } => {
590                let mode = if *must_break {
591                    Mode::Break
592                } else if fits(
593                    opts.line_width as isize - col as isize,
594                    &cmds,
595                    Cmd {
596                        indent: cmd.indent,
597                        mode: Mode::Flat,
598                        doc: content,
599                    },
600                    opts,
601                ) {
602                    Mode::Flat
603                } else {
604                    Mode::Break
605                };
606                cmds.push(Cmd {
607                    indent: cmd.indent,
608                    mode,
609                    doc: content,
610                });
611            }
612            Doc::Line(kind) => {
613                let newline = match cmd.mode {
614                    Mode::Break => true,
615                    Mode::Flat => match kind {
616                        LineKind::Hard => true,
617                        LineKind::Space => {
618                            out.push(' ');
619                            col += 1;
620                            false
621                        }
622                        LineKind::Soft => false,
623                    },
624                };
625                if newline {
626                    // Before breaking, emit any deferred line suffixes on this line, then
627                    // reprocess the newline with an empty buffer.
628                    if !line_suffixes.is_empty() {
629                        cmds.push(cmd);
630                        while let Some(suffix) = line_suffixes.pop() {
631                            cmds.push(suffix);
632                        }
633                        continue;
634                    }
635                    out.push('\n');
636                    for _ in 0..cmd.indent {
637                        out.push(' ');
638                    }
639                    col = cmd.indent;
640                }
641            }
642        }
643    }
644
645    finalize(out)
646}
647
648/// Trim trailing whitespace from each line, drop leading/trailing blank lines, and ensure a single
649/// trailing newline. This keeps output stable under re-formatting regardless of verbatim spans.
650///
651/// Single streaming pass: each line is `trim_end`ed; leading blank lines are skipped, and interior
652/// blank lines are buffered as `pending_blanks` so any run of them that turns out to be trailing is
653/// dropped rather than emitted. Equivalent to the former collect-join-trim, without the temporary
654/// `Vec<&str>` or the intermediate joined `String`.
655fn finalize(raw: String) -> String {
656    let mut out = String::with_capacity(raw.len());
657    let mut started = false;
658    let mut pending_blanks = 0usize;
659    for line in raw.lines() {
660        let line = line.trim_end();
661        if line.is_empty() {
662            if started {
663                pending_blanks += 1;
664            }
665            // leading blank lines are dropped entirely
666            continue;
667        }
668        if started {
669            // one separator for the previous content line, plus any buffered interior blanks
670            out.push('\n');
671            for _ in 0..pending_blanks {
672                out.push('\n');
673            }
674        }
675        pending_blanks = 0;
676        out.push_str(line);
677        started = true;
678    }
679    if started {
680        out.push('\n');
681    }
682    out
683}
684
685#[cfg(test)]
686mod tests {
687    use super::*;
688
689    fn p(doc: &Doc, width: usize) -> String {
690        print(
691            doc,
692            &PrintOptions {
693                line_width: width,
694                indent_width: 4,
695            },
696        )
697    }
698
699    #[test]
700    fn text_and_concat() {
701        let doc = concat(vec![text("a"), space(), text("b")]);
702        assert_eq!(p(&doc, 80), "a b\n");
703    }
704
705    #[test]
706    fn cjk_text_is_measured_double_width() {
707        assert_eq!(text_width("abc"), 3);
708        assert_eq!(text_width("長芋"), 4); // two wide chars
709        assert_eq!(text_width("a長"), 3); // mixed
710                                          // Flat width is 長(2) + space(1) + 芋(2) = 5: fits at 5, breaks at 4.
711        let doc = group(concat(vec![text("長"), line(), text("芋")]));
712        assert_eq!(p(&doc, 4), "長\n芋\n");
713        assert_eq!(p(&doc, 5), "長 芋\n");
714    }
715
716    #[test]
717    fn group_stays_flat_when_it_fits() {
718        let doc = group(concat(vec![text("a"), line(), text("b")]));
719        assert_eq!(p(&doc, 80), "a b\n");
720    }
721
722    #[test]
723    fn group_breaks_when_too_wide() {
724        let doc = group(concat(vec![
725            text("aaaa"),
726            indent(concat(vec![line(), text("bbbb")])),
727        ]));
728        assert_eq!(p(&doc, 5), "aaaa\n    bbbb\n");
729    }
730
731    #[test]
732    fn hard_line_forces_enclosing_group_to_break() {
733        let doc = group(concat(vec![text("a"), hard_line(), text("b")]));
734        assert_eq!(p(&doc, 80), "a\nb\n");
735    }
736
737    #[test]
738    fn soft_line_is_nothing_when_flat() {
739        let doc = group(concat(vec![text("a"), soft_line(), text("b")]));
740        assert_eq!(p(&doc, 80), "ab\n");
741    }
742
743    #[test]
744    fn join_interleaves_separator() {
745        let doc = join(text(", "), vec![text("a"), text("b"), text("c")]);
746        assert_eq!(p(&doc, 80), "a, b, c\n");
747    }
748
749    #[test]
750    fn best_fitting_picks_first_candidate_that_fits() {
751        let doc = best_fitting(vec![
752            text("short"),
753            concat(vec![text("fallback"), hard_line(), text("layout")]),
754        ]);
755        assert_eq!(p(&doc, 10), "short\n");
756    }
757
758    #[test]
759    fn best_fitting_falls_back_to_last_candidate() {
760        let doc = best_fitting(vec![
761            text("too-wide"),
762            concat(vec![text("fallback"), hard_line(), text("layout")]),
763        ]);
764        assert_eq!(p(&doc, 4), "fallback\nlayout\n");
765    }
766
767    #[test]
768    fn nested_indent_accumulates() {
769        let doc = concat(vec![
770            text("a"),
771            indent(concat(vec![
772                hard_line(),
773                text("b"),
774                indent(concat(vec![hard_line(), text("c")])),
775            ])),
776        ]);
777        assert_eq!(p(&doc, 80), "a\n    b\n        c\n");
778    }
779
780    #[test]
781    fn expanded_group_breaks_even_when_it_fits() {
782        let doc = group_expanded(concat(vec![
783            text("("),
784            indent(concat(vec![soft_line(), text("a")])),
785            soft_line(),
786            text(")"),
787        ]));
788        assert_eq!(p(&doc, 80), "(\n    a\n)\n");
789    }
790
791    #[test]
792    fn expanded_inner_group_forces_outer_to_break() {
793        // An exploded inner collection can't sit inline, so the outer group breaks too: the soft
794        // line before the inner group becomes a newline.
795        let inner = group_expanded(concat(vec![
796            text("("),
797            indent(concat(vec![soft_line(), text("x")])),
798            soft_line(),
799            text(")"),
800        ]));
801        let outer = group(concat(vec![text("a"), line(), inner]));
802        assert_eq!(p(&outer, 80), "a\n(\n    x\n)\n");
803    }
804
805    #[test]
806    fn line_suffix_defers_content_to_the_end_of_the_line() {
807        // The suffix is emitted before the hard line, even though it appears first in the concat.
808        let doc = concat(vec![
809            line_suffix(text(" -- note")),
810            text("code"),
811            hard_line(),
812            text("next"),
813        ]);
814        assert_eq!(p(&doc, 80), "code -- note\nnext\n");
815    }
816
817    #[test]
818    fn line_suffix_flushes_at_end_of_document() {
819        let doc = concat(vec![line_suffix(text(" -- note")), text("code")]);
820        assert_eq!(p(&doc, 80), "code -- note\n");
821    }
822
823    #[test]
824    fn break_parent_forces_its_group_to_break() {
825        let doc = group(concat(vec![
826            text("a"),
827            break_parent(),
828            indent(concat(vec![line(), text("b")])),
829        ]));
830        assert_eq!(p(&doc, 80), "a\n    b\n");
831    }
832
833    #[test]
834    fn empty_document_prints_nothing() {
835        assert_eq!(p(&empty(), 80), "");
836    }
837
838    #[test]
839    fn trailing_whitespace_is_trimmed() {
840        // A broken group whose indented line is immediately followed by a hard line must not leave
841        // spaces dangling on the blank line.
842        let doc = concat(vec![text("a"), indent(hard_line()), hard_line(), text("b")]);
843        assert_eq!(p(&doc, 80), "a\n\nb\n");
844    }
845
846    #[test]
847    fn text_caches_its_display_width() {
848        // The width folded in by `text` must equal a fresh measurement, for both ASCII (fast path)
849        // and wide CJK input — this is what the fit/break decision now reads instead of re-scanning.
850        for s in ["select", "a, b, c", "長芋", "a長b", ""] {
851            match text(s) {
852                Doc::Text(stored, width) => {
853                    assert_eq!(stored, s);
854                    assert_eq!(width, text_width(s));
855                }
856                other => panic!("text() should build a Text node, got {other:?}"),
857            }
858        }
859        // The ASCII fast path agrees with the per-char fold.
860        assert_eq!(
861            text_width("SELECT a, b"),
862            "SELECT a, b".chars().map(char_width).sum::<usize>()
863        );
864    }
865
866    #[test]
867    fn group_precomputes_its_forced_break() {
868        // A plain group over flat content stays optional; a group wrapping a hard line is marked
869        // must-break at construction, which is what the printer reads in O(1).
870        match group(concat(vec![text("a"), line(), text("b")])) {
871            Doc::Group { must_break, .. } => assert!(!must_break),
872            other => panic!("expected a group, got {other:?}"),
873        }
874        match group(concat(vec![text("a"), hard_line(), text("b")])) {
875            Doc::Group { must_break, .. } => assert!(must_break),
876            other => panic!("expected a group, got {other:?}"),
877        }
878        // group_expanded is always must-break.
879        match group_expanded(text("x")) {
880            Doc::Group { must_break, .. } => assert!(must_break),
881            other => panic!("expected a group, got {other:?}"),
882        }
883    }
884
885    #[test]
886    fn finalize_drops_leading_and_trailing_blank_lines_but_keeps_interior() {
887        // Leading and trailing blank lines vanish; an interior blank-line run is preserved; every
888        // line is right-trimmed. Equivalent to the previous collect/join/trim implementation.
889        let raw = String::from("\n  \na  \n\n\nb \n\n  \n");
890        assert_eq!(finalize(raw), "a\n\n\nb\n");
891        assert_eq!(finalize(String::new()), "");
892        assert_eq!(finalize(String::from("   \n  ")), "");
893        assert_eq!(finalize(String::from("only")), "only\n");
894    }
895
896    #[test]
897    fn source_code_slice_single_line_behaves_like_text() {
898        let doc = concat(vec![text("a "), source_code_slice("b = c"), text(" d")]);
899        assert_eq!(p(&doc, 80), "a b = c d\n");
900    }
901
902    #[test]
903    fn source_code_slice_caches_its_last_line_width() {
904        match source_code_slice("alpha\nbeta\nlong-tail") {
905            Doc::SourceCodeSlice(s, width) => {
906                assert_eq!(s, "alpha\nbeta\nlong-tail");
907                assert_eq!(width, "long-tail".len());
908            }
909            other => panic!("expected a source slice, got {other:?}"),
910        }
911        match source_code_slice("x\n") {
912            Doc::SourceCodeSlice(_, width) => assert_eq!(width, 0),
913            other => panic!("expected a source slice, got {other:?}"),
914        }
915    }
916
917    #[test]
918    fn source_code_slice_reindents_continuation_lines_under_indent() {
919        let doc = concat(vec![
920            text("header"),
921            indent(concat(vec![hard_line(), source_code_slice("v1\nv2\nv3")])),
922            hard_line(),
923            text("footer"),
924        ]);
925        assert_eq!(p(&doc, 80), "header\n    v1\n    v2\n    v3\nfooter\n");
926    }
927
928    #[test]
929    fn source_code_slice_tail_width_participates_in_fits() {
930        let doc = group(concat(vec![
931            source_code_slice("aa\nbbbb"),
932            line(),
933            text("c"),
934        ]));
935        assert_eq!(p(&doc, 6), "aa\nbbbb c\n");
936        assert_eq!(p(&doc, 5), "aa\nbbbb\nc\n");
937    }
938
939    #[test]
940    fn if_group_breaks_selects_layout_by_group_mode() {
941        let doc = group(concat(vec![
942            text("a"),
943            if_group_breaks(text(","), empty()),
944            line(),
945            text("b"),
946        ]));
947        assert_eq!(p(&doc, 80), "a b\n");
948        assert_eq!(p(&doc, 1), "a,\nb\n");
949    }
950
951    #[test]
952    fn if_group_breaks_flat_arm_forced_break_propagates_but_broken_arm_does_not() {
953        assert!(has_forced_break(&if_group_breaks(empty(), hard_line())));
954        assert!(!has_forced_break(&if_group_breaks(hard_line(), empty())));
955    }
956
957    #[test]
958    fn block_indent_frames_content_on_its_own_indented_line() {
959        let doc = concat(vec![text("{"), block_indent(text("body")), text("}")]);
960        assert_eq!(p(&doc, 80), "{\n    body\n}\n");
961    }
962
963    #[test]
964    fn doc_write_composes_format_values_in_order() {
965        let mut f = DocBuffer::new();
966        doc_write!(f, [text("a"), text(" = "), text("b")]);
967        assert_eq!(p(&group(f.finish()), 80), "a = b\n");
968    }
969
970    #[test]
971    fn format_rule_decouples_layout_from_data() {
972        struct AssignRule;
973        impl FormatRule<str> for AssignRule {
974            fn fmt(&self, item: &str, buffer: &mut DocBuffer) {
975                doc_write!(
976                    buffer,
977                    [text("x => \""), text(item.to_string()), text("\"")]
978                );
979            }
980        }
981
982        let mut f = DocBuffer::new();
983        f.write_with_rule(&AssignRule, "v");
984        assert_eq!(p(&f.finish(), 80), "x => \"v\"\n");
985    }
986}