Skip to main content

tstring_format_doc/
lib.rs

1use unicode_width::UnicodeWidthStr;
2
3#[derive(Clone, Debug, PartialEq, Eq)]
4pub enum Doc {
5    Text(String),
6    Concat(Vec<Doc>),
7    Indent(Box<Doc>),
8    Line,
9    SoftLine,
10    HardLine,
11    Group(Box<Doc>),
12}
13
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub struct RenderOptions {
16    pub line_length: usize,
17    pub indent_width: usize,
18}
19
20impl Default for RenderOptions {
21    fn default() -> Self {
22        Self {
23            line_length: 80,
24            indent_width: 2,
25        }
26    }
27}
28
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30enum Mode {
31    Flat,
32    Break,
33}
34
35#[derive(Clone, Debug)]
36struct Command<'a> {
37    indent: usize,
38    mode: Mode,
39    doc: &'a Doc,
40}
41
42impl Doc {
43    #[must_use]
44    pub fn text(value: impl Into<String>) -> Self {
45        Self::Text(value.into())
46    }
47
48    #[must_use]
49    pub fn concat(parts: Vec<Doc>) -> Self {
50        Self::Concat(parts)
51    }
52
53    #[must_use]
54    pub fn indent(self) -> Self {
55        Self::Indent(Box::new(self))
56    }
57
58    #[must_use]
59    pub fn group(self) -> Self {
60        Self::Group(Box::new(self))
61    }
62
63    #[must_use]
64    pub fn line() -> Self {
65        Self::Line
66    }
67
68    #[must_use]
69    pub fn soft_line() -> Self {
70        Self::SoftLine
71    }
72
73    #[must_use]
74    pub fn hard_line() -> Self {
75        Self::HardLine
76    }
77}
78
79#[must_use]
80pub fn render(doc: &Doc, options: RenderOptions) -> String {
81    let mut out = String::new();
82    let mut width = 0usize;
83    let mut stack = vec![Command {
84        indent: 0,
85        mode: Mode::Break,
86        doc,
87    }];
88
89    while let Some(command) = stack.pop() {
90        match command.doc {
91            Doc::Text(text) => {
92                out.push_str(text);
93                width = width_after_text(width, text);
94            }
95            Doc::Concat(parts) => {
96                for part in parts.iter().rev() {
97                    stack.push(Command {
98                        indent: command.indent,
99                        mode: command.mode,
100                        doc: part,
101                    });
102                }
103            }
104            Doc::Indent(inner) => {
105                stack.push(Command {
106                    indent: command.indent + options.indent_width,
107                    mode: command.mode,
108                    doc: inner,
109                });
110            }
111            Doc::Line => match command.mode {
112                Mode::Flat => {
113                    out.push(' ');
114                    width += 1;
115                }
116                Mode::Break => {
117                    push_newline(&mut out, command.indent);
118                    width = command.indent;
119                }
120            },
121            Doc::SoftLine => match command.mode {
122                Mode::Flat => {}
123                Mode::Break => {
124                    push_newline(&mut out, command.indent);
125                    width = command.indent;
126                }
127            },
128            Doc::HardLine => {
129                push_newline(&mut out, command.indent);
130                width = command.indent;
131            }
132            Doc::Group(inner) => {
133                let next_mode = match command.mode {
134                    Mode::Flat => Mode::Flat,
135                    Mode::Break => {
136                        let remaining = options.line_length.saturating_sub(width);
137                        if fits(
138                            remaining,
139                            vec![Command {
140                                indent: command.indent,
141                                mode: Mode::Flat,
142                                doc: inner,
143                            }],
144                        ) {
145                            Mode::Flat
146                        } else {
147                            Mode::Break
148                        }
149                    }
150                };
151                stack.push(Command {
152                    indent: command.indent,
153                    mode: next_mode,
154                    doc: inner,
155                });
156            }
157        }
158    }
159
160    out
161}
162
163#[must_use]
164pub fn flat_width(doc: &Doc) -> Option<usize> {
165    let mut remaining = usize::MAX;
166    if fits(
167        remaining,
168        vec![Command {
169            indent: 0,
170            mode: Mode::Flat,
171            doc,
172        }],
173    ) {
174        accumulate_flat_width(doc, &mut remaining)?;
175        Some(usize::MAX - remaining)
176    } else {
177        None
178    }
179}
180
181#[must_use]
182pub fn has_forced_break(doc: &Doc) -> bool {
183    match doc {
184        Doc::Text(text) => text.contains('\n') || text.contains('\r'),
185        Doc::Concat(parts) => parts.iter().any(has_forced_break),
186        Doc::Indent(inner) | Doc::Group(inner) => has_forced_break(inner),
187        Doc::Line | Doc::SoftLine => false,
188        Doc::HardLine => true,
189    }
190}
191
192fn fits(mut remaining: usize, mut stack: Vec<Command<'_>>) -> bool {
193    while let Some(command) = stack.pop() {
194        match command.doc {
195            Doc::Text(text) => {
196                let Some(width) = text_flat_width(text) else {
197                    return false;
198                };
199                if width > remaining {
200                    return false;
201                }
202                remaining -= width;
203            }
204            Doc::Concat(parts) => {
205                for part in parts.iter().rev() {
206                    stack.push(Command {
207                        indent: command.indent,
208                        mode: command.mode,
209                        doc: part,
210                    });
211                }
212            }
213            Doc::Indent(inner) => stack.push(Command {
214                indent: command.indent,
215                mode: command.mode,
216                doc: inner,
217            }),
218            Doc::Line => {
219                if remaining == 0 {
220                    return false;
221                }
222                remaining -= 1;
223            }
224            Doc::SoftLine => {}
225            Doc::HardLine => return false,
226            Doc::Group(inner) => stack.push(Command {
227                indent: command.indent,
228                mode: Mode::Flat,
229                doc: inner,
230            }),
231        }
232    }
233
234    true
235}
236
237fn accumulate_flat_width(doc: &Doc, remaining: &mut usize) -> Option<()> {
238    match doc {
239        Doc::Text(text) => {
240            *remaining -= text_flat_width(text)?;
241        }
242        Doc::Concat(parts) => {
243            for part in parts {
244                accumulate_flat_width(part, remaining)?;
245            }
246        }
247        Doc::Indent(inner) | Doc::Group(inner) => {
248            accumulate_flat_width(inner, remaining)?;
249        }
250        Doc::Line => *remaining -= 1,
251        Doc::SoftLine => {}
252        Doc::HardLine => return None,
253    }
254    Some(())
255}
256
257fn push_newline(out: &mut String, indent: usize) {
258    out.push('\n');
259    for _ in 0..indent {
260        out.push(' ');
261    }
262}
263
264fn width_after_text(current_width: usize, text: &str) -> usize {
265    match text.rsplit_once('\n') {
266        Some((_, tail)) => tail.width(),
267        None => current_width + text.width(),
268    }
269}
270
271fn text_flat_width(text: &str) -> Option<usize> {
272    if text.contains('\n') || text.contains('\r') {
273        return None;
274    }
275    Some(text.width())
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn groups_break_when_they_do_not_fit() {
284        let doc = Doc::concat(vec![
285            Doc::text("<div"),
286            Doc::concat(vec![Doc::line(), Doc::text("class=\"hello\"")]).indent(),
287            Doc::soft_line(),
288            Doc::text(">"),
289        ])
290        .group();
291
292        assert_eq!(
293            render(
294                &doc,
295                RenderOptions {
296                    line_length: 10,
297                    indent_width: 2,
298                }
299            ),
300            "<div\n  class=\"hello\"\n>"
301        );
302    }
303
304    #[test]
305    fn text_with_newline_forces_break() {
306        let doc = Doc::concat(vec![Doc::text("{\n  value\n}"), Doc::text("</div>")]);
307        assert!(has_forced_break(&doc));
308        assert_eq!(flat_width(&doc), None);
309    }
310}