From 5189ddb3d5e1b0b95c3b79c8e9956f7706da0c3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alja=C5=BE=20Mur=20Er=C5=BEen?=
 <aljazerzen@users.noreply.github.com>
Date: Sat, 11 May 2024 22:21:41 +0200
Subject: [PATCH] refactor: cleanup union parsing (#4462)

---
 prqlc/prqlc-parser/src/lexer.rs            |  1 +
 prqlc/prqlc-parser/src/types.rs            | 47 +++++++++++++---------
 prqlc/prqlc/tests/integration/resolving.rs | 15 ++++---
 3 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/prqlc/prqlc-parser/src/lexer.rs b/prqlc/prqlc-parser/src/lexer.rs
index f84c8c7b937b..10cf73d09929 100644
--- a/prqlc/prqlc-parser/src/lexer.rs
+++ b/prqlc/prqlc-parser/src/lexer.rs
@@ -108,6 +108,7 @@ pub fn lex_token() -> impl Parser<char, Token, Error = Cheap<char>> {
         just("internal"),
         just("func"),
         just("import"),
+        just("enum"),
     ))
     .then_ignore(end_expr())
     .map(|x| x.to_string())
diff --git a/prqlc/prqlc-parser/src/types.rs b/prqlc/prqlc-parser/src/types.rs
index 8a53ae1463aa..3650a3774140 100644
--- a/prqlc/prqlc-parser/src/types.rs
+++ b/prqlc/prqlc-parser/src/types.rs
@@ -81,25 +81,32 @@ pub fn type_expr() -> impl Parser<TokenKind, Ty, Error = PError> {
             .map(TyKind::Tuple)
             .labelled("tuple");
 
-        let union_parenthesized = ident_part()
-            .then_ignore(ctrl('='))
-            .or_not()
-            .then(nested_type_expr.clone())
-            .padded_by(new_line().repeated())
-            .separated_by(just(TokenKind::Or))
-            .allow_trailing()
-            .then_ignore(new_line().repeated())
-            .delimited_by(ctrl('('), ctrl(')'))
-            .recover_with(nested_delimiters(
-                TokenKind::Control('('),
-                TokenKind::Control(')'),
-                [
-                    (TokenKind::Control('{'), TokenKind::Control('}')),
-                    (TokenKind::Control('('), TokenKind::Control(')')),
-                    (TokenKind::Control('['), TokenKind::Control(']')),
-                ],
-                |_| vec![],
-            ))
+        let enum_ = keyword("enum")
+            .ignore_then(
+                ident_part()
+                    .then(ctrl('=').ignore_then(nested_type_expr.clone()).or_not())
+                    .map(|(name, ty)| {
+                        (
+                            Some(name),
+                            ty.unwrap_or_else(|| Ty::new(TyKind::Tuple(vec![]))),
+                        )
+                    })
+                    .padded_by(new_line().repeated())
+                    .separated_by(ctrl(','))
+                    .allow_trailing()
+                    .then_ignore(new_line().repeated())
+                    .delimited_by(ctrl('{'), ctrl('}'))
+                    .recover_with(nested_delimiters(
+                        TokenKind::Control('{'),
+                        TokenKind::Control('}'),
+                        [
+                            (TokenKind::Control('{'), TokenKind::Control('}')),
+                            (TokenKind::Control('('), TokenKind::Control(')')),
+                            (TokenKind::Control('['), TokenKind::Control(']')),
+                        ],
+                        |_| vec![],
+                    )),
+            )
             .map(TyKind::Union)
             .labelled("union");
 
@@ -120,7 +127,7 @@ pub fn type_expr() -> impl Parser<TokenKind, Ty, Error = PError> {
             .map(TyKind::Array)
             .labelled("array");
 
-        let term = choice((basic, ident, func, tuple, array, union_parenthesized))
+        let term = choice((basic, ident, func, tuple, array, enum_))
             .map_with_span(into_ty)
             .boxed();
 
diff --git a/prqlc/prqlc/tests/integration/resolving.rs b/prqlc/prqlc/tests/integration/resolving.rs
index be1a36291534..00ae1b759d60 100644
--- a/prqlc/prqlc/tests/integration/resolving.rs
+++ b/prqlc/prqlc/tests/integration/resolving.rs
@@ -59,9 +59,9 @@ fn resolve_types_01() {
 #[test]
 fn resolve_types_02() {
     assert_snapshot!(resolve(r#"
-    type A = int || ()
+    type A = int || {}
     "#).unwrap(), @r###"
-    type A = int
+    type A = int || {}
     "###)
 }
 
@@ -78,16 +78,15 @@ fn resolve_types_03() {
 fn resolve_types_04() {
     assert_snapshot!(resolve(
         r#"
-    type Status = (
-        Paid = () ||
-        Unpaid = float ||
-        Canceled = {reason = text, cancelled_at = timestamp} ||
-    )
+    type Status = enum {
+        Paid = {},
+        Unpaid = float,
+        Canceled = {reason = text, cancelled_at = timestamp},
+    }
     "#,
     )
     .unwrap(), @r###"
     type Status = (
-      Paid = () ||
       Unpaid = float ||
       {reason = text, cancelled_at = timestamp} ||
     )