1use crate::config::FormatConfig;
7use crate::doc::Doc;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11enum Mode {
12 Flat,
14 Break,
16}
17
18#[derive(Debug, Clone)]
20struct PrintCommand<'a> {
21 indent: usize,
23 mode: Mode,
25 doc: &'a Doc,
27}
28
29pub struct Printer {
31 config: FormatConfig,
32}
33
34impl Printer {
35 pub fn new(config: FormatConfig) -> Self {
37 Self { config }
38 }
39
40 pub fn print(&self, doc: &Doc) -> String {
42 let mut output = String::new();
43 let mut pos = 0; let mut stack = vec![PrintCommand {
45 indent: 0,
46 mode: Mode::Break,
47 doc,
48 }];
49
50 while let Some(cmd) = stack.pop() {
51 match cmd.doc {
52 Doc::Nil => {}
53
54 Doc::Text(s) => {
55 output.push_str(s);
56 pos += s.chars().count();
57 }
58
59 Doc::Line => {
60 if cmd.mode == Mode::Flat {
61 output.push(' ');
62 pos += 1;
63 } else {
64 output.push('\n');
65 output.push_str(&self.indent_string(cmd.indent));
66 pos = cmd.indent * self.config.indent_width;
67 }
68 }
69
70 Doc::SoftLine => {
71 if cmd.mode == Mode::Flat {
72 } else {
74 output.push('\n');
75 output.push_str(&self.indent_string(cmd.indent));
76 pos = cmd.indent * self.config.indent_width;
77 }
78 }
79
80 Doc::HardLine => {
81 output.push('\n');
82 output.push_str(&self.indent_string(cmd.indent));
83 pos = cmd.indent * self.config.indent_width;
84 }
85
86 Doc::Indent(inner) => {
87 stack.push(PrintCommand {
88 indent: cmd.indent + 1,
89 mode: cmd.mode,
90 doc: inner,
91 });
92 }
93
94 Doc::Concat(left, right) => {
95 stack.push(PrintCommand {
97 indent: cmd.indent,
98 mode: cmd.mode,
99 doc: right,
100 });
101 stack.push(PrintCommand {
102 indent: cmd.indent,
103 mode: cmd.mode,
104 doc: left,
105 });
106 }
107
108 Doc::Group(inner) => {
109 if self.fits(cmd.indent, pos, inner) {
111 stack.push(PrintCommand {
112 indent: cmd.indent,
113 mode: Mode::Flat,
114 doc: inner,
115 });
116 } else {
117 stack.push(PrintCommand {
118 indent: cmd.indent,
119 mode: Mode::Break,
120 doc: inner,
121 });
122 }
123 }
124
125 Doc::IfBreak { flat, broken } => {
126 let chosen = if cmd.mode == Mode::Flat { flat } else { broken };
127 stack.push(PrintCommand {
128 indent: cmd.indent,
129 mode: cmd.mode,
130 doc: chosen,
131 });
132 }
133 }
134 }
135
136 output
137 }
138
139 fn fits(&self, indent: usize, current_pos: usize, doc: &Doc) -> bool {
141 let remaining = self.config.max_width.saturating_sub(current_pos);
142 self.fits_within(remaining, indent, doc)
143 }
144
145 fn fits_within(&self, mut remaining: usize, indent: usize, doc: &Doc) -> bool {
147 let mut stack = vec![(indent, doc)];
148
149 while let Some((ind, d)) = stack.pop() {
150 match d {
151 Doc::Nil => {}
152
153 Doc::Text(s) => {
154 let len = s.chars().count();
155 if len > remaining {
156 return false;
157 }
158 remaining -= len;
159 }
160
161 Doc::Line => {
162 if remaining == 0 {
164 return false;
165 }
166 remaining -= 1;
167 }
168
169 Doc::SoftLine => {
170 }
172
173 Doc::HardLine => {
174 return false;
176 }
177
178 Doc::Indent(inner) => {
179 stack.push((ind + 1, inner));
180 }
181
182 Doc::Concat(left, right) => {
183 stack.push((ind, right));
184 stack.push((ind, left));
185 }
186
187 Doc::Group(inner) => {
188 stack.push((ind, inner));
190 }
191
192 Doc::IfBreak { flat, .. } => {
193 stack.push((ind, flat));
195 }
196 }
197 }
198
199 true
200 }
201
202 fn indent_string(&self, level: usize) -> String {
204 let width = level * self.config.indent_width;
205 if self.config.use_tabs {
206 "\t".repeat(level)
207 } else {
208 " ".repeat(width)
209 }
210 }
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216
217 fn print(doc: &Doc) -> String {
218 Printer::new(FormatConfig::default()).print(doc)
219 }
220
221 fn print_width(doc: &Doc, width: usize) -> String {
222 Printer::new(FormatConfig {
223 max_width: width,
224 ..Default::default()
225 })
226 .print(doc)
227 }
228
229 #[test]
230 fn test_text() {
231 let doc = Doc::text("hello");
232 assert_eq!(print(&doc), "hello");
233 }
234
235 #[test]
236 fn test_concat() {
237 let doc = Doc::text("hello").concat(Doc::text(" world"));
238 assert_eq!(print(&doc), "hello world");
239 }
240
241 #[test]
242 fn test_line_in_break_mode() {
243 let doc = Doc::text("a").concat(Doc::line()).concat(Doc::text("b"));
244 assert_eq!(print(&doc), "a\nb");
246 }
247
248 #[test]
249 fn test_group_fits() {
250 let doc = Doc::group(Doc::text("a").concat(Doc::line()).concat(Doc::text("b")));
251 assert_eq!(print_width(&doc, 80), "a b");
253 }
254
255 #[test]
256 fn test_group_breaks() {
257 let doc = Doc::group(
258 Doc::text("hello")
259 .concat(Doc::line())
260 .concat(Doc::text("world")),
261 );
262 assert_eq!(print_width(&doc, 5), "hello\nworld");
264 }
265
266 #[test]
267 fn test_indent() {
268 let doc = Doc::text("a").concat(Doc::indent(Doc::hardline().concat(Doc::text("b"))));
270 assert_eq!(print(&doc), "a\n b");
271 }
272
273 #[test]
274 fn test_softline_flat() {
275 let doc = Doc::group(
276 Doc::text("[")
277 .concat(Doc::softline())
278 .concat(Doc::text("]")),
279 );
280 assert_eq!(print_width(&doc, 80), "[]");
282 }
283
284 #[test]
285 fn test_softline_break() {
286 let doc = Doc::group(
287 Doc::text("[")
288 .concat(Doc::softline())
289 .concat(Doc::text("very_long_content"))
290 .concat(Doc::softline())
291 .concat(Doc::text("]")),
292 );
293 assert_eq!(print_width(&doc, 10), "[\nvery_long_content\n]");
295 }
296
297 #[test]
298 fn test_hardline_forces_break() {
299 let doc = Doc::group(
300 Doc::text("a")
301 .concat(Doc::hardline())
302 .concat(Doc::text("b")),
303 );
304 assert_eq!(print_width(&doc, 80), "a\nb");
306 }
307
308 #[test]
309 fn test_if_break() {
310 let trailing_comma = Doc::if_break(Doc::Nil, Doc::text(","));
311
312 let doc = Doc::group(
313 Doc::text("[")
314 .concat(Doc::softline())
315 .concat(Doc::text("a"))
316 .concat(trailing_comma.clone())
317 .concat(Doc::softline())
318 .concat(Doc::text("]")),
319 );
320
321 assert_eq!(print_width(&doc, 80), "[a]");
323
324 let doc = Doc::group(
326 Doc::text("[")
327 .concat(Doc::softline())
328 .concat(Doc::text("very_long_item"))
329 .concat(trailing_comma)
330 .concat(Doc::softline())
331 .concat(Doc::text("]")),
332 );
333 assert_eq!(print_width(&doc, 10), "[\nvery_long_item,\n]");
334 }
335
336 #[test]
337 fn test_nested_indent() {
338 let inner = Doc::indent(Doc::hardline().concat(Doc::text("c")));
340 let outer = Doc::indent(Doc::hardline().concat(Doc::text("b")).concat(inner));
341 let doc = Doc::text("a").concat(outer);
342
343 assert_eq!(print(&doc), "a\n b\n c");
344 }
345}