From f88955517ebf55885fe4a1246c1512060ed049da Mon Sep 17 00:00:00 2001
From: Jeremy Wall <jeremy@marzhillstudios.com>
Date: Wed, 22 May 2019 19:05:48 -0500
Subject: [PATCH] DEV: Handle TRACE, fail, and format expressions with embedded
 comments.

---
 src/ast/printer/mod.rs  | 19 +++++++++++++-
 src/ast/printer/test.rs | 56 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/src/ast/printer/mod.rs b/src/ast/printer/mod.rs
index 7d0f1c4..956653f 100644
--- a/src/ast/printer/mod.rs
+++ b/src/ast/printer/mod.rs
@@ -282,10 +282,16 @@ where
             }
             Expression::Debug(_def) => {
                 self.w.write("TRACE ".as_bytes())?;
+                if self.has_comment(_def.expr.pos().line) {
+                    self.w.write("\n".as_bytes())?;
+                }
                 self.render_expr(&_def.expr)?;
             }
             Expression::Fail(_def) => {
                 self.w.write("fail ".as_bytes())?;
+                if self.has_comment(_def.message.pos().line) {
+                    self.w.write("\n".as_bytes())?;
+                }
                 self.render_expr(&_def.message)?;
             }
             Expression::Format(_def) => {
@@ -293,13 +299,24 @@ where
                 write!(self.w, " % ")?;
                 match _def.args {
                     FormatArgs::Single(ref e) => {
+                        if self.has_comment(e.pos().line) {
+                            self.w.write("\n".as_bytes())?;
+                        }
                         self.render_expr(e)?;
                     }
                     FormatArgs::List(ref es) => {
                         self.w.write("(\n".as_bytes())?;
                         self.curr_indent += self.indent_size;
                         let indent = self.make_indent();
-                        let mut prefix = "";
+                        let mut prefix = if es
+                            .first()
+                            .and_then(|e| Some(self.has_comment(e.pos().line)))
+                            .unwrap_or(false)
+                        {
+                            "\n"
+                        } else {
+                            ""
+                        };
                         for e in es.iter() {
                             write!(self.w, "{}{}", prefix, indent)?;
                             self.render_expr(e)?;
diff --git a/src/ast/printer/test.rs b/src/ast/printer/test.rs
index 85a8ea6..dd619df 100644
--- a/src/ast/printer/test.rs
+++ b/src/ast/printer/test.rs
@@ -506,3 +506,59 @@ fn test_call_expression_with_embedded_comment_in_args() {
         format!("{}\n", input.trim())
     );
 }
+
+#[test]
+fn test_copy_expression_with_embedded_comment_in_args() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// a comment\nmyfunc{\n  foo = arg1,\n  // another comment\n  bar = arg2,\n};";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(
+        String::from_utf8(buffer).unwrap(),
+        format!("{}\n", input.trim())
+    );
+}
+
+#[test]
+fn test_trace_expression_with_embedded_comment() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// a comment\nTRACE \n// another comment\nfoo;";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(
+        String::from_utf8(buffer).unwrap(),
+        format!("{}\n", input.trim())
+    );
+}
+
+#[test]
+fn test_fail_expression_with_embedded_comment() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// a comment\nfail \n// another comment\nfoo;";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(
+        String::from_utf8(buffer).unwrap(),
+        format!("{}\n", input.trim())
+    );
+}
+
+#[test]
+fn test_format_expression_with_embedded_comment() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// a comment\n\"@(item.bar)\" % \n// another comment\nfoo;";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(
+        String::from_utf8(buffer).unwrap(),
+        format!("{}\n", input.trim())
+    );
+}