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