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/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"
+        "#)
+    }
 }

prqlc/prqlc/src/codegen/ast.rs

@@ -438,6 +438,18 @@ impl WriteSource for pr::Stmt {
                     }
                 }
             },
+            pr::StmtKind::TypeDef(pr::TypeDef {
+                name,
+                value:
+                    pr::Ty {
+                        kind: pr::TyKind::Enum(enum_tuple),
+                        ..
+                    },
+            }) => {
+                r += opt.consume(&format!("enum {} ", name))?;
+                r += &enum_tuple.write(opt)?;
+                r += "\n";
+            }
             pr::StmtKind::TypeDef(type_def) => {
                 r += opt.consume(&format!("type {}", type_def.name))?;
                 r += opt.consume(" = ")?;

prqlc/prqlc/src/codegen/types.rs

@@ -58,6 +58,7 @@ impl WriteSource for pr::TyKind {
                 r += &func.return_ty.as_deref().write(opt)?;
                 Some(r)
             }
+            Enum(tuple) => Some(format!("enum {}", tuple.write(opt)?)),
         }
     }
 }

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

@@ -348,6 +348,7 @@ pub fn fold_type<T: ?Sized + PlFold>(fold: &mut T, ty: Ty) -> Result<Ty> {
                 .transpose()?,
             ),
             TyKind::Ident(_) | TyKind::Primitive(_) => ty.kind,
+            TyKind::Enum(_) => ty.kind,
         },
         span: ty.span,
         name: ty.name,

prqlc/prqlc/src/semantic/resolver/functions.rs

@@ -6,7 +6,7 @@ use itertools::Itertools;
 use super::Resolver;
 use crate::ir::decl::{Decl, DeclKind, Module};
 use crate::ir::pl::*;
-use crate::pr::{Ty, TyFunc};
+use crate::pr::{Ty, TyFunc, TyKind};
 use crate::semantic::resolver::types;
 use crate::semantic::{NS_PARAM, NS_THAT, NS_THIS};
 use crate::Result;
@@ -307,7 +307,7 @@ impl Resolver<'_> {
                     let mut fields_new = Vec::with_capacity(fields.len());
                     self.in_flight_tuple_aliases.push(Vec::new());
                     for field in fields {
-                        let field = self.fold_within_namespace(field, &param.name)?;
+                        let field = self.fold_within_namespace(field, param)?;
 
                         // add aliased columns into scope
                         if let Some(alias) = field.alias.clone() {
@@ -361,7 +361,7 @@ impl Resolver<'_> {
         param: &FuncParam,
         func_name: &Option<Ident>,
     ) -> Result<Result<Expr, Expr>> {
-        let mut arg = self.fold_within_namespace(arg, &param.name)?;
+        let mut arg = self.fold_within_namespace(arg, param)?;
 
         // don't validate types of unresolved exprs
         if arg.id.is_some() {
@@ -387,12 +387,41 @@ impl Resolver<'_> {
         Ok(Ok(arg))
     }
 
-    fn fold_within_namespace(&mut self, expr: Expr, param_name: &str) -> Result<Expr> {
-        let prev_namespace = self.default_namespace.take();
+    fn fold_within_namespace(&mut self, expr: Expr, param: &FuncParam) -> Result<Expr> {
+        let param_name = &param.name;
 
         if param_name.starts_with("noresolve.") {
             return Ok(expr);
-        } else if let Some((ns, _)) = param_name.split_once('.') {
+        }
+
+        if let Some(Ty {
+            name: Some(name),
+            kind: TyKind::Enum(enum_expr),
+            ..
+        }) = &param.ty
+        {
+            if let crate::pr::Expr {
+                kind: crate::pr::ExprKind::Tuple(enum_tuple),
+                ..
+            } = &**enum_expr
+            {
+                // function parameters with a declared enum type will
+                // automatically have that namespace module searched first
+                self.current_module_path.push(name.to_string());
+                let res = self.fold_expr(expr).push_hint(format!(
+                    "perhaps you meant one of: {}",
+                    enum_tuple
+                        .iter()
+                        .filter_map(|expr| expr.alias.clone())
+                        .join(", ")
+                ));
+                self.current_module_path.pop();
+                return res;
+            }
+        }
+
+        let prev_namespace = self.default_namespace.take();
+        if let Some((ns, _)) = param_name.split_once('.') {
             self.default_namespace = Some(ns.to_string());
         } else {
             self.default_namespace = None;