chore: fix wrapping issue for constraints

noir-lang · Nov 27, 2023 · b3b4cca · b3b4cca
1 parent 70ccb7e
commit b3b4cca
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 15 deletions.
diff --git a/tooling/nargo_fmt/src/visitor/stmt.rs b/tooling/nargo_fmt/src/visitor/stmt.rs
@@ -4,7 +4,7 @@ use noirc_frontend::{
     ConstrainKind, ConstrainStatement, ExpressionKind, ForRange, Statement, StatementKind,
 };
 
-use crate::rewrite;
+use crate::{rewrite, visitor::expr::wrap_exprs};
 
 use super::ExpressionType;
 
@@ -33,30 +33,44 @@ impl super::FmtVisitor<'_> {
                     self.push_rewrite(format!("{let_str} {expr_str};"), span);
                 }
                 StatementKind::Constrain(ConstrainStatement(expr, message, kind)) => {
+                    let mut nested_shape = self.shape();
+                    let shape = nested_shape;
+
+                    nested_shape.indent.block_indent(self.config);
+
                     let message =
                         message.map_or(String::new(), |message| format!(", \"{message}\""));
-                    let constrain = match kind {
+
+                    let (callee, args) = match kind {
                         ConstrainKind::Assert => {
-                            let assertion = rewrite::sub_expr(self, self.shape(), expr);
+                            let assertion = rewrite::sub_expr(self, nested_shape, expr);
+                            let args = format!("{assertion}{message}");
 
-                            format!("assert({assertion}{message});")
+                            ("assert", args)
                         }
                         ConstrainKind::AssertEq => {
                             if let ExpressionKind::Infix(infix) = expr.kind {
-                                let lhs = rewrite::sub_expr(self, self.shape(), infix.lhs);
-                                let rhs = rewrite::sub_expr(self, self.shape(), infix.rhs);
+                                let lhs = rewrite::sub_expr(self, nested_shape, infix.lhs);
+                                let rhs = rewrite::sub_expr(self, nested_shape, infix.rhs);
 
-                                format!("assert_eq({lhs}, {rhs}{message});")
+                                let args = format!("{lhs}, {rhs}{message}");
+
+                                ("assert_eq", args)
                             } else {
                                 unreachable!()
                             }
                         }
                         ConstrainKind::Constrain => {
                             let expr = rewrite::sub_expr(self, self.shape(), expr);
-                            format!("constrain {expr};")
+                            let constrain = format!("constrain {expr};");
+                            self.push_rewrite(constrain, span);
+                            return;
                         }
                     };
 
+                    let args = wrap_exprs("(", ")", args, nested_shape, shape, true);
+                    let constrain = format!("{callee}{args};");
+
                     self.push_rewrite(constrain, span);
                 }
                 StatementKind::For(for_stmt) => {

diff --git a/tooling/nargo_fmt/tests/expected/call.nr b/tooling/nargo_fmt/tests/expected/call.nr
@@ -39,10 +39,12 @@ fn foo() {
 
     assert(x == y);
 
-    assert(p4_affine.eq(
-        Gaffine::new(
-            6890855772600357754907169075114257697580319025794532037257385534741338397365,
-            4338620300185947561074059802482547481416142213883829469920100239455078257889
+    assert(
+        p4_affine.eq(
+            Gaffine::new(
+                6890855772600357754907169075114257697580319025794532037257385534741338397365,
+                4338620300185947561074059802482547481416142213883829469920100239455078257889
+            )
         )
-    ));
+    );
 }
diff --git a/tooling/nargo_fmt/tests/expected/infix.nr b/tooling/nargo_fmt/tests/expected/infix.nr
@@ -9,10 +9,12 @@ fn foo() {
 }
 
 fn big() {
-    assert(bjj_affine.contains(bjj_affine.gen)
+    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));
+        & bjj_affine.contains(p5_affine)
+    );
 }