chore(fmt): refactor the way we handle shapes in the formatter

tooling/nargo_fmt/src/rewrite.rs

@@ -1,5 +1,9 @@
 mod array;
+mod expr;
 mod infix;
+mod parenthesized;
 
 pub(crate) use array::rewrite as array;
+pub(crate) use expr::{rewrite as expr, rewrite_sub_expr as sub_expr};
 pub(crate) use infix::rewrite as infix;
+pub(crate) use parenthesized::rewrite as parenthesized;

tooling/nargo_fmt/src/rewrite/array.rs

@@ -26,7 +26,7 @@ pub(crate) fn rewrite(mut visitor: FmtVisitor, array: Vec<Expression>, array_spa
         let end = item_span.start();
 
         let leading = visitor.slice(start..end).trim_matches(pattern);
-        let item = visitor.format_sub_expr(item);
+        let item = super::sub_expr(&visitor, visitor.shape(), item);
         let next_start = items.peek().map_or(end_position, |expr| expr.span.start());
         let trailing = visitor.slice(item_span.end()..next_start);
         let offset = trailing

tooling/nargo_fmt/src/rewrite/expr.rs

@@ -0,0 +1,144 @@
+use noirc_frontend::{token::Token, ArrayLiteral, Expression, ExpressionKind, Literal, UnaryOp};
+
+use crate::visitor::{
+    expr::{format_brackets, format_parens},
+    ExpressionType, FmtVisitor, Indent, Shape,
+};
+
+pub(crate) fn rewrite_sub_expr(
+    visitor: &FmtVisitor,
+    shape: Shape,
+    expression: Expression,
+) -> String {
+    rewrite(visitor, expression, ExpressionType::SubExpression, shape)
+}
+
+pub(crate) fn rewrite(
+    visitor: &FmtVisitor,
+    Expression { kind, span }: Expression,
+    expr_type: ExpressionType,
+    shape: Shape,
+) -> String {
+    match kind {
+        ExpressionKind::Block(block) => {
+            let mut visitor = visitor.fork();
+            visitor.visit_block(block, span);
+            visitor.finish()
+        }
+        ExpressionKind::Prefix(prefix) => {
+            let op = match prefix.operator {
+                UnaryOp::Minus => "-",
+                UnaryOp::Not => "!",
+                UnaryOp::MutableReference => "&mut ",
+                UnaryOp::Dereference { implicitly_added } => {
+                    if implicitly_added {
+                        ""
+                    } else {
+                        "*"
+                    }
+                }
+            };
+
+            format!("{op}{}", rewrite_sub_expr(visitor, shape, prefix.rhs))
+        }
+        ExpressionKind::Cast(cast) => {
+            format!("{} as {}", rewrite_sub_expr(visitor, shape, cast.lhs), cast.r#type)
+        }
+        kind @ ExpressionKind::Infix(_) => {
+            super::infix(visitor.fork(), Expression { kind, span }, shape)
+        }
+        ExpressionKind::Call(call_expr) => {
+            let args_span =
+                visitor.span_before(call_expr.func.span.end()..span.end(), Token::LeftParen);
+
+            let callee = rewrite_sub_expr(visitor, shape, *call_expr.func);
+            let args = format_parens(
+                visitor.config.fn_call_width.into(),
+                visitor.fork(),
+                shape,
+                false,
+                call_expr.arguments,
+                args_span,
+                true,
+            );
+
+            format!("{callee}{args}")
+        }
+        ExpressionKind::MethodCall(method_call_expr) => {
+            let args_span = visitor.span_before(
+                method_call_expr.method_name.span().end()..span.end(),
+                Token::LeftParen,
+            );
+
+            let object = rewrite_sub_expr(visitor, shape, method_call_expr.object);
+            let method = method_call_expr.method_name.to_string();
+            let args = format_parens(
+                visitor.config.fn_call_width.into(),
+                visitor.fork(),
+                shape,
+                false,
+                method_call_expr.arguments,
+                args_span,
+                true,
+            );
+
+            format!("{object}.{method}{args}")
+        }
+        ExpressionKind::MemberAccess(member_access_expr) => {
+            let lhs_str = rewrite_sub_expr(visitor, shape, member_access_expr.lhs);
+            format!("{}.{}", lhs_str, member_access_expr.rhs)
+        }
+        ExpressionKind::Index(index_expr) => {
+            let index_span = visitor
+                .span_before(index_expr.collection.span.end()..span.end(), Token::LeftBracket);
+
+            let collection = rewrite_sub_expr(visitor, shape, index_expr.collection);
+            let index = format_brackets(visitor.fork(), false, vec![index_expr.index], index_span);
+
+            format!("{collection}{index}")
+        }
+        ExpressionKind::Tuple(exprs) => {
+            format_parens(None, visitor.fork(), shape, exprs.len() == 1, exprs, span, false)
+        }
+        ExpressionKind::Literal(literal) => match literal {
+            Literal::Integer(_) | Literal::Bool(_) | Literal::Str(_) | Literal::FmtStr(_) => {
+                visitor.slice(span).to_string()
+            }
+            Literal::Array(ArrayLiteral::Repeated { repeated_element, length }) => {
+                let repeated = rewrite_sub_expr(visitor, shape, *repeated_element);
+                let length = rewrite_sub_expr(visitor, shape, *length);
+
+                format!("[{repeated}; {length}]")
+            }
+            Literal::Array(ArrayLiteral::Standard(exprs)) => {
+                super::array(visitor.fork(), exprs, span)
+            }
+            Literal::Unit => "()".to_string(),
+        },
+        ExpressionKind::Parenthesized(sub_expr) => {
+            super::parenthesized(visitor, shape, span, *sub_expr)
+        }
+        ExpressionKind::Constructor(constructor) => {
+            let type_name = visitor.slice(span.start()..constructor.type_name.span().end());
+            let fields_span = visitor
+                .span_before(constructor.type_name.span().end()..span.end(), Token::LeftBrace);
+
+            visitor.format_struct_lit(type_name, fields_span, *constructor)
+        }
+        ExpressionKind::If(if_expr) => {
+            let allow_single_line = expr_type == ExpressionType::SubExpression;
+
+            if allow_single_line {
+                let mut visitor = visitor.fork();
+                visitor.indent = Indent::default();
+                if let Some(line) = visitor.format_if_single_line(*if_expr.clone()) {
+                    return line;
+                }
+            }
+
+            visitor.format_if(*if_expr)
+        }
+        ExpressionKind::Lambda(_) | ExpressionKind::Variable(_) => visitor.slice(span).to_string(),
+        ExpressionKind::Error => unreachable!(),
+    }
+}

tooling/nargo_fmt/src/rewrite/infix.rs

@@ -3,8 +3,9 @@ use std::iter::zip;
 use noirc_frontend::{Expression, ExpressionKind};
 
 use crate::{
+    rewrite,
     utils::{first_line_width, is_single_line},
-    visitor::{ExpressionType, FmtVisitor, Shape},
+    visitor::{FmtVisitor, Shape},
 };
 
 pub(crate) fn rewrite(visitor: FmtVisitor, expr: Expression, shape: Shape) -> String {
@@ -16,9 +17,9 @@ pub(crate) fn rewrite(visitor: FmtVisitor, expr: Expression, shape: Shape) -> St
 
             format!(
                 "{} {} {}",
-                visitor.format_sub_expr(infix.lhs),
+                rewrite::sub_expr(&visitor, shape, infix.lhs),
                 infix.operator.contents.as_string(),
-                visitor.format_sub_expr(infix.rhs)
+                rewrite::sub_expr(&visitor, shape, infix.rhs)
             )
         }
     }
@@ -87,10 +88,10 @@ pub(crate) fn flatten(
             }
             _ => {
                 let rewrite = if result.is_empty() {
-                    visitor.format_expr(node.clone(), ExpressionType::SubExpression)
+                    rewrite::sub_expr(&visitor, visitor.shape(), node.clone())
                 } else {
                     visitor.indent.block_indent(visitor.config);
-                    visitor.format_expr(node.clone(), ExpressionType::SubExpression)
+                    rewrite::sub_expr(&visitor, visitor.shape(), node.clone())
                 };
 
                 result.push(rewrite);

tooling/nargo_fmt/src/rewrite/parenthesized.rs

@@ -0,0 +1,67 @@
+use noirc_frontend::{hir::resolution::errors::Span, Expression, ExpressionKind};
+
+use crate::visitor::{FmtVisitor, Shape};
+
+pub(crate) fn rewrite(
+    visitor: &FmtVisitor<'_>,
+    shape: Shape,
+    mut span: Span,
+    mut sub_expr: Expression,
+) -> String {
+    let remove_nested_parens = visitor.config.remove_nested_parens;
+
+    let mut leading;
+    let mut trailing;
+
+    loop {
+        let leading_span = span.start() + 1..sub_expr.span.start();
+        let trailing_span = sub_expr.span.end()..span.end() - 1;
+
+        leading = visitor.format_comment(leading_span.into());
+        trailing = visitor.format_comment(trailing_span.into());
+
+        if let ExpressionKind::Parenthesized(ref sub_sub_expr) = sub_expr.kind {
+            if remove_nested_parens && leading.is_empty() && trailing.is_empty() {
+                span = sub_expr.span;
+                sub_expr = *sub_sub_expr.clone();
+                continue;
+            }
+        }
+
+        break;
+    }
+
+    if !leading.contains("//") && !trailing.contains("//") {
+        let sub_expr = super::sub_expr(visitor, shape, sub_expr);
+        format!("({leading}{sub_expr}{trailing})")
+    } else {
+        let mut visitor = visitor.fork();
+
+        let indent = visitor.indent.to_string_with_newline();
+        visitor.indent.block_indent(visitor.config);
+        let nested_indent = visitor.indent.to_string_with_newline();
+
+        let sub_expr = super::sub_expr(&visitor, shape, sub_expr);
+
+        let mut result = String::new();
+        result.push('(');
+
+        if !leading.is_empty() {
+            result.push_str(&nested_indent);
+            result.push_str(&leading);
+        }
+
+        result.push_str(&nested_indent);
+        result.push_str(&sub_expr);
+
+        if !trailing.is_empty() {
+            result.push_str(&nested_indent);
+            result.push_str(&trailing);
+        }
+
+        result.push_str(&indent);
+        result.push(')');
+
+        result
+    }
+}

tooling/nargo_fmt/src/utils.rs

@@ -1,4 +1,5 @@
-use crate::visitor::FmtVisitor;
+use crate::rewrite;
+use crate::visitor::{FmtVisitor, Shape};
 use noirc_frontend::hir::resolution::errors::Span;
 use noirc_frontend::lexer::Lexer;
 use noirc_frontend::token::Token;
@@ -40,15 +41,22 @@ impl Expr {
 
 pub(crate) struct Exprs<'me, T> {
     pub(crate) visitor: &'me FmtVisitor<'me>,
+    shape: Shape,
     pub(crate) elements: std::iter::Peekable<std::vec::IntoIter<T>>,
     pub(crate) last_position: u32,
     pub(crate) end_position: u32,
 }
 
 impl<'me, T: Item> Exprs<'me, T> {
-    pub(crate) fn new(visitor: &'me FmtVisitor<'me>, span: Span, elements: Vec<T>) -> Self {
+    pub(crate) fn new(
+        visitor: &'me FmtVisitor<'me>,
+        shape: Shape,
+        span: Span,
+        elements: Vec<T>,
+    ) -> Self {
         Self {
             visitor,
+            shape,
             last_position: span.start() + 1, /*(*/
             end_position: span.end() - 1,    /*)*/
             elements: elements.into_iter().peekable(),
@@ -70,7 +78,7 @@ impl<T: Item> Iterator for Exprs<'_, T> {
         let next_start = self.elements.peek().map_or(self.end_position, |expr| expr.start());
 
         let (leading, different_line) = self.leading(start, end);
-        let expr = element.format(self.visitor);
+        let expr = element.format(self.visitor, self.shape);
         let trailing = self.trailing(element_span.end(), next_start, is_last);
 
         Expr { leading, value: expr, trailing, different_line }.into()
@@ -196,7 +204,7 @@ pub(crate) fn count_newlines(slice: &str) -> usize {
 pub(crate) trait Item {
     fn span(&self) -> Span;
 
-    fn format(self, visitor: &FmtVisitor) -> String;
+    fn format(self, visitor: &FmtVisitor, shape: Shape) -> String;
 
     fn start(&self) -> u32 {
         self.span().start()
@@ -212,8 +220,8 @@ impl Item for Expression {
         self.span
     }
 
-    fn format(self, visitor: &FmtVisitor) -> String {
-        visitor.format_sub_expr(self)
+    fn format(self, visitor: &FmtVisitor, shape: Shape) -> String {
+        rewrite::sub_expr(visitor, shape, self)
     }
 }
 
@@ -223,11 +231,11 @@ impl Item for (Ident, Expression) {
         (name.span().start()..value.span.end()).into()
     }
 
-    fn format(self, visitor: &FmtVisitor) -> String {
+    fn format(self, visitor: &FmtVisitor, shape: Shape) -> String {
         let (name, expr) = self;
 
         let name = name.0.contents;
-        let expr = visitor.format_sub_expr(expr);
+        let expr = rewrite::sub_expr(visitor, shape, expr);
 
         if name == expr {
             name
@@ -242,7 +250,7 @@ impl Item for Param {
         self.span
     }
 
-    fn format(self, visitor: &FmtVisitor) -> String {
+    fn format(self, visitor: &FmtVisitor, _shape: Shape) -> String {
         let visibility = match self.visibility {
             Visibility::Public => "pub ",
             Visibility::Private => "",
@@ -259,7 +267,7 @@ impl Item for Ident {
         self.span()
     }
 
-    fn format(self, visitor: &FmtVisitor) -> String {
+    fn format(self, visitor: &FmtVisitor, _shape: Shape) -> String {
         visitor.slice(self.span()).into()
     }
 }
@@ -268,10 +276,22 @@ pub(crate) fn first_line_width(exprs: &str) -> usize {
     exprs.lines().next().map_or(0, |line: &str| line.chars().count())
 }
 
+pub(crate) fn last_line_width(s: &str) -> usize {
+    s.rsplit('\n').next().unwrap_or("").chars().count()
+}
+
 pub(crate) fn is_single_line(s: &str) -> bool {
     !s.chars().any(|c| c == '\n')
 }
 
 pub(crate) fn last_line_contains_single_line_comment(s: &str) -> bool {
     s.lines().last().map_or(false, |line| line.contains("//"))
 }
+
+pub(crate) fn last_line_used_width(s: &str, offset: usize) -> usize {
+    if s.contains('\n') {
+        last_line_width(s)
+    } else {
+        offset + s.chars().count()
+    }
+}