pretty_print/render/
mod.rs

1use crate::{BufferWrite, PrettyTree};
2use alloc::{rc::Rc, vec, vec::Vec};
3use color_ansi::AnsiStyle;
4use core::fmt::{Debug, Display, Formatter};
5
6#[cfg(feature = "std")]
7pub mod write_io;
8
9pub mod write_fmt;
10
11/// Trait representing the operations necessary to render a document
12pub trait Render {
13    /// The type of the output
14    type Error;
15
16    /// Write a string to the output
17    fn write_str(&mut self, s: &str) -> Result<usize, Self::Error>;
18
19    /// Write a character to the output
20    fn write_str_all(&mut self, mut s: &str) -> Result<(), Self::Error> {
21        while !s.is_empty() {
22            let count = self.write_str(s)?;
23            s = &s[count..];
24        }
25        Ok(())
26    }
27
28    /// Write a character to the output
29    fn fail_doc(&self) -> Self::Error;
30}
31
32/// The given text, which must not contain line breaks.
33#[derive(Debug)]
34pub struct PrettyFormatter<'a> {
35    tree: &'a PrettyTree,
36    width: usize,
37}
38
39impl<'a> Display for PrettyFormatter<'a> {
40    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
41        self.tree.render_fmt(self.width, f)
42    }
43}
44
45impl PrettyTree {
46    /// Returns a value which implements `std::fmt::Display`
47    ///
48    /// ```
49    /// use pretty::{BoxDoc, Doc};
50    /// let doc =
51    ///     BoxDoc::<()>::group(BoxDoc::text("hello").append(Doc::line()).append(Doc::text("world")));
52    /// assert_eq!(format!("{}", doc.pretty(80)), "hello world");
53    /// ```
54    #[inline]
55    pub fn pretty(&self, width: usize) -> PrettyFormatter {
56        PrettyFormatter { tree: self, width }
57    }
58}
59
60/// Trait representing the operations necessary to write an annotated document.
61pub trait RenderAnnotated: Render {
62    /// Push an annotation onto the stack
63    fn push_annotation(&mut self, annotation: Rc<AnsiStyle>) -> Result<(), Self::Error>;
64    /// Pop an annotation from the stack
65    fn pop_annotation(&mut self) -> Result<(), Self::Error>;
66}
67
68#[derive(Debug)]
69enum Annotation<A> {
70    Push(Rc<A>),
71    Pop,
72}
73
74macro_rules! make_spaces {
75    () => { "" };
76    ($s: tt $($t: tt)*) => { concat!("          ", make_spaces!($($t)*)) };
77}
78
79pub(crate) const SPACES: &str = make_spaces!(,,,,,,,,,,);
80
81fn append_docs2(ldoc: Rc<PrettyTree>, rdoc: Rc<PrettyTree>, mut consumer: impl FnMut(Rc<PrettyTree>)) -> Rc<PrettyTree> {
82    let d = append_docs(rdoc, &mut consumer);
83    consumer(d);
84    append_docs(ldoc, &mut consumer)
85}
86
87fn append_docs(mut doc: Rc<PrettyTree>, consumer: &mut impl FnMut(Rc<PrettyTree>)) -> Rc<PrettyTree> {
88    loop {
89        // Since appended documents often appear in sequence on the left side we
90        // gain a slight performance increase by batching these pushes (avoiding
91        // to push and directly pop `Append` documents)
92        match doc.as_ref() {
93            PrettyTree::Append { lhs, rhs } => {
94                let d = append_docs(rhs.clone(), consumer);
95                consumer(d);
96                doc = lhs.clone();
97            }
98            _ => return doc,
99        }
100    }
101}
102
103pub fn best<W>(doc: Rc<PrettyTree>, width: usize, out: &mut W) -> Result<(), W::Error>
104where
105    W: RenderAnnotated,
106    W: ?Sized,
107{
108    Best {
109        pos: 0,
110        back_cmds: vec![RenderCommand { indent: 0, mode: Mode::Break, node: doc }],
111        front_cmds: vec![],
112        annotation_levels: vec![],
113        width,
114    }
115    .best(0, out)?;
116
117    Ok(())
118}
119
120#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
121enum Mode {
122    Break,
123    Flat,
124}
125
126struct RenderCommand {
127    indent: usize,
128    mode: Mode,
129    node: Rc<PrettyTree>,
130}
131
132fn write_newline<W>(ind: usize, out: &mut W) -> Result<(), W::Error>
133where
134    W: ?Sized + Render,
135{
136    out.write_str_all("\n")?;
137    write_spaces(ind, out)
138}
139
140fn write_spaces<W>(spaces: usize, out: &mut W) -> Result<(), W::Error>
141where
142    W: ?Sized + Render,
143{
144    let mut inserted = 0;
145    while inserted < spaces {
146        let insert = core::cmp::min(SPACES.len(), spaces - inserted);
147        inserted += out.write_str(&SPACES[..insert])?;
148    }
149
150    Ok(())
151}
152
153struct Best {
154    pos: usize,
155    back_cmds: Vec<RenderCommand>,
156    front_cmds: Vec<Rc<PrettyTree>>,
157    annotation_levels: Vec<usize>,
158    width: usize,
159}
160
161impl Best {
162    fn fitting(&mut self, next: Rc<PrettyTree>, mut pos: usize, ind: usize) -> bool {
163        let mut bidx = self.back_cmds.len();
164        self.front_cmds.clear(); // clear from previous calls from best
165        self.front_cmds.push(next);
166        let mut mode = Mode::Flat;
167
168        loop {
169            let mut doc = match self.front_cmds.pop() {
170                None => {
171                    if bidx == 0 {
172                        // All commands have been processed
173                        return true;
174                    }
175                    else {
176                        bidx -= 1;
177                        mode = Mode::Break;
178                        self.back_cmds[bidx].node.clone()
179                    }
180                }
181                Some(cmd) => cmd,
182            };
183
184            loop {
185                match doc.as_ref() {
186                    PrettyTree::Nil => {}
187                    PrettyTree::Append { lhs, rhs } => {
188                        doc = append_docs2(lhs.clone(), rhs.clone(), |send| self.front_cmds.push(send));
189                        continue;
190                    }
191                    // Newlines inside the group makes it not fit, but those outside lets it
192                    // fit on the current line
193                    PrettyTree::Hardline => return mode == Mode::Break,
194                    PrettyTree::RenderLength { length: len, body: _ } => {
195                        pos += len;
196                        if pos > self.width {
197                            return false;
198                        }
199                    }
200                    PrettyTree::StaticText(str) => {
201                        pos += str.len();
202                        if pos > self.width {
203                            return false;
204                        }
205                    }
206                    PrettyTree::Text(ref str) => {
207                        pos += str.len();
208                        if pos > self.width {
209                            return false;
210                        }
211                    }
212                    PrettyTree::MaybeInline { block: flat, inline } => {
213                        doc = match mode {
214                            Mode::Break => flat.clone(),
215                            Mode::Flat => inline.clone(),
216                        };
217                        continue;
218                    }
219
220                    PrettyTree::Column { invoke: function } => {
221                        doc = Rc::new(function(pos));
222                        continue;
223                    }
224                    PrettyTree::Nesting { invoke: function } => {
225                        doc = Rc::new(function(ind));
226                        continue;
227                    }
228                    PrettyTree::Nest { space: _, doc: next }
229                    | PrettyTree::Group { items: next }
230                    | PrettyTree::Annotated { style: _, body: next }
231                    | PrettyTree::Union { lhs: _, rhs: next } => {
232                        doc = next.clone();
233                        continue;
234                    }
235                    PrettyTree::Fail => return false,
236                }
237                break;
238            }
239        }
240    }
241
242    fn best<W>(&mut self, top: usize, out: &mut W) -> Result<bool, W::Error>
243    where
244        W: RenderAnnotated,
245        W: ?Sized,
246    {
247        let mut fits = true;
248
249        while top < self.back_cmds.len() {
250            let mut cmd = self.back_cmds.pop().unwrap();
251            loop {
252                let RenderCommand { indent: ind, mode, node } = cmd;
253                match node.as_ref() {
254                    PrettyTree::Nil => {}
255                    PrettyTree::Append { lhs, rhs } => {
256                        cmd.node = append_docs2(lhs.clone(), rhs.clone(), |send| {
257                            self.back_cmds.push(RenderCommand { indent: ind, mode, node: send })
258                        });
259                        continue;
260                    }
261                    PrettyTree::MaybeInline { block, inline } => {
262                        cmd.node = match mode {
263                            Mode::Break => block.clone(),
264                            Mode::Flat => inline.clone(),
265                        };
266                        continue;
267                    }
268                    PrettyTree::Group { items } => {
269                        match mode {
270                            Mode::Break if self.fitting(items.clone(), self.pos, ind) => {
271                                cmd.mode = Mode::Flat;
272                            }
273                            _ => {}
274                        }
275                        cmd.node = items.clone();
276                        continue;
277                    }
278                    PrettyTree::Nest { space, doc } => {
279                        // Once https://doc.rust-lang.org/std/primitive.usize.html#method.saturating_add_signed is stable
280                        // this can be replaced
281                        let new_ind = if *space >= 0 {
282                            ind.saturating_add(*space as usize)
283                        }
284                        else {
285                            ind.saturating_sub(space.unsigned_abs())
286                        };
287                        cmd = RenderCommand { indent: new_ind, mode, node: doc.clone() };
288                        continue;
289                    }
290                    PrettyTree::Hardline => {
291                        // The next document may have different indentation so we should use it if we can
292                        match self.back_cmds.pop() {
293                            Some(next) => {
294                                write_newline(next.indent, out)?;
295                                self.pos = next.indent;
296                                cmd = next;
297                                continue;
298                            }
299                            None => {
300                                write_newline(ind, out)?;
301                                self.pos = ind;
302                            }
303                        }
304                    }
305                    PrettyTree::RenderLength { length: len, body: doc } => match doc.as_ref() {
306                        PrettyTree::Text(s) => {
307                            out.write_str_all(s)?;
308                            self.pos += len;
309                            fits &= self.pos <= self.width;
310                        }
311                        PrettyTree::StaticText(s) => {
312                            out.write_str_all(s)?;
313                            self.pos += len;
314                            fits &= self.pos <= self.width;
315                        }
316                        _ => unreachable!(),
317                    },
318                    PrettyTree::Text(ref s) => {
319                        out.write_str_all(s)?;
320                        self.pos += s.len();
321                        fits &= self.pos <= self.width;
322                    }
323                    PrettyTree::StaticText(s) => {
324                        out.write_str_all(s)?;
325                        self.pos += s.len();
326                        fits &= self.pos <= self.width;
327                    }
328                    PrettyTree::Annotated { style: color, body: doc } => {
329                        out.push_annotation(color.clone())?;
330                        self.annotation_levels.push(self.back_cmds.len());
331                        cmd.node = doc.clone();
332                        continue;
333                    }
334                    PrettyTree::Union { lhs: left, rhs: right } => {
335                        let pos = self.pos;
336                        let annotation_levels = self.annotation_levels.len();
337                        let bcmds = self.back_cmds.len();
338
339                        self.back_cmds.push(RenderCommand { indent: ind, mode, node: left.clone() });
340
341                        let mut buffer = BufferWrite::new(0);
342
343                        match self.best(bcmds, &mut buffer) {
344                            Ok(true) => buffer.render(out)?,
345                            Ok(false) | Err(_) => {
346                                self.pos = pos;
347                                self.back_cmds.truncate(bcmds);
348                                self.annotation_levels.truncate(annotation_levels);
349                                cmd.node = right.clone();
350                                continue;
351                            }
352                        }
353                    }
354                    PrettyTree::Column { invoke: column } => {
355                        cmd.node = Rc::new(column(self.pos));
356                        continue;
357                    }
358                    PrettyTree::Nesting { invoke: nesting } => {
359                        cmd.node = Rc::new(nesting(self.pos));
360                        continue;
361                    }
362                    PrettyTree::Fail => return Err(out.fail_doc()),
363                }
364
365                break;
366            }
367            while self.annotation_levels.last() == Some(&self.back_cmds.len()) {
368                self.annotation_levels.pop();
369                out.pop_annotation()?;
370            }
371        }
372        Ok(fits)
373    }
374}