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
11pub trait Render {
13 type Error;
15
16 fn write_str(&mut self, s: &str) -> Result<usize, Self::Error>;
18
19 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 fn fail_doc(&self) -> Self::Error;
30}
31
32#[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 #[inline]
55 pub fn pretty(&self, width: usize) -> PrettyFormatter {
56 PrettyFormatter { tree: self, width }
57 }
58}
59
60pub trait RenderAnnotated: Render {
62 fn push_annotation(&mut self, annotation: Rc<AnsiStyle>) -> Result<(), Self::Error>;
64 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 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(); 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 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 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 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 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}