diff --git a/tooling/nargo_fmt/src/visitor/stmt.rs b/tooling/nargo_fmt/src/visitor/stmt.rs index c27b7911d03..eabdb1a150a 100644 --- 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 index 7824ba37089..de78d7c4edb 100644 --- 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 index cbc73045fe3..228dfdf68c4 100644 --- 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) + ); }