chore: enhance code formatting for If expressions (#3246)

tooling/nargo_fmt/src/config.rs

@@ -47,6 +47,7 @@ config! {
     remove_nested_parens: bool, true, "Remove nested parens";
     short_array_element_width_threshold: usize, 10, "Width threshold for an array element to be considered short";
     array_width: usize, 100, "Maximum width of an array literal before falling back to vertical formatting";
+    single_line_if_else_max_width: usize, 100, "Maximum line length for single line if-else expressions";
 }
 
 impl Config {

tooling/nargo_fmt/src/utils.rs

@@ -218,7 +218,7 @@ impl Item for Expression {
     }
 
     fn format(self, visitor: &FmtVisitor) -> String {
-        visitor.format_expr(self)
+        visitor.format_subexpr(self)
     }
 }
 
@@ -232,7 +232,7 @@ impl Item for (Ident, Expression) {
         let (name, expr) = self;
 
         let name = name.0.contents;
-        let expr = visitor.format_expr(expr);
+        let expr = visitor.format_subexpr(expr);
 
         if name == expr {
             name

tooling/nargo_fmt/src/visitor.rs

@@ -164,7 +164,7 @@ impl<'me> FmtVisitor<'me> {
     }
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, Default)]
 struct Indent {
     block_indent: usize,
 }
@@ -197,3 +197,9 @@ struct Shape {
     width: usize,
     indent: Indent,
 }
+
+#[derive(PartialEq, Eq, Debug)]
+pub(crate) enum ExpressionType {
+    Statement,
+    SubExpression,
+}

tooling/nargo_fmt/src/visitor/expr.rs

@@ -1,26 +1,35 @@
 use noirc_frontend::{
     hir::resolution::errors::Span, token::Token, ArrayLiteral, BlockExpression,
-    ConstructorExpression, Expression, ExpressionKind, Literal, Statement, UnaryOp,
+    ConstructorExpression, Expression, ExpressionKind, IfExpression, Literal, Statement,
+    StatementKind, UnaryOp,
 };
 
-use super::{FmtVisitor, Shape};
+use super::{ExpressionType, FmtVisitor, Indent, Shape};
 use crate::{
     utils::{self, Expr, FindToken, Item},
     Config,
 };
 
 impl FmtVisitor<'_> {
-    pub(crate) fn visit_expr(&mut self, expr: Expression) {
+    pub(crate) fn visit_expr(&mut self, expr: Expression, expr_type: ExpressionType) {
         let span = expr.span;
 
-        let rewrite = self.format_expr(expr);
+        let rewrite = self.format_expr(expr, expr_type);
         let rewrite = utils::recover_comment_removed(self.slice(span), rewrite);
         self.push_rewrite(rewrite, span);
 
         self.last_position = span.end();
     }
 
-    pub(crate) fn format_expr(&self, Expression { kind, mut span }: Expression) -> String {
+    pub(crate) fn format_subexpr(&self, expression: Expression) -> String {
+        self.format_expr(expression, ExpressionType::SubExpression)
+    }
+
+    pub(crate) fn format_expr(
+        &self,
+        Expression { kind, mut span }: Expression,
+        expr_type: ExpressionType,
+    ) -> String {
         match kind {
             ExpressionKind::Block(block) => {
                 let mut visitor = self.fork();
@@ -41,24 +50,24 @@ impl FmtVisitor<'_> {
                     }
                 };
 
-                format!("{op}{}", self.format_expr(prefix.rhs))
+                format!("{op}{}", self.format_subexpr(prefix.rhs))
             }
             ExpressionKind::Cast(cast) => {
-                format!("{} as {}", self.format_expr(cast.lhs), cast.r#type)
+                format!("{} as {}", self.format_subexpr(cast.lhs), cast.r#type)
             }
             ExpressionKind::Infix(infix) => {
                 format!(
                     "{} {} {}",
-                    self.format_expr(infix.lhs),
+                    self.format_subexpr(infix.lhs),
                     infix.operator.contents.as_string(),
-                    self.format_expr(infix.rhs)
+                    self.format_subexpr(infix.rhs)
                 )
             }
             ExpressionKind::Call(call_expr) => {
                 let args_span =
                     self.span_before(call_expr.func.span.end()..span.end(), Token::LeftParen);
 
-                let callee = self.format_expr(*call_expr.func);
+                let callee = self.format_subexpr(*call_expr.func);
                 let args = format_parens(self.fork(), false, call_expr.arguments, args_span);
 
                 format!("{callee}{args}")
@@ -69,21 +78,21 @@ impl FmtVisitor<'_> {
                     Token::LeftParen,
                 );
 
-                let object = self.format_expr(method_call_expr.object);
+                let object = self.format_subexpr(method_call_expr.object);
                 let method = method_call_expr.method_name.to_string();
                 let args = format_parens(self.fork(), false, method_call_expr.arguments, args_span);
 
                 format!("{object}.{method}{args}")
             }
             ExpressionKind::MemberAccess(member_access_expr) => {
-                let lhs_str = self.format_expr(member_access_expr.lhs);
+                let lhs_str = self.format_subexpr(member_access_expr.lhs);
                 format!("{}.{}", lhs_str, member_access_expr.rhs)
             }
             ExpressionKind::Index(index_expr) => {
                 let index_span = self
                     .span_before(index_expr.collection.span.end()..span.end(), Token::LeftBracket);
 
-                let collection = self.format_expr(index_expr.collection);
+                let collection = self.format_subexpr(index_expr.collection);
                 let index = format_brackets(self.fork(), false, vec![index_expr.index], index_span);
 
                 format!("{collection}{index}")
@@ -96,8 +105,8 @@ impl FmtVisitor<'_> {
                     self.slice(span).to_string()
                 }
                 Literal::Array(ArrayLiteral::Repeated { repeated_element, length }) => {
-                    let repeated = self.format_expr(*repeated_element);
-                    let length = self.format_expr(*length);
+                    let repeated = self.format_subexpr(*repeated_element);
+                    let length = self.format_subexpr(*length);
 
                     format!("[{repeated}; {length}]")
                 }
@@ -131,7 +140,7 @@ impl FmtVisitor<'_> {
                 }
 
                 if !leading.contains("//") && !trailing.contains("//") {
-                    let sub_expr = self.format_expr(*sub_expr);
+                    let sub_expr = self.format_subexpr(*sub_expr);
                     format!("({leading}{sub_expr}{trailing})")
                 } else {
                     let mut visitor = self.fork();
@@ -140,7 +149,7 @@ impl FmtVisitor<'_> {
                     visitor.indent.block_indent(self.config);
                     let nested_indent = visitor.indent.to_string_with_newline();
 
-                    let sub_expr = visitor.format_expr(*sub_expr);
+                    let sub_expr = visitor.format_subexpr(*sub_expr);
 
                     let mut result = String::new();
                     result.push('(');
@@ -171,11 +180,66 @@ impl FmtVisitor<'_> {
 
                 self.format_struct_lit(type_name, fields_span, *constructor)
             }
-            // TODO:
-            _expr => self.slice(span).to_string(),
+            ExpressionKind::If(if_expr) => {
+                let allow_single_line = expr_type == ExpressionType::SubExpression;
+
+                if allow_single_line {
+                    let mut visitor = self.fork();
+                    visitor.indent = Indent::default();
+                    if let Some(line) = visitor.format_if_single_line(*if_expr.clone()) {
+                        return line;
+                    }
+                }
+
+                self.format_if(*if_expr)
+            }
+            _ => self.slice(span).to_string(),
         }
     }
 
+    fn format_if(&self, if_expr: IfExpression) -> String {
+        let condition_str = self.format_subexpr(if_expr.condition);
+        let consequence_str = self.format_subexpr(if_expr.consequence);
+
+        let mut result = format!("if {condition_str} {consequence_str}");
+
+        if let Some(alternative) = if_expr.alternative {
+            let alternative = if let Some(ExpressionKind::If(if_expr)) =
+                extract_simple_expr(alternative.clone()).map(|expr| expr.kind)
+            {
+                self.format_if(*if_expr)
+            } else {
+                self.format_expr(alternative, ExpressionType::Statement)
+            };
+
+            result.push_str(" else ");
+            result.push_str(&alternative);
+        };
+
+        result
+    }
+
+    fn format_if_single_line(&self, if_expr: IfExpression) -> Option<String> {
+        let condition_str = self.format_subexpr(if_expr.condition);
+        let consequence_str = self.format_subexpr(extract_simple_expr(if_expr.consequence)?);
+
+        let if_str = if let Some(alternative) = if_expr.alternative {
+            let alternative_str = if let Some(ExpressionKind::If(_)) =
+                extract_simple_expr(alternative.clone()).map(|expr| expr.kind)
+            {
+                return None;
+            } else {
+                self.format_expr(extract_simple_expr(alternative)?, ExpressionType::Statement)
+            };
+
+            format!("if {} {{ {} }} else {{ {} }}", condition_str, consequence_str, alternative_str)
+        } else {
+            format!("if {{{}}} {{{}}}", condition_str, consequence_str)
+        };
+
+        (if_str.len() <= self.config.single_line_if_else_max_width).then_some(if_str)
+    }
+
     fn format_struct_lit(
         &self,
         type_name: &str,
@@ -515,3 +579,15 @@ fn has_single_line_comment(slice: &str) -> bool {
 fn no_long_exprs(exprs: &[Expr], max_width: usize) -> bool {
     exprs.iter().all(|expr| expr.value.len() <= max_width)
 }
+
+fn extract_simple_expr(expr: Expression) -> Option<Expression> {
+    if let ExpressionKind::Block(mut block) = expr.kind {
+        if block.len() == 1 {
+            if let StatementKind::Expression(expr) = block.pop().unwrap() {
+                return expr.into();
+            }
+        }
+    }
+
+    None
+}

tooling/nargo_fmt/src/visitor/stmt.rs

@@ -1,18 +1,30 @@
+use std::iter::zip;
+
 use noirc_frontend::{Statement, StatementKind};
 
+use super::ExpressionType;
+
 impl super::FmtVisitor<'_> {
     pub(crate) fn visit_stmts(&mut self, stmts: Vec<Statement>) {
-        for Statement { kind, span } in stmts {
+        let len = stmts.len();
+
+        for (Statement { kind, span }, index) in zip(stmts, 1..) {
+            let is_last = index == len;
+
             match kind {
-                StatementKind::Expression(expr) => self.visit_expr(expr),
+                StatementKind::Expression(expr) => self.visit_expr(
+                    expr,
+                    if is_last { ExpressionType::SubExpression } else { ExpressionType::Statement },
+                ),
                 StatementKind::Semi(expr) => {
-                    self.visit_expr(expr);
+                    self.visit_expr(expr, ExpressionType::Statement);
                     self.push_str(";");
                 }
                 StatementKind::Let(let_stmt) => {
                     let let_str =
                         self.slice(span.start()..let_stmt.expression.span.start()).trim_end();
-                    let expr_str = self.format_expr(let_stmt.expression);
+                    let expr_str =
+                        self.format_expr(let_stmt.expression, ExpressionType::SubExpression);
 
                     self.push_rewrite(format!("{let_str} {expr_str};"), span);
                 }

tooling/nargo_fmt/tests/expected/expr.nr

@@ -94,3 +94,29 @@ fn parenthesized() {
 fn constructor() {
     Point { x: 5, y: 10 };
 }
+
+fn if_expr() {
+    if true {
+        println("Hello :D");
+    }
+}
+
+fn return_if_expr() {
+    if true { 42 } else { 40 + 2 }
+}
+
+fn return_if_expr() {
+    if true {
+        42
+    };
+
+    if true { 42 } else { 40 + 2 }
+}
+
+fn if_if() {
+    if cond {
+        some();
+    } else {
+        none();
+    }.bar().baz();
+}

tooling/nargo_fmt/tests/expected/if.nr

@@ -0,0 +1,40 @@
+fn main() {
+    if false {
+        ();
+        ();
+    }
+
+    if false // lone if comment
+    {
+        ();
+        ();
+    }
+
+    let a = if 0 > 1 { 0 } else { 0 };
+
+    if true {
+        ();
+    } else if false {
+        ();
+        ();
+    } else {
+        ();
+        ();
+        ();
+    }
+
+    if true // else-if-chain if comment
+    {
+        ();
+    }
+    else if false // else-if-chain else-if comment
+    {
+        ();
+        ();
+    } else // else-if-chain else comment
+    {
+        ();
+        ();
+        ();
+    }
+}

tooling/nargo_fmt/tests/expected/nested_if_else.nr

@@ -1,3 +1,9 @@
 fn nested_if_else() {
-    if false { 1 } else if false { 2 } else { 3 }
+    if false {
+        1
+    } else if false {
+        2
+    } else {
+        3
+    }
 }