1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

/// This crate contains a utility struct for generating indented code strings.
///
/// The `CodeBuilder` struct accepts lines of code as input, and writes them to a string buffer.
/// The code can then be finalized with `result`, which returns an indented code string.

use std::cmp::max;

#[derive(Debug)]
pub struct CodeBuilder {
    code: String,
    indent_level: i32,
    indent_size: i32,
}

impl CodeBuilder {
    /// Adds a single line of code to this code builder, formatting it based on previous code.
    pub fn add_line<S>(&mut self, line: S)
        where S: AsRef<str>
    {
        let line = line.as_ref().trim();
        let indent_change = (line.matches("{").count() as i32) - (line.matches("}").count() as i32);
        let new_indent_level = max(0, self.indent_level + indent_change);

        // Lines starting with '}' should be de-indented even if they contain '{' after; in
        // addition, lines ending with ':' are typically labels, e.g., in LLVM.
        let this_line_indent = if line.starts_with("}") || line.ends_with(":") {
            self.indent_level - 1
        } else {
            self.indent_level
        };

        // self.code.push_str(this_line_indent.as_ref());
        for _ in 0..this_line_indent * self.indent_size {
            self.code.push(' ');
        }
        self.code.push_str(line);
        self.code.push_str("\n");

        self.indent_level = new_indent_level;
    }

    /// Adds one or more lines (split by "\n") to this code builder.
    pub fn add<S>(&mut self, code: S)
        where S: AsRef<str>
    {
        for l in code.as_ref().lines() {
            self.add_line(l);
        }
    }

    /// Adds the result of another CodeBuilder.
    pub fn add_code(&mut self, builder: &CodeBuilder) {
        self.code.push_str(builder.code.as_ref());
    }

    /// Returns the code in this code builder so far.
    pub fn result(&self) -> &str {
        self.code.as_str()
    }

    /// Returns a new CodeBuilder.
    pub fn new() -> CodeBuilder {
        CodeBuilder::with_indent_size(2)
    }

    /// Returns a new CodeBuilder with the given indent size.
    pub fn with_indent_size(indent_size: i32) -> CodeBuilder {
        CodeBuilder {
            code: String::new(),
            indent_level: 0,
            indent_size: indent_size,
        }
    }

    /// Returns a formatted string using the CodeBuilder.
    pub fn format<S>(indent_size: i32, code: S) -> String
        where S: AsRef<str>
    {
        let mut c = CodeBuilder::with_indent_size(indent_size);
        c.add(code);
        c.code
    }
}

#[test]
fn code_builder_basic() {
    let input = "
class A {
blahblah;
}";
    let expected = "
class A {
  blahblah;
}
";
    assert_eq!(CodeBuilder::format(2, input), expected);

    let input = "
class A {
if (c) {
duh
     }
}";
    let expected = "
class A {
  if (c) {
    duh
  }
}
";
    assert_eq!(CodeBuilder::format(2, input), expected);

    let input = "
class A {
if (c) { duh }
}";
    let expected = "
class A {
  if (c) { duh }
}
";
    assert_eq!(CodeBuilder::format(2, input), expected);

    let input = "
class A {
if (c) { duh }
myLabel:
blah
}";
    let expected = "
class A {
  if (c) { duh }
myLabel:
  blah
}
";
    assert_eq!(CodeBuilder::format(2, input), expected);
}