baobao_codegen/
code_builder.rs1use crate::Indent;
4
5#[derive(Debug, Clone)]
23pub struct CodeBuilder {
24 indent_level: usize,
25 indent: Indent,
26 buffer: String,
27}
28
29impl CodeBuilder {
30 pub fn new(indent: Indent) -> Self {
32 Self {
33 indent_level: 0,
34 indent,
35 buffer: String::new(),
36 }
37 }
38
39 pub fn rust() -> Self {
41 Self::new(Indent::RUST)
42 }
43
44 pub fn typescript() -> Self {
46 Self::new(Indent::TYPESCRIPT)
47 }
48
49 pub fn go() -> Self {
51 Self::new(Indent::GO)
52 }
53
54 pub fn line(mut self, s: &str) -> Self {
56 self.write_indent();
57 self.buffer.push_str(s);
58 self.buffer.push('\n');
59 self
60 }
61
62 pub fn blank(mut self) -> Self {
64 self.buffer.push('\n');
65 self
66 }
67
68 pub fn raw(mut self, s: &str) -> Self {
70 self.buffer.push_str(s);
71 self
72 }
73
74 pub fn indent(mut self) -> Self {
76 self.indent_level += 1;
77 self
78 }
79
80 pub fn dedent(mut self) -> Self {
82 self.indent_level = self.indent_level.saturating_sub(1);
83 self
84 }
85
86 pub fn block<F>(self, header: &str, f: F) -> Self
100 where
101 F: FnOnce(Self) -> Self,
102 {
103 let builder = self.line(header).indent();
104 f(builder).dedent()
105 }
106
107 pub fn block_with_close<F>(self, header: &str, close: &str, f: F) -> Self
121 where
122 F: FnOnce(Self) -> Self,
123 {
124 let builder = self.line(header).indent();
125 f(builder).dedent().line(close)
126 }
127
128 pub fn doc(mut self, prefix: &str, text: &str) -> Self {
130 self.write_indent();
131 self.buffer.push_str(prefix);
132 self.buffer.push(' ');
133 self.buffer.push_str(text);
134 self.buffer.push('\n');
135 self
136 }
137
138 pub fn rust_doc(self, text: &str) -> Self {
140 self.doc("///", text)
141 }
142
143 pub fn jsdoc(mut self, text: &str) -> Self {
145 self.write_indent();
146 self.buffer.push_str("/** ");
147 self.buffer.push_str(text);
148 self.buffer.push_str(" */\n");
149 self
150 }
151
152 pub fn when<F>(self, condition: bool, f: F) -> Self
154 where
155 F: FnOnce(Self) -> Self,
156 {
157 if condition { f(self) } else { self }
158 }
159
160 pub fn each<T, I, F>(mut self, items: I, f: F) -> Self
162 where
163 I: IntoIterator<Item = T>,
164 F: Fn(Self, T) -> Self,
165 {
166 for item in items {
167 self = f(self, item);
168 }
169 self
170 }
171
172 pub fn current_indent(&self) -> usize {
174 self.indent_level
175 }
176
177 pub fn build(self) -> String {
179 self.buffer
180 }
181
182 pub fn as_str(&self) -> &str {
184 &self.buffer
185 }
186
187 fn write_indent(&mut self) {
188 for _ in 0..self.indent_level {
189 self.buffer.push_str(self.indent.as_str());
190 }
191 }
192}
193
194impl Default for CodeBuilder {
195 fn default() -> Self {
196 Self::rust()
197 }
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203
204 #[test]
205 fn test_basic_line() {
206 let code = CodeBuilder::rust().line("let x = 1;").build();
207 assert_eq!(code, "let x = 1;\n");
208 }
209
210 #[test]
211 fn test_indentation() {
212 let code = CodeBuilder::rust()
213 .line("fn main() {")
214 .indent()
215 .line("println!(\"Hello\");")
216 .dedent()
217 .line("}")
218 .build();
219
220 assert_eq!(code, "fn main() {\n println!(\"Hello\");\n}\n");
221 }
222
223 #[test]
224 fn test_block() {
225 let code = CodeBuilder::rust()
226 .block_with_close("impl Foo {", "}", |b| b.line("fn bar(&self) {}"))
227 .build();
228
229 assert_eq!(code, "impl Foo {\n fn bar(&self) {}\n}\n");
230 }
231
232 #[test]
233 fn test_blank_line() {
234 let code = CodeBuilder::rust()
235 .line("use std::io;")
236 .blank()
237 .line("fn main() {}")
238 .build();
239
240 assert_eq!(code, "use std::io;\n\nfn main() {}\n");
241 }
242
243 #[test]
244 fn test_doc_comment() {
245 let code = CodeBuilder::rust()
246 .rust_doc("A test function")
247 .line("fn test() {}")
248 .build();
249
250 assert_eq!(code, "/// A test function\nfn test() {}\n");
251 }
252
253 #[test]
254 fn test_conditional() {
255 let with_debug = CodeBuilder::rust()
256 .when(true, |b| b.line("#[derive(Debug)]"))
257 .line("struct Foo;")
258 .build();
259
260 let without_debug = CodeBuilder::rust()
261 .when(false, |b| b.line("#[derive(Debug)]"))
262 .line("struct Foo;")
263 .build();
264
265 assert_eq!(with_debug, "#[derive(Debug)]\nstruct Foo;\n");
266 assert_eq!(without_debug, "struct Foo;\n");
267 }
268
269 #[test]
270 fn test_each() {
271 let code = CodeBuilder::rust()
272 .line("enum Color {")
273 .indent()
274 .each(["Red", "Green", "Blue"], |b, color| {
275 b.line(&format!("{},", color))
276 })
277 .dedent()
278 .line("}")
279 .build();
280
281 assert_eq!(code, "enum Color {\n Red,\n Green,\n Blue,\n}\n");
282 }
283
284 #[test]
285 fn test_typescript_indent() {
286 let code = CodeBuilder::typescript()
287 .line("function foo() {")
288 .indent()
289 .line("return 1;")
290 .dedent()
291 .line("}")
292 .build();
293
294 assert_eq!(code, "function foo() {\n return 1;\n}\n");
295 }
296}