refactor: block indentation

1 files changed, 25 insertions, 9 deletions

diff --git a/src/main.rs b/src/main.rs
index 14e3f22..c28914f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -24,6 +24,26 @@ fn write_to_file<P: AsRef<Path>>(path: P, content: &str) -> std::io::Result<()>
     Ok(())
 }
 
+/* Indent block of code */
+fn indent<'a>(input: Cow<'a,str>, indent: usize) -> Cow<'a,str> {
+    if indent > 0 {
+        let prefix = format!("{:indent$}", "");
+        let mut output = String::with_capacity(input.len() + indent*input.lines().count());
+        input.lines().enumerate().for_each(|(i,line)| {
+            if i > 0 {
+                output.push('\n');
+            }
+            if !line.is_empty() {
+                output.push_str(&prefix);
+                output.push_str(line);
+            }
+        });
+        Cow::Owned(output)
+    } else {
+        input
+    }
+}
+
 /*
  * Here are some notes on the following function
  *
@@ -68,15 +88,11 @@ fn build(blocks: &Blocks) {
         // We can use it to recursively build blocks of code until no more substitutions are
         // necessary (i.e., `replace_all` returns a `Borrowed`).
         while let Cow::Owned(step) = MACRO.replace_all(&code, |caps: &Captures| {
-            let indent = caps[1].len();
-            blocks.get(&caps[2])
-                .expect("Block not present")
-                .lines()
-                .map(|l| format!("{:indent$}{}", "", l) )
-                .collect::<Vec<_>>()
-                .join("\n")
-            }
-        ) {
+            indent(
+                blocks.get(&caps[2]).expect("Block not present").clone(),
+                caps[1].len()
+            )
+        }) {
             code = Cow::from(step);
         }
         write_to_file(k, &code).expect("Unable to write to file");