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: Option<char>,
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    fn check_stream(&mut self) {
215        while self.right_total - self.left_total > self.space {
216            if *self.scan_stack.front().unwrap() == self.buf.index_of_first() {
217                self.scan_stack.pop_front().unwrap();
218                self.buf.first_mut().size = SIZE_INFINITY;
219            }
220
221            self.advance_left();
222
223            if self.buf.is_empty() {
224                break;
225            }
226        }
227    }
228
229    fn advance_left(&mut self) {
230        while self.buf.first().size >= 0 {
231            let left = self.buf.pop_first();
232
233            match left.token {
234                Token::String(string) => {
235                    self.left_total += left.size;
236                    self.print_string(string);
237                }
238                Token::Break(token) => {
239                    self.left_total += token.blank_space as isize;
240                    self.print_break(token, left.size);
241                }
242                Token::Begin(token) => self.print_begin(token, left.size),
243                Token::End => self.print_end(),
244            }
245
246            if self.buf.is_empty() {
247                break;
248            }
249        }
250    }
251
252    fn check_stack(&mut self, mut depth: usize) {
253        while let Some(&index) = self.scan_stack.back() {
254            let entry = &mut self.buf[index];
255            match entry.token {
256                Token::Begin(_) => {
257                    if depth == 0 {
258                        break;
259                    }
260                    self.scan_stack.pop_back().unwrap();
261                    entry.size += self.right_total;
262                    depth -= 1;
263                }
264                Token::End => {
265                    self.scan_stack.pop_back().unwrap();
266                    entry.size = 1;
267                    depth += 1;
268                }
269                Token::Break(_) => {
270                    self.scan_stack.pop_back().unwrap();
271                    entry.size += self.right_total;
272                    if depth == 0 {
273                        break;
274                    }
275                }
276                Token::String(_) => unreachable!(),
277            }
278        }
279    }
280
281    fn get_top(&self) -> PrintFrame {
282        const OUTER: PrintFrame = PrintFrame::Broken(0, Breaks::Inconsistent);
283        self.print_stack.last().map_or(OUTER, PrintFrame::clone)
284    }
285
286    fn print_begin(&mut self, token: BeginToken, size: isize) {
287        if cfg!(prettyplease_debug) {
288            self.out.push(match token.breaks {
289                Breaks::Consistent => '«',
290                Breaks::Inconsistent => '‹',
291            });
292            if cfg!(prettyplease_debug_indent) {
293                self.out
294                    .extend(token.offset.to_string().chars().map(|ch| match ch {
295                        '0'..='9' => ['₀', '₁', '₂', '₃', '₄', '₅', '₆', '₇', '₈', '₉']
296                            [(ch as u8 - b'0') as usize],
297                        '-' => '₋',
298                        _ => unreachable!(),
299                    }));
300            }
301        }
302        if size > self.space {
303            self.print_stack
304                .push(PrintFrame::Broken(self.indent, token.breaks));
305            self.indent = usize::try_from(self.indent as isize + token.offset).unwrap();
306        } else {
307            self.print_stack.push(PrintFrame::Fits(token.breaks));
308        }
309    }
310
311    fn print_end(&mut self) {
312        let breaks = match self.print_stack.pop().unwrap() {
313            PrintFrame::Broken(indent, breaks) => {
314                self.indent = indent;
315                breaks
316            }
317            PrintFrame::Fits(breaks) => breaks,
318        };
319        if cfg!(prettyplease_debug) {
320            self.out.push(match breaks {
321                Breaks::Consistent => '»',
322                Breaks::Inconsistent => '›',
323            });
324        }
325    }
326
327    fn print_break(&mut self, token: BreakToken, size: isize) {
328        let fits = token.never_break
329            || match self.get_top() {
330                PrintFrame::Fits(..) => true,
331                PrintFrame::Broken(.., Breaks::Consistent) => false,
332                PrintFrame::Broken(.., Breaks::Inconsistent) => size <= self.space,
333            };
334        if fits {
335            self.pending_indentation += token.blank_space;
336            self.space -= token.blank_space as isize;
337            if let Some(no_break) = token.no_break {
338                self.out.push(no_break);
339                self.space -= no_break.len_utf8() as isize;
340            }
341            if cfg!(prettyplease_debug) {
342                self.out.push('·');
343            }
344        } else {
345            if let Some(pre_break) = token.pre_break {
346                self.print_indent();
347                self.out.push(pre_break);
348            }
349            if cfg!(prettyplease_debug) {
350                self.out.push('·');
351            }
352            self.out.push('\n');
353            let indent = self.indent as isize + token.offset;
354            self.pending_indentation = usize::try_from(indent).unwrap();
355            self.space = cmp::max(MARGIN - indent, MIN_SPACE);
356            if let Some(post_break) = token.post_break {
357                self.print_indent();
358                self.out.push(post_break);
359                self.space -= post_break.len_utf8() as isize;
360            }
361        }
362    }
363
364    fn print_string(&mut self, string: Cow<'static, str>) {
365        self.print_indent();
366        self.out.push_str(&string);
367        self.space -= string.len() as isize;
368    }
369
370    fn print_indent(&mut self) {
371        self.out.reserve(self.pending_indentation);
372        self.out
373            .extend(iter::repeat(' ').take(self.pending_indentation));
374        self.pending_indentation = 0;
375    }
376}