1 use std::fmt::{self, Write};
2 use std::ops::Deref;
3 
4 /// Helper structure to maintain indentation automatically when printing.
5 #[derive(Default)]
6 pub struct Source {
7     s: String,
8     indent: usize,
9 }
10 
11 impl Source {
push_str(&mut self, src: &str)12     pub fn push_str(&mut self, src: &str) {
13         let lines = src.lines().collect::<Vec<_>>();
14         for (i, line) in lines.iter().enumerate() {
15             let trimmed = line.trim();
16             if trimmed.starts_with('}') && self.s.ends_with("  ") {
17                 self.s.pop();
18                 self.s.pop();
19             }
20             self.s.push_str(if lines.len() == 1 {
21                 line
22             } else {
23                 line.trim_start()
24             });
25             if trimmed.ends_with('{') {
26                 self.indent += 1;
27             }
28             if trimmed.starts_with('}') {
29                 // Note that a `saturating_sub` is used here to prevent a panic
30                 // here in the case of invalid code being generated in debug
31                 // mode. It's typically easier to debug those issues through
32                 // looking at the source code rather than getting a panic.
33                 self.indent = self.indent.saturating_sub(1);
34             }
35             if i != lines.len() - 1 || src.ends_with('\n') {
36                 self.newline();
37             }
38         }
39     }
40 
indent(&mut self, amt: usize)41     pub fn indent(&mut self, amt: usize) {
42         self.indent += amt;
43     }
44 
deindent(&mut self, amt: usize)45     pub fn deindent(&mut self, amt: usize) {
46         self.indent -= amt;
47     }
48 
newline(&mut self)49     fn newline(&mut self) {
50         self.s.push('\n');
51         for _ in 0..self.indent {
52             self.s.push_str("  ");
53         }
54     }
55 
as_mut_string(&mut self) -> &mut String56     pub fn as_mut_string(&mut self) -> &mut String {
57         &mut self.s
58     }
59 }
60 
61 impl Write for Source {
write_str(&mut self, s: &str) -> fmt::Result62     fn write_str(&mut self, s: &str) -> fmt::Result {
63         self.push_str(s);
64         Ok(())
65     }
66 }
67 
68 impl Deref for Source {
69     type Target = str;
deref(&self) -> &str70     fn deref(&self) -> &str {
71         &self.s
72     }
73 }
74 
75 impl From<Source> for String {
from(s: Source) -> String76     fn from(s: Source) -> String {
77         s.s
78     }
79 }
80 
81 #[cfg(test)]
82 mod tests {
83     use super::Source;
84 
85     #[test]
simple_append()86     fn simple_append() {
87         let mut s = Source::default();
88         s.push_str("x");
89         assert_eq!(s.s, "x");
90         s.push_str("y");
91         assert_eq!(s.s, "xy");
92         s.push_str("z ");
93         assert_eq!(s.s, "xyz ");
94         s.push_str(" a ");
95         assert_eq!(s.s, "xyz  a ");
96         s.push_str("\na");
97         assert_eq!(s.s, "xyz  a \na");
98     }
99 
100     #[test]
newline_remap()101     fn newline_remap() {
102         let mut s = Source::default();
103         s.push_str("function() {\n");
104         s.push_str("y\n");
105         s.push_str("}\n");
106         assert_eq!(s.s, "function() {\n  y\n}\n");
107     }
108 
109     #[test]
if_else()110     fn if_else() {
111         let mut s = Source::default();
112         s.push_str("if() {\n");
113         s.push_str("y\n");
114         s.push_str("} else if () {\n");
115         s.push_str("z\n");
116         s.push_str("}\n");
117         assert_eq!(s.s, "if() {\n  y\n} else if () {\n  z\n}\n");
118     }
119 
120     #[test]
trim_ws()121     fn trim_ws() {
122         let mut s = Source::default();
123         s.push_str(
124             "function() {
125                 x
126         }",
127         );
128         assert_eq!(s.s, "function() {\n  x\n}");
129     }
130 }
131