chore: format infix expression (#3451)

tooling/nargo_fmt/src/lib.rs

@@ -20,6 +20,7 @@
 /// in both placement and content during the formatting process.
 mod config;
 pub mod errors;
+mod rewrite;
 mod utils;
 mod visitor;
 

tooling/nargo_fmt/src/rewrite.rs

@@ -0,0 +1,3 @@
+mod infix;
+
+pub(crate) use infix::rewrite as infix;

tooling/nargo_fmt/src/rewrite/infix.rs

@@ -0,0 +1,112 @@
+use std::iter::zip;
+
+use noirc_frontend::{Expression, ExpressionKind};
+
+use crate::{
+    utils::{first_line_width, is_single_line},
+    visitor::{ExpressionType, FmtVisitor, Shape},
+};
+
+pub(crate) fn rewrite(visitor: FmtVisitor, expr: Expression, shape: Shape) -> String {
+    match flatten(visitor.fork(), &expr) {
+        Some((exprs, separators)) => rewrite_single_line(shape, &exprs, &separators)
+            .unwrap_or_else(|| rewrite_multiline(visitor, &exprs, &separators)),
+        None => {
+            let ExpressionKind::Infix(infix) = expr.kind else { unreachable!() };
+
+            format!(
+                "{} {} {}",
+                visitor.format_sub_expr(infix.lhs),
+                infix.operator.contents.as_string(),
+                visitor.format_sub_expr(infix.rhs)
+            )
+        }
+    }
+}
+
+fn rewrite_single_line(shape: Shape, exprs: &[String], separators: &[String]) -> Option<String> {
+    let mut result = String::new();
+
+    for (rewrite, separator) in zip(exprs, separators) {
+        if !is_single_line(rewrite) || result.len() > shape.width {
+            return None;
+        }
+
+        result.push_str(rewrite);
+        result.push(' ');
+        result.push_str(separator);
+        result.push(' ');
+    }
+
+    let last = exprs.last().unwrap();
+    result.push_str(last);
+
+    if first_line_width(&result) > shape.width {
+        return None;
+    }
+
+    result.into()
+}
+
+fn rewrite_multiline(visitor: FmtVisitor, exprs: &[String], separators: &[String]) -> String {
+    let mut visitor = visitor.fork();
+    visitor.indent.block_indent(visitor.config);
+    let indent_str = visitor.indent.to_string_with_newline();
+
+    let mut result = exprs[0].clone();
+
+    for (rewrite, separator) in exprs[1..].iter().zip(separators.iter()) {
+        result.push_str(&indent_str);
+        result.push_str(separator);
+        result.push(' ');
+        result.push_str(rewrite);
+    }
+
+    result
+}
+
+pub(crate) fn flatten(
+    mut visitor: FmtVisitor,
+    mut node: &Expression,
+) -> Option<(Vec<String>, Vec<String>)> {
+    let top_operator = match node.kind {
+        ExpressionKind::Infix(ref infix) => infix.operator.contents,
+        _ => return None,
+    };
+
+    let mut result = Vec::new();
+
+    let mut stack: Vec<&Expression> = Vec::new();
+    let mut separators = Vec::new();
+
+    loop {
+        match &node.kind {
+            ExpressionKind::Infix(infix) if top_operator == infix.operator.contents => {
+                stack.push(node);
+                node = &infix.lhs;
+            }
+            _ => {
+                let rewrite = if result.is_empty() {
+                    visitor.format_expr(node.clone(), ExpressionType::SubExpression)
+                } else {
+                    visitor.indent.block_indent(visitor.config);
+                    visitor.format_expr(node.clone(), ExpressionType::SubExpression)
+                };
+
+                result.push(rewrite);
+
+                let Some(pop) = stack.pop() else { break; };
+
+                match &pop.kind {
+                    ExpressionKind::Infix(infix) => {
+                        separators.push(infix.operator.contents.to_string());
+                        node = &infix.rhs;
+                    }
+                    _ => unreachable!(),
+                }
+            }
+        }
+    }
+
+    (result, separators).into()
+}

tooling/nargo_fmt/src/utils.rs

@@ -236,3 +236,11 @@ impl Item for (Ident, Expression) {
         }
     }
 }
+
+pub(crate) fn first_line_width(exprs: &str) -> usize {
+    exprs.lines().next().map_or(0, |line: &str| line.chars().count())
+}
+
+pub(crate) fn is_single_line(s: &str) -> bool {
+    !s.chars().any(|c| c == '\n')
+}

tooling/nargo_fmt/src/visitor.rs

@@ -10,10 +10,10 @@ use crate::{
 };
 
 pub(crate) struct FmtVisitor<'me> {
-    config: &'me Config,
+    pub(crate) config: &'me Config,
     buffer: String,
     pub(crate) source: &'me str,
-    indent: Indent,
+    pub(crate) indent: Indent,
     last_position: u32,
 }
 
@@ -245,37 +245,37 @@ impl<'me> FmtVisitor<'me> {
 }
 
 #[derive(Clone, Copy, Debug, Default)]
-struct Indent {
+pub(crate) struct Indent {
     block_indent: usize,
 }
 
 impl Indent {
-    fn width(&self) -> usize {
+    pub(crate) fn width(&self) -> usize {
         self.block_indent
     }
 
-    fn block_indent(&mut self, config: &Config) {
+    pub(crate) fn block_indent(&mut self, config: &Config) {
         self.block_indent += config.tab_spaces;
     }
 
-    fn block_unindent(&mut self, config: &Config) {
+    pub(crate) fn block_unindent(&mut self, config: &Config) {
         self.block_indent -= config.tab_spaces;
     }
 
-    fn to_string_with_newline(self) -> String {
+    pub(crate) fn to_string_with_newline(self) -> String {
         "\n".to_string() + &self.to_string()
     }
 
     #[allow(clippy::inherent_to_string)]
-    fn to_string(self) -> String {
+    pub(crate) fn to_string(self) -> String {
         " ".repeat(self.block_indent)
     }
 }
 
 #[derive(Clone, Copy, Debug)]
-struct Shape {
-    width: usize,
-    indent: Indent,
+pub(crate) struct Shape {
+    pub(crate) width: usize,
+    pub(crate) indent: Indent,
 }
 
 #[derive(PartialEq, Eq, Debug)]

tooling/nargo_fmt/src/visitor/expr.rs

@@ -6,7 +6,8 @@ use noirc_frontend::{
 
 use super::{ExpressionType, FmtVisitor, Indent, Shape};
 use crate::{
-    utils::{self, Expr, FindToken, Item},
+    rewrite,
+    utils::{self, first_line_width, Expr, FindToken, Item},
     Config,
 };
 
@@ -52,13 +53,9 @@ impl FmtVisitor<'_> {
             ExpressionKind::Cast(cast) => {
                 format!("{} as {}", self.format_sub_expr(cast.lhs), cast.r#type)
             }
-            ExpressionKind::Infix(infix) => {
-                format!(
-                    "{} {} {}",
-                    self.format_sub_expr(infix.lhs),
-                    infix.operator.contents.as_string(),
-                    self.format_sub_expr(infix.rhs)
-                )
+            kind @ ExpressionKind::Infix(_) => {
+                let shape = self.shape();
+                rewrite::infix(self.fork(), Expression { kind, span }, shape)
             }
             ExpressionKind::Call(call_expr) => {
                 let args_span =
@@ -493,7 +490,7 @@ fn wrap_exprs(
     nested_shape: Shape,
     shape: Shape,
 ) -> String {
-    let first_line_width = exprs.lines().next().map_or(0, |line| line.chars().count());
+    let first_line_width = first_line_width(&exprs);
 
     if first_line_width <= shape.width {
         let allow_trailing_newline = exprs

tooling/nargo_fmt/tests/expected/expr.nr

@@ -101,9 +101,9 @@ fn parenthesized() {
 
 fn parenthesized() {
     value + (
-        // line
-        x as Field
-    )
+            // line
+            x as Field
+        )
 }
 
 fn constructor() {

tooling/nargo_fmt/tests/expected/infix.nr

@@ -7,3 +7,12 @@ fn foo() {
     40/*test*/ + 2 == 42;
     40 + 2/*test*/ == 42;
 }
+
+fn big() {
+    assert(bjj_affine.contains(bjj_affine.gen)
+        & bjj_affine.contains(p1_affine)
+        & bjj_affine.contains(p2_affine)
+        & bjj_affine.contains(p3_affine)
+        & bjj_affine.contains(p4_affine)
+        & bjj_affine.contains(p5_affine));
+}

tooling/nargo_fmt/tests/input/infix.nr

@@ -7,3 +7,7 @@ fn foo() {
 40/*test*/ + 2 == 42;
                 40 + 2/*test*/ == 42;
 }
+
+fn big() {
+    assert(bjj_affine.contains(bjj_affine.gen) & bjj_affine.contains(p1_affine) & bjj_affine.contains(p2_affine) & bjj_affine.contains(p3_affine) & bjj_affine.contains(p4_affine) & bjj_affine.contains(p5_affine));
+}