ast2str_lib/builder.rs
1//! This module defines the data structures and functions used to draw trees.
2use std::borrow::Cow;
3
4use crate::{AstToStr, Symbols};
5
6/// A builder struct for formatting AST nodes.
7#[derive(Debug)]
8pub struct TreeBuilder<'a, 's> {
9 /// The name of the node.
10 name: Cow<'a, str>,
11 /// The list of this node's children.
12 children: Vec<String>,
13 /// The symbols to use for drawing the tree.
14 symbols: &'s dyn Symbols,
15}
16
17impl<'a, 's> TreeBuilder<'a, 's> {
18 /// Creates a new [`TreeBuilder`] instance with the given node name and [`Symbols`].
19 ///
20 /// [`TreeBuilder`]: struct.TreeBuilder.html
21 /// [`Symbols`]: trait.Symbols.html
22 pub fn new<S: Into<Cow<'a, str>>>(name: S, symbols: &'s dyn Symbols) -> Self {
23 Self {
24 name: name.into(),
25 children: Vec::new(),
26 symbols,
27 }
28 }
29
30 /// Creates a new [`TreeBuilder`] configured to use [`DefaultSymbols`].
31 ///
32 /// [`TreeBuilder`]: struct.TreeBuilder.html
33 /// [`DefaultSymbols`]: struct.DefaultSymbols.html
34 pub fn with_default_symbols<S: Into<Cow<'a, str>>>(name: S) -> Self {
35 Self::new(name, &crate::DefaultSymbols)
36 }
37
38 /// Adds a new child with the given name to the tree, recursively calling [`AstToStr::ast_to_str_impl`] on the value.
39 ///
40 /// # Example
41 /// ```
42 /// # use ast2str_lib as ast2str;
43 /// use ast2str::{builder::TreeBuilder};
44 ///
45 /// let mut tree = TreeBuilder::with_default_symbols("Expr");
46 /// tree = tree.field("kind", &"ExprKind::Literal");
47 /// tree = tree.field("span", &(3..10));
48 /// assert_eq!(
49 /// tree.build(),
50 /// r#"
51 /// Expr
52 /// ├─kind: "ExprKind::Literal"
53 /// ╰─span: 3..10
54 /// "#.trim());
55 /// ```
56 ///
57 /// [`AstToStr::ast_to_str_impl`]: trait.AstToStr.html#tymethod.ast_to_str_impl
58 pub fn field<A: AstToStr>(mut self, name: &str, value: &A) -> Self {
59 self.add_format(name.into(), value.ast_to_str_impl(self.symbols));
60 self
61 }
62
63 /// Adds a new child with the given name to the tree, formatting its value as [`Display`] surrounded with backticks.
64 ///
65 /// # Example
66 /// ```
67 /// # use ast2str_lib as ast2str;
68 /// use ast2str::{builder::TreeBuilder};
69 ///
70 /// let mut tree = TreeBuilder::with_default_symbols("Token");
71 /// tree = tree.quoted("lexeme", "\"a string\"");
72 /// tree = tree.field("span", &(0..8));
73 /// assert_eq!(
74 /// tree.build(),
75 /// r#"
76 /// Token
77 /// ├─lexeme: `"a string"`
78 /// ╰─span: 0..8
79 /// "#.trim());
80 /// ```
81 ///
82 /// [`Display`]: https://doc.rust-lang.org/std/fmt/trait.Display.html
83 pub fn quoted<S: std::fmt::Display>(mut self, name: &str, value: S) -> Self {
84 self.add_format(name.into(), format!("`{}`", value));
85 self
86 }
87
88 /// Adds a new child with the given name to the tree, formatting its value as [`Display`].
89 ///
90 /// # Example
91 /// ```
92 /// # use ast2str_lib as ast2str;
93 /// use ast2str::{builder::TreeBuilder};
94 ///
95 /// let mut tree = TreeBuilder::with_default_symbols("Variable");
96 /// tree = tree.display("name", "x");
97 /// assert_eq!(
98 /// tree.build(),
99 /// r#"
100 /// Variable
101 /// ╰─name: x
102 /// "#.trim());
103 /// ```
104 ///
105 /// [`Display`]: https://doc.rust-lang.org/std/fmt/trait.Display.html
106 pub fn display<S: std::fmt::Display>(mut self, name: &str, value: S) -> Self {
107 self.add_format(name.into(), value.to_string());
108 self
109 }
110
111 /// Adds a new child with the given name to the tree, formatting its value as [`Debug`].
112 ///
113 /// # Example
114 /// ```
115 /// # use ast2str_lib as ast2str;
116 /// use ast2str::{builder::TreeBuilder};
117 ///
118 /// let mut tree = TreeBuilder::with_default_symbols("Binary");
119 /// tree = tree.field("left", &1);
120 /// tree = tree.debug("operator", &std::cmp::Ordering::Less);
121 /// tree = tree.field("right", &3);
122 /// assert_eq!(
123 /// tree.build(),
124 /// r#"
125 /// Binary
126 /// ├─left: 1
127 /// ├─operator: Less
128 /// ╰─right: 3
129 /// "#.trim());
130 /// ```
131 ///
132 /// [`Debug`]: https://doc.rust-lang.org/std/fmt/trait.Debug.html]
133 pub fn debug<S: std::fmt::Debug>(mut self, name: &str, value: S) -> Self {
134 self.children.push(format!("{}: {:?}", name, value));
135 self
136 }
137
138 /// Attempts to add a new child to the tree with [`field`] if the option is [`Some`], recursively calling [`AstToStr::ast_to_str_impl`] on its value.
139 /// If the option is `None`, falls back to [`display`] with the given default as its value.
140 ///
141 /// # Example
142 /// ```
143 /// # use ast2str_lib as ast2str;
144 /// use ast2str::{builder::TreeBuilder};
145 ///
146 /// let mut tree = TreeBuilder::with_default_symbols("VarDecl");
147 /// tree = tree.quoted("name", "x");
148 /// tree = tree.option("initializer", "<None>", &None::<i32>);
149 /// assert_eq!(
150 /// tree.build(),
151 /// r#"
152 /// VarDecl
153 /// ├─name: `x`
154 /// ╰─initializer: <None>
155 /// "#.trim());
156 ///
157 /// let mut tree = TreeBuilder::with_default_symbols("VarDecl");
158 /// tree = tree.quoted("name", "x");
159 /// tree = tree.option("initializer", "<None>", &Some(7));
160 /// assert_eq!(
161 /// tree.build(),
162 /// r#"
163 /// VarDecl
164 /// ├─name: `x`
165 /// ╰─initializer: 7
166 /// "#.trim());
167 /// ```
168 ///
169 /// [`AstToStr::ast_to_str_impl`]: trait.AstToStr.html#tymethod.ast_to_str_impl
170 /// [`field`]: struct.TreeBuilder.html#method.field
171 /// [`display`]: struct.TreeBuilder.html#method.display
172 pub fn option<A: AstToStr>(self, field: &str, default: &str, option: &Option<A>) -> Self {
173 if let Some(child) = option {
174 self.field(field, child)
175 } else {
176 self.display(field, default)
177 }
178 }
179
180 /// Adds the given collection of items to the tree as a child, recursively calling [`AstToStr::ast_to_str_impl`] on each item.
181 /// If the collection is empty, [`Symbols::missing_items_symbol`] will be displayed; otherwise, [`Symbols::item_list_symbol`].
182 ///
183 /// # Example
184 /// ```
185 /// # use ast2str_lib as ast2str;
186 /// use ast2str::{builder::TreeBuilder};
187 ///
188 /// let mut tree = TreeBuilder::with_default_symbols("List");
189 /// tree = tree.list("non_empty", &[1usize, 2, 3]);
190 /// tree = tree.list("empty", &Vec::<usize>::new());
191 /// assert_eq!(
192 /// tree.build(),
193 /// r#"
194 /// List
195 /// ├─non_empty=↓
196 /// │ ├─1
197 /// │ ├─2
198 /// │ ╰─3
199 /// ╰─empty=✕
200 /// "#.trim());
201 /// ```
202 ///
203 /// [`AstToStr::ast_to_str_impl`]: trait.AstToStr.html#tymethod.ast_to_str_impl
204 /// [`Symbols::missing_items_symbol`]: trait.Symbols.html#tymethod.missing_items_symbol
205 /// [`Symbols::item_list_symbol`]: trait.Symbols.html#tymethod.item_list_symbol
206 /// [`DefaultSymbols`]: struct.DefaultSymbols.html
207 pub fn list<'b, S: AstToStr + 'static>(
208 mut self,
209 name: &str,
210 collection: impl IntoIterator<Item = &'b S>,
211 ) -> Self {
212 self.children.push(print_ast_list_generic(
213 name,
214 collection,
215 |s| s.ast_to_str_impl(self.symbols),
216 self.symbols,
217 ));
218 self
219 }
220
221 /// Just like [`list`], but allows the user to pass a function, applying it to every item.
222 ///
223 /// # Example
224 /// ```
225 /// # use ast2str_lib as ast2str;
226 /// use ast2str::{builder::TreeBuilder};
227 ///
228 /// let mut tree = TreeBuilder::with_default_symbols("List");
229 /// tree = tree.list_map("values", &[1, 2, 3], |x| x * x);
230 /// assert_eq!(
231 /// tree.build(),
232 /// r#"
233 /// List
234 /// ╰─values=↓
235 /// ├─1
236 /// ├─4
237 /// ╰─9
238 /// "#.trim());
239 /// ```
240 /// [`list`]: struct.TreeBuilder.html#method.list
241 pub fn list_map<T, A: AstToStr>(
242 mut self,
243 name: &str,
244 collection: impl IntoIterator<Item = T>,
245 f: impl Fn(T) -> A,
246 ) -> Self {
247 self.children.push(print_ast_list_generic(
248 name,
249 collection,
250 |t| f(t).ast_to_str_impl(self.symbols),
251 self.symbols,
252 ));
253 self
254 }
255
256 /// Adds the given value to the tree as a child.
257 ///
258 /// # Example
259 /// ```
260 /// # use ast2str_lib as ast2str;
261 /// use ast2str::{builder::TreeBuilder};
262 ///
263 /// let mut tree = TreeBuilder::with_default_symbols("root");
264 /// tree = tree.add_child(std::f64::consts::PI);
265 /// tree = tree.add_child("any ToString value works");
266 /// assert_eq!(
267 /// tree.build(),
268 /// r#"
269 /// root
270 /// ├─3.141592653589793
271 /// ╰─any ToString value works
272 /// "#.trim());
273 /// ```
274 /// [`list`]: struct.TreeBuilder.html#method.list
275 pub fn add_child<S: ToString>(mut self, child: S) -> Self {
276 self.children.push(child.to_string());
277 self
278 }
279
280 /// Consumes the builder, formatting the node and its children as a tree. See the other methods for examples.
281 pub fn build(self) -> String {
282 format(&self.name, self.children, self.symbols)
283 }
284
285 fn add_format(&mut self, key: String, value: String) {
286 let key = if value.starts_with('=') {
287 key
288 } else {
289 format!("{}: ", key)
290 };
291 self.children.push(format!("{}{}", key, value));
292 }
293}
294
295// This enum should be private, but it has been public for a while. Although unlikely, but someone could be
296// importing it in their code.
297#[doc(hidden)]
298#[derive(Debug, Clone, Copy, PartialEq, Eq)]
299pub enum TreeIndent {
300 Trunk,
301 Branch,
302}
303
304impl TreeIndent {
305 #[inline]
306 pub fn is_trunk(&self) -> bool {
307 *self == TreeIndent::Trunk
308 }
309}
310
311/// A function used by [`TreeBuilder::list`] and [`TreeBuilder::list_map`] for list formatting.
312/// Also see [`print_ast_list_without_node_name`].
313pub fn print_ast_list_generic<T>(
314 node: &str,
315 collection: impl IntoIterator<Item = T>,
316 f: impl Fn(T) -> String,
317 s: &dyn Symbols,
318) -> String {
319 format!(
320 "{}{}",
321 node,
322 print_ast_list_without_node_name(collection, f, s)
323 )
324}
325
326/// A lower-level function behind [`print_ast_list_generic`] that _actually_ does the formatting.
327/// Useful for manually implementing [`AstToStr`] for [`Vec`]-like types.
328///
329/// # Example
330/// ```
331/// # use ast2str_lib as ast2str;
332/// use ast2str::{AstToStr, Symbols, builder::print_ast_list_without_node_name};
333///
334/// struct Wrapper(Vec<i32>);
335/// impl<'a> IntoIterator for &'a Wrapper {
336/// type Item = &'a i32;
337/// type IntoIter = std::slice::Iter<'a, i32>;
338/// fn into_iter(self) -> Self::IntoIter {
339/// self.0.iter()
340/// }
341/// }
342
343/// impl AstToStr for Wrapper {
344/// fn ast_to_str_impl(&self, symbols: &dyn Symbols) -> String {
345/// print_ast_list_without_node_name(self, |x| x.ast_to_str_impl(symbols), symbols)
346/// }
347/// }
348///
349/// let wrapper = Wrapper(vec![1, 2, 3]);
350/// assert_eq!(wrapper.ast_to_str(), r#"
351/// =↓
352/// ├─1
353/// ├─2
354/// ╰─3
355/// "#.trim());
356/// ```
357pub fn print_ast_list_without_node_name<T>(
358 collection: impl IntoIterator<Item = T>,
359 f: impl Fn(T) -> String,
360 s: &dyn Symbols,
361) -> String {
362 let mut collection = collection.into_iter().peekable();
363 let symbol = if collection.peek().is_none() {
364 s.missing_items_symbol()
365 } else {
366 s.item_list_symbol()
367 };
368 format(&format!("={}", symbol)[..], collection.map(f).collect(), s)
369}
370
371/// Very inefficiently formats the given node and children into a tree by indenting every line.
372///
373// # Example
374/// ```
375/// # use ast2str_lib as ast2str;
376/// use ast2str::{DefaultSymbols, builder::format};
377///
378/// let output = format(
379/// "A",
380/// vec![
381/// format("B", vec!["b1".to_string(), "b2".to_string(), "b3".to_string()], &DefaultSymbols),
382/// format("C", vec![format("D", vec!["d1".to_string(), "d2".to_string()], &DefaultSymbols)], &DefaultSymbols)
383/// ],
384/// &DefaultSymbols
385/// );
386/// assert_eq!(output, r#"
387/// A
388/// ├─B
389/// │ ├─b1
390/// │ ├─b2
391/// │ ╰─b3
392/// ╰─C
393/// ╰─D
394/// ├─d1
395/// ╰─d2
396/// "#.trim());
397/// ```
398pub fn format(tree: &str, children: Vec<String>, symbols: &dyn Symbols) -> String {
399 let mut new_tree = Vec::with_capacity(children.len());
400 new_tree.push(tree.to_string());
401
402 if !children.is_empty() {
403 for child in &children[0..children.len() - 1] {
404 new_tree.push(indent(child, TreeIndent::Trunk, symbols));
405 }
406
407 new_tree.push(indent(
408 children.last().unwrap(),
409 TreeIndent::Branch,
410 symbols,
411 ));
412 }
413
414 new_tree.join("\n")
415}
416
417fn indent(tree: &str, kind: TreeIndent, s: &dyn Symbols) -> String {
418 let tree = tree.split('\n').collect::<Vec<&str>>();
419 let mut new_tree = vec![];
420
421 // Handle the root first
422 let wood = if kind.is_trunk() {
423 format!("{}{}", s.right_branch(), s.horizontal_bar())
424 } else {
425 format!("{}{}", s.left_bottom_corner(), s.horizontal_bar())
426 };
427 new_tree.push(format!("{}{}", wood, tree[0]));
428
429 let indent = s.indent();
430
431 if tree.len() > 1 {
432 for child in tree[1..tree.len()].iter() {
433 let wood = if kind.is_trunk() {
434 format!("{}{}", s.vertical_bar(), indent)
435 } else {
436 indent.repeat(2)
437 };
438
439 new_tree.push(format!("{}{}", wood, child));
440 }
441 }
442
443 new_tree.join("\n")
444}
445
446#[cfg(test)]
447pub mod tests {
448 use super::*;
449 use crate::AstToStr;
450 use pretty_assertions::assert_eq;
451
452 #[derive(Debug)]
453 enum BinOpKind {
454 Plus,
455 Minus,
456 }
457 struct Token {
458 lexeme: &'static str,
459 }
460
461 struct If {
462 condition: Box<Expr>,
463 then: Vec<Expr>,
464 otherwise: Option<Vec<Expr>>,
465 }
466
467 enum Expr {
468 Binary(Box<Expr>, BinOpKind, Box<Expr>),
469 Variable(Token),
470 If(If),
471 }
472
473 impl AstToStr for Token {
474 fn ast_to_str_impl(&self, _: &dyn crate::Symbols) -> String {
475 self.lexeme.to_string()
476 }
477 }
478
479 impl AstToStr for Expr {
480 fn ast_to_str_impl(&self, s: &dyn crate::Symbols) -> String {
481 match self {
482 Expr::Binary(l, op, r) => TreeBuilder::new("Expr::Binary", s)
483 .field("left", l)
484 .debug("op", op)
485 .field("right", r)
486 .build(),
487 Expr::Variable(v) => TreeBuilder::new("Expr::Variable", s)
488 .quoted("name", &v.lexeme)
489 .build(),
490 Expr::If(r#if) => TreeBuilder::new("Expr::If", s)
491 .field("condition", &r#if.condition)
492 .list("then", r#if.then.iter())
493 .option("otherwise", "------", &r#if.otherwise)
494 .build(),
495 }
496 }
497 }
498
499 #[test]
500 fn test_builder() {
501 let ast = Expr::If(If {
502 condition: Box::new(Expr::Binary(
503 Box::new(Expr::Variable(Token { lexeme: "a" })),
504 BinOpKind::Plus,
505 Box::new(Expr::Variable(Token { lexeme: "b" })),
506 )),
507 then: vec![Expr::Binary(
508 Box::new(Expr::Variable(Token { lexeme: "c" })),
509 BinOpKind::Minus,
510 Box::new(Expr::Variable(Token { lexeme: "d" })),
511 )],
512 otherwise: Some(vec![]),
513 });
514
515 assert_eq!(
516 ast.ast_to_str(),
517 r#"Expr::If
518├─condition: Expr::Binary
519│ ├─left: Expr::Variable
520│ │ ╰─name: `a`
521│ ├─op: Plus
522│ ╰─right: Expr::Variable
523│ ╰─name: `b`
524├─then=↓
525│ ╰─Expr::Binary
526│ ├─left: Expr::Variable
527│ │ ╰─name: `c`
528│ ├─op: Minus
529│ ╰─right: Expr::Variable
530│ ╰─name: `d`
531╰─otherwise=✕"#
532 );
533 }
534
535 #[test]
536 fn test_builder_with_custom_symbols() {
537 let ast = Expr::If(If {
538 condition: Box::new(Expr::Binary(
539 Box::new(Expr::Variable(Token { lexeme: "a" })),
540 BinOpKind::Plus,
541 Box::new(Expr::Variable(Token { lexeme: "b" })),
542 )),
543 then: vec![Expr::Binary(
544 Box::new(Expr::Variable(Token { lexeme: "c" })),
545 BinOpKind::Minus,
546 Box::new(Expr::Variable(Token { lexeme: "d" })),
547 )],
548 otherwise: Some(vec![]),
549 });
550
551 println!("{}", ast.ast_to_str_impl(&crate::TestSymbols));
552 assert_eq!(
553 ast.ast_to_str_impl(&crate::TestSymbols),
554 r#"Expr::If
555 condition: Expr::Binary
556 left: Expr::Variable
557 name: `a`
558 op: Plus
559 right: Expr::Variable
560 name: `b`
561 then=
562 Expr::Binary
563 left: Expr::Variable
564 name: `c`
565 op: Minus
566 right: Expr::Variable
567 name: `d`
568 otherwise="#
569 );
570 }
571}