DEV: Handle comments embedded in binary expressions and call expressions.

src/ast/printer/mod.rs

@@ -123,6 +123,15 @@ where
         Ok(())
     }
 
+    fn has_comment(&self, line: usize) -> bool {
+        if line > self.last_line {
+            if let Some(next_comment_line) = self.comment_group_lines.last() {
+                return *next_comment_line < line;
+            }
+        }
+        false
+    }
+
     fn render_list_def(&mut self, def: &ListDef) -> std::io::Result<()> {
         write!(self.w, "[")?;
         self.curr_indent += self.indent_size;
@@ -209,13 +218,14 @@ where
     }
 
     pub fn render_expr(&mut self, expr: &Expression) -> std::io::Result<()> {
+        self.render_comment_if_needed(expr.pos().line)?;
         match expr {
             Expression::Binary(_def) => {
                 let op = match _def.kind {
                     BinaryExprType::AND => " && ",
                     BinaryExprType::OR => " || ",
                     BinaryExprType::DOT => ".",
-                    BinaryExprType::Equal => " = ",
+                    BinaryExprType::Equal => " == ",
                     BinaryExprType::NotEqual => " != ",
                     BinaryExprType::GTEqual => " >= ",
                     BinaryExprType::LTEqual => " <= ",
@@ -231,8 +241,14 @@ where
                     BinaryExprType::REMatch => " ~ ",
                     BinaryExprType::NotREMatch => " !~ ",
                 };
+                let right_line = _def.right.pos().line;
                 self.render_expr(&_def.left)?;
                 self.w.write(op.as_bytes())?;
+                if self.has_comment(right_line) {
+                    // if we'll be rendering a comment then we should
+                    // add a new line here
+                    self.w.write("\n".as_bytes())?;
+                }
                 self.render_expr(&_def.right)?;
             }
             Expression::Call(_def) => {
@@ -245,6 +261,7 @@ where
                     write!(self.w, "\n")?;
                 }
                 for e in _def.arglist.iter() {
+                    self.render_comment_if_needed(e.pos().line)?;
                     if has_args {
                         write!(self.w, "{}", indent)?;
                     }

src/ast/printer/test.rs

@@ -464,3 +464,45 @@ fn test_list_expression_with_embedded_comment() {
         format!("{}\n", input.trim())
     );
 }
+
+#[test]
+fn test_binary_expression_with_embedded_comment() {
+    let mut comment_map = BTreeMap::new();
+    let input = "true == \n// false is not true\nfalse;";
+    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_empty_call_expression_with_comment() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// a comment\nmyfunc();";
+    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_call_expression_with_embedded_comment_in_args() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// a comment\nmyfunc(\n  arg1,\n  // another comment\n  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())
+    );
+}