feat: generic type arguments (#4122)

Author: aljazerzen

prqlc/prqlc-ast/src/expr.rs

@@ -102,6 +102,9 @@ pub struct Func {
 
     /// Named function parameters.
     pub named_params: Vec<FuncParam>,
+
+    /// Generic type arguments within this function.
+    pub generic_type_params: Vec<GenericTypeParam>,
 }
 
 #[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
@@ -114,6 +117,17 @@ pub struct FuncParam {
     pub default_value: Option<Box<Expr>>,
 }
 
+#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
+pub struct GenericTypeParam {
+    /// Assigned name of this generic type argument.
+    pub name: String,
+
+    /// Possible values of this type argument.
+    /// For a given instance of this function, the argument must be
+    /// exactly one of types in the domain.
+    pub domain: Vec<Ty>,
+}
+
 /// A value and a series of functions that are to be applied to that value one after another.
 #[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
 pub struct Pipeline {

prqlc/prqlc-ast/src/types.rs

@@ -45,6 +45,9 @@ pub enum TyKind {
 
     /// Type that is the largest subtype of `base` while not a subtype of `exclude`.
     Difference { base: Box<Ty>, exclude: Box<Ty> },
+
+    /// A generic argument. Contains id of the function call node and generic type param name.
+    GenericArg((usize, String)),
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, EnumAsInner)]

prqlc/prqlc-parser/src/expr.rs

@@ -333,24 +333,34 @@ where
         .map(ExprKind::Internal)
         .map_with_span(into_expr);
 
+    let generic_args = ident_part()
+        .then_ignore(ctrl(':'))
+        .then(type_expr().separated_by(ctrl('|')))
+        .map(|(name, domain)| GenericTypeParam { name, domain })
+        .separated_by(ctrl(','))
+        .at_least(1)
+        .delimited_by(ctrl('<'), ctrl('>'))
+        .or_not()
+        .map(|x| x.unwrap_or_default());
+
     choice((
         // func
-        keyword("func").ignore_then(
+        keyword("func").ignore_then(generic_args).then(
             param
                 .clone()
                 .separated_by(new_line().repeated())
                 .allow_leading()
                 .allow_trailing(),
         ),
         // plain
-        param.repeated(),
+        param.repeated().map(|params| (Vec::new(), params)),
     ))
     .then_ignore(just(Token::ArrowThin))
     // return type
     .then(type_expr().delimited_by(ctrl('<'), ctrl('>')).or_not())
     // body
     .then(choice((internal, func_call(expr))))
-    .map(|((params, return_ty), body)| {
+    .map(|(((generic_type_params, params), return_ty), body)| {
         let (pos, name) = params
             .into_iter()
             .map(|((name, ty), default_value)| FuncParam {
@@ -366,6 +376,7 @@ where
 
             body: Box::new(body),
             return_ty,
+            generic_type_params,
         })
     })
     .map(ExprKind::Func)

prqlc/prqlc-parser/src/test.rs

@@ -1065,6 +1065,7 @@ fn test_function() {
               - name: x
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-26"
     "###);
     assert_yaml_snapshot!(parse_single("let identity = x ->  x\n").unwrap()
@@ -1083,6 +1084,7 @@ fn test_function() {
               - name: x
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-22"
     "###);
     assert_yaml_snapshot!(parse_single("let plus_one = x ->  (x + 1)\n").unwrap()
@@ -1107,6 +1109,7 @@ fn test_function() {
               - name: x
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-28"
     "###);
     assert_yaml_snapshot!(parse_single("let plus_one = x ->  x + 1\n").unwrap()
@@ -1131,6 +1134,7 @@ fn test_function() {
               - name: x
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-26"
     "###);
 
@@ -1174,6 +1178,7 @@ fn test_function() {
               - name: x
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-51"
     "###);
 
@@ -1192,6 +1197,7 @@ fn test_function() {
               - name: return_constant
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-28"
     "###);
 
@@ -1217,6 +1223,7 @@ fn test_function() {
               - name: X
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-28"
     "###);
 
@@ -1273,6 +1280,7 @@ fn test_function() {
               - name: x
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:13-147"
     "###);
 
@@ -1301,6 +1309,7 @@ fn test_function() {
                 default_value:
                   Ident:
                     - a
+            generic_type_params: []
       span: "0:0-27"
     "###);
 }
@@ -1686,6 +1695,7 @@ fn test_inline_pipeline() {
               - name: x
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:0-37"
     "###);
 }
@@ -2417,6 +2427,7 @@ fn test_annotation() {
               - name: b
                 default_value: ~
             named_params: []
+            generic_type_params: []
       span: "0:9-61"
       annotations:
         - expr:

prqlc/prqlc/src/codegen/ast.rs

@@ -137,6 +137,23 @@ impl WriteSource for ExprKind {
             }
             Func(c) => {
                 let mut r = "func ".to_string();
+                if !c.generic_type_params.is_empty() {
+                    r += opt.consume("<")?;
+                    for generic_param in &c.generic_type_params {
+                        r += opt.consume(&write_ident_part(&generic_param.name))?;
+                        r += opt.consume(": ")?;
+                        r += &opt.consume(
+                            SeparatedExprs {
+                                exprs: &generic_param.domain,
+                                inline: " | ",
+                                line_end: "|",
+                            }
+                            .write(opt.clone())?,
+                        )?;
+                    }
+                    r += opt.consume("> ")?;
+                }
+
                 for param in &c.params {
                     r += opt.consume(&write_ident_part(&param.name))?;
                     r += opt.consume(" ")?;

prqlc/prqlc/src/codegen/types.rs

@@ -83,6 +83,7 @@ impl WriteSource for TyKind {
                 let exclude = exclude.write(opt.clone())?;
                 Some(format!("{base} - {exclude}"))
             }
+            GenericArg(_) => Some("?".to_string()),
         }
     }
 }

prqlc/prqlc/src/ir/pl/expr.rs

@@ -5,7 +5,7 @@ use enum_as_inner::EnumAsInner;
 use serde::{Deserialize, Serialize};
 
 use prqlc_ast::expr::generic;
-use prqlc_ast::{Ident, Literal, Span, Ty};
+use prqlc_ast::{GenericTypeParam, Ident, Literal, Span, Ty};
 
 use crate::codegen::write_ty;
 
@@ -114,6 +114,9 @@ pub struct Func {
     /// Named function parameters.
     pub named_params: Vec<FuncParam>,
 
+    /// Generic type arguments within this function.
+    pub generic_type_params: Vec<GenericTypeParam>,
+
     /// Arguments that have already been provided.
     pub args: Vec<Expr>,
 

prqlc/prqlc/src/ir/pl/fold.rs

@@ -350,7 +350,11 @@ pub fn fold_type<T: ?Sized + PlFold>(fold: &mut T, ty: Ty) -> Result<Ty> {
                 base: Box::new(fold.fold_type(*base)?),
                 exclude: Box::new(fold.fold_type(*exclude)?),
             },
-            TyKind::Any | TyKind::Ident(_) | TyKind::Primitive(_) | TyKind::Singleton(_) => ty.kind,
+            TyKind::Any
+            | TyKind::Ident(_)
+            | TyKind::Primitive(_)
+            | TyKind::Singleton(_)
+            | TyKind::GenericArg(_) => ty.kind,
         },
         span: ty.span,
         name: ty.name,

prqlc/prqlc/src/semantic/ast_expand.rs

@@ -44,6 +44,7 @@ pub fn expand_expr(expr: Expr) -> Result<pl::Expr> {
                 name_hint: None,
                 args: Vec::new(),
                 env: HashMap::new(),
+                generic_type_params: v.generic_type_params,
             }
             .into(),
         ),
@@ -295,6 +296,7 @@ fn restrict_expr_kind(value: pl::ExprKind) -> ExprKind {
                     body: restrict_expr_box(v.body),
                     params: restrict_func_params(v.params),
                     named_params: restrict_func_params(v.named_params),
+                    generic_type_params: v.generic_type_params,
                 }
                 .into(),
             );

prqlc/prqlc/src/semantic/eval.rs

@@ -425,6 +425,7 @@ fn new_func(name: &str, params: &[&str]) -> Expr {
         named_params: Default::default(),
         args: Default::default(),
         env: Default::default(),
+        generic_type_params: Default::default(),
     }));
     Expr {
         alias: Some(name.to_string()),