From 6c26bdb68e223cab6426f7bccaa605118541059e Mon Sep 17 00:00:00 2001 From: kek kek kek Date: Wed, 8 Nov 2023 07:50:29 -0800 Subject: [PATCH] chore: format infix expression (#3451) --- tooling/nargo_fmt/src/lib.rs | 1 + tooling/nargo_fmt/src/rewrite.rs | 3 + tooling/nargo_fmt/src/rewrite/infix.rs | 112 ++++++++++++++++++++++ tooling/nargo_fmt/src/utils.rs | 8 ++ tooling/nargo_fmt/src/visitor.rs | 22 ++--- tooling/nargo_fmt/src/visitor/expr.rs | 15 ++- tooling/nargo_fmt/tests/expected/expr.nr | 6 +- tooling/nargo_fmt/tests/expected/infix.nr | 9 ++ tooling/nargo_fmt/tests/input/infix.nr | 4 + 9 files changed, 157 insertions(+), 23 deletions(-) create mode 100644 tooling/nargo_fmt/src/rewrite.rs create mode 100644 tooling/nargo_fmt/src/rewrite/infix.rs diff --git a/tooling/nargo_fmt/src/lib.rs b/tooling/nargo_fmt/src/lib.rs index 39e978aa306..d731934c3c3 100644 --- a/tooling/nargo_fmt/src/lib.rs +++ b/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; diff --git a/tooling/nargo_fmt/src/rewrite.rs b/tooling/nargo_fmt/src/rewrite.rs new file mode 100644 index 00000000000..15ef3a1af33 --- /dev/null +++ b/tooling/nargo_fmt/src/rewrite.rs @@ -0,0 +1,3 @@ +mod infix; + +pub(crate) use infix::rewrite as infix; diff --git a/tooling/nargo_fmt/src/rewrite/infix.rs b/tooling/nargo_fmt/src/rewrite/infix.rs new file mode 100644 index 00000000000..0fbfa07a841 --- /dev/null +++ b/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 { + 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, Vec)> { + 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() +} diff --git a/tooling/nargo_fmt/src/utils.rs b/tooling/nargo_fmt/src/utils.rs index ca7afecd3cf..84f532678ff 100644 --- a/tooling/nargo_fmt/src/utils.rs +++ b/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') +} diff --git a/tooling/nargo_fmt/src/visitor.rs b/tooling/nargo_fmt/src/visitor.rs index 4f83b3e8e2c..2c520fbad73 100644 --- a/tooling/nargo_fmt/src/visitor.rs +++ b/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)] diff --git a/tooling/nargo_fmt/src/visitor/expr.rs b/tooling/nargo_fmt/src/visitor/expr.rs index 68606b8826d..d8e0c6d9d31 100644 --- a/tooling/nargo_fmt/src/visitor/expr.rs +++ b/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 diff --git a/tooling/nargo_fmt/tests/expected/expr.nr b/tooling/nargo_fmt/tests/expected/expr.nr index d7b3be74fe2..03a26835ee3 100644 --- a/tooling/nargo_fmt/tests/expected/expr.nr +++ b/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() { diff --git a/tooling/nargo_fmt/tests/expected/infix.nr b/tooling/nargo_fmt/tests/expected/infix.nr index 0688781ba48..cbc73045fe3 100644 --- a/tooling/nargo_fmt/tests/expected/infix.nr +++ b/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)); +} diff --git a/tooling/nargo_fmt/tests/input/infix.nr b/tooling/nargo_fmt/tests/input/infix.nr index b16ccd02982..059e58c6b64 100644 --- a/tooling/nargo_fmt/tests/input/infix.nr +++ b/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)); +}