Clean-up and some more minor improvements

partiql · Jul 31, 2024 · 63b0023 · 63b0023
1 parent 88f61bd
commit 63b0023
Show file tree

Hide file tree

Showing 12 changed files with 206 additions and 223 deletions.
diff --git a/partiql-ast/src/pretty.rs b/partiql-ast/src/pretty.rs
@@ -213,14 +213,24 @@ impl PrettyDoc for Select {
         D::Doc: Clone,
         A: Clone,
     {
+        fn format<'b, C, D, A>(child: &'b C, arena: &'b D) -> DocBuilder<'b, D, A>
+        where
+            D: DocAllocator<'b, A>,
+            D::Doc: Clone,
+            A: Clone,
+            C: PrettyDoc,
+        {
+            child.pretty_doc(arena).group()
+        }
+
         fn delegate<'b, C, D, A>(child: &'b Option<C>, arena: &'b D) -> Option<DocBuilder<'b, D, A>>
         where
             D: DocAllocator<'b, A>,
             D::Doc: Clone,
             A: Clone,
             C: PrettyDoc,
         {
-            child.as_ref().map(|inner| inner.pretty_doc(arena).group())
+            child.as_ref().map(|inner| format(inner, arena))
         }
 
         let Select {
@@ -232,25 +242,32 @@ impl PrettyDoc for Select {
             group_by,
             having,
         } = self;
-        let clauses = [
-            Some(project.pretty_doc(arena).group()),
+        let mut clauses = [
+            Some(format(project, arena)),
             delegate(exclude, arena),
-            from.as_ref().map(|inner| inner.pretty_doc(arena).group()),
-            from_let
-                .as_ref()
-                .map(|inner| inner.pretty_doc(arena).group()),
-            where_clause
-                .as_ref()
-                .map(|inner| inner.pretty_doc(arena).group()),
-            group_by
-                .as_ref()
-                .map(|inner| inner.pretty_doc(arena).group()),
-            having.as_ref().map(|inner| inner.pretty_doc(arena).group()),
+            delegate(from, arena),
+            delegate(from_let, arena),
+            delegate(where_clause, arena),
+            delegate(group_by, arena),
+            delegate(having, arena),
         ]
         .into_iter()
         .flatten();
 
-        arena.intersperse(clauses, arena.softline()).group()
+        let mut result = arena.nil();
+        let separator = arena.line();
+        if let Some(first) = clauses.next() {
+            let mut curr = first;
+
+            for clause in clauses {
+                result = result.append(curr.append(separator.clone()).group());
+                curr = clause;
+            }
+
+            result = result.append(curr);
+        }
+
+        result
     }
 }
 
@@ -401,6 +418,7 @@ impl PrettyDoc for Expr {
                 unreachable!();
             }
         }
+        .group()
     }
 }
 
@@ -576,8 +594,6 @@ impl PrettyDoc for BinOp {
             arena.softline()
         };
         let expr = arena.intersperse([lhs, op, rhs], sep).group();
-        //let paren_expr = [arena.text("("), expr, arena.text(")")];
-        //arena.concat(paren_expr).group()
         pretty_parenthesized_doc(expr, arena).group()
     }
 }
@@ -704,24 +720,7 @@ impl PrettyDoc for SimpleCase {
         } = self;
 
         let search = expr.pretty_doc(arena);
-
-        let branches = cases
-            .iter()
-            .map(|ExprPair { first, second }| {
-                let kw_when = arena.text("WHEN");
-                let test = first.pretty_doc(arena);
-                let kw_then = arena.text("THEN");
-                let then = second.pretty_doc(arena);
-                arena
-                    .intersperse([kw_when, test, kw_then, then], arena.space())
-                    .group()
-            })
-            .chain(
-                default
-                    .iter()
-                    .map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)).group()),
-            );
-
+        let branches = case_branches(arena, cases, default);
         pretty_seq_doc(branches, "CASE", Some(search), "END", " ", arena)
     }
 }
@@ -735,27 +734,39 @@ impl PrettyDoc for SearchedCase {
     {
         let SearchedCase { cases, default } = self;
 
-        let branches = cases
-            .iter()
-            .map(|ExprPair { first, second }| {
-                let kw_when = arena.text("WHEN");
-                let test = first.pretty_doc(arena);
-                let kw_then = arena.text("THEN");
-                let then = second.pretty_doc(arena);
-                arena
-                    .intersperse([kw_when, test, kw_then, then], arena.space())
-                    .group()
-            })
-            .chain(
-                default
-                    .iter()
-                    .map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)).group()),
-            );
-
+        let branches = case_branches(arena, cases, default);
         pretty_seq_doc(branches, "CASE", None, "END", " ", arena)
     }
 }
 
