baobao_codegen/builder/
code_builder.rs1use super::{CodeFragment, Indent, Renderable};
4
5#[derive(Debug, Clone)]
41pub struct CodeBuilder {
42 indent_level: usize,
43 indent: Indent,
44 buffer: String,
45}
46
47impl CodeBuilder {
48 pub fn new(indent: Indent) -> Self {
50 Self {
51 indent_level: 0,
52 indent,
53 buffer: String::new(),
54 }
55 }
56
57 pub fn rust() -> Self {
59 Self::new(Indent::RUST)
60 }
61
62 pub fn typescript() -> Self {
64 Self::new(Indent::TYPESCRIPT)
65 }
66
67 pub fn go() -> Self {
69 Self::new(Indent::GO)
70 }
71
72 pub fn push_line(&mut self, s: &str) -> &mut Self {
78 self.write_indent();
79 self.buffer.push_str(s);
80 self.buffer.push('\n');
81 self
82 }
83
84 pub fn push_blank(&mut self) -> &mut Self {
86 self.buffer.push('\n');
87 self
88 }
89
90 pub fn push_raw(&mut self, s: &str) -> &mut Self {
92 self.buffer.push_str(s);
93 self
94 }
95
96 pub fn push_indent(&mut self) -> &mut Self {
98 self.indent_level += 1;
99 self
100 }
101
102 pub fn push_dedent(&mut self) -> &mut Self {
104 self.indent_level = self.indent_level.saturating_sub(1);
105 self
106 }
107
108 pub fn push_jsdoc(&mut self, text: &str) -> &mut Self {
110 self.write_indent();
111 self.buffer.push_str("/** ");
112 self.buffer.push_str(text);
113 self.buffer.push_str(" */\n");
114 self
115 }
116
117 pub fn push_rust_doc(&mut self, text: &str) -> &mut Self {
119 self.write_indent();
120 self.buffer.push_str("/// ");
121 self.buffer.push_str(text);
122 self.buffer.push('\n');
123 self
124 }
125
126 pub fn emit(&mut self, node: &impl Renderable) -> &mut Self {
130 for fragment in node.to_fragments() {
131 self.apply_fragment(fragment);
132 }
133 self
134 }
135
136 pub fn apply_fragment(&mut self, fragment: CodeFragment) {
138 match fragment {
139 CodeFragment::Line(s) => {
140 self.push_line(&s);
141 }
142 CodeFragment::Blank => {
143 self.push_blank();
144 }
145 CodeFragment::Raw(s) => {
146 self.push_raw(&s);
147 }
148 CodeFragment::Block {
149 header,
150 body,
151 close,
152 } => {
153 self.push_line(&header);
154 self.push_indent();
155 for f in body {
156 self.apply_fragment(f);
157 }
158 self.push_dedent();
159 if let Some(c) = close {
160 self.push_line(&c);
161 }
162 }
163 CodeFragment::Indent(fragments) => {
164 self.push_indent();
165 for f in fragments {
166 self.apply_fragment(f);
167 }
168 self.push_dedent();
169 }
170 CodeFragment::Sequence(fragments) => {
171 for f in fragments {
172 self.apply_fragment(f);
173 }
174 }
175 CodeFragment::JsDoc(text) => {
176 self.push_jsdoc(&text);
177 }
178 CodeFragment::RustDoc(text) => {
179 self.push_rust_doc(&text);
180 }
181 }
182 }
183
184 pub fn line(mut self, s: &str) -> Self {
190 self.push_line(s);
191 self
192 }
193
194 pub fn blank(mut self) -> Self {
196 self.push_blank();
197 self
198 }
199
200 pub fn raw(mut self, s: &str) -> Self {
202 self.push_raw(s);
203 self
204 }
205
206 pub fn indent(mut self) -> Self {
208 self.push_indent();
209 self
210 }
211
212 pub fn dedent(mut self) -> Self {
214 self.push_dedent();
215 self
216 }
217
218 pub fn block<F>(self, header: &str, f: F) -> Self
232 where
233 F: FnOnce(Self) -> Self,
234 {
235 let builder = self.line(header).indent();
236 f(builder).dedent()
237 }
238
239 pub fn block_with_close<F>(self, header: &str, close: &str, f: F) -> Self
253 where
254 F: FnOnce(Self) -> Self,
255 {
256 let builder = self.line(header).indent();
257 f(builder).dedent().line(close)
258 }
259
260 pub fn doc(mut self, prefix: &str, text: &str) -> Self {
262 self.write_indent();
263 self.buffer.push_str(prefix);
264 self.buffer.push(' ');
265 self.buffer.push_str(text);
266 self.buffer.push('\n');
267 self
268 }
269
270 pub fn rust_doc(self, text: &str) -> Self {
272 self.doc("///", text)
273 }
274
275 pub fn jsdoc(mut self, text: &str) -> Self {
277 self.push_jsdoc(text);
278 self
279 }
280
281 pub fn when<F>(self, condition: bool, f: F) -> Self
283 where
284 F: FnOnce(Self) -> Self,
285 {
286 if condition { f(self) } else { self }
287 }
288
289 pub fn each<T, I, F>(mut self, items: I, f: F) -> Self
291 where
292 I: IntoIterator<Item = T>,
293 F: Fn(Self, T) -> Self,
294 {
295 for item in items {
296 self = f(self, item);
297 }
298 self
299 }
300
301 pub fn current_indent(&self) -> usize {
303 self.indent_level
304 }
305
306 pub fn build(self) -> String {
308 self.buffer
309 }
310
311 pub fn as_str(&self) -> &str {
313 &self.buffer
314 }
315
316 fn write_indent(&mut self) {
317 for _ in 0..self.indent_level {
318 self.buffer.push_str(self.indent.as_str());
319 }
320 }
321}
322
323impl Default for CodeBuilder {
324 fn default() -> Self {
325 Self::rust()
326 }
327}
328
329#[cfg(test)]
330mod tests {
331 use super::*;
332
333 #[test]
338 fn test_basic_line() {
339 let code = CodeBuilder::rust().line("let x = 1;").build();
340 assert_eq!(code, "let x = 1;\n");
341 }
342
343 #[test]
344 fn test_indentation() {
345 let code = CodeBuilder::rust()
346 .line("fn main() {")
347 .indent()
348 .line("println!(\"Hello\");")
349 .dedent()
350 .line("}")
351 .build();
352
353 assert_eq!(code, "fn main() {\n println!(\"Hello\");\n}\n");
354 }
355
356 #[test]
357 fn test_block() {
358 let code = CodeBuilder::rust()
359 .block_with_close("impl Foo {", "}", |b| b.line("fn bar(&self) {}"))
360 .build();
361
362 assert_eq!(code, "impl Foo {\n fn bar(&self) {}\n}\n");
363 }
364
365 #[test]
366 fn test_blank_line() {
367 let code = CodeBuilder::rust()
368 .line("use std::io;")
369 .blank()
370 .line("fn main() {}")
371 .build();
372
373 assert_eq!(code, "use std::io;\n\nfn main() {}\n");
374 }
375
376 #[test]
377 fn test_doc_comment() {
378 let code = CodeBuilder::rust()
379 .rust_doc("A test function")
380 .line("fn test() {}")
381 .build();
382
383 assert_eq!(code, "/// A test function\nfn test() {}\n");
384 }
385
386 #[test]
387 fn test_conditional() {
388 let with_debug = CodeBuilder::rust()
389 .when(true, |b| b.line("#[derive(Debug)]"))
390 .line("struct Foo;")
391 .build();
392
393 let without_debug = CodeBuilder::rust()
394 .when(false, |b| b.line("#[derive(Debug)]"))
395 .line("struct Foo;")
396 .build();
397
398 assert_eq!(with_debug, "#[derive(Debug)]\nstruct Foo;\n");
399 assert_eq!(without_debug, "struct Foo;\n");
400 }
401
402 #[test]
403 fn test_each() {
404 let code = CodeBuilder::rust()
405 .line("enum Color {")
406 .indent()
407 .each(["Red", "Green", "Blue"], |b, color| {
408 b.line(&format!("{},", color))
409 })
410 .dedent()
411 .line("}")
412 .build();
413
414 assert_eq!(code, "enum Color {\n Red,\n Green,\n Blue,\n}\n");
415 }
416
417 #[test]
418 fn test_typescript_indent() {
419 let code = CodeBuilder::typescript()
420 .line("function foo() {")
421 .indent()
422 .line("return 1;")
423 .dedent()
424 .line("}")
425 .build();
426
427 assert_eq!(code, "function foo() {\n return 1;\n}\n");
428 }
429
430 #[test]
435 fn test_mutable_api_basic() {
436 let mut builder = CodeBuilder::rust();
437 builder
438 .push_line("let x = 1;")
439 .push_blank()
440 .push_line("let y = 2;");
441 assert_eq!(builder.build(), "let x = 1;\n\nlet y = 2;\n");
442 }
443
444 #[test]
445 fn test_mutable_api_indentation() {
446 let mut builder = CodeBuilder::typescript();
447 builder
448 .push_line("function foo() {")
449 .push_indent()
450 .push_line("return 1;")
451 .push_dedent()
452 .push_line("}");
453 assert_eq!(builder.build(), "function foo() {\n return 1;\n}\n");
454 }
455
456 #[test]
457 fn test_emit_with_fragments() {
458 struct SimpleNode;
459 impl Renderable for SimpleNode {
460 fn to_fragments(&self) -> Vec<CodeFragment> {
461 vec![
462 CodeFragment::Line("// comment".to_string()),
463 CodeFragment::Line("let x = 1;".to_string()),
464 ]
465 }
466 }
467
468 let mut builder = CodeBuilder::rust();
469 builder.emit(&SimpleNode);
470 assert_eq!(builder.build(), "// comment\nlet x = 1;\n");
471 }
472
473 #[test]
474 fn test_emit_block_fragment() {
475 struct BlockNode;
476 impl Renderable for BlockNode {
477 fn to_fragments(&self) -> Vec<CodeFragment> {
478 vec![CodeFragment::Block {
479 header: "fn main() {".to_string(),
480 body: vec![CodeFragment::Line("println!(\"Hello\");".to_string())],
481 close: Some("}".to_string()),
482 }]
483 }
484 }
485
486 let mut builder = CodeBuilder::rust();
487 builder.emit(&BlockNode);
488 assert_eq!(
489 builder.build(),
490 "fn main() {\n println!(\"Hello\");\n}\n"
491 );
492 }
493
494 #[test]
495 fn test_emit_jsdoc_fragment() {
496 struct DocNode;
497 impl Renderable for DocNode {
498 fn to_fragments(&self) -> Vec<CodeFragment> {
499 vec![
500 CodeFragment::JsDoc("A function".to_string()),
501 CodeFragment::Line("function foo() {}".to_string()),
502 ]
503 }
504 }
505
506 let mut builder = CodeBuilder::typescript();
507 builder.emit(&DocNode);
508 assert_eq!(builder.build(), "/** A function */\nfunction foo() {}\n");
509 }
510}