prettyplease/
algorithm.rs

1// Adapted from https://github.com/rust-lang/rust/blob/1.57.0/compiler/rustc_ast_pretty/src/pp.rs.
2// See "Algorithm notes" in the crate-level rustdoc.
3
4use crate::ring::RingBuffer;
5use crate::{MARGIN, MIN_SPACE};
6use std::borrow::Cow;
7use std::cmp;
8use std::collections::VecDeque;
9use std::iter;
10
11#[derive(Clone, Copy, PartialEq)]
12pub enum Breaks {
13    Consistent,
14    Inconsistent,
15}
16
17#[derive(Clone, Copy, Default)]
18pub struct BreakToken {
19    pub offset: isize,
20    pub blank_space: usize,
21    pub pre_break: Option<char>,
22    pub post_break: &'static str,
23    pub no_break: Option<char>,
24    pub if_nonempty: bool,
25    pub never_break: bool,
26}
27
28#[derive(Clone, Copy)]
29pub struct BeginToken {
30    pub offset: isize,
31    pub breaks: Breaks,
32}
33
34#[derive(Clone)]
35pub enum Token {
36    String(Cow<'static, str>),
37    Break(BreakToken),
38    Begin(BeginToken),
39    End,
40}
41
42#[derive(Copy, Clone)]
43enum PrintFrame {
44    Fits(Breaks),
45    Broken(usize, Breaks),
46}
47
48pub const SIZE_INFINITY: isize = 0xffff;
49
50pub struct Printer {
51    out: String,
52    // Number of spaces left on line
53    space: isize,
54    // Ring-buffer of tokens and calculated sizes
55    buf: RingBuffer<BufEntry>,
56    // Total size of tokens already printed
57    left_total: isize,
58    // Total size of tokens enqueued, including printed and not yet printed
59    right_total: isize,
60    // Holds the ring-buffer index of the Begin that started the current block,
61    // possibly with the most recent Break after that Begin (if there is any) on
62    // top of it. Values are pushed and popped on the back of the queue using it
63    // like stack, and elsewhere old values are popped from the front of the
64    // queue as they become irrelevant due to the primary ring-buffer advancing.
65    scan_stack: VecDeque<usize>,
66    // Stack of blocks-in-progress being flushed by print
67    print_stack: Vec<PrintFrame>,
68    // Level of indentation of current line
69    indent: usize,
70    // Buffered indentation to avoid writing trailing whitespace
71    pending_indentation: usize,
72}
73
74#[derive(Clone)]
75struct BufEntry {
76    token: Token,
77    size: isize,
78}
79
80impl Printer {
81    pub fn new() -> Self {
82        Printer {
83            out: String::new(),
84            space: MARGIN,
85            buf: RingBuffer::new(),
86            left_total: 0,
87            right_total: 0,
88            scan_stack: VecDeque::new(),
89            print_stack: Vec::new(),
90            indent: 0,
91            pending_indentation: 0,
92        }
93    }
94
95    pub fn eof(mut self) -> String {
96        if !self.scan_stack.is_empty() {
97            self.check_stack(0);
98            self.advance_left();
99        }
100        self.out
101    }
102
103    pub fn scan_begin(&mut self, token: BeginToken) {
104        if self.scan_stack.is_empty() {
105            self.left_total = 1;
106            self.right_total = 1;
107            self.buf.clear();
108        }
109        let right = self.buf.push(BufEntry {
110            token: Token::Begin(token),
111            size: -self.right_total,
112        });
113        self.scan_stack.push_back(right);
114    }
115
116    pub fn scan_end(&mut self) {
117        if self.scan_stack.is_empty() {
118            self.print_end();
119        } else {
120            if !self.buf.is_empty() {
121                if let Token::Break(break_token) = self.buf.last().token {
122                    if self.buf.len() >= 2 {
123                        if let Token::Begin(_) = self.buf.second_last().token {
124                            self.buf.pop_last();
125                            self.buf.pop_last();
126                            self.scan_stack.pop_back();
127                            self.scan_stack.pop_back();
128                            self.right_total -= break_token.blank_space as isize;
129                            return;
130                        }
131                    }
132                    if break_token.if_nonempty {
133                        self.buf.pop_last();
134                        self.scan_stack.pop_back();
135                        self.right_total -= break_token.blank_space as isize;
136                    }
137                }
138            }
139            let right = self.buf.push(BufEntry {
140                token: Token::End,
141                size: -1,
142            });
143            self.scan_stack.push_back(right);
144        }
145    }
146
147    pub fn scan_break(&mut self, token: BreakToken) {
148        if self.scan_stack.is_empty() {
149            self.left_total = 1;
150            self.right_total = 1;
151            self.buf.clear();
152        } else {
153            self.check_stack(0);
154        }
155        let right = self.buf.push(BufEntry {
156            token: Token::Break(token),
157            size: -self.right_total,
158        });
159        self.scan_stack.push_back(right);
160        self.right_total += token.blank_space as isize;
161    }
162
163    pub fn scan_string(&mut self, string: Cow<'static, str>) {
164        if self.scan_stack.is_empty() {
165            self.print_string(string);
166        } else {
167            let len = string.len() as isize;
168            self.buf.push(BufEntry {
169                token: Token::String(string),
170                size: len,
171            });
172            self.right_total += len;
173            self.check_stream();
174        }
175    }
176
177    pub fn offset(&mut self, offset: isize) {
178        match &mut self.buf.last_mut().token {
179            Token::Break(token) => token.offset += offset,
180            Token::Begin(_) => {}
181            Token::String(_) | Token::End => unreachable!(),
182        }
183    }
184
185    pub fn end_with_max_width(&mut self, max: isize) {
186        let mut depth = 1;
187        for &index in self.scan_stack.iter().rev() {
188            let entry = &self.buf[index];
189            match entry.token {
190                Token::Begin(_) => {
191                    depth -= 1;
192                    if depth == 0 {
193                        if entry.size < 0 {
194                            let actual_width = entry.size + self.right_total;
195                            if actual_width > max {
196                                self.buf.push(BufEntry {
197                                    token: Token::String(Cow::Borrowed("")),
198                                    size: SIZE_INFINITY,
199                                });
200                                self.right_total += SIZE_INFINITY;
201                            }
202                        }
203                        break;
204                    }
205                }
206                Token::End => depth += 1,
207                Token::Break(_) => {}
208                Token::String(_) => unreachable!(),
209            }
210        }
211        self.scan_end();
212    }
213
214    pub fn ends_with(&self, ch: char) -> bool {
215        for i in self.buf.index_range().rev() {
216            if let Token::String(token) = &self.buf[i].token {
217                return token.ends_with(ch);
218            }
219        }
220        self.out.ends_with(ch)
221    }
222
223    fn check_stream(&mut self) {
224        while self.right_total - self.left_total > self.space {
225            if *self.scan_stack.front().unwrap() == self.buf.index_range().start {
226                self.scan_stack.pop_front().unwrap();
227                self.buf.first_mut().size = SIZE_INFINITY;
228            }
229
230            self.advance_left();
231
232            if self.buf.is_empty() {
233                break;
234            }
235        }
236    }
237
238    fn advance_left(&mut self) {
239        while self.buf.first().size >= 0 {
240            let left = self.buf.pop_first();
241
242            match left.token {
243                Token::String(string) => {
244                    self.left_total += left.size;
245                    self.print_string(string);
246                }
247                Token::Break(token) => {
248                    self.left_total += token.blank_space as isize;
249                    self.print_break(token, left.size);
250                }
251                Token::Begin(token) => self.print_begin(token, left.size),
252                Token::End => self.print_end(),
253            }
254
255            if self.buf.is_empty() {
256                break;
257            }
258        }
259    }
260
261    fn check_stack(&mut self, mut depth: usize) {
262        while let Some(&index) = self.scan_stack.back() {
263            let entry = &mut self.buf[index];
264            match entry.token {
265                Token::Begin(_) => {
266                    if depth == 0 {
267                        break;
268                    }
269                    self.scan_stack.pop_back().unwrap();
270                    entry.size += self.right_total;
271                    depth -= 1;
272                }
273                Token::End => {
274                    self.scan_stack.pop_back().unwrap();
275                    entry.size = 1;
276                    depth += 1;
277                }
278                Token::Break(_) => {
279                    self.scan_stack.pop_back().unwrap();
280                    entry.size += self.right_total;
281                    if depth == 0 {
282                        break;
283                    }
284                }
285                Token::String(_) => unreachable!(),
286            }
287        }
288    }
289
290    fn get_top(&self) -> PrintFrame {
291        const OUTER: PrintFrame = PrintFrame::Broken(0, Breaks::Inconsistent);
292        self.print_stack.last().map_or(OUTER, PrintFrame::clone)
293    }
294
295    fn print_begin(&mut self, token: BeginToken, size: isize) {
296        if cfg!(prettyplease_debug) {
297            self.out.push(match token.breaks {
298                Breaks::Consistent => '«',
299                Breaks::Inconsistent => '‹',
300            });
301            if cfg!(prettyplease_debug_indent) {
302                self.out
303                    .extend(token.offset.to_string().chars().map(|ch| match ch {
304                        '0'..='9' => ['₀', '₁', '₂', '₃', '₄', '₅', '₆', '₇', '₈', '₉']
305                            [(ch as u8 - b'0') as usize],
306                        '-' => '₋',
307                        _ => unreachable!(),
308                    }));
309            }
310        }
311        if size > self.space {
312            self.print_stack
313                .push(PrintFrame::Broken(self.indent, token.breaks));
314            self.indent = usize::try_from(self.indent as isize + token.offset).unwrap();
315        } else {
316            self.print_stack.push(PrintFrame::Fits(token.breaks));
317        }
318    }
319
320    fn print_end(&mut self) {
321        let breaks = match self.print_stack.pop().unwrap() {
322            PrintFrame::Broken(indent, breaks) => {
323                self.indent = indent;
324                breaks
325            }
326            PrintFrame::Fits(breaks) => breaks,
327        };
328        if cfg!(prettyplease_debug) {
329            self.out.push(match breaks {
330                Breaks::Consistent => '»',
331                Breaks::Inconsistent => '›',
332            });
333        }
334    }
335
336    fn print_break(&mut self, token: BreakToken, size: isize) {
337        let fits = token.never_break
338            || match self.get_top() {
339                PrintFrame::Fits(..) => true,
340                PrintFrame::Broken(.., Breaks::Consistent) => false,
341                PrintFrame::Broken(.., Breaks::Inconsistent) => size <= self.space,
342            };
343        if fits {
344            self.pending_indentation += token.blank_space;
345            self.space -= token.blank_space as isize;
346            if let Some(no_break) = token.no_break {
347                self.out.push(no_break);
348                self.space -= no_break.len_utf8() as isize;
349            }
350            if cfg!(prettyplease_debug) {
351                self.out.push('·');
352            }
353        } else {
354            if let Some(pre_break) = token.pre_break {
355                self.print_indent();
356                self.out.push(pre_break);
357            }
358            if cfg!(prettyplease_debug) {
359                self.out.push('·');
360            }
361            self.out.push('\n');
362            let indent = self.indent as isize + token.offset;
363            self.pending_indentation = usize::try_from(indent).unwrap();
364            self.space = cmp::max(MARGIN - indent, MIN_SPACE);
365            if !token.post_break.is_empty() {
366                self.print_indent();
367                self.out.push_str(token.post_break);
368                self.space -= token.post_break.len() as isize;
369            }
370        }
371    }
372
373    fn print_string(&mut self, string: Cow<'static, str>) {
374        self.print_indent();
375        self.out.push_str(&string);
376        self.space -= string.len() as isize;
377    }
378
379    fn print_indent(&mut self) {
380        self.out.reserve(self.pending_indentation);
381        self.out
382            .extend(iter::repeat(' ').take(self.pending_indentation));
383        self.pending_indentation = 0;
384    }
385}