Fix formatting conditionals with comments

sway-ast/src/expr/mod.rs

@@ -406,6 +406,13 @@ impl Spanned for MatchBranchKind {
 pub struct CodeBlockContents {
     pub statements: Vec<Statement>,
     pub final_expr_opt: Option<Box<Expr>>,
+    pub span: Span,
+}
+
+impl Spanned for CodeBlockContents {
+    fn span(&self) -> Span {
+        self.span.clone()
+    }
 }
 
 #[derive(Clone, Debug, Serialize)]

sway-parse/src/attribute.rs

@@ -237,6 +237,7 @@ mod tests {
                   inner: Nil,
                   span: (178, 180),
                 ))),
+                span: (161, 193),
               ),
               span: (160, 194),
             ),

sway-parse/src/expr/mod.rs

@@ -139,6 +139,7 @@ impl ParseToEnd for CodeBlockContents {
         let code_block_contents = CodeBlockContents {
             statements,
             final_expr_opt,
+            span: parser.full_span().clone(),
         };
         Ok((code_block_contents, consumed))
     }

sway-parse/src/module.rs

@@ -132,6 +132,7 @@ mod tests {
                         inner: Nil,
                         span: (70, 72),
                       ))),
+                      span: (53, 85),
                     ),
                     span: (52, 86),
                   ),

sway-parse/src/parser.rs

@@ -149,6 +149,10 @@ impl<'a, 'e> Parser<'a, 'e> {
         }
         Ok(())
     }
+
+    pub fn full_span(&self) -> &Span {
+        &self.full_span
+    }
 }
 
 pub struct Peeker<'a> {

sway-parse/tests/noop_script.rs

@@ -55,6 +55,7 @@ fn noop_script_file() {
                     inner: Nil,
                     span: (48, 50),
                   ))),
+                  span: (39, 57),
                 ),
                 span: (38, 58),
               ),

swayfmt/src/items/item_fn/tests.rs

@@ -88,3 +88,35 @@ fmt_test_item!(  fn_nested_if_lets
     };
 }"
 );
+
+fmt_test_item!(  fn_conditional_with_comment
+"fn conditional_with_comment() {
+    if true {
+        // comment here
+    }
+}",
+intermediate_whitespace
+"fn conditional_with_comment() {
+    if true {
+        // comment here
+    }
+}"
+);
+
+fmt_test_item!(  fn_conditional_with_comment_and_else
+"fn conditional_with_comment() {
+    if true {
+        // if
+    } else {
+        // else
+    }
+}",
+intermediate_whitespace
+"fn conditional_with_comment() {
+    if true {
+        // if
+    } else {
+        // else
+    }
+}"
+);

swayfmt/src/utils/language/expr/code_block.rs

@@ -1,4 +1,5 @@
 use crate::{
+    comments::write_comments,
     config::items::ItemBraceStyle,
     formatter::{shape::LineStyle, *},
     utils::{
@@ -8,7 +9,7 @@ use crate::{
 };
 use std::fmt::Write;
 use sway_ast::CodeBlockContents;
-use sway_types::ast::Delimiter;
+use sway_types::{ast::Delimiter, Spanned};
 
 impl Format for CodeBlockContents {
     fn format(
@@ -52,6 +53,15 @@ impl Format for CodeBlockContents {
                     }
                 }
             }
+        } else {
+            let comments: bool = write_comments(
+                formatted_code,
+                self.span().start()..self.span().end(),
+                formatter,
+            )?;
+            if !comments {
+                formatter.shape.block_unindent(&formatter.config);
+            }
         }
 
         Ok(())

swayfmt/src/utils/language/expr/conditional.rs

@@ -142,16 +142,24 @@ fn format_then_block(
     formatted_code: &mut FormattedCode,
     formatter: &mut Formatter,
 ) -> Result<(), FormatterError> {
+    IfExpr::open_curly_brace(formatted_code, formatter)?;
+
     if !if_expr.then_block.get().statements.is_empty()
         || if_expr.then_block.get().final_expr_opt.is_some()
     {
-        IfExpr::open_curly_brace(formatted_code, formatter)?;
         if_expr.then_block.get().format(formatted_code, formatter)?;
-        if if_expr.else_opt.is_none() {
-            IfExpr::close_curly_brace(formatted_code, formatter)?;
-        }
     } else {
-        write!(formatted_code, " {{}}")?;
+        let comments = write_comments(
+            formatted_code,
+            if_expr.then_block.span().start()..if_expr.then_block.span().end(),
+            formatter,
+        )?;
+        if !comments {
+            formatter.shape.block_unindent(&formatter.config);
+        }
+    }
+    if if_expr.else_opt.is_none() {
+        IfExpr::close_curly_brace(formatted_code, formatter)?;
     }
 
     Ok(())