pretty/
render.rs

Help
1use std::{cmp, fmt, io};
2
3#[cfg(feature = "termcolor")]
4use termcolor::{ColorSpec, WriteColor};
5
6use crate::{Doc, DocPtr};
7
8/// Trait representing the operations necessary to render a document
9pub trait Render {
10    type Error;
11
12    fn write_str(&mut self, s: &str) -> Result<usize, Self::Error>;
13
14    fn write_str_all(&mut self, mut s: &str) -> Result<(), Self::Error> {
15        while !s.is_empty() {
16            let count = self.write_str(s)?;
17            s = &s[count..];
18        }
19        Ok(())
20    }
21
22    fn fail_doc(&self) -> Self::Error;
23}
24
25/// Writes to something implementing `std::io::Write`
26pub struct IoWrite<W> {
27    upstream: W,
28}
29
30impl<W> IoWrite<W> {
31    pub fn new(upstream: W) -> IoWrite<W> {
32        IoWrite { upstream }
33    }
34}
35
36impl<W> Render for IoWrite<W>
37where
38    W: io::Write,
39{
40    type Error = io::Error;
41
42    fn write_str(&mut self, s: &str) -> io::Result<usize> {
43        self.upstream.write(s.as_bytes())
44    }
45
46    fn write_str_all(&mut self, s: &str) -> io::Result<()> {
47        self.upstream.write_all(s.as_bytes())
48    }
49
50    fn fail_doc(&self) -> Self::Error {
51        io::Error::new(io::ErrorKind::Other, "Document failed to render")
52    }
53}
54
55/// Writes to something implementing `std::fmt::Write`
56pub struct FmtWrite<W> {
57    upstream: W,
58}
59
60impl<W> FmtWrite<W> {
61    pub fn new(upstream: W) -> FmtWrite<W> {
62        FmtWrite { upstream }
63    }
64}
65
66impl<W> Render for FmtWrite<W>
67where
68    W: fmt::Write,
69{
70    type Error = fmt::Error;
71
72    fn write_str(&mut self, s: &str) -> Result<usize, fmt::Error> {
73        self.write_str_all(s).map(|_| s.len())
74    }
75
76    fn write_str_all(&mut self, s: &str) -> fmt::Result {
77        self.upstream.write_str(s)
78    }
79
80    fn fail_doc(&self) -> Self::Error {
81        fmt::Error
82    }
83}
84
85/// Trait representing the operations necessary to write an annotated document.
86pub trait RenderAnnotated<'a, A>: Render {
87    fn push_annotation(&mut self, annotation: &'a A) -> Result<(), Self::Error>;
88    fn pop_annotation(&mut self) -> Result<(), Self::Error>;
89}
90
91impl<A, W> RenderAnnotated<'_, A> for IoWrite<W>
92where
93    W: io::Write,
94{
95    fn push_annotation(&mut self, _: &A) -> Result<(), Self::Error> {
96        Ok(())
97    }
98
99    fn pop_annotation(&mut self) -> Result<(), Self::Error> {
100        Ok(())
101    }
102}
103
104impl<A, W> RenderAnnotated<'_, A> for FmtWrite<W>
105where
106    W: fmt::Write,
107{
108    fn push_annotation(&mut self, _: &A) -> Result<(), Self::Error> {
109        Ok(())
110    }
111
112    fn pop_annotation(&mut self) -> Result<(), Self::Error> {
113        Ok(())
114    }
115}
116
117#[cfg(feature = "termcolor")]
118pub struct TermColored<W> {
119    color_stack: Vec<ColorSpec>,
120    upstream: W,
121}
122
123#[cfg(feature = "termcolor")]
124impl<W> TermColored<W> {
125    pub fn new(upstream: W) -> TermColored<W> {
126        TermColored {
127            color_stack: Vec::new(),
128            upstream,
129        }
130    }
131}
132
133#[cfg(feature = "termcolor")]
134impl<W> Render for TermColored<W>
135where
136    W: io::Write,
137{
138    type Error = io::Error;
139
140    fn write_str(&mut self, s: &str) -> io::Result<usize> {
141        self.upstream.write(s.as_bytes())
142    }
143
144    fn write_str_all(&mut self, s: &str) -> io::Result<()> {
145        self.upstream.write_all(s.as_bytes())
146    }
147
148    fn fail_doc(&self) -> Self::Error {
149        io::Error::new(io::ErrorKind::Other, "Document failed to render")
150    }
151}
152
153#[cfg(feature = "termcolor")]
154impl<W> RenderAnnotated<'_, ColorSpec> for TermColored<W>
155where
156    W: WriteColor,
157{
158    fn push_annotation(&mut self, color: &ColorSpec) -> Result<(), Self::Error> {
159        self.color_stack.push(color.clone());
160        self.upstream.set_color(color)
161    }
162
163    fn pop_annotation(&mut self) -> Result<(), Self::Error> {
164        self.color_stack.pop();
165        match self.color_stack.last() {
166            Some(previous) => self.upstream.set_color(previous),
167            None => self.upstream.reset(),
168        }
169    }
170}
171
172enum Annotation<'a, A> {
173    Push(&'a A),
174    Pop,
175}
176
177struct BufferWrite<'a, A> {
178    buffer: String,
179    annotations: Vec<(usize, Annotation<'a, A>)>,
180}
181
182impl<'a, A> BufferWrite<'a, A> {
183    fn new() -> Self {
184        BufferWrite {
185            buffer: String::new(),
186            annotations: Vec::new(),
187        }
188    }
189
190    fn render<W>(&mut self, render: &mut W) -> Result<(), W::Error>
191    where
192        W: RenderAnnotated<'a, A>,
193        W: ?Sized,
194    {
195        let mut start = 0;
196        for (end, annotation) in &self.annotations {
197            let s = &self.buffer[start..*end];
198            if !s.is_empty() {
199                render.write_str_all(s)?;
200            }
201            start = *end;
202            match annotation {
203                Annotation::Push(a) => render.push_annotation(a)?,
204                Annotation::Pop => render.pop_annotation()?,
205            }
206        }
207        let s = &self.buffer[start..];
208        if !s.is_empty() {
209            render.write_str_all(s)?;
210        }
211        Ok(())
212    }
213}
214
215impl<A> Render for BufferWrite<'_, A> {
216    type Error = ();
217
218    fn write_str(&mut self, s: &str) -> Result<usize, Self::Error> {
219        self.buffer.push_str(s);
220        Ok(s.len())
221    }
222
223    fn write_str_all(&mut self, s: &str) -> Result<(), Self::Error> {
224        self.buffer.push_str(s);
225        Ok(())
226    }
227
228    fn fail_doc(&self) -> Self::Error {}
229}
230
231impl<'a, A> RenderAnnotated<'a, A> for BufferWrite<'a, A> {
232    fn push_annotation(&mut self, a: &'a A) -> Result<(), Self::Error> {
233        self.annotations
234            .push((self.buffer.len(), Annotation::Push(a)));
235        Ok(())
236    }
237
238    fn pop_annotation(&mut self) -> Result<(), Self::Error> {
239        self.annotations.push((self.buffer.len(), Annotation::Pop));
240        Ok(())
241    }
242}
243
244macro_rules! make_spaces {
245    () => { "" };
246    ($s: tt $($t: tt)*) => { concat!("          ", make_spaces!($($t)*)) };
247}
248
249pub(crate) const SPACES: &str = make_spaces!(,,,,,,,,,,);
250
251fn append_docs2<'a, 'd, T, A>(
252    ldoc: &'d Doc<'a, T, A>,
253    rdoc: &'d Doc<'a, T, A>,
254    mut consumer: impl FnMut(&'d Doc<'a, T, A>),
255) -> &'d Doc<'a, T, A>
256where
257    T: DocPtr<'a, A>,
258{
259    let d = append_docs(rdoc, &mut consumer);
260    consumer(d);
261    append_docs(ldoc, &mut consumer)
262}
263
264fn append_docs<'a, 'd, T, A>(
265    mut doc: &'d Doc<'a, T, A>,
266    consumer: &mut impl FnMut(&'d Doc<'a, T, A>),
267) -> &'d Doc<'a, T, A>
268where
269    T: DocPtr<'a, A>,
270{
271    loop {
272        // Since appended documents often appear in sequence on the left side we
273        // gain a slight performance increase by batching these pushes (avoiding
274        // to push and directly pop `Append` documents)
275        match doc {
276            Doc::Append(l, r) => {
277                consumer(r);
278                doc = l;
279            }
280            _ => return doc,
281        }
282    }
283}
284
285pub fn best<'a, W, T, A>(doc: &Doc<'a, T, A>, width: usize, out: &mut W) -> Result<(), W::Error>
286where
287    T: DocPtr<'a, A> + 'a,
288    for<'b> W: RenderAnnotated<'b, A>,
289    W: ?Sized,
290{
291    let temp_arena = &typed_arena::Arena::new();
292    Best {
293        pos: 0,
294        bcmds: vec![(0, Mode::Break, doc)],
295        fcmds: vec![],
296        annotation_levels: vec![],
297        width,
298        temp_arena,
299    }
300    .best(0, out)?;
301
302    Ok(())
303}
304
305#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
306enum Mode {
307    Break,
308    Flat,
309}
310
311type Cmd<'d, 'a, T, A> = (usize, Mode, &'d Doc<'a, T, A>);
312
313fn write_newline<W>(ind: usize, out: &mut W) -> Result<(), W::Error>
314where
315    W: ?Sized + Render,
316{
317    out.write_str_all("\n")?;
318    write_spaces(ind, out)
319}
320
321fn write_spaces<W>(spaces: usize, out: &mut W) -> Result<(), W::Error>
322where
323    W: ?Sized + Render,
324{
325    let mut inserted = 0;
326    while inserted < spaces {
327        let insert = cmp::min(SPACES.len(), spaces - inserted);
328        inserted += out.write_str(&SPACES[..insert])?;
329    }
330
331    Ok(())
332}
333
334struct Best<'d, 'a, T, A>
335where
336    T: DocPtr<'a, A> + 'a,
337{
338    pos: usize,
339    bcmds: Vec<Cmd<'d, 'a, T, A>>,
340    fcmds: Vec<&'d Doc<'a, T, A>>,
341    annotation_levels: Vec<usize>,
342    width: usize,
343    temp_arena: &'d typed_arena::Arena<T>,
344}
345
346impl<'d, 'a, T, A> Best<'d, 'a, T, A>
347where
348    T: DocPtr<'a, A> + 'a,
349{
350    fn fitting(&mut self, next: &'d Doc<'a, T, A>, mut pos: usize, ind: usize) -> bool {
351        let mut bidx = self.bcmds.len();
352        self.fcmds.clear(); // clear from previous calls from best
353        self.fcmds.push(next);
354
355        let mut mode = Mode::Flat;
356        loop {
357            let mut doc = match self.fcmds.pop() {
358                None => {
359                    if bidx == 0 {
360                        // All commands have been processed
361                        return true;
362                    } else {
363                        bidx -= 1;
364                        mode = Mode::Break;
365                        self.bcmds[bidx].2
366                    }
367                }
368                Some(cmd) => cmd,
369            };
370
371            loop {
372                match *doc {
373                    Doc::Nil => {}
374                    Doc::Append(ref ldoc, ref rdoc) => {
375                        doc = append_docs2(ldoc, rdoc, |doc| self.fcmds.push(doc));
376                        continue;
377                    }
378                    // Newlines inside the group makes it not fit, but those outside lets it
379                    // fit on the current line
380                    Doc::Hardline => return mode == Mode::Break,
381                    Doc::RenderLen(len, _) => {
382                        pos += len;
383                        if pos > self.width {
384                            return false;
385                        }
386                    }
387                    Doc::BorrowedText(str) => {
388                        pos += str.len();
389                        if pos > self.width {
390                            return false;
391                        }
392                    }
393                    Doc::OwnedText(ref str) => {
394                        pos += str.len();
395                        if pos > self.width {
396                            return false;
397                        }
398                    }
399                    Doc::SmallText(ref str) => {
400                        pos += str.len();
401                        if pos > self.width {
402                            return false;
403                        }
404                    }
405                    Doc::FlatAlt(ref b, ref f) => {
406                        doc = match mode {
407                            Mode::Break => b,
408                            Mode::Flat => f,
409                        };
410                        continue;
411                    }
412
413                    Doc::Column(ref f) => {
414                        doc = self.temp_arena.alloc(f(pos));
415                        continue;
416                    }
417                    Doc::Nesting(ref f) => {
418                        doc = self.temp_arena.alloc(f(ind));
419                        continue;
420                    }
421                    Doc::Nest(_, ref next)
422                    | Doc::Group(ref next)
423                    | Doc::Annotated(_, ref next)
424                    | Doc::Union(_, ref next) => {
425                        doc = next;
426                        continue;
427                    }
428                    Doc::Fail => return false,
429                }
430                break;
431            }
432        }
433    }
434
435    fn best<W>(&mut self, top: usize, out: &mut W) -> Result<bool, W::Error>
436    where
437        W: RenderAnnotated<'d, A>,
438        W: ?Sized,
439    {
440        let mut fits = true;
441
442        while top < self.bcmds.len() {
443            let mut cmd = self.bcmds.pop().unwrap();
444            loop {
445                let (ind, mode, doc) = cmd;
446                match *doc {
447                    Doc::Nil => {}
448                    Doc::Append(ref ldoc, ref rdoc) => {
449                        cmd.2 = append_docs2(ldoc, rdoc, |doc| self.bcmds.push((ind, mode, doc)));
450                        continue;
451                    }
452                    Doc::FlatAlt(ref b, ref f) => {
453                        cmd.2 = match mode {
454                            Mode::Break => b,
455                            Mode::Flat => f,
456                        };
457                        continue;
458                    }
459                    Doc::Group(ref doc) => {
460                        if let Mode::Break = mode {
461                            if self.fitting(doc, self.pos, ind) {
462                                cmd.1 = Mode::Flat;
463                            }
464                        }
465                        cmd.2 = doc;
466                        continue;
467                    }
468                    Doc::Nest(off, ref doc) => {
469                        // Once https://doc.rust-lang.org/std/primitive.usize.html#method.saturating_add_signed is stable
470                        // this can be replaced
471                        let new_ind = if off >= 0 {
472                            ind.saturating_add(off as usize)
473                        } else {
474                            ind.saturating_sub(off.unsigned_abs())
475                        };
476                        cmd = (new_ind, mode, doc);
477                        continue;
478                    }
479                    Doc::Hardline => {
480                        // The next document may have different indentation so we should use it if
481                        // we can
482                        if let Some(next) = self.bcmds.pop() {
483                            write_newline(next.0, out)?;
484                            self.pos = next.0;
485                            cmd = next;
486                            continue;
487                        } else {
488                            write_newline(ind, out)?;
489                            self.pos = ind;
490                        }
491                    }
492                    Doc::RenderLen(len, ref doc) => match **doc {
493                        Doc::OwnedText(ref s) => {
494                            out.write_str_all(s)?;
495                            self.pos += len;
496                            fits &= self.pos <= self.width;
497                        }
498                        Doc::BorrowedText(s) => {
499                            out.write_str_all(s)?;
500                            self.pos += len;
501                            fits &= self.pos <= self.width;
502                        }
503                        Doc::SmallText(ref s) => {
504                            out.write_str_all(s)?;
505                            self.pos += len;
506                            fits &= self.pos <= self.width;
507                        }
508                        _ => unreachable!(),
509                    },
510                    Doc::OwnedText(ref s) => {
511                        out.write_str_all(s)?;
512                        self.pos += s.len();
513                        fits &= self.pos <= self.width;
514                    }
515                    Doc::BorrowedText(s) => {
516                        out.write_str_all(s)?;
517                        self.pos += s.len();
518                        fits &= self.pos <= self.width;
519                    }
520                    Doc::SmallText(ref s) => {
521                        out.write_str_all(s)?;
522                        self.pos += s.len();
523                        fits &= self.pos <= self.width;
524                    }
525                    Doc::Annotated(ref ann, ref doc) => {
526                        out.push_annotation(ann)?;
527                        self.annotation_levels.push(self.bcmds.len());
528                        cmd.2 = doc;
529                        continue;
530                    }
531                    Doc::Union(ref l, ref r) => {
532                        let pos = self.pos;
533                        let annotation_levels = self.annotation_levels.len();
534                        let bcmds = self.bcmds.len();
535
536                        self.bcmds.push((ind, mode, l));
537
538                        let mut buffer = BufferWrite::new();
539
540                        match self.best(bcmds, &mut buffer) {
541                            Ok(true) => buffer.render(out)?,
542                            Ok(false) | Err(()) => {
543                                self.pos = pos;
544                                self.bcmds.truncate(bcmds);
545                                self.annotation_levels.truncate(annotation_levels);
546                                cmd.2 = r;
547                                continue;
548                            }
549                        }
550                    }
551                    Doc::Column(ref f) => {
552                        cmd.2 = self.temp_arena.alloc(f(self.pos));
553                        continue;
554                    }
555                    Doc::Nesting(ref f) => {
556                        cmd.2 = self.temp_arena.alloc(f(ind));
557                        continue;
558                    }
559                    Doc::Fail => return Err(out.fail_doc()),
560                }
561
562                break;
563            }
564            while self.annotation_levels.last() == Some(&self.bcmds.len()) {
565                self.annotation_levels.pop();
566                out.pop_annotation()?;
567            }
568        }
569
570        Ok(fits)
571    }
572}
pretty/render.rs

pretty/
render.rs