+fn case_branches<'b, D, A>(
+    arena: &'b D,
+    cases: &'b Vec<ExprPair>,
+    default: &'b Option<Box<Expr>>,
+) -> impl Iterator<Item = DocBuilder<'b, D, A>>
+where
+    D: DocAllocator<'b, A>,
+    D::Doc: Clone,
+    A: Clone + 'b,
+{
+    cases
+        .iter()
+        .map(|ExprPair { first, second }| {
+            let kw_when = arena.text("WHEN");
+            let test = first.pretty_doc(arena);
+            let kw_then = arena.text("THEN");
+            let then = second.pretty_doc(arena);
+            arena
+                .intersperse([kw_when, test, kw_then, then], arena.space())
+                .group()
+        })
+        .chain(
+            default
+                .iter()
+                .map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)).group()),
+        )
+}
+
 impl PrettyDoc for Struct {
     fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
     where
@@ -1280,10 +1291,13 @@ where
     } else {
         start
     };
-    let body = arena.line().append(arena.intersperse(seq, sep)).group();
+    let body = arena
+        .line()
+        .append(arena.intersperse(seq, sep))
+        .append(arena.line())
+        .group();
     start
         .append(body.nest(MINOR_NEST_INDENT))
-        .append(arena.line())
         .append(end)
         .group()
 }

diff --git a/partiql/tests/snapshots/pretty__pretty.snap b/partiql/tests/snapshots/pretty__pretty.snap
@@ -1,5 +1,5 @@
 ---
-source: partiql-ast/tests/pretty.rs
+source: partiql/tests/pretty.rs
 expression: doc
 ---
 ========================================================================================================================================================================================================
@@ -24,7 +24,8 @@ SELECT foo, bar, baz, thud.*, grunt.a[*].b[2].*, count(1) AS n FROM <<
     { 'foo': 'baz' },
     { 'foo': 'bar' },
     { 'foo': 'baz' }
-  >> GROUP BY foo ORDER BY n DESC
+  >>
+GROUP BY foo ORDER BY n DESC
 
 ------------------------------------------------------------------------------------------------------------------------
 SELECT foo, bar, baz, thud.*, grunt.a[*].b[2].*, count(1) AS n FROM <<
@@ -35,7 +36,8 @@ SELECT foo, bar, baz, thud.*, grunt.a[*].b[2].*, count(1) AS n FROM <<
     { 'foo': 'baz' },
     { 'foo': 'bar' },
     { 'foo': 'baz' }
-  >> GROUP BY foo ORDER BY n DESC
+  >>
+GROUP BY foo ORDER BY n DESC
 
 --------------------------------------------------------------------------------
 SELECT foo, bar, baz, thud.*, grunt.a[*].b[2].*, count(1) AS n FROM <<
@@ -46,7 +48,8 @@ SELECT foo, bar, baz, thud.*, grunt.a[*].b[2].*, count(1) AS n FROM <<
     { 'foo': 'baz' },
     { 'foo': 'bar' },
     { 'foo': 'baz' }
-  >> GROUP BY foo ORDER BY n DESC
+  >>
+GROUP BY foo ORDER BY n DESC
 
 ----------------------------------------
 SELECT foo, bar, baz, thud.*,
@@ -64,12 +67,14 @@ FROM <<
     { 'foo': 'baz' },
     { 'foo': 'bar' },
     { 'foo': 'baz' }
-  >> GROUP BY foo ORDER BY n DESC
+  >>
+GROUP BY foo ORDER BY n DESC
 
 ------------------------------
 SELECT foo, bar, baz, thud.*,
   grunt.a[*].b[2].*,
-  count(1) AS n FROM <<
+  count(1) AS n
+FROM <<
     {
       'foo': 'foo',
       'x': 9,
@@ -82,8 +87,8 @@ SELECT foo, bar, baz, thud.*,
     { 'foo': 'baz' },
     { 'foo': 'bar' },
     { 'foo': 'baz' }
-  >> GROUP BY foo
-ORDER BY n DESC
+  >>
+GROUP BY foo ORDER BY n DESC
 
 --------------------
 SELECT foo, bar,
@@ -97,18 +102,24 @@ FROM <<
       'y': 5,
       z: -11
     },
-    { 'foo': 'bar'
+    {
+      'foo': 'bar'
     },
-    { 'foo': 'qux'
+    {
+      'foo': 'qux'
     },
-    { 'foo': 'bar'
+    {
+      'foo': 'bar'
     },
-    { 'foo': 'baz'
+    {
+      'foo': 'baz'
     },
-    { 'foo': 'bar'
+    {
+      'foo': 'bar'
     },
     { 'foo': 'baz' }
-  >> GROUP BY foo
+  >>
+GROUP BY foo
 ORDER BY n DESC
 
 ----------

diff --git a/partiql/tests/snapshots/pretty__pretty2.snap b/partiql/tests/snapshots/pretty__pretty2.snap
@@ -16,21 +16,23 @@ ORDER BY n DESC
 
 ------------------------------------------------------------------------------------------------------------------------
 SELECT foo, bar, baz, thud, grunt, count(1) AS n FROM (SELECT * FROM table1) WHERE ((bar BETWEEN 3 AND 25 AND
-    NOT (baz LIKE 'example%')) OR (foo.a.b[*] IS MISSING)) GROUP BY foo ORDER BY n DESC
+    NOT (baz LIKE 'example%')) OR (foo.a.b[*] IS MISSING))
+GROUP BY foo ORDER BY n DESC
 
 --------------------------------------------------------------------------------
 SELECT foo, bar, baz, thud, grunt, count(1) AS n FROM (SELECT * FROM table1)
 WHERE ((bar BETWEEN 3 AND 25 AND NOT (baz LIKE 'example%')) OR
-  (foo.a.b[*] IS MISSING)) GROUP BY foo ORDER BY n DESC
+  (foo.a.b[*] IS MISSING))
+GROUP BY foo ORDER BY n DESC
 
 ----------------------------------------
 SELECT foo, bar, baz, thud, grunt,
-  count(1) AS n FROM (SELECT *
-        FROM table1)
+  count(1) AS n
+FROM (SELECT * FROM table1)
 WHERE ((bar BETWEEN 3 AND 25 AND
     NOT (baz LIKE 'example%')) OR
-  (foo.a.b[*] IS MISSING)) GROUP BY foo
-ORDER BY n DESC
+  (foo.a.b[*] IS MISSING))
+GROUP BY foo ORDER BY n DESC
 
 ------------------------------
 SELECT foo, bar, baz, thud,

diff --git a/partiql/tests/snapshots/pretty__pretty_case_1.snap b/partiql/tests/snapshots/pretty__pretty_case_1.snap
@@ -23,15 +23,17 @@ SELECT VALUE CASE
     WHEN ((x + 1) = f) THEN 'TWO' 
     WHEN (((x + 1) > d) AND ((x + 1) < 100)) THEN '>= THREE < 100' 
     ELSE '?'
-  END FROM << -1.0000, i, f, d, 100, NULL, MISSING >> AS x
+  END
+FROM << -1.0000, i, f, d, 100, NULL, MISSING >> AS x
 
 --------------------------------------------------------------------------------
 SELECT VALUE CASE
     WHEN ((x + 1) < i) THEN '< ONE' 
     WHEN ((x + 1) = f) THEN 'TWO' 
     WHEN (((x + 1) > d) AND ((x + 1) < 100)) THEN '>= THREE < 100' 
     ELSE '?'
-  END FROM << -1.0000, i, f, d, 100, NULL, MISSING >> AS x
+  END
+FROM << -1.0000, i, f, d, 100, NULL, MISSING >> AS x
 
 ----------------------------------------
 SELECT VALUE CASE
@@ -40,7 +42,8 @@ SELECT VALUE CASE
     WHEN (((x + 1) > d) AND
     ((x + 1) < 100)) THEN '>= THREE < 100' 
     ELSE '?'
-  END FROM <<
+  END
+FROM <<
     -1.0000,
     i,
     f,
@@ -57,7 +60,8 @@ SELECT VALUE CASE
     WHEN (((x + 1) > d) AND
     ((x + 1) < 100)) THEN '>= THREE < 100' 
     ELSE '?'
-  END FROM <<
+  END
+FROM <<
     -1.0000,
     i,
     f,
@@ -75,7 +79,8 @@ SELECT VALUE CASE
     AND
     ((x + 1) < 100)) THEN '>= THREE < 100' 
     ELSE '?'
-  END FROM <<
+  END
+FROM <<
     -1.0000,
     i,
     f,