Function types (#6891)

* Change Fn to fn for function types

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Support args and return types in function types

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Use fancy function types in the docs

Signed-off-by: Nick Cameron <nrc@ncameron.org>

---------

Signed-off-by: Nick Cameron <nrc@ncameron.org>

docs/kcl-std/functions/std-array-map.md

@@ -10,7 +10,7 @@ Apply a function to every element of a list.
 ```kcl
 map(
   @array: [any],
-  f: Fn,
+  f: fn(any): any,
 ): [any]
 ```
 
@@ -22,7 +22,7 @@ Given a list like `[a, b, c]`, and a function like `f`, returns
 | Name | Type | Description | Required |
 |----------|------|-------------|----------|
 | `array` | [`[any]`](/docs/kcl-std/types/std-types-any) | Input array. The output array is this input array, but every element has had the function `f` run on it. | Yes |
-| `f` | [`Fn`](/docs/kcl-std/types/std-types-Fn) | A function. The output array is just the input array, but `f` has been run on every item. | Yes |
+| `f` | [`fn(any): any`](/docs/kcl-std/types/std-types-fn) | A function. The output array is just the input array, but `f` has been run on every item. | Yes |
 
 ### Returns
 

docs/kcl-std/functions/std-array-reduce.md

@@ -11,7 +11,7 @@ layout: manual
 reduce(
   @array: [any],
   initial: any,
-  f: Fn,
+  f: fn(any, accum: any): any,
 ): [any]
 ```
 
@@ -24,7 +24,7 @@ using the previous value and the element.
 |----------|------|-------------|----------|
 | `array` | [`[any]`](/docs/kcl-std/types/std-types-any) | Each element of this array gets run through the function `f`, combined with the previous output from `f`, and then used for the next run. | Yes |
 | `initial` | [`any`](/docs/kcl-std/types/std-types-any) | The first time `f` is run, it will be called with the first item of `array` and this initial starting value. | Yes |
-| `f` | [`Fn`](/docs/kcl-std/types/std-types-Fn) | Run once per item in the input `array`. This function takes an item from the array, and the previous output from `f` (or `initial` on the very first run). The final time `f` is run, its output is returned as the final output from `reduce`. | Yes |
+| `f` | [`fn(any, accum: any): any`](/docs/kcl-std/types/std-types-fn) | Run once per item in the input `array`. This function takes an item from the array, and the previous output from `f` (or `initial` on the very first run). The final time `f` is run, its output is returned as the final output from `reduce`. | Yes |
 
 ### Returns
 

docs/kcl-std/index.md

@@ -140,13 +140,13 @@ See also the [types overview](/docs/kcl-lang/types)
 
 * [**Primitive types**](/docs/kcl-lang/types)
   * [`End`](/docs/kcl-lang/types#End)
-  * [`Fn`](/docs/kcl-std/types/std-types-Fn)
   * [`ImportedGeometry`](/docs/kcl-std/types/std-types-ImportedGeometry)
   * [`Start`](/docs/kcl-lang/types#Start)
   * [`TagDeclarator`](/docs/kcl-lang/types#TagDeclarator)
   * [`TagIdentifier`](/docs/kcl-lang/types#TagIdentifier)
   * [`any`](/docs/kcl-std/types/std-types-any)
   * [`bool`](/docs/kcl-std/types/std-types-bool)
+  * [`fn`](/docs/kcl-std/types/std-types-fn)
   * [`number`](/docs/kcl-std/types/std-types-number)
   * [`string`](/docs/kcl-std/types/std-types-string)
   * [`tag`](/docs/kcl-std/types/std-types-tag)

docs/kcl-std/modules/std-types.md

@@ -17,7 +17,6 @@ Types can (optionally) be used to describe a function's arguments and returned v
 * [`Axis3d`](/docs/kcl-std/types/std-types-Axis3d)
 * [`Edge`](/docs/kcl-std/types/std-types-Edge)
 * [`Face`](/docs/kcl-std/types/std-types-Face)
-* [`Fn`](/docs/kcl-std/types/std-types-Fn)
 * [`Helix`](/docs/kcl-std/types/std-types-Helix)
 * [`ImportedGeometry`](/docs/kcl-std/types/std-types-ImportedGeometry)
 * [`Plane`](/docs/kcl-std/types/std-types-Plane)
@@ -27,6 +26,7 @@ Types can (optionally) be used to describe a function's arguments and returned v
 * [`Solid`](/docs/kcl-std/types/std-types-Solid)
 * [`any`](/docs/kcl-std/types/std-types-any)
 * [`bool`](/docs/kcl-std/types/std-types-bool)
+* [`fn`](/docs/kcl-std/types/std-types-fn)
 * [`number`](/docs/kcl-std/types/std-types-number)
 * [`string`](/docs/kcl-std/types/std-types-string)
 * [`tag`](/docs/kcl-std/types/std-types-tag)

docs/kcl-std/types/std-types-fn.md

@@ -1,5 +1,5 @@
 ---
-title: "Fn"
+title: "fn"
 subtitle: "Type in std::types"
 excerpt: "The type of any function in KCL."
 layout: manual

rust/kcl-lib/src/docs/gen_std_tests.rs

@@ -674,6 +674,8 @@ fn cleanup_type_string(input: &str, fmt_for_text: bool) -> String {
 
             if fmt_for_text && ty.starts_with("number") {
                 format!("[{prefix}{ty}{suffix}](/docs/kcl-std/types/std-types-number)")
+            } else if fmt_for_text && ty.starts_with("fn") {
+                format!("[{prefix}{ty}{suffix}](/docs/kcl-std/types/std-types-fn)")
             } else if fmt_for_text && SPECIAL_TYPES.contains(&ty) {
                 format!("[{prefix}{ty}{suffix}](/docs/kcl-lang/types#{ty})")
             } else if fmt_for_text && DECLARED_TYPES.contains(&ty) {

rust/kcl-lib/src/docs/mod.rs

@@ -39,7 +39,7 @@ const DECLARED_TYPES: [&str; 17] = [
     "Axis2d",
     "Axis3d",
     "ImportedGeometry",
-    "Fn",
+    "fn",
 ];
 
 lazy_static::lazy_static! {

rust/kcl-lib/src/execution/types.rs

@@ -183,7 +183,7 @@ impl RuntimeType {
             AstPrimitiveType::Named(name) => Self::from_alias(&name.name, exec_state, source_range)?,
             AstPrimitiveType::Tag => RuntimeType::Primitive(PrimitiveType::Tag),
             AstPrimitiveType::ImportedGeometry => RuntimeType::Primitive(PrimitiveType::ImportedGeometry),
-            AstPrimitiveType::Function => RuntimeType::Primitive(PrimitiveType::Function),
+            AstPrimitiveType::Function(_) => RuntimeType::Primitive(PrimitiveType::Function),
         })
     }
 

rust/kcl-lib/src/parsing/ast/digest.rs

@@ -2,8 +2,8 @@ use sha2::{Digest as DigestTrait, Sha256};
 
 use crate::parsing::ast::types::{
     Annotation, ArrayExpression, ArrayRangeExpression, AscribedExpression, BinaryExpression, BinaryPart, BodyItem,
-    CallExpressionKw, DefaultParamVal, ElseIf, Expr, ExpressionStatement, FunctionExpression, Identifier, IfExpression,
-    ImportItem, ImportSelector, ImportStatement, ItemVisibility, KclNone, LabelledExpression, Literal,
+    CallExpressionKw, DefaultParamVal, ElseIf, Expr, ExpressionStatement, FunctionExpression, FunctionType, Identifier,
+    IfExpression, ImportItem, ImportSelector, ImportStatement, ItemVisibility, KclNone, LabelledExpression, Literal,
     LiteralIdentifier, LiteralValue, MemberExpression, MemberObject, Name, ObjectExpression, ObjectProperty, Parameter,
     PipeExpression, PipeSubstitution, PrimitiveType, Program, ReturnStatement, TagDeclarator, Type, TypeDeclaration,
     UnaryExpression, VariableDeclaration, VariableDeclarator, VariableKind,
@@ -233,13 +233,28 @@ impl PrimitiveType {
             PrimitiveType::Boolean => hasher.update(b"bool"),
             PrimitiveType::Tag => hasher.update(b"tag"),
             PrimitiveType::ImportedGeometry => hasher.update(b"ImportedGeometry"),
-            PrimitiveType::Function => hasher.update(b"Fn"),
+            PrimitiveType::Function(f) => hasher.update(f.compute_digest()),
         }
 
         hasher.finalize().into()
     }
 }
 
+impl FunctionType {
+    compute_digest!(|slf, hasher| {
+        if let Some(u) = &mut slf.unnamed_arg {
+            hasher.update(u.compute_digest());
+        }
+        slf.named_args.iter_mut().for_each(|(a, t)| {
+            a.compute_digest();
+            t.compute_digest();
+        });
+        if let Some(r) = &mut slf.return_type {
+            hasher.update(r.compute_digest());
+        }
+    });
+}
+
 impl Parameter {
     compute_digest!(|slf, hasher| {
         hasher.update(slf.identifier.compute_digest());

rust/kcl-lib/src/parsing/ast/types/mod.rs

@@ -3199,8 +3199,8 @@ pub enum PrimitiveType {
     Tag,
     /// Imported from other CAD system.
     ImportedGeometry,
-    /// `Fn`, type of functions.
-    Function,
+    /// `fn`, type of functions.
+    Function(FunctionType),
     /// An identifier used as a type (not really a primitive type, but whatever).
     Named(Node<Identifier>),
 }
@@ -3215,7 +3215,6 @@ impl PrimitiveType {
             ("number", None) => Some(PrimitiveType::Number(NumericSuffix::None)),
             ("number", Some(s)) => Some(PrimitiveType::Number(s)),
             ("ImportedGeometry", None) => Some(PrimitiveType::ImportedGeometry),
-            ("Fn", None) => Some(PrimitiveType::Function),
             _ => None,
         }
     }
@@ -3236,12 +3235,57 @@ impl fmt::Display for PrimitiveType {
             PrimitiveType::Boolean => write!(f, "bool"),
             PrimitiveType::Tag => write!(f, "tag"),
             PrimitiveType::ImportedGeometry => write!(f, "ImportedGeometry"),
-            PrimitiveType::Function => write!(f, "Fn"),
+            PrimitiveType::Function(t) => {
+                write!(f, "fn")?;
+                if t.unnamed_arg.is_some() || !t.named_args.is_empty() || t.return_type.is_some() {
+                    write!(f, "(")?;
+                    if let Some(u) = &t.unnamed_arg {
+                        write!(f, "{u}")?;
+                        if !t.named_args.is_empty() {
+                            write!(f, ", ")?;
+                        }
+                    }
+                    for (i, (a, t)) in t.named_args.iter().enumerate() {
+                        if i != 0 {
+                            write!(f, ", ")?;
+                        }
+                        write!(f, "{}: {t}", a.name)?;
+                    }
+                    write!(f, ")")?;
+                    if let Some(r) = &t.return_type {
+                        write!(f, ": {r}")?;
+                    }
+                }
+                Ok(())
+            }
             PrimitiveType::Named(n) => write!(f, "{}", n.name),
         }
     }
 }
 
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
+#[ts(export)]
+pub struct FunctionType {
+    pub unnamed_arg: Option<BoxNode<Type>>,
+    pub named_args: Vec<(Node<Identifier>, Node<Type>)>,
+    pub return_type: Option<BoxNode<Type>>,
+
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    #[ts(optional)]
+    pub digest: Option<Digest>,
+}
+
+impl FunctionType {
+    pub fn empty_fn_type() -> Self {
+        FunctionType {
+            unnamed_arg: None,
+            named_args: Vec::new(),
+            return_type: None,
+            digest: None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
 #[ts(export)]
 #[serde(tag = "type")]

rust/kcl-lib/src/parsing/parser.rs

@@ -25,11 +25,11 @@ use crate::{
         ast::types::{
             Annotation, ArrayExpression, ArrayRangeExpression, BinaryExpression, BinaryOperator, BinaryPart, BodyItem,
             BoxNode, CallExpressionKw, CommentStyle, DefaultParamVal, ElseIf, Expr, ExpressionStatement,
-            FunctionExpression, Identifier, IfExpression, ImportItem, ImportSelector, ImportStatement, ItemVisibility,
-            LabeledArg, Literal, LiteralIdentifier, LiteralValue, MemberExpression, MemberObject, Name, Node, NodeList,
-            NonCodeMeta, NonCodeNode, NonCodeValue, ObjectExpression, ObjectProperty, Parameter, PipeExpression,
-            PipeSubstitution, PrimitiveType, Program, ReturnStatement, Shebang, TagDeclarator, Type, TypeDeclaration,
-            UnaryExpression, UnaryOperator, VariableDeclaration, VariableDeclarator, VariableKind,
+            FunctionExpression, FunctionType, Identifier, IfExpression, ImportItem, ImportSelector, ImportStatement,
+            ItemVisibility, LabeledArg, Literal, LiteralIdentifier, LiteralValue, MemberExpression, MemberObject, Name,
+            Node, NodeList, NonCodeMeta, NonCodeNode, NonCodeValue, ObjectExpression, ObjectProperty, Parameter,
+            PipeExpression, PipeSubstitution, PrimitiveType, Program, ReturnStatement, Shebang, TagDeclarator, Type,
+            TypeDeclaration, UnaryExpression, UnaryOperator, VariableDeclaration, VariableDeclarator, VariableKind,
         },
         math::BinaryExpressionToken,
         token::{Token, TokenSlice, TokenType},
@@ -1261,7 +1261,7 @@ fn function_decl(i: &mut TokenSlice) -> PResult<Node<FunctionExpression>> {
     fn return_type(i: &mut TokenSlice) -> PResult<Node<Type>> {
         colon(i)?;
         ignore_whitespace(i);
-        argument_type(i)
+        type_(i)
     }
 
     let open = open_paren(i)?;
@@ -2013,7 +2013,7 @@ fn expression_but_not_pipe(i: &mut TokenSlice) -> PResult<Expr> {
     .context(expected("a KCL value"))
     .parse_next(i)?;
 
-    let ty = opt((colon, opt(whitespace), argument_type)).parse_next(i)?;
+    let ty = opt((colon, opt(whitespace), type_)).parse_next(i)?;
     if let Some((_, _, ty)) = ty {
         expr = Expr::AscribedExpression(Box::new(AscribedExpression::new(expr, ty)))
     }
@@ -2083,7 +2083,7 @@ fn possible_operands(i: &mut TokenSlice) -> PResult<Expr> {
     ))
     .parse_next(i)?;
 
-    let ty = opt((colon, opt(whitespace), argument_type)).parse_next(i)?;
+    let ty = opt((colon, opt(whitespace), type_)).parse_next(i)?;
     if let Some((_, _, ty)) = ty {
         expr = Expr::AscribedExpression(Box::new(AscribedExpression::new(expr, ty)))
     }
@@ -2233,7 +2233,21 @@ fn ty_decl(i: &mut TokenSlice) -> PResult<BoxNode<TypeDeclaration>> {
     let start = visibility_token.map(|t| t.start).unwrap_or_else(|| decl_token.start);
     whitespace(i)?;
 
-    let name = identifier(i)?;
+    let name = alt((
+        fun.map(|t| {
+            Node::new(
+                Identifier {
+                    name: "fn".to_owned(),
+                    digest: None,
+                },
+                t.start,
+                t.end,
+                t.module_id,
+            )
+        }),
+        identifier,
+    ))
+    .parse_next(i)?;
     let mut end = name.end;
 
     let args = if peek((opt(whitespace), open_paren)).parse_next(i).is_ok() {
@@ -2253,7 +2267,7 @@ fn ty_decl(i: &mut TokenSlice) -> PResult<BoxNode<TypeDeclaration>> {
         ignore_whitespace(i);
         equals(i)?;
         ignore_whitespace(i);
-        let ty = argument_type(i)?;
+        let ty = type_(i)?;
 
         ParseContext::warn(CompilationError::err(
             ty.as_source_range(),
@@ -2755,12 +2769,8 @@ fn labeled_argument(i: &mut TokenSlice) -> PResult<LabeledArg> {
         .parse_next(i)
 }
 
-/// A type of a function argument.
-/// This can be:
-/// - a primitive type, e.g. `number` or `string` or `bool`
-/// - an array type, e.g. `[number]` or `[string]` or `[bool]`
-/// - an object type, e.g. `{x: number, y: number}` or `{name: string, age: number}`
-fn argument_type(i: &mut TokenSlice) -> PResult<Node<Type>> {
+/// Parse a type in various positions.
+fn type_(i: &mut TokenSlice) -> PResult<Node<Type>> {
     let type_ = alt((
         // Object types
         // TODO it is buggy to treat object fields like parameters since the parameters parser assumes a terminating `)`.
@@ -2800,14 +2810,69 @@ fn argument_type(i: &mut TokenSlice) -> PResult<Node<Type>> {
 }
 
 fn primitive_type(i: &mut TokenSlice) -> PResult<Node<PrimitiveType>> {
-    let ident = identifier(i)?;
+    alt((
+        // A function type: `fn` (`(` type?, (id: type,)* `)` (`:` type)?)?
+        (
+            fun,
+            opt((
+                // `(` type?, (id: type,)* `)`
+                delimited(
+                    open_paren,
+                    opt(alt((
+                        // type, (id: type,)+
+                        (
+                            type_,
+                            comma,
+                            opt(whitespace),
+                            separated(
+                                1..,
+                                (identifier, colon, opt(whitespace), type_).map(|(id, _, _, ty)| (id, ty)),
+                                comma_sep,
+                            ),
+                        )
+                            .map(|(t, _, _, args)| (Some(t), args)),
+                        // (id: type,)+
+                        separated(
+                            1..,
+                            (identifier, colon, opt(whitespace), type_).map(|(id, _, _, ty)| (id, ty)),
+                            comma_sep,
+                        )
+                        .map(|args| (None, args)),
+                        // type
+                        type_.map(|t| (Some(t), Vec::new())),
+                    ))),
+                    close_paren,
+                ),
+                // `:` type
+                opt((colon, opt(whitespace), type_)),
+            )),
+        )
+            .map(|(t, tys)| {
+                let mut ft = FunctionType::empty_fn_type();
 
-    let suffix = opt(delimited(open_paren, uom_for_type, close_paren)).parse_next(i)?;
+                if let Some((args, ret)) = tys {
+                    if let Some((unnamed, named)) = args {
+                        if let Some(unnamed) = unnamed {
+                            ft.unnamed_arg = Some(Box::new(unnamed));
+                        }
+                        ft.named_args = named;
+                    }
+                    if let Some((_, _, ty)) = ret {
+                        ft.return_type = Some(Box::new(ty));
+                    }
+                }
 
+                Node::new(PrimitiveType::Function(ft), t.start, t.end, t.module_id)
+            }),
+        // A named type, possibly with a numeric suffix.
+        (identifier, opt(delimited(open_paren, uom_for_type, close_paren))).map(|(ident, suffix)| {
             let mut result = Node::new(PrimitiveType::Boolean, ident.start, ident.end, ident.module_id);
-    result.inner = PrimitiveType::primitive_from_str(&ident.name, suffix).unwrap_or(PrimitiveType::Named(ident));
-
-    Ok(result)
+            result.inner =
+                PrimitiveType::primitive_from_str(&ident.name, suffix).unwrap_or(PrimitiveType::Named(ident));
+            result
+        }),
+    ))
+    .parse_next(i)
 }
 
 fn array_type(i: &mut TokenSlice) -> PResult<Node<Type>> {
@@ -2817,7 +2882,7 @@ fn array_type(i: &mut TokenSlice) -> PResult<Node<Type>> {
     }
 
     open_bracket(i)?;
-    let ty = argument_type(i)?;
+    let ty = type_(i)?;
     let len = opt((
         semi_colon,
         opt_whitespace,
@@ -2905,7 +2970,7 @@ fn parameter(i: &mut TokenSlice) -> PResult<ParamDescription> {
         any.verify(|token: &Token| !matches!(token.token_type, TokenType::Brace) || token.value != ")"),
         opt(question_mark),
         opt(whitespace),
-        opt((colon, opt(whitespace), argument_type).map(|tup| tup.2)),
+        opt((colon, opt(whitespace), type_).map(|tup| tup.2)),
         opt(whitespace),
         opt((equals, opt(whitespace), literal).map(|(_, _, literal)| literal)),
     )
@@ -4777,6 +4842,20 @@ let myBox = box(p=[0,0], h=-3, l=-16, w=-10)
         assert_no_err(some_program_string);
     }
     #[test]
+    fn parse_function_types() {
+        let code = r#"foo = x: fn
+foo = x: fn(number)
+fn foo(x: fn(): number): fn { return 0 }
+fn foo(x: fn(a, b: number(mm), c: d): number(Angle)): fn { return 0 }
+type fn
+type foo = fn
+type foo = fn(a: string, b: { f: fn(): any })
+type foo = fn([fn])
+type foo = fn(fn, f: fn(number(_))): [fn([any]): string]
+    "#;
+        assert_no_err(code);
+    }
+    #[test]
     fn test_parse_tag_starting_with_bang() {
         let some_program_string = r#"startSketchOn(XY)
     |> startProfile(at = [0, 0])

rust/kcl-lib/src/unparser.rs

@@ -2575,6 +2575,27 @@ sketch002 = startSketchOn({
         assert_eq!(actual, input);
     }
 
+    #[test]
+    fn recast_function_types() {
+        let input = r#"foo = x: fn
+foo = x: fn(number)
+fn foo(x: fn(): number): fn {
+  return 0
+}
+fn foo(x: fn(a, b: number(mm), c: d): number(Angle)): fn {
+  return 0
+}
+type fn
+type foo = fn
+type foo = fn(a: string, b: { f: fn(): any })
+type foo = fn([fn])
+type foo = fn(fn, f: fn(number(_))): [fn([any]): string]
+"#;
+        let ast = crate::parsing::top_level_parse(input).unwrap();
+        let actual = ast.recast(&FormatOptions::new(), 0);
+        assert_eq!(actual, input);
+    }
+
     #[test]
     fn unparse_call_inside_function_args_multiple_lines() {
         let input = r#"fn foo() {

rust/kcl-lib/std/array.kcl

@@ -40,7 +40,7 @@ export fn map(
   /// Input array. The output array is this input array, but every element has had the function `f` run on it.
   @array: [any],
   /// A function. The output array is just the input array, but `f` has been run on every item.
-  f: Fn,
+  f: fn(any): any,
 ): [any] {}
 
 /// Take a starting value. Then, for each element of an array, calculate the next value,
@@ -132,7 +132,7 @@ export fn reduce(
   /// The first time `f` is run, it will be called with the first item of `array` and this initial starting value.
   initial: any,
   /// Run once per item in the input `array`. This function takes an item from the array, and the previous output from `f` (or `initial` on the very first run). The final time `f` is run, its output is returned as the final output from `reduce`.
-  f: Fn,
+  f: fn(any, accum: any): any,
 ): [any] {}
 
 /// Append an element to the end of an array.

rust/kcl-lib/std/types.kcl

@@ -169,7 +169,7 @@ export type ImportedGeometry
 
 /// The type of any function in KCL.
 @(impl = primitive)
-export type Fn
+export type fn
 
 /// An abstract plane.
 ///

rust/kcl-lib/tests/argument_error/execution_error.snap

@@ -4,7 +4,8 @@ description: Error from executing argument_error.kcl
 ---
 KCL Semantic error
 
-  × semantic: f requires a value with type `Fn`, but found array (list)
+  × semantic: f requires a value with type `fn(any): any`, but found array
+  │ (list)
    ╭─[5:1]
  4 │ 
  5 │ map(f, f = [0, 1])
@@ -14,7 +15,8 @@ KCL Semantic error
    ╰────
   ╰─▶ KCL Semantic error
       
-        × semantic: f requires a value with type `Fn`, but found array (list)
+        × semantic: f requires a value with type `fn(any): any`, but found
+        │ array (list)
          ╭─[5:12]
        4 │
        5 │ map(f, f = [0, 1])

rust/kcl-lib/tests/kcl_samples/walkie-talkie/artifact_graph_flowchart.snap.md

@@ -709,7 +709,6 @@ flowchart LR
   115 --- 179
   115 x--> 228
   115 --- 256
-  115 x--> 307
   115 --- 308
   164 <--x 116
   116 --- 184