1use crate::{BufferWrite, PrettyTree, Text};
2use alloc::{rc::Rc, vec, vec::Vec};
3use color_ansi::AnsiStyle;
4use core::{
5 fmt::{Debug, Display, Formatter},
6 slice,
7};
8
9#[cfg(feature = "std")]
10pub mod write_io;
11
12pub mod write_fmt;
13
14pub trait Render<'a, T> {
16 type Error;
18
19 fn write_all(&mut self, s: &[T]) -> Result<(), Self::Error>;
21
22 fn fail_doc(&self) -> Self::Error;
24}
25
26pub struct PrettyFormatter<'b, 'a, T> {
28 tree: &'b PrettyTree<'a, T>,
29 width: usize,
30}
31
32impl<'a, 'b, T: Text<'a> + Debug> Debug for PrettyFormatter<'b, 'a, T> {
33 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
34 f.debug_struct("PrettyFormatter").field("tree", &self.tree).field("width", &self.width).finish()
35 }
36}
37
38impl<'a, 'b, T: Text<'a>> Display for PrettyFormatter<'b, 'a, T> {
39 fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
40 self.tree.render_fmt(self.width, f)
41 }
42}
43
44impl<'a, T> PrettyTree<'a, T> {
45 #[inline]
54 pub fn pretty(&self, width: usize) -> PrettyFormatter<'_, 'a, T> {
55 PrettyFormatter { tree: self, width }
56 }
57}
58
59pub trait RenderAnnotated<'a, T>: Render<'a, T> {
61 fn push_annotation(&mut self, annotation: Rc<AnsiStyle>) -> Result<(), Self::Error>;
63 fn pop_annotation(&mut self) -> Result<(), Self::Error>;
65}
66
67#[derive(Debug)]
68enum Annotation<A> {
69 Push(Rc<A>),
70 Pop,
71}
72
73macro_rules! make_spaces {
74 () => { "" };
75 ($s: tt $($t: tt)*) => { concat!(" ", make_spaces!($($t)*)) };
76}
77
78pub(crate) const SPACES: &str = make_spaces!(,,,,,,,,,,);
79
80fn append_docs2<'a, T>(
81 ldoc: Rc<PrettyTree<'a, T>>,
82 rdoc: Rc<PrettyTree<'a, T>>,
83 mut consumer: impl FnMut(Rc<PrettyTree<'a, T>>),
84) -> Rc<PrettyTree<'a, T>> {
85 let d = append_docs(rdoc, &mut consumer);
86 consumer(d);
87 append_docs(ldoc, &mut consumer)
88}
89
90fn append_docs<'a, T>(
91 mut doc: Rc<PrettyTree<'a, T>>,
92 consumer: &mut impl FnMut(Rc<PrettyTree<'a, T>>),
93) -> Rc<PrettyTree<'a, T>> {
94 loop {
95 match doc.as_ref() {
99 PrettyTree::Append { lhs, rhs } => {
100 let d = append_docs(rhs.clone(), consumer);
101 consumer(d);
102 doc = lhs.clone();
103 }
104 _ => return doc,
105 }
106 }
107}
108
109pub fn best<'a, W, T: Text<'a>>(doc: Rc<PrettyTree<'a, T>>, width: usize, out: &mut W) -> Result<(), W::Error>
110where
111 W: RenderAnnotated<'a, T>,
112 W: ?Sized,
113{
114 Best {
115 pos: 0,
116 back_cmds: vec![RenderCommand { indent: 0, mode: Mode::Break, node: doc }],
117 front_cmds: vec![],
118 annotation_levels: vec![],
119 width,
120 }
121 .best(0, out)?;
122
123 Ok(())
124}
125
126#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
127enum Mode {
128 Break,
129 Flat,
130}
131
132struct RenderCommand<'a, T> {
133 indent: usize,
134 mode: Mode,
135 node: Rc<PrettyTree<'a, T>>,
136}
137
138fn write_newline<'a, W, T: Text<'a>>(ind: usize, out: &mut W) -> Result<(), W::Error>
139where
140 W: ?Sized + Render<'a, T>,
141{
142 out.write_all(&[T::newline()])?;
143 write_spaces(ind, out)
144}
145
146fn write_spaces<'a, W, T: Text<'a>>(spaces: usize, out: &mut W) -> Result<(), W::Error>
147where
148 W: ?Sized + Render<'a, T>,
149{
150 let mut inserted = 0;
151 while inserted < spaces {
152 let insert = core::cmp::min(SPACES.len(), spaces - inserted);
153 out.write_all(&[T::from_static_spaces(&SPACES[..insert])])?;
154 inserted += insert;
155 }
156
157 Ok(())
158}
159
160struct Best<'a, T> {
161 pos: usize,
162 back_cmds: Vec<RenderCommand<'a, T>>,
163 front_cmds: Vec<Rc<PrettyTree<'a, T>>>,
164 annotation_levels: Vec<usize>,
165 width: usize,
166}
167
168impl<'a, T: Text<'a>> Best<'a, T> {
169 fn fitting(&mut self, next: Rc<PrettyTree<'a, T>>, mut pos: usize, ind: usize) -> bool {
170 let mut bidx = self.back_cmds.len();
171 self.front_cmds.clear(); self.front_cmds.push(next);
173 let mut mode = Mode::Flat;
174
175 loop {
176 let mut doc = match self.front_cmds.pop() {
177 None => {
178 if bidx == 0 {
179 return true;
181 } else {
182 bidx -= 1;
183 mode = Mode::Break;
184 self.back_cmds[bidx].node.clone()
185 }
186 }
187 Some(cmd) => cmd,
188 };
189
190 loop {
191 match doc.as_ref() {
192 PrettyTree::Nil => {}
193 PrettyTree::Append { lhs, rhs } => {
194 doc = append_docs2(lhs.clone(), rhs.clone(), |send| self.front_cmds.push(send));
195 continue;
196 }
197 PrettyTree::Hardline => return mode == Mode::Break,
200 PrettyTree::RenderLength { length: len, body: _ } => {
201 pos += len;
202 if pos > self.width {
203 return false;
204 }
205 }
206 PrettyTree::Text(ref s) => {
207 pos += s.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<'a, T>,
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 let new_ind = if *space >= 0 {
282 ind.saturating_add(*space as usize)
283 } else {
284 ind.saturating_sub(space.unsigned_abs())
285 };
286 cmd = RenderCommand { indent: new_ind, mode, node: doc.clone() };
287 continue;
288 }
289 PrettyTree::Hardline => {
290 match self.back_cmds.pop() {
292 Some(next) => {
293 write_newline(next.indent, out)?;
294 self.pos = next.indent;
295 cmd = next;
296 continue;
297 }
298 None => {
299 write_newline(ind, out)?;
300 self.pos = ind;
301 }
302 }
303 }
304 PrettyTree::RenderLength { length: len, body: doc } => match doc.as_ref() {
305 PrettyTree::Text(s) => {
306 out.write_all(slice::from_ref(s))?;
307 self.pos += len;
308 fits &= self.pos <= self.width;
309 }
310 _ => unreachable!(),
311 },
312 PrettyTree::Text(ref s) => {
313 out.write_all(slice::from_ref(s))?;
314 self.pos += s.len();
315 fits &= self.pos <= self.width;
316 }
317 PrettyTree::Annotated { style: color, body: doc } => {
318 out.push_annotation(color.clone())?;
319 self.annotation_levels.push(self.back_cmds.len());
320 cmd.node = doc.clone();
321 continue;
322 }
323 PrettyTree::Union { lhs: left, rhs: right } => {
324 let pos = self.pos;
325 let annotation_levels = self.annotation_levels.len();
326 let bcmds = self.back_cmds.len();
327
328 self.back_cmds.push(RenderCommand { indent: ind, mode, node: left.clone() });
329
330 let mut buffer = BufferWrite::new(0);
331
332 match self.best(bcmds, &mut buffer) {
333 Ok(true) => buffer.render(out)?,
334 Ok(false) | Err(_) => {
335 self.pos = pos;
336 self.back_cmds.truncate(bcmds);
337 self.annotation_levels.truncate(annotation_levels);
338 cmd.node = right.clone();
339 continue;
340 }
341 }
342 }
343 PrettyTree::Column { invoke: column } => {
344 cmd.node = Rc::new(column(self.pos));
345 continue;
346 }
347 PrettyTree::Nesting { invoke: nesting } => {
348 cmd.node = Rc::new(nesting(self.pos));
349 continue;
350 }
351 PrettyTree::Fail => return Err(out.fail_doc()),
352 }
353
354 break;
355 }
356 while self.annotation_levels.last() == Some(&self.back_cmds.len()) {
357 self.annotation_levels.pop();
358 out.pop_annotation()?;
359 }
360 }
361 Ok(fits)
362 }
363}