feat: Syntax for environment types now works with generics

crates/nargo_cli/tests/compile_success_empty/closure_explicit_types/src/main.nr

@@ -3,49 +3,60 @@ fn ret_normal_lambda1() -> fn() -> Field {
     || 10
 }
 
-// explicitly specified empty capture group
-fn ret_normal_lambda2() -> fn[]() -> Field {
-    || 20
-}
-
 // return lamda that captures a thing
-fn ret_closure1() -> fn[Field]() -> Field {
+fn ret_closure1() -> fn[(Field,)]() -> Field {
     let x = 20;
     || x + 10
 }
 
 // return lamda that captures two things
-fn ret_closure2() -> fn[Field,Field]() -> Field {
+fn ret_closure2() -> fn[(Field,Field)]() -> Field {
     let x = 20;
     let y = 10;
     || x + y + 10
 }
 
 // return lamda that captures two things with different types
-fn ret_closure3() -> fn[u32,u64]() -> u64 {
+fn ret_closure3() -> fn[(u32,u64)]() -> u64 {
     let x: u32 = 20;
     let y: u64 = 10;
     || x as u64 + y + 10
 }
 
 // accepts closure that has 1 thing in its env, calls it and returns the result
-fn accepts_closure1(f: fn[Field]() -> Field) -> Field {
+fn accepts_closure1(f: fn[(Field,)]() -> Field) -> Field {
     f()
 }
 
 // accepts closure that has 1 thing in its env and returns it
-fn accepts_closure2(f: fn[Field]() -> Field) -> fn[Field]() -> Field {
+fn accepts_closure2(f: fn[(Field,)]() -> Field) -> fn[(Field,)]() -> Field {
     f
 }
 
 // accepts closure with different types in the capture group
-fn accepts_closure3(f: fn[u32, u64]() -> u64) -> u64 {
+fn accepts_closure3(f: fn[(u32, u64)]() -> u64) -> u64 {
     f()
 }
 
+// generic over closure environments
+fn add_results<Env1, Env2>(f1: fn[Env1]() -> Field, f2: fn[Env2]() -> Field) -> Field {
+    f1() + f2()
+}
+
+// a *really* generic function
+fn map<T, U, N, Env>(arr: [T; N], f: fn[Env](T) -> U) -> [U; N] {
+    let first_elem = f(arr[0]);
+    let mut ret = [first_elem; N];
+
+    for i in 1 .. N {
+        ret[i] = f(arr[i]);
+    }
+
+    ret
+}
+
 fn main() {
     assert(ret_normal_lambda1()() == 10);
-    assert(ret_normal_lambda2()() == 20);
     assert(ret_closure1()() == 30);
     assert(ret_closure2()() == 40);
     assert(ret_closure3()() == 40);
@@ -57,4 +68,13 @@ fn main() {
     let y: u32 = 30;
     let z: u64 = 40;
     assert(accepts_closure3(|| y as u64 + z) == 70);
+
+    let w = 50;
+    assert(add_results(|| 100, || x ) == 150);
+    assert(add_results(|| x + 100, || w + x ) == 250);
+
+    let arr = [1,2,3,4];
+
+    assert(map(arr, |n| n + 1) == [2, 3, 4, 5]);
+    assert(map(arr, |n| n + x) == [51, 52, 53, 54]);
 }

crates/noirc_frontend/src/ast/function.rs

@@ -1,8 +1,10 @@
 use std::fmt::Display;
 
+use noirc_errors::Span;
+
 use crate::{token::Attribute, FunctionReturnType, Ident, Pattern, Visibility};
 
-use super::{FunctionDefinition, UnresolvedType};
+use super::{FunctionDefinition, UnresolvedType, UnresolvedTypeData};
 
 // A NoirFunction can be either a foreign low level function or a function definition
 // A closure / function definition will be stored under a name, so we do not differentiate between their variants
@@ -42,7 +44,9 @@ impl NoirFunction {
 
     pub fn return_type(&self) -> UnresolvedType {
         match &self.def.return_type {
-            FunctionReturnType::Default(_) => UnresolvedType::Unit,
+            FunctionReturnType::Default(_) => {
+                UnresolvedType::without_span(UnresolvedTypeData::Unit)
+            }
             FunctionReturnType::Ty(ty, _) => ty.clone(),
         }
     }
@@ -67,6 +71,9 @@ impl NoirFunction {
     pub fn number_of_statements(&self) -> usize {
         self.def.body.0.len()
     }
+    pub fn span(&self) -> Span {
+        self.def.span
+    }
 
     pub fn foreign(&self) -> Option<&FunctionDefinition> {
         match &self.kind {

crates/noirc_frontend/src/ast/mod.rs

@@ -31,7 +31,7 @@ use iter_extended::vecmap;
 /// require name resolution to resolve any type names used
 /// for structs within, but are otherwise identical to Types.
 #[derive(Debug, PartialEq, Eq, Clone, Hash)]
-pub enum UnresolvedType {
+pub enum UnresolvedTypeData {
     FieldElement,
     Array(Option<UnresolvedTypeExpression>, Box<UnresolvedType>), // [4]Witness = Array(4, Witness)
     Integer(Signedness, u32),                                     // u32 = Integer(unsigned, 32)
@@ -60,6 +60,16 @@ pub enum UnresolvedType {
     Error,
 }
 
+#[derive(Debug, PartialEq, Eq, Clone, Hash)]
+pub struct UnresolvedType {
+    pub typ: UnresolvedTypeData,
+
+    // The span is None in the cases where the User omitted a type:
+    //  fn Foo() {}  --- return type is UnresolvedType::Unit without a span
+    //  let x = 100; --- type is UnresolvedType::Unspecified without a span
+    pub span: Option<Span>,
+}
+
 /// The precursor to TypeExpression, this is the type that the parser allows
 /// to be used in the length position of an array type. Only constants, variables,
 /// and numeric binary operators are allowed here.
@@ -76,14 +86,14 @@ pub enum UnresolvedTypeExpression {
 }
 
 impl Recoverable for UnresolvedType {
-    fn error(_: Span) -> Self {
-        UnresolvedType::Error
+    fn error(span: Span) -> Self {
+        UnresolvedType { typ: UnresolvedTypeData::Error, span: Some(span) }
     }
 }
 
-impl std::fmt::Display for UnresolvedType {
+impl std::fmt::Display for UnresolvedTypeData {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        use UnresolvedType::*;
+        use UnresolvedTypeData::*;
         match self {
             FieldElement => write!(f, "Field"),
             Array(len, typ) => match len {
@@ -95,7 +105,7 @@ impl std::fmt::Display for UnresolvedType {
                 Signedness::Unsigned => write!(f, "u{num_bits}"),
             },
             Named(s, args) => {
-                let args = vecmap(args, ToString::to_string);
+                let args = vecmap(args, |arg| ToString::to_string(&arg.typ));
                 if args.is_empty() {
                     write!(f, "{s}")
                 } else {
@@ -116,12 +126,12 @@ impl std::fmt::Display for UnresolvedType {
             Function(args, ret, env) => {
                 let args = vecmap(args, ToString::to_string);
 
-                match &**env {
-                    UnresolvedType::Unit => {
+                match &env.as_ref().typ {
+                    UnresolvedTypeData::Unit => {
                         write!(f, "fn({}) -> {ret}", args.join(", "))
                     }
-                    UnresolvedType::Tuple(env_types) => {
-                        let env_types = vecmap(env_types, ToString::to_string);
+                    UnresolvedTypeData::Tuple(env_types) => {
+                        let env_types = vecmap(env_types, |arg| ToString::to_string(&arg.typ));
                         write!(f, "fn[{}]({}) -> {ret}", env_types.join(", "), args.join(", "))
                     }
                     _ => unreachable!(),
@@ -135,6 +145,12 @@ impl std::fmt::Display for UnresolvedType {
     }
 }
 
+impl std::fmt::Display for UnresolvedType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.typ.fmt(f)
+    }
+}
+
 impl std::fmt::Display for UnresolvedTypeExpression {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -148,13 +164,27 @@ impl std::fmt::Display for UnresolvedTypeExpression {
 }
 
 impl UnresolvedType {
-    pub fn from_int_token(token: IntType) -> UnresolvedType {
-        use {IntType::*, UnresolvedType::Integer};
+    pub fn without_span(typ: UnresolvedTypeData) -> UnresolvedType {
+        UnresolvedType { typ, span: None }
+    }
+
+    pub fn unspecified() -> UnresolvedType {
+        UnresolvedType { typ: UnresolvedTypeData::Unspecified, span: None }
+    }
+}
+
+impl UnresolvedTypeData {
+    pub fn from_int_token(token: IntType) -> UnresolvedTypeData {
+        use {IntType::*, UnresolvedTypeData::Integer};
         match token {
             Signed(num_bits) => Integer(Signedness::Signed, num_bits),
             Unsigned(num_bits) => Integer(Signedness::Unsigned, num_bits),
         }
     }
+
+    pub fn with_span(&self, span: Span) -> UnresolvedType {
+        UnresolvedType { typ: self.clone(), span: Some(span) }
+    }
 }
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]

crates/noirc_frontend/src/hir/resolution/errors.rs

@@ -74,6 +74,8 @@ pub enum ResolverError {
     ContractFunctionInternalInNormalFunction { span: Span },
     #[error("Numeric constants should be printed without formatting braces")]
     NumericConstantInFormatString { name: String, span: Span },
+    #[error("Closure environment must be a tuple or unit type")]
+    InvalidClosureEnvironment { typ: Type, span: Span },
 }
 
 impl ResolverError {
@@ -283,6 +285,9 @@ impl From<ResolverError> for Diagnostic {
                 "Numeric constants should be printed without formatting braces".to_string(),
                 span,
             ),
+            ResolverError::InvalidClosureEnvironment { span, typ } => Diagnostic::simple_error(
+                format!("{typ} is not a valid closure environment type"),
+                "Closure environment must be a tuple or unit type".to_string(), span),
         }
     }
 }

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

@@ -36,7 +36,8 @@ use crate::{
 use crate::{
     ArrayLiteral, ContractFunctionType, Distinctness, Generics, LValue, NoirStruct, NoirTypeAlias,
     Path, Pattern, Shared, StructType, Trait, Type, TypeAliasType, TypeBinding, TypeVariable,
-    UnaryOp, UnresolvedGenerics, UnresolvedType, UnresolvedTypeExpression, Visibility, ERROR_IDENT,
+    UnaryOp, UnresolvedGenerics, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
+    Visibility, ERROR_IDENT,
 };
 use fm::FileId;
 use iter_extended::vecmap;
@@ -331,9 +332,11 @@ impl<'a> Resolver<'a> {
     /// Translates an UnresolvedType into a Type and appends any
     /// freshly created TypeVariables created to new_variables.
     fn resolve_type_inner(&mut self, typ: UnresolvedType, new_variables: &mut Generics) -> Type {
-        match typ {
-            UnresolvedType::FieldElement => Type::FieldElement,
-            UnresolvedType::Array(size, elem) => {
+        use UnresolvedTypeData::*;
+
+        match typ.typ {
+            FieldElement => Type::FieldElement,
+            Array(size, elem) => {
                 let elem = Box::new(self.resolve_type_inner(*elem, new_variables));
                 let size = if size.is_none() {
                     Type::NotConstant
@@ -342,32 +345,50 @@ impl<'a> Resolver<'a> {
                 };
                 Type::Array(Box::new(size), elem)
             }
-            UnresolvedType::Expression(expr) => self.convert_expression_type(expr),
-            UnresolvedType::Integer(sign, bits) => Type::Integer(sign, bits),
-            UnresolvedType::Bool => Type::Bool,
-            UnresolvedType::String(size) => {
+            Expression(expr) => self.convert_expression_type(expr),
+            Integer(sign, bits) => Type::Integer(sign, bits),
+            Bool => Type::Bool,
+            String(size) => {
                 let resolved_size = self.resolve_array_size(size, new_variables);
                 Type::String(Box::new(resolved_size))
             }
-            UnresolvedType::FormatString(size, fields) => {
+            FormatString(size, fields) => {
                 let resolved_size = self.convert_expression_type(size);
                 let fields = self.resolve_type_inner(*fields, new_variables);
                 Type::FmtString(Box::new(resolved_size), Box::new(fields))
             }
-            UnresolvedType::Unit => Type::Unit,
-            UnresolvedType::Unspecified => Type::Error,
-            UnresolvedType::Error => Type::Error,
-            UnresolvedType::Named(path, args) => self.resolve_named_type(path, args, new_variables),
-            UnresolvedType::Tuple(fields) => {
+            Unit => Type::Unit,
+            Unspecified => Type::Error,
+            Error => Type::Error,
+            Named(path, args) => self.resolve_named_type(path, args, new_variables),
+            Tuple(fields) => {
                 Type::Tuple(vecmap(fields, |field| self.resolve_type_inner(field, new_variables)))
             }
-            UnresolvedType::Function(args, ret, env) => {
+            Function(args, ret, env) => {
                 let args = vecmap(args, |arg| self.resolve_type_inner(arg, new_variables));
                 let ret = Box::new(self.resolve_type_inner(*ret, new_variables));
+
+                // expect() here is valid, because the only places we don't have a span are omitted types
+                // e.g. a function without return type implicitly has a spanless UnresolvedType::Unit return type
+                // To get an invalid env type, the user must explicitly specify the type, which will have a span
+                let env_span = env.span.expect("Unexpected missing span for closure environment type");
+
                 let env = Box::new(self.resolve_type_inner(*env, new_variables));
-                Type::Function(args, ret, env)
+
+                match *env {
+                    Type::Unit | Type::Tuple(_) | Type::NamedGeneric(_, _) => {
+                        Type::Function(args, ret, env)
+                    }
+                    _ => {
+                        self.push_err(ResolverError::InvalidClosureEnvironment {
+                            typ: *env,
+                            span: env_span,
+                        });
+                        Type::Error
+                    }
+                }
             }
-            UnresolvedType::MutableReference(element) => {
+            MutableReference(element) => {
                 Type::MutableReference(Box::new(self.resolve_type_inner(*element, new_variables)))
             }
         }
@@ -576,9 +597,9 @@ impl<'a> Resolver<'a> {
     /// Translates a (possibly Unspecified) UnresolvedType to a Type.
     /// Any UnresolvedType::Unspecified encountered are replaced with fresh type variables.
     fn resolve_inferred_type(&mut self, typ: UnresolvedType) -> Type {
-        match typ {
-            UnresolvedType::Unspecified => self.interner.next_type_variable(),
-            other => self.resolve_type_inner(other, &mut vec![]),
+        match &typ.typ {
+            UnresolvedTypeData::Unspecified => self.interner.next_type_variable(),
+            _ => self.resolve_type_inner(typ, &mut vec![]),
         }
     }
 

crates/noirc_frontend/src/parser/mod.rs

@@ -403,7 +403,7 @@ impl ForRange {
                 // let fresh1 = array;
                 let let_array = Statement::Let(LetStatement {
                     pattern: Pattern::Identifier(array_ident.clone()),
-                    r#type: UnresolvedType::Unspecified,
+                    r#type: UnresolvedType::unspecified(),
                     expression: array,
                 });
 
@@ -436,7 +436,7 @@ impl ForRange {
                 // let elem = array[i];
                 let let_elem = Statement::Let(LetStatement {
                     pattern: Pattern::Identifier(identifier),
-                    r#type: UnresolvedType::Unspecified,
+                    r#type: UnresolvedType::unspecified(),
                     expression: Expression::new(loop_element, array_span),
                 });