feat: Add simple enumerations to type system

CHANGELOG.md

@@ -6,6 +6,9 @@
 
 **Features**:
 
+- Add simple enumeration type, and dedicated support for name resolution when an
+  enum is used as a function parameter. (@kgutwin, #6104)
+
 **Fixes**:
 
 **Documentation**:

prqlc/bindings/prqlc-python/src/lib.rs

@@ -232,15 +232,15 @@ mod test {
     fn debug_prql_lineage() {
         assert_snapshot!(
             debug::prql_lineage(r#"from a | select { beta, gamma }"#).unwrap(),
-            @r###"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":123,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":124,"target_name":null}}],"inputs":[{"id":121,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":121,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":126},{"id":123,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[121],"parent":125},{"id":124,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[121],"parent":125},{"id":125,"kind":"Tuple","span":"1:16-31","children":[123,124],"parent":126},{"id":126,"kind":"TransformCall: Select","span":"1:9-31","children":[121,125]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"###
+            @r#"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":127,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":128,"target_name":null}}],"inputs":[{"id":125,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":125,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":130},{"id":127,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[125],"parent":129},{"id":128,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[125],"parent":129},{"id":129,"kind":"Tuple","span":"1:16-31","children":[127,128],"parent":130},{"id":130,"kind":"TransformCall: Select","span":"1:9-31","children":[125,129]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"#
         );
     }
 
     #[test]
     fn debug_pl_to_lineage() {
         assert_snapshot!(
             prql_to_pl(r#"from a | select { beta, gamma }"#).and_then(|x| debug::pl_to_lineage(&x)).unwrap(),
-            @r###"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":123,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":124,"target_name":null}}],"inputs":[{"id":121,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":121,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":126},{"id":123,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[121],"parent":125},{"id":124,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[121],"parent":125},{"id":125,"kind":"Tuple","span":"1:16-31","children":[123,124],"parent":126},{"id":126,"kind":"TransformCall: Select","span":"1:9-31","children":[121,125]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"###
+            @r#"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":127,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":128,"target_name":null}}],"inputs":[{"id":125,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":125,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":130},{"id":127,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[125],"parent":129},{"id":128,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[125],"parent":129},{"id":129,"kind":"Tuple","span":"1:16-31","children":[127,128],"parent":130},{"id":130,"kind":"TransformCall: Select","span":"1:9-31","children":[125,129]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"#
         );
     }
 }

prqlc/prqlc-parser/src/parser/expr.rs

@@ -173,6 +173,41 @@ where
     })
 }
 
+pub(crate) fn literal_tuple<'a, I>() -> impl Parser<'a, I, Expr, ParserError<'a>> + Clone + 'a
+where
+    I: Input<'a, Token = lr::Token, Span = Span> + BorrowInput<'a>,
+{
+    use chumsky::recovery::{skip_then_retry_until, via_parser};
+
+    sequence(maybe_aliased(expr()).validate(|expr, extra, emit| {
+        let span = extra.span();
+
+        if expr.alias.is_none() {
+            emit.emit(Rich::custom(span, "must specify an alias for this value"));
+        }
+
+        match expr.kind {
+            ExprKind::Literal(_) => (),
+            _ => emit.emit(Rich::custom(extra.span(), "expected a literal value")),
+        };
+
+        expr
+    }))
+    .delimited_by(
+        ctrl('{'),
+        ctrl('}')
+            .recover_with(via_parser(end()))
+            .recover_with(skip_then_retry_until(
+                any_ref().ignored(),
+                ctrl('}').ignored().or(ctrl(',').ignored()).or(end()),
+            )),
+    )
+    .map(ExprKind::Tuple)
+    .labelled("literal tuple")
+    .map_with(|kind, extra| ExprKind::into_expr(kind, extra.span()))
+    .boxed()
+}
+
 fn tuple<'a, I>(
     nested_expr: impl Parser<'a, I, Expr, ParserError<'a>> + Clone + 'a,
 ) -> impl Parser<'a, I, ExprKind, ParserError<'a>> + Clone + 'a

prqlc/prqlc-parser/src/parser/pr/types.rs

@@ -3,6 +3,7 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use strum::AsRefStr;
 
+use crate::parser::pr::expr::Expr;
 use crate::parser::pr::ident::Ident;
 use crate::span::Span;
 
@@ -32,6 +33,9 @@ pub enum TyKind {
 
     /// Type of functions with defined params and return types.
     Function(Option<TyFunc>),
+
+    /// Simple enumeration type
+    Enum(Box<Expr>),
 }
 
 impl TyKind {

prqlc/prqlc-parser/src/parser/stmt.rs

@@ -6,7 +6,7 @@ use chumsky::prelude::*;
 use itertools::Itertools;
 use semver::VersionReq;
 
-use super::expr::{expr, expr_call, ident, pipeline};
+use super::expr::{expr, expr_call, ident, literal_tuple, pipeline};
 use super::{ctrl, ident_part, into_stmt, keyword, new_line, pipe, with_doc_comment};
 use crate::lexer::lr;
 use crate::lexer::lr::{Literal, TokenKind};
@@ -74,7 +74,13 @@ where
         let stmt_kind = new_line()
             .repeated()
             .collect::<Vec<_>>()
-            .ignore_then(choice((module_def, type_def(), import_def(), var_def())));
+            .ignore_then(choice((
+                module_def,
+                type_def(),
+                enum_def(),
+                import_def(),
+                var_def(),
+            )));
 
         // Currently doc comments need to be before the annotation; probably
         // should relax this?
@@ -226,12 +232,32 @@ where
     I: Input<'a, Token = lr::Token, Span = Span> + BorrowInput<'a> + chumsky::input::ValueInput<'a>,
 {
     keyword("type")
-        .ignore_then(ident_part())
-        .then(ctrl('=').ignore_then(type_expr()))
-        .map(|(name, value)| StmtKind::TypeDef(TypeDef { name, value }))
+        .ignore_then(
+            ident_part()
+                .then(ctrl('=').ignore_then(type_expr()))
+                .map(|(name, value)| StmtKind::TypeDef(TypeDef { name, value })),
+        )
         .labelled("type definition")
 }
 
+fn enum_def<'a, I>() -> impl Parser<'a, I, StmtKind, ParserError<'a>> + Clone
+where
+    I: Input<'a, Token = lr::Token, Span = Span> + BorrowInput<'a> + chumsky::input::ValueInput<'a>,
+{
+    keyword("enum")
+        .ignore_then(ident_part().then(literal_tuple()).map(|(name, value)| {
+            StmtKind::TypeDef(TypeDef {
+                name: name.clone(),
+                value: Ty {
+                    span: value.span,
+                    kind: TyKind::Enum(Box::new(value)),
+                    name: Some(name),
+                },
+            })
+        }))
+        .labelled("enum definition")
+}
+
 fn import_def<'a, I>() -> impl Parser<'a, I, StmtKind, ParserError<'a>> + Clone
 where
     I: Input<'a, Token = lr::Token, Span = Span> + BorrowInput<'a> + chumsky::input::ValueInput<'a>,
@@ -397,7 +423,7 @@ mod tests {
                     0:0-73,
                 ),
                 reason: Simple(
-                    "Expected one of import statement, module definition, new line, pipeline, something else, type definition or variable definition, but didn't find anything before the end.",
+                    "Expected one of enum definition, import statement, module definition, new line, pipeline, something else, type definition or variable definition, but didn't find anything before the end.",
                 ),
                 hints: [],
                 code: None,
@@ -698,4 +724,152 @@ mod tests {
           span: "0:0-139"
         "#);
     }
+
+    #[test]
+    fn enums() {
+        assert_yaml_snapshot!(parse_module_contents(r#"
+        enum foo { First = 0, Second = 1 }
+        "#).unwrap(), @r#"
+        - TypeDef:
+            name: foo
+            value:
+              kind:
+                Enum:
+                  Tuple:
+                    - Literal:
+                        Integer: 0
+                      span: "0:28-29"
+                      alias: First
+                    - Literal:
+                        Integer: 1
+                      span: "0:40-41"
+                      alias: Second
+                  span: "0:18-43"
+              span: "0:18-43"
+              name: foo
+          span: "0:0-43"
+        "#)
+    }
+
+    #[test]
+    fn enum_must_be_tuple_literal() {
+        assert_debug_snapshot!(parse_module_contents(r#"
+        enum foo
+        "#).unwrap_err(), @r#"
+        [
+            Error {
+                kind: Error,
+                span: Some(
+                    0:17-18,
+                ),
+                reason: Expected {
+                    who: None,
+                    expected: "literal tuple",
+                    found: "new line",
+                },
+                hints: [],
+                code: None,
+            },
+        ]
+        "#);
+
+        assert_debug_snapshot!(parse_module_contents(r#"
+        enum foo 4
+        "#).unwrap_err(), @r#"
+        [
+            Error {
+                kind: Error,
+                span: Some(
+                    0:18-19,
+                ),
+                reason: Expected {
+                    who: None,
+                    expected: "literal tuple",
+                    found: "4",
+                },
+                hints: [],
+                code: None,
+            },
+        ]
+        "#);
+
+        assert_debug_snapshot!(parse_module_contents(r#"
+        enum foo { "First", "Second" }
+        "#).unwrap_err(), @r#"
+        [
+            Error {
+                kind: Error,
+                span: Some(
+                    0:20-27,
+                ),
+                reason: Simple(
+                    "must specify an alias for this value",
+                ),
+                hints: [],
+                code: None,
+            },
+            Error {
+                kind: Error,
+                span: Some(
+                    0:29-37,
+                ),
+                reason: Simple(
+                    "must specify an alias for this value",
+                ),
+                hints: [],
+                code: None,
+            },
+        ]
+        "#);
+
+        assert_debug_snapshot!(parse_module_contents(r#"
+        enum foo { First = 4 + 5 }
+        "#).unwrap_err(), @r#"
+        [
+            Error {
+                kind: Error,
+                span: Some(
+                    0:20-33,
+                ),
+                reason: Simple(
+                    "expected a literal value",
+                ),
+                hints: [],
+                code: None,
+            },
+        ]
+        "#)
+    }
+
+    #[test]
+    fn func_named_arg_type() {
+        assert_yaml_snapshot!(parse_module_contents(r#"
+        let f = func x <int>:0 -> x
+        "#).unwrap(), @r#"
+        - VarDef:
+            kind: Let
+            name: f
+            value:
+              Func:
+                return_ty: ~
+                body:
+                  Ident:
+                    - x
+                  span: "0:35-36"
+                params: []
+                named_params:
+                  - name: x
+                    ty:
+                      kind:
+                        Primitive: Int
+                      span: "0:25-28"
+                      name: ~
+                    default_value:
+                      Literal:
+                        Integer: 0
+                      span: "0:30-31"
+              span: "0:17-36"
+          span: "0:0-36"
+        "#)
+    }
 }