chore: type formatting (#3618)

* chore: type formatting

* chore: clippy fix

* chore: use map_with_span instead of spanned

Co-authored-by: jfecher <jfecher11@gmail.com>

* chore: fix compile error

* apply feedback

* chore: apply apply feedback

* chore: fix tests

---------

Co-authored-by: jfecher <jfecher11@gmail.com>

compiler/noirc_errors/src/position.rs

@@ -65,6 +65,10 @@ impl Span {
         Span::inclusive(start, start)
     }
 
+    pub fn empty(position: u32) -> Span {
+        Span::from(position..position)
+    }
+
     #[must_use]
     pub fn merge(self, other: Span) -> Span {
         Span(self.0.merge(other.0))

compiler/noirc_frontend/src/ast/mod.rs

@@ -41,6 +41,8 @@ pub enum UnresolvedTypeData {
     FormatString(UnresolvedTypeExpression, Box<UnresolvedType>),
     Unit,
 
+    Parenthesized(Box<UnresolvedType>),
+
     /// A Named UnresolvedType can be a struct type or a type variable
     Named(Path, Vec<UnresolvedType>),
 
@@ -152,6 +154,7 @@ impl std::fmt::Display for UnresolvedTypeData {
             Unit => write!(f, "()"),
             Error => write!(f, "error"),
             Unspecified => write!(f, "unspecified"),
+            Parenthesized(typ) => write!(f, "({typ})"),
         }
     }
 }

compiler/noirc_frontend/src/ast/statement.rs

@@ -198,7 +198,11 @@ impl From<Ident> for Expression {
     fn from(i: Ident) -> Expression {
         Expression {
             span: i.0.span(),
-            kind: ExpressionKind::Variable(Path { segments: vec![i], kind: PathKind::Plain }),
+            kind: ExpressionKind::Variable(Path {
+                span: i.span(),
+                segments: vec![i],
+                kind: PathKind::Plain,
+            }),
         }
     }
 }
@@ -311,6 +315,7 @@ impl UseTree {
 pub struct Path {
     pub segments: Vec<Ident>,
     pub kind: PathKind,
+    pub span: Span,
 }
 
 impl Path {
@@ -330,18 +335,11 @@ impl Path {
     }
 
     pub fn from_ident(name: Ident) -> Path {
-        Path { segments: vec![name], kind: PathKind::Plain }
+        Path { span: name.span(), segments: vec![name], kind: PathKind::Plain }
     }
 
     pub fn span(&self) -> Span {
-        let mut segments = self.segments.iter();
-        let first_segment = segments.next().expect("ice : cannot have an empty path");
-        let mut span = first_segment.0.span();
-
-        for segment in segments {
-            span = span.merge(segment.0.span());
-        }
-        span
+        self.span
     }
 
     pub fn last_segment(&self) -> Ident {
@@ -545,8 +543,11 @@ impl ForRange {
 
                 // array.len()
                 let segments = vec![array_ident];
-                let array_ident =
-                    ExpressionKind::Variable(Path { segments, kind: PathKind::Plain });
+                let array_ident = ExpressionKind::Variable(Path {
+                    segments,
+                    kind: PathKind::Plain,
+                    span: array_span,
+                });
 
                 let end_range = ExpressionKind::MethodCall(Box::new(MethodCallExpression {
                     object: Expression::new(array_ident.clone(), array_span),
@@ -561,8 +562,11 @@ impl ForRange {
 
                 // array[i]
                 let segments = vec![Ident::new(index_name, array_span)];
-                let index_ident =
-                    ExpressionKind::Variable(Path { segments, kind: PathKind::Plain });
+                let index_ident = ExpressionKind::Variable(Path {
+                    segments,
+                    kind: PathKind::Plain,
+                    span: array_span,
+                });
 
                 let loop_element = ExpressionKind::Index(Box::new(IndexExpression {
                     collection: Expression::new(array_ident, array_span),

compiler/noirc_frontend/src/hir/resolution/import.rs

@@ -1,5 +1,5 @@
 use iter_extended::partition_results;
-use noirc_errors::CustomDiagnostic;
+use noirc_errors::{CustomDiagnostic, Span};
 
 use crate::graph::CrateId;
 use std::collections::BTreeMap;
@@ -202,7 +202,11 @@ fn resolve_external_dep(
     // Create an import directive for the dependency crate
     let path_without_crate_name = &path[1..]; // XXX: This will panic if the path is of the form `use dep::std` Ideal algorithm will not distinguish between crate and module
 
-    let path = Path { segments: path_without_crate_name.to_vec(), kind: PathKind::Plain };
+    let path = Path {
+        segments: path_without_crate_name.to_vec(),
+        kind: PathKind::Plain,
+        span: Span::default(),
+    };
     let dep_directive =
         ImportDirective { module_id: dep_module.local_id, path, alias: directive.alias.clone() };
 

compiler/noirc_frontend/src/hir/resolution/resolver.rs

@@ -441,6 +441,7 @@ impl<'a> Resolver<'a> {
             MutableReference(element) => {
                 Type::MutableReference(Box::new(self.resolve_type_inner(*element, new_variables)))
             }
+            Parenthesized(typ) => self.resolve_type_inner(*typ, new_variables),
         }
     }
 
@@ -1787,6 +1788,7 @@ impl<'a> Resolver<'a> {
                     self.verify_type_valid_for_program_input(element);
                 }
             }
+            UnresolvedTypeData::Parenthesized(typ) => self.verify_type_valid_for_program_input(typ),
         }
     }
 

compiler/noirc_frontend/src/parser/parser.rs

@@ -726,21 +726,21 @@ fn token_kind(token_kind: TokenKind) -> impl NoirParser<Token> {
 
 fn path() -> impl NoirParser<Path> {
     let idents = || ident().separated_by(just(Token::DoubleColon)).at_least(1);
-    let make_path = |kind| move |segments| Path { segments, kind };
+    let make_path = |kind| move |segments, span| Path { segments, kind, span };
 
     let prefix = |key| keyword(key).ignore_then(just(Token::DoubleColon));
-    let path_kind = |key, kind| prefix(key).ignore_then(idents()).map(make_path(kind));
+    let path_kind = |key, kind| prefix(key).ignore_then(idents()).map_with_span(make_path(kind));
 
     choice((
         path_kind(Keyword::Crate, PathKind::Crate),
         path_kind(Keyword::Dep, PathKind::Dep),
-        idents().map(make_path(PathKind::Plain)),
+        idents().map_with_span(make_path(PathKind::Plain)),
     ))
 }
 
 fn empty_path() -> impl NoirParser<Path> {
-    let make_path = |kind| move |_| Path { segments: Vec::new(), kind };
-    let path_kind = |key, kind| keyword(key).map(make_path(kind));
+    let make_path = |kind| move |_, span| Path { segments: Vec::new(), kind, span };
+    let path_kind = |key, kind| keyword(key).map_with_span(make_path(kind));
 
     choice((path_kind(Keyword::Crate, PathKind::Crate), path_kind(Keyword::Dep, PathKind::Dep)))
 }
@@ -1015,13 +1015,24 @@ fn parse_type_inner(
         named_type(recursive_type_parser.clone()),
         named_trait(recursive_type_parser.clone()),
         array_type(recursive_type_parser.clone()),
-        recursive_type_parser.clone().delimited_by(just(Token::LeftParen), just(Token::RightParen)),
+        parenthesized_type(recursive_type_parser.clone()),
         tuple_type(recursive_type_parser.clone()),
         function_type(recursive_type_parser.clone()),
         mutable_reference_type(recursive_type_parser),
     ))
 }
 
+fn parenthesized_type(
+    recursive_type_parser: impl NoirParser<UnresolvedType>,
+) -> impl NoirParser<UnresolvedType> {
+    recursive_type_parser
+        .delimited_by(just(Token::LeftParen), just(Token::RightParen))
+        .map_with_span(|typ, span| UnresolvedType {
+            typ: UnresolvedTypeData::Parenthesized(Box::new(typ)),
+            span: span.into(),
+        })
+}
+
 fn optional_visibility() -> impl NoirParser<Visibility> {
     keyword(Keyword::Pub).or_not().map(|opt| match opt {
         Some(_) => Visibility::Public,
@@ -1177,7 +1188,9 @@ where
         .ignore_then(type_parser.clone())
         .then_ignore(just(Token::RightBracket))
         .or_not()
-        .map_with_span(|t, span| t.unwrap_or_else(|| UnresolvedTypeData::Unit.with_span(span)));
+        .map_with_span(|t, span| {
+            t.unwrap_or_else(|| UnresolvedTypeData::Unit.with_span(Span::empty(span.end())))
+        });
 
     keyword(Keyword::Fn)
         .ignore_then(env)

tooling/nargo_fmt/build.rs

@@ -46,8 +46,8 @@ fn generate_formatter_tests(test_file: &mut File, test_data_dir: &Path) {
             .collect::<Vec<_>>()
             .join("\n");
 
-        let output_source_path = outputs_dir.join(file_name);
-        let output_source = std::fs::read_to_string(output_source_path).unwrap();
+        let output_source_path = outputs_dir.join(file_name).display().to_string();
+        let output_source = std::fs::read_to_string(output_source_path.clone()).unwrap();
 
         write!(
             test_file,
@@ -63,6 +63,7 @@ fn format_{test_name}() {{
     let config = nargo_fmt::Config::of("{config}").unwrap();
     let fmt_text = nargo_fmt::format(&input, parsed_module, &config);
 
+    std::fs::write("{output_source_path}", fmt_text.clone());
 
     similar_asserts::assert_eq!(fmt_text, expected_output);
 }}

tooling/nargo_fmt/src/rewrite.rs

@@ -2,8 +2,10 @@ mod array;
 mod expr;
 mod infix;
 mod parenthesized;
+mod typ;
 
 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;
+pub(crate) use typ::rewrite as typ;

tooling/nargo_fmt/src/rewrite/typ.rs

@@ -0,0 +1,70 @@
+use noirc_frontend::{UnresolvedType, UnresolvedTypeData};
+
+use crate::{
+    utils::span_is_empty,
+    visitor::{FmtVisitor, Shape},
+};
+
+pub(crate) fn rewrite(visitor: &FmtVisitor, _shape: Shape, typ: UnresolvedType) -> String {
+    match typ.typ {
+        UnresolvedTypeData::Array(length, element) => {
+            let typ = rewrite(visitor, _shape, *element);
+            if let Some(length) = length {
+                let length = visitor.slice(length.span());
+                format!("[{typ}; {length}]")
+            } else {
+                format!("[{typ}]")
+            }
+        }
+        UnresolvedTypeData::Parenthesized(typ) => {
+            let typ = rewrite(visitor, _shape, *typ);
+            format!("({typ})")
+        }
+        UnresolvedTypeData::MutableReference(typ) => {
+            let typ = rewrite(visitor, _shape, *typ);
+            format!("&mut {typ}")
+        }
+        UnresolvedTypeData::Tuple(mut types) => {
+            if types.len() == 1 {
+                let typ = types.pop().unwrap();
+                let typ = rewrite(visitor, _shape, typ);
+
+                format!("({typ},)")
+            } else {
+                let types: Vec<_> =
+                    types.into_iter().map(|typ| rewrite(visitor, _shape, typ)).collect();
+                let types = types.join(", ");
+                format!("({types})")
+            }
+        }
+        UnresolvedTypeData::Function(args, return_type, env) => {
+            let env = if span_is_empty(env.span.unwrap()) {
+                "".into()
+            } else {
+                let ty = rewrite(visitor, _shape, *env);
+                format!("[{ty}]")
+            };
+
+            let args = args
+                .into_iter()
+                .map(|arg| rewrite(visitor, _shape, arg))
+                .collect::<Vec<_>>()
+                .join(", ");
+
+            let return_type = rewrite(visitor, _shape, *return_type);
+
+            format!("fn{env}({args}) -> {return_type}")
+        }
+        UnresolvedTypeData::Unspecified => todo!(),
+        UnresolvedTypeData::FieldElement
+        | UnresolvedTypeData::Integer(_, _)
+        | UnresolvedTypeData::Bool
+        | UnresolvedTypeData::Named(_, _)
+        | UnresolvedTypeData::Unit
+        | UnresolvedTypeData::Expression(_)
+        | UnresolvedTypeData::String(_)
+        | UnresolvedTypeData::FormatString(_, _)
+        | UnresolvedTypeData::TraitAsType(_, _) => visitor.slice(typ.span.unwrap()).into(),
+        UnresolvedTypeData::Error => unreachable!(),
+    }
+}

tooling/nargo_fmt/src/utils.rs

@@ -250,13 +250,13 @@ impl Item for Param {
         self.span
     }
 
-    fn format(self, visitor: &FmtVisitor, _shape: Shape) -> String {
+    fn format(self, visitor: &FmtVisitor, shape: Shape) -> String {
         let visibility = match self.visibility {
             Visibility::Public => "pub ",
             Visibility::Private => "",
         };
         let pattern = visitor.slice(self.pattern.span());
-        let ty = visitor.slice(self.typ.span.unwrap());
+        let ty = rewrite::typ(visitor, shape, self.typ);
 
         format!("{pattern}: {visibility}{ty}")
     }
@@ -295,3 +295,7 @@ pub(crate) fn last_line_used_width(s: &str, offset: usize) -> usize {
         offset + s.chars().count()
     }
 }
+
+pub(crate) fn span_is_empty(span: Span) -> bool {
+    span.start() == span.end()
+}

tooling/nargo_fmt/src/visitor/item.rs

@@ -6,6 +6,7 @@ use noirc_frontend::{
 };
 
 use crate::{
+    rewrite,
     utils::{last_line_contains_single_line_comment, last_line_used_width, FindToken},
     visitor::expr::{format_seq, NewlineMode},
 };
@@ -122,7 +123,8 @@ impl super::FmtVisitor<'_> {
                 result.push_str("pub ");
             }
 
-            result.push_str(self.slice(span));
+            let typ = rewrite::typ(self, self.shape(), func.return_type());
+            result.push_str(&typ);
 
             let slice = self.slice(span.end()..func_span.start());
             if !slice.trim().is_empty() {

tooling/nargo_fmt/tests/expected/fn.nr

@@ -36,7 +36,21 @@ fn apply_binary_field_op<N>(
     registers: &mut Registers<N>
 ) -> bool {}
 
-fn main() -> distinct pub [Field;2] {}
+fn main() -> distinct pub [Field; 2] {}
+
+fn ret_normal_lambda1() -> ((fn() -> Field)) {}
+
+fn ret_normal_lambda1() -> fn() -> Field {}
+
+fn ret_closure1() -> fn[(Field,)]() -> Field {}
+
+fn ret_closure2() -> fn[(Field, Field)]() -> Field {}
+
+fn ret_closure3() -> fn[(u32, u64)]() -> u64 {}
+
+fn make_counter() -> fn[(&mut Field,)]() -> Field {}
+
+fn get_some<Env>(generator: fn[Env]() -> Field) -> [Field; 5] {}
 
 fn main(
     message: [u8; 10],
@@ -45,3 +59,5 @@ fn main(
     pub_key_y: Field,
     signature: [u8; 64]
 ) {}
+
+pub fn from_baz(x: [Field; crate::foo::MAGIC_NUMBER]) {}

tooling/nargo_fmt/tests/input/fn.nr

@@ -25,6 +25,22 @@ fn apply_binary_field_op<N>(lhs: RegisterIndex, rhs: RegisterIndex, result: Regi
 
 fn main() -> distinct pub [Field;2] {}
 
+fn ret_normal_lambda1() -> ((fn() -> Field)) {}
+
+fn ret_normal_lambda1() -> fn() -> Field {}
+
+fn ret_closure1() -> fn[(Field,)]() -> Field {}
+
+fn ret_closure2() -> fn[(Field,Field)]() -> Field {}
+
+fn ret_closure3() -> fn[(u32,u64)]() -> u64 {}
+
+fn make_counter() -> fn[(&mut Field,)]() -> Field {}
+
+fn get_some<Env>(generator: fn[Env]() -> Field) -> [Field;5] {}
+
 fn main(
     message: [u8; 10], message_field: Field, pub_key_x: Field, pub_key_y: Field, signature: [u8; 64]
 ) {}
+
+pub fn from_baz(x: [Field; crate::foo::MAGIC_NUMBER]) {}