use crate::codegen::common::{emit_doc, DocCommentStyle};
#[derive(Debug, Clone)]
pub struct CodeWriter {
buf: String,
depth: usize,
unit: String,
}
impl CodeWriter {
pub fn new(unit: impl Into<String>) -> Self {
Self {
buf: String::new(),
depth: 0,
unit: unit.into(),
}
}
pub fn four_space() -> Self {
Self::new(" ")
}
pub fn two_space() -> Self {
Self::new(" ")
}
pub fn tabs() -> Self {
Self::new("\t")
}
pub fn depth(&self) -> usize {
self.depth
}
#[must_use]
pub fn with_depth(mut self, depth: usize) -> Self {
self.depth = depth;
self
}
pub fn indent_str(&self) -> String {
self.unit.repeat(self.depth)
}
pub fn indent(&mut self) -> &mut Self {
self.depth += 1;
self
}
pub fn dedent(&mut self) -> &mut Self {
self.depth = self.depth.saturating_sub(1);
self
}
pub fn line(&mut self, s: impl AsRef<str>) -> &mut Self {
let s = s.ref_str();
if s.is_empty() {
self.buf.push('\n');
} else {
self.buf.push_str(&self.unit.repeat(self.depth));
self.buf.push_str(s);
self.buf.push('\n');
}
self
}
pub fn blank(&mut self) -> &mut Self {
self.buf.push('\n');
self
}
pub fn raw(&mut self, s: impl AsRef<str>) -> &mut Self {
self.buf.push_str(s.ref_str());
self
}
pub fn block_raw(&mut self, s: impl AsRef<str>) -> &mut Self {
let s = s.ref_str();
if s.is_empty() {
return self;
}
let prefix = self.unit.repeat(self.depth);
let ends_with_newline = s.ends_with('\n');
let mut lines = s.split('\n').peekable();
while let Some(line) = lines.next() {
let is_last = lines.peek().is_none();
if is_last && line.is_empty() && ends_with_newline {
break;
}
if line.is_empty() {
self.buf.push('\n');
} else {
self.buf.push_str(&prefix);
self.buf.push_str(line);
self.buf.push('\n');
}
}
self
}
pub fn scope(&mut self, f: impl FnOnce(&mut Self)) -> &mut Self {
self.indent();
f(self);
self.dedent();
self
}
pub fn block(
&mut self,
open: impl AsRef<str>,
close: impl AsRef<str>,
f: impl FnOnce(&mut Self),
) -> &mut Self {
self.line(open);
self.scope(f);
self.line(close);
self
}
pub fn doc(&mut self, doc: &Option<String>, style: DocCommentStyle) -> &mut Self {
let prefix = self.unit.repeat(self.depth);
emit_doc(&mut self.buf, doc, &prefix, style);
self
}
pub fn as_str(&self) -> &str {
&self.buf
}
pub fn is_empty(&self) -> bool {
self.buf.is_empty()
}
pub fn finish(self) -> String {
self.buf
}
}
trait RefStr {
fn ref_str(&self) -> &str;
}
impl<T: AsRef<str>> RefStr for T {
fn ref_str(&self) -> &str {
self.as_ref()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn line_indents_and_newline_terminates() {
let mut w = CodeWriter::four_space();
w.line("a");
w.indent();
w.line("b");
w.dedent();
w.line("c");
assert_eq!(w.finish(), "a\n b\nc\n");
}
#[test]
fn empty_line_is_bare_newline() {
let mut w = CodeWriter::four_space();
w.indent();
w.line("");
w.blank();
w.line("x");
assert_eq!(w.finish(), "\n\n x\n");
}
#[test]
fn scope_restores_depth() {
let mut w = CodeWriter::two_space();
w.line("outer");
w.scope(|w| {
w.line("inner");
w.scope(|w| {
w.line("deepest");
});
w.line("inner again");
});
w.line("outer again");
assert_eq!(
w.finish(),
"outer\n inner\n deepest\n inner again\nouter again\n"
);
}
#[test]
fn block_brackets_body() {
let mut w = CodeWriter::four_space();
w.block("if (x) {", "}", |w| {
w.line("do_a();");
w.line("do_b();");
});
assert_eq!(w.finish(), "if (x) {\n do_a();\n do_b();\n}\n");
}
#[test]
fn nested_blocks() {
let mut w = CodeWriter::four_space();
w.block("class A:", "", |w| {
w.block("def f(self):", "", |w| {
w.line("pass");
});
});
assert_eq!(w.finish(), "class A:\n def f(self):\n pass\n\n\n");
}
#[test]
fn raw_appends_verbatim() {
let mut w = CodeWriter::four_space();
w.indent();
w.raw("no-indent");
w.raw(" continues");
assert_eq!(w.finish(), "no-indent continues");
}
#[test]
fn block_raw_reindents_relative_structure() {
let mut w = CodeWriter::four_space();
w.indent();
w.block_raw("def foo():\n return 1\n");
assert_eq!(w.finish(), " def foo():\n return 1\n");
}
#[test]
fn block_raw_preserves_blank_lines_without_trailing_ws() {
let mut w = CodeWriter::two_space();
w.indent();
w.block_raw("a\n\nb\n");
assert_eq!(w.finish(), " a\n\n b\n");
}
#[test]
fn block_raw_without_trailing_newline() {
let mut w = CodeWriter::four_space();
w.block_raw("one\ntwo");
assert_eq!(w.finish(), "one\ntwo\n");
}
#[test]
fn block_raw_empty_is_noop() {
let mut w = CodeWriter::four_space();
w.block_raw("");
assert!(w.is_empty());
}
#[test]
fn doc_uses_current_indent() {
let mut w = CodeWriter::four_space();
w.indent();
w.doc(&Some("Hello.".to_string()), DocCommentStyle::TripleSlash);
w.line("fn f() {}");
assert_eq!(w.finish(), " /// Hello.\n fn f() {}\n");
}
#[test]
fn doc_none_is_noop() {
let mut w = CodeWriter::four_space();
w.doc(&None, DocCommentStyle::Hash);
assert!(w.is_empty());
}
#[test]
fn accepts_string_and_str() {
let mut w = CodeWriter::four_space();
let owned = String::from("owned");
w.line(&owned);
w.line("borrowed");
w.line(format!("fmt {}", 1));
assert_eq!(w.finish(), "owned\nborrowed\nfmt 1\n");
}
#[test]
fn indent_str_reflects_depth() {
let mut w = CodeWriter::new(" ");
assert_eq!(w.indent_str(), "");
w.indent().indent();
assert_eq!(w.indent_str(), " ");
assert_eq!(w.depth(), 2);
}
#[test]
fn with_depth_seeds_initial_indentation() {
let mut w = CodeWriter::four_space().with_depth(2);
w.line("if x:");
w.scope(|w| {
w.line("pass");
});
assert_eq!(w.finish(), " if x:\n pass\n");
}
#[test]
fn dedent_saturates_at_zero() {
let mut w = CodeWriter::four_space();
w.dedent().dedent();
w.line("x");
assert_eq!(w.finish(), "x\n");
}
}