From a264233d4c06f5fd865c57ff2f91a4c540e52b90 Mon Sep 17 00:00:00 2001 From: Josh Pschorr Date: Mon, 8 Jul 2024 12:35:05 -0700 Subject: [PATCH] Initial partial AST pretty-printer --- CHANGELOG.md | 5 + .../src/ast_to_dot.rs | 5 +- partiql-ast-passes/src/name_resolver.rs | 3 +- partiql-ast/Cargo.toml | 14 +- partiql-ast/src/ast.rs | 9 +- partiql-ast/src/lib.rs | 3 + partiql-ast/src/pretty.rs | 1003 +++++++++++++++++ partiql-ast/tests/common.rs | 50 + partiql-logical-planner/src/lower.rs | 19 +- partiql-parser/src/parse/partiql.lalrpop | 37 +- 10 files changed, 1102 insertions(+), 46 deletions(-) create mode 100644 partiql-ast/src/pretty.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index d04b324e..2f96f2c8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,10 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed +- *BREAKING:* partiql-ast: changed modeling of `GroupByExpr` `strategy` field to be an `Option` +- *BREAKING:* partiql-ast: changed modeling of `PathStep` to split `PathExpr` to `PathIndex` (e.g., `[2]`) and `PathProject` (e.g., `.a`) +- *BREAKING:* partiql-ast: changed modeling of `PathStep` to rename `PathWildcard` to `PathForEach` (for `[*]`) ### Added +- partiql-ast: Pretty-printing of AST ### Fixed +- Minor documentation issues ## [0.8.0] ### Changed diff --git a/extension/partiql-extension-visualize/src/ast_to_dot.rs b/extension/partiql-extension-visualize/src/ast_to_dot.rs index 3d80a3fe..3c641eb9 100644 --- a/extension/partiql-extension-visualize/src/ast_to_dot.rs +++ b/extension/partiql-extension-visualize/src/ast_to_dot.rs @@ -631,8 +631,9 @@ impl ToDot for AstToDot { impl ToDot for AstToDot { fn to_dot(&mut self, out: &mut Scope<'_, '_>, ast: &ast::PathStep) -> Targets { match &ast { - ast::PathStep::PathExpr(e) => self.to_dot(out, e), - ast::PathStep::PathWildCard => vec![out.node_auto_labelled("*").id()], + ast::PathStep::PathProject(e) => self.to_dot(out, e), + ast::PathStep::PathIndex(e) => self.to_dot(out, e), + ast::PathStep::PathForEach => vec![out.node_auto_labelled("*").id()], ast::PathStep::PathUnpivot => vec![out.node_auto_labelled("Unpivot").id()], } } diff --git a/partiql-ast-passes/src/name_resolver.rs b/partiql-ast-passes/src/name_resolver.rs index 5ce5b09c..e6994668 100644 --- a/partiql-ast-passes/src/name_resolver.rs +++ b/partiql-ast-passes/src/name_resolver.rs @@ -514,7 +514,8 @@ fn infer_alias(expr: &ast::Expr) -> Option { match expr { ast::Expr::VarRef(ast::AstNode { node, .. }) => Some(node.name.clone()), ast::Expr::Path(ast::AstNode { node, .. }) => match node.steps.last() { - Some(ast::PathStep::PathExpr(expr)) => infer_alias(&expr.index), + Some(ast::PathStep::PathProject(expr)) => infer_alias(&expr.index), + Some(ast::PathStep::PathIndex(expr)) => infer_alias(&expr.index), _ => None, }, _ => None, diff --git a/partiql-ast/Cargo.toml b/partiql-ast/Cargo.toml index 7e0828f7..c7416968 100644 --- a/partiql-ast/Cargo.toml +++ b/partiql-ast/Cargo.toml @@ -9,8 +9,8 @@ readme = "../README.md" keywords = ["sql", "ast", "query", "compilers", "interpreters"] categories = ["database", "compilers", "ast-implementations"] exclude = [ - "**/.git/**", - "**/.github/**", + "**/.git/**", + "**/.github/**", ] version.workspace = true edition.workspace = true @@ -23,16 +23,18 @@ bench = false indexmap = { version = "1.9", default-features = false } rust_decimal = { version = "1.25.0", default-features = false, features = ["std"] } serde = { version = "1.*", features = ["derive"], optional = true } +pretty = "0.12" [dev-dependencies] +partiql-parser = { path = "../partiql-parser", version = "0.8" } [features] default = [] serde = [ - "dep:serde", - "rust_decimal/serde-with-str", - "rust_decimal/serde", - "indexmap/serde", + "dep:serde", + "rust_decimal/serde-with-str", + "rust_decimal/serde", + "indexmap/serde", ] [dependencies.partiql-ast-macros] diff --git a/partiql-ast/src/ast.rs b/partiql-ast/src/ast.rs index 7b712a4e..2d5b1860 100644 --- a/partiql-ast/src/ast.rs +++ b/partiql-ast/src/ast.rs @@ -644,7 +644,7 @@ pub enum CallArg { PositionalType(Type), /// named argument to a function call (e.g., the `"from" : 2` in `substring(a, "from":2)` Named(CallArgNamed), - /// E.g. `AS: VARCHAR` in `CAST('abc' AS VARCHAR` + /// E.g. `AS: VARCHAR` in `CAST('abc' AS VARCHAR)` NamedType(CallArgNamedType), } @@ -676,9 +676,10 @@ pub struct Path { #[derive(Visit, Clone, Debug, PartialEq)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub enum PathStep { - PathExpr(PathExpr), + PathProject(PathExpr), + PathIndex(PathExpr), #[visit(skip)] - PathWildCard, + PathForEach, #[visit(skip)] PathUnpivot, } @@ -789,7 +790,7 @@ pub enum JoinSpec { #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct GroupByExpr { #[visit(skip)] - pub strategy: GroupingStrategy, + pub strategy: Option, pub keys: Vec>, #[visit(skip)] pub group_as_alias: Option, diff --git a/partiql-ast/src/lib.rs b/partiql-ast/src/lib.rs index 6fb8ca12..30f517da 100644 --- a/partiql-ast/src/lib.rs +++ b/partiql-ast/src/lib.rs @@ -8,4 +8,7 @@ //! This API is currently unstable and subject to change. pub mod ast; + +pub mod pretty; + pub mod visit; diff --git a/partiql-ast/src/pretty.rs b/partiql-ast/src/pretty.rs new file mode 100644 index 00000000..6f63824e --- /dev/null +++ b/partiql-ast/src/pretty.rs @@ -0,0 +1,1003 @@ +use crate::ast::*; +use pretty::{Arena, BoxAllocator, DocAllocator, DocBuilder, IoWrite, Pretty, RcDoc}; +use std::io; +use std::io::Write; +const MINOR_NEST_INDENT: isize = 2; +const SUBQUERY_INDENT: isize = 6; + +pub(crate) trait PrettyDoc { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone; +} + +pub trait ToPretty { + fn to_pretty_string(&self, width: usize) -> io::Result { + let mut out = Vec::new(); + self.to_pretty(width, &mut out)?; + // TODO error instead of unwrap + Ok(String::from_utf8(out).unwrap()) + } + + /// Writes a rendered document to a `std::io::Write` object. + #[inline] + fn to_pretty(&self, width: usize, out: &mut W) -> io::Result<()> + where + W: ?Sized + io::Write; +} + +impl ToPretty for AstNode +where + T: PrettyDoc, +{ + fn to_pretty(&self, width: usize, out: &mut W) -> io::Result<()> + where + W: ?Sized + Write, + { + let allocator = Arena::new(); + let DocBuilder(alloc, doc) = self.node.pretty_doc::<_, ()>(&allocator); + doc.render(width, out) + } +} + +impl PrettyDoc for AstNode +where + T: PrettyDoc, +{ + #[inline] + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + self.node.pretty_doc(allocator) + } +} + +impl PrettyDoc for Box +where + T: PrettyDoc, +{ + #[inline] + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + self.as_ref().pretty_doc(allocator) + } +} + +impl PrettyDoc for str { + #[inline] + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + allocator.concat(["'", self, "'"]) + } +} + +impl PrettyDoc for rust_decimal::Decimal { + #[inline] + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + allocator.text(self.to_string()) + } +} + +impl PrettyDoc for TopLevelQuery { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + // TODO With + self.query.pretty_doc(allocator) + } +} + +impl PrettyDoc for Query { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let Query { + set, + order_by, + limit_offset, + } = self; + + let clauses = [ + Some(set.pretty_doc(allocator)), + order_by.as_ref().map(|inner| inner.pretty_doc(allocator)), + limit_offset + .as_ref() + .map(|inner| inner.pretty_doc(allocator)), + ] + .into_iter() + .filter_map(|x| x); + + allocator.intersperse(clauses, allocator.softline()).group() + } +} + +impl PrettyDoc for QuerySet { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + match self { + QuerySet::BagOp(_) => { + todo!() + } + QuerySet::Select(sel) => sel.pretty_doc(allocator), + QuerySet::Expr(_) => { + todo!() + } + QuerySet::Values(_) => { + todo!() + } + QuerySet::Table(_) => { + todo!() + } + } + } +} + +impl PrettyDoc for Select { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let Select { + project, + from, + from_let, + where_clause, + group_by, + having, + } = self; + let clauses = [ + Some(project.pretty_doc(allocator).group()), + from.as_ref() + .map(|inner| inner.pretty_doc(allocator).group()), + from_let + .as_ref() + .map(|inner| inner.pretty_doc(allocator).group()), + where_clause + .as_ref() + .map(|inner| inner.pretty_doc(allocator).group()), + group_by + .as_ref() + .map(|inner| inner.pretty_doc(allocator).group()), + having + .as_ref() + .map(|inner| inner.pretty_doc(allocator).group()), + ] + .into_iter() + .filter_map(|x| x); + + allocator.intersperse(clauses, allocator.softline()).group() + } +} + +impl PrettyDoc for Projection { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + if self.setq.is_some() { + todo!("project SetQuantifier") + } + self.kind.pretty_doc(allocator) + } +} + +impl PrettyDoc for ProjectionKind { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + match self { + ProjectionKind::ProjectStar => allocator.text("SELECT *"), + ProjectionKind::ProjectList(l) => { + let list = pretty_list(l, allocator); + allocator + .text("SELECT") + .append(allocator.softline()) + .append(list) + } + + ProjectionKind::ProjectPivot(_) => { + todo!() + } + ProjectionKind::ProjectValue(_) => { + todo!() + } + } + .group() + } +} + +impl PrettyDoc for ProjectItem { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + match self { + ProjectItem::ProjectAll(_) => { + todo!() + } + ProjectItem::ProjectExpr(e) => e.pretty_doc(allocator), + } + } +} + +impl PrettyDoc for ProjectExpr { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + pretty_source_as_alias(&self.expr, self.as_alias.as_ref(), allocator) + .unwrap_or_else(|| self.expr.pretty_doc(allocator)) + } +} + +impl PrettyDoc for Expr { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + match self { + Expr::Lit(inner) => inner.pretty_doc(allocator), + Expr::VarRef(inner) => inner.pretty_doc(allocator), + Expr::BinOp(inner) => inner.pretty_doc(allocator), + Expr::UniOp(inner) => inner.pretty_doc(allocator), + Expr::Like(inner) => inner.pretty_doc(allocator), + Expr::Between(inner) => inner.pretty_doc(allocator), + Expr::In(_) => { + todo!() + } + Expr::Case(_) => { + todo!() + } + Expr::Struct(inner) => inner.pretty_doc(allocator), + Expr::Bag(inner) => inner.pretty_doc(allocator), + Expr::List(inner) => inner.pretty_doc(allocator), + Expr::Sexp(inner) => inner.pretty_doc(allocator), + Expr::Path(inner) => inner.pretty_doc(allocator), + Expr::Call(inner) => inner.pretty_doc(allocator), + + Expr::CallAgg(inner) => inner.pretty_doc(allocator), + + Expr::Query(inner) => { + let inner = inner.pretty_doc(allocator).group(); + allocator + .text("(") + .append(inner.nest(SUBQUERY_INDENT)) + .append(allocator.text(")")) + } + Expr::Error => { + unreachable!(); + } + } + } +} + +impl PrettyDoc for Path { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let Path { root, steps } = self; + let mut path = root.pretty_doc(allocator); + for step in steps { + path = path.append(match step { + PathStep::PathProject(e) => { + allocator.text(".").append(e.index.pretty_doc(allocator)) + } + PathStep::PathIndex(e) => allocator + .text("[") + .append(e.index.pretty_doc(allocator)) + .append(allocator.text("]")), + PathStep::PathForEach => allocator.text("[*]"), + PathStep::PathUnpivot => allocator.text(".*"), + }); + } + + path + } +} + +impl PrettyDoc for VarRef { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let name = self.name.pretty_doc(allocator); + match self.qualifier { + ScopeQualifier::Unqualified => name, + ScopeQualifier::Qualified => allocator.text("@").append(name).group(), + } + } +} + +impl PrettyDoc for Lit { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + match self { + Lit::Null => allocator.text("NULL"), + Lit::Missing => allocator.text("MISSING"), + Lit::Int8Lit(inner) => allocator.text(inner.to_string()), + Lit::Int16Lit(inner) => allocator.text(inner.to_string()), + Lit::Int32Lit(inner) => allocator.text(inner.to_string()), + Lit::Int64Lit(inner) => allocator.text(inner.to_string()), + Lit::DecimalLit(inner) => inner.pretty_doc(allocator), + Lit::NumericLit(inner) => inner.pretty_doc(allocator), + Lit::RealLit(inner) => allocator.text(inner.to_string()), + Lit::FloatLit(inner) => allocator.text(inner.to_string()), + Lit::DoubleLit(inner) => allocator.text(inner.to_string()), + Lit::BoolLit(inner) => allocator.text(inner.to_string()), + Lit::IonStringLit(inner) => inner.pretty_doc(allocator), + Lit::CharStringLit(inner) => inner.pretty_doc(allocator), + Lit::NationalCharStringLit(inner) => inner.pretty_doc(allocator), + Lit::BitStringLit(inner) => inner.pretty_doc(allocator), + Lit::HexStringLit(inner) => inner.pretty_doc(allocator), + Lit::StructLit(inner) => inner.pretty_doc(allocator), + Lit::BagLit(inner) => inner.pretty_doc(allocator), + Lit::ListLit(inner) => inner.pretty_doc(allocator), + Lit::TypedLit(s, t) => { + todo!("typed literal"); + } + } + } +} + +impl PrettyDoc for BinOp { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let BinOp { kind, lhs, rhs } = self; + let (nest, sym) = match kind { + BinOpKind::Add => (0, "+"), + BinOpKind::Div => (0, "/"), + BinOpKind::Exp => (0, "^"), + BinOpKind::Mod => (0, "%"), + BinOpKind::Mul => (0, "*"), + BinOpKind::Sub => (0, "-"), + BinOpKind::And => (MINOR_NEST_INDENT, "AND"), + BinOpKind::Or => (MINOR_NEST_INDENT, "OR"), + BinOpKind::Concat => (0, "||"), + BinOpKind::Eq => (0, "="), + BinOpKind::Gt => (0, ">"), + BinOpKind::Gte => (0, ">="), + BinOpKind::Lt => (0, "<"), + BinOpKind::Lte => (0, "<="), + BinOpKind::Ne => (0, "<>"), + BinOpKind::Is => (0, "IS"), + }; + let op = allocator.text(sym); + let lhs = lhs.pretty_doc(allocator).nest(nest); + let rhs = rhs.pretty_doc(allocator).nest(nest); + let sep = allocator.space(); + let expr = allocator.intersperse([lhs, op, rhs], sep).group(); + let paren_expr = [allocator.text("("), expr, allocator.text(")")]; + allocator.concat(paren_expr).group() + } +} + +impl PrettyDoc for UniOp { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + // TODO NOT LIKE, NOT IN, NOT BETWEEN? + let UniOp { kind, expr } = self; + let (sym, paren) = match kind { + UniOpKind::Pos => ("+", false), + UniOpKind::Neg => ("-", false), + UniOpKind::Not => ("NOT ", true), + }; + let op = allocator.text(sym); + let expr = expr.pretty_doc(allocator); + if paren { + let open = allocator.text("("); + let close = allocator.text(")"); + allocator.concat([op, open, expr, close]).group() + } else { + allocator.concat([op, expr]).group() + } + } +} + +impl PrettyDoc for Like { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let Like { + value, + pattern, + escape, + } = self; + + let sep = allocator.space(); + let value = value.pretty_doc(allocator); + let kw_like = allocator.text("LIKE"); + let pattern = pattern.pretty_doc(allocator); + if let Some(escape) = escape { + let kw_esc = allocator.text("ESCAPE"); + let escape = escape.pretty_doc(allocator); + allocator.intersperse([value, kw_like, pattern, kw_esc, escape], sep) + } else { + allocator.intersperse([value, kw_like, pattern], sep) + } + .group() + } +} + +impl PrettyDoc for Between { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let Between { value, from, to } = self; + + let value = value.pretty_doc(allocator); + let kw_b = allocator.text("BETWEEN"); + let kw_a = allocator.text("AND"); + let from = from.pretty_doc(allocator); + let to = to.pretty_doc(allocator); + let sep = allocator.space(); + let expr = allocator + .intersperse([value, kw_b, from, kw_a, to], sep) + .group(); + expr.group() + } +} + +impl PrettyDoc for Struct { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let wrapped = self.fields.iter().map(|p| unsafe { + let x: &'b StructExprPair = std::mem::transmute(p); + x + }); + pretty_seq(wrapped, "{", "}", ",", allocator) + } +} + +pub struct StructExprPair(pub ExprPair); + +impl PrettyDoc for StructExprPair { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let k = self.0.first.pretty_doc(allocator); + let v = self.0.second.pretty_doc(allocator); + let sep = allocator.text(": "); + + k.append(sep).group().append(v).group() + } +} + +impl PrettyDoc for Bag { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + pretty_seq(&self.values, "<<", ">>", ",", allocator) + } +} + +impl PrettyDoc for List { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + pretty_seq(&self.values, "[", "]", ",", allocator) + } +} + +impl PrettyDoc for Sexp { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + todo!("remove s-expr from ast?"); + } +} + +impl PrettyDoc for Call { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let name = self.func_name.pretty_doc(allocator); + let list = pretty_list(&self.args, allocator); + name.append(allocator.text("(")) + .append(list.nest(MINOR_NEST_INDENT)) + .append(allocator.text(")")) + } +} + +impl PrettyDoc for CallAgg { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let name = self.func_name.pretty_doc(allocator); + let list = pretty_list(&self.args, allocator); + name.append(allocator.text("(")) + .append(list.nest(MINOR_NEST_INDENT)) + .append(allocator.text(")")) + } +} + +impl PrettyDoc for CallArg { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + match self { + CallArg::Star() => allocator.text("*"), + CallArg::Positional(arg) => arg.pretty_doc(allocator), + CallArg::PositionalType(_) => { + todo!() + } + CallArg::Named(_) => { + todo!() + } + CallArg::NamedType(_) => { + todo!() + } + } + } +} + +impl PrettyDoc for SymbolPrimitive { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let sym = allocator.text(self.value.as_str()); + match self.case { + CaseSensitivity::CaseSensitive => allocator + .text("\"") + .append(sym) + .append(allocator.text("\"")), + CaseSensitivity::CaseInsensitive => sym, + } + } +} + +impl PrettyDoc for FromClause { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + self.source.pretty_doc(allocator) + } +} + +impl PrettyDoc for FromSource { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let doc = match self { + FromSource::FromLet(fl) => fl.pretty_doc(allocator), + FromSource::Join(_) => { + todo!() + } + }; + + allocator + .text("FROM") + .append(allocator.space()) + .append(doc) + .group() + } +} + +impl PrettyDoc for FromLet { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let FromLet { + expr, + kind, + as_alias, + at_alias, + by_alias, + } = self; + + let expr = expr.pretty_doc(allocator); + let as_alias = pretty_as_alias(as_alias.as_ref(), allocator); + let at_alias = pretty_at_alias(at_alias.as_ref(), allocator); + let by_alias = pretty_by_alias(by_alias.as_ref(), allocator); + let aliases: Vec<_> = [as_alias, at_alias, by_alias] + .into_iter() + .filter_map(|x| x) + .collect(); + + let clause = match kind { + FromLetKind::Scan => expr, + FromLetKind::Unpivot => allocator + .text("UNPIVOT") + .append(allocator.space()) + .append(expr), + }; + + if aliases.is_empty() { + clause + } else { + let spec = allocator.intersperse(aliases, allocator.text("#")).group(); + clause.append(allocator.space()).append(spec) + } + .group() + } +} + +impl PrettyDoc for Let { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + todo!() + } +} + +impl PrettyDoc for WhereClause { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let expr = self.expr.pretty_doc(allocator).nest(MINOR_NEST_INDENT); + allocator + .text("WHERE") + .append(allocator.space()) + .append(expr) + .group() + } +} + +impl PrettyDoc for GroupByExpr { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let GroupByExpr { + strategy, + keys, + group_as_alias, + } = self; + + let mut doc = match strategy { + None => allocator.text("GROUP"), + Some(GroupingStrategy::GroupFull) => allocator.text("GROUP ALL"), + Some(GroupingStrategy::GroupPartial) => allocator.text("GROUP PARTIAL"), + }; + + if !keys.is_empty() { + doc = doc + .append(allocator.space()) + .append(allocator.text("BY")) + .group(); + doc = doc + .append(allocator.softline()) + .append(pretty_list(keys, allocator)); + } + + match group_as_alias { + None => doc, + Some(gas) => { + let gas = pretty_source_as_alias("GROUP", Some(gas), allocator); + doc.append(gas) + } + } + .group() + } +} + +impl PrettyDoc for GroupKey { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + pretty_source_as_alias(&self.expr, self.as_alias.as_ref(), allocator) + .unwrap_or_else(|| self.expr.pretty_doc(allocator)) + } +} + +impl PrettyDoc for HavingClause { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + todo!() + } +} + +impl PrettyDoc for OrderByExpr { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + if self.sort_specs.is_empty() { + allocator.text("ORDER BY PRESERVE") + } else { + allocator + .text("ORDER BY") + .append(allocator.space()) + .append(pretty_list(&self.sort_specs, allocator)) + } + .group() + } +} + +impl PrettyDoc for SortSpec { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + let SortSpec { + expr, + ordering_spec, + null_ordering_spec, + } = self; + let mut doc = expr.pretty_doc(allocator); + if let Some(os) = ordering_spec { + let os = allocator.space().append(os.pretty_doc(allocator)).group(); + doc = doc.append(os) + }; + if let Some(nos) = null_ordering_spec { + let nos = allocator.space().append(nos.pretty_doc(allocator)).group(); + doc = doc.append(nos) + }; + + doc.group() + } +} + +impl PrettyDoc for OrderingSpec { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + allocator.text(match self { + OrderingSpec::Asc => "ASC", + OrderingSpec::Desc => "DESC", + }) + } +} + +impl PrettyDoc for NullOrderingSpec { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + allocator.text(match self { + NullOrderingSpec::First => "NULLS FIRST", + NullOrderingSpec::Last => "NULLS LAST", + }) + } +} + +impl PrettyDoc for LimitOffsetClause { + fn pretty_doc<'b, D, A>(&'b self, allocator: &'b D) -> DocBuilder<'b, D, A> + where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, + { + allocator.text("LIMIT OFFSET") + } +} + +fn pretty_seq<'i, 'b, I, P, D, A>( + list: I, + start: &'static str, + end: &'static str, + sep: &'static str, + allocator: &'b D, +) -> DocBuilder<'b, D, A> +where + I: IntoIterator, + P: PrettyDoc + 'b, + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, +{ + let start = allocator.text(start); + let end = allocator.text(end); + let sep = allocator.text(sep).append(allocator.line()); + let seq = list.into_iter().map(|l| l.pretty_doc(allocator)); + let body = allocator + .line() + .append(allocator.intersperse(seq, sep)) + .group(); + start + .append(body.nest(MINOR_NEST_INDENT)) + .append(allocator.line()) + .append(end) + .group() +} + +fn pretty_list<'b, P, D, A>(list: &'b Vec

, allocator: &'b D) -> DocBuilder<'b, D, A> +where + P: PrettyDoc, + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, +{ + let sep = allocator.text(",").append(allocator.softline()); + allocator + .intersperse(list.iter().map(|l| l.pretty_doc(allocator)), sep) + .nest(MINOR_NEST_INDENT) + .group() +} + +fn pretty_alias_helper<'b, D, A>( + kw: &'static str, + sym: Option<&'b SymbolPrimitive>, + allocator: &'b D, +) -> Option> +where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, +{ + sym.map(|sym| { + allocator + .space() + .append(allocator.text(kw)) + .append(allocator.space()) + .append(sym.pretty_doc(allocator)) + .group() + }) +} + +fn pretty_source_as_alias<'b, S, D, A>( + source: &'b S, + as_alias: Option<&'b SymbolPrimitive>, + allocator: &'b D, +) -> Option> +where + S: PrettyDoc + ?Sized, + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, +{ + pretty_as_alias(as_alias, allocator).map(|alias| { + let expr = source.pretty_doc(allocator); + allocator.concat([expr, alias]).group() + }) +} + +fn pretty_as_alias<'b, D, A>( + sym: Option<&'b SymbolPrimitive>, + allocator: &'b D, +) -> Option> +where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, +{ + pretty_alias_helper("AS", sym, allocator) +} + +fn pretty_at_alias<'b, D, A>( + sym: Option<&'b SymbolPrimitive>, + allocator: &'b D, +) -> Option> +where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, +{ + pretty_alias_helper("AT", sym, allocator) +} + +fn pretty_by_alias<'b, D, A>( + sym: Option<&'b SymbolPrimitive>, + allocator: &'b D, +) -> Option> +where + D: DocAllocator<'b, A>, + D::Doc: Clone, + A: Clone, +{ + pretty_alias_helper("BY", sym, allocator) +} diff --git a/partiql-ast/tests/common.rs b/partiql-ast/tests/common.rs index 60a4158b..e729d5f5 100644 --- a/partiql-ast/tests/common.rs +++ b/partiql-ast/tests/common.rs @@ -1,3 +1,53 @@ +use partiql_ast::pretty::ToPretty; +use partiql_parser::ParserResult; + pub fn setup() { // setup test code goes here } + +#[track_caller] +#[inline] +fn parse(statement: &str) -> ParserResult<'_> { + partiql_parser::Parser::default().parse(statement) +} + +#[test] +fn pretty() { + let query = "select foo,bar, baz,thud.*,grunt.a[*].b[2].*, count(1) as n from + << + { 'foo': 'foo', 'x': 9, 'y':5, z:-11 }, + { 'foo': 'bar' }, + { 'foo': 'qux' }, + { 'foo': 'bar' }, + { 'foo': 'baz' }, + { 'foo': 'bar' }, + { 'foo': 'baz' } + >> group by foo order by n desc"; + + let res = parse(query); + assert!(res.is_ok()); + let res = res.unwrap(); + + for w in [180, 120, 80, 40, 30, 20, 10] { + println!("{:- Visitor<'ast> for AstToLogical<'a> { Traverse::Continue } - fn enter_path_step(&mut self, _path_step: &'ast PathStep) -> Traverse { - if let PathStep::PathExpr(expr) = _path_step { + fn enter_path_step(&mut self, path_step: &'ast PathStep) -> Traverse { + if let PathStep::PathIndex(expr) | PathStep::PathProject(expr) = path_step { self.enter_env(); match *(expr.index) { Expr::VarRef(_) => { @@ -1401,9 +1401,9 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> { Traverse::Continue } - fn exit_path_step(&mut self, _path_step: &'ast PathStep) -> Traverse { - let step = match _path_step { - PathStep::PathExpr(_s) => { + fn exit_path_step(&mut self, path_step: &'ast PathStep) -> Traverse { + let step = match path_step { + PathStep::PathProject(_s) | PathStep::PathIndex(_s) => { let mut env = self.exit_env(); eq_or_fault!(self, env.len(), 1, "env.len() != 1"); @@ -1425,8 +1425,8 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> { } } } - PathStep::PathWildCard => { - not_yet_implemented_fault!(self, "PathStep::PathWildCard".to_string()); + PathStep::PathForEach => { + not_yet_implemented_fault!(self, "PathStep::PathForEach".to_string()); } PathStep::PathUnpivot => { not_yet_implemented_fault!(self, "PathStep::PathUnpivot".to_string()); @@ -1627,8 +1627,9 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> { .map(|SymbolPrimitive { value, case: _ }| value.clone()); let strategy = match _group_by_expr.strategy { - GroupingStrategy::GroupFull => logical::GroupingStrategy::GroupFull, - GroupingStrategy::GroupPartial => logical::GroupingStrategy::GroupPartial, + None => logical::GroupingStrategy::GroupFull, + Some(GroupingStrategy::GroupFull) => logical::GroupingStrategy::GroupFull, + Some(GroupingStrategy::GroupPartial) => logical::GroupingStrategy::GroupPartial, }; // What follows is an approach to implement section 11.2.1 of the PartiQL spec diff --git a/partiql-parser/src/parse/partiql.lalrpop b/partiql-parser/src/parse/partiql.lalrpop index 65b350e9..d23828be 100644 --- a/partiql-parser/src/parse/partiql.lalrpop +++ b/partiql-parser/src/parse/partiql.lalrpop @@ -223,31 +223,26 @@ FwsClauses: ast::AstNode = { SelectClause: ast::AstNode = { "SELECT" "*" => state.node(ast::Projection { kind: ast::ProjectionKind::ProjectStar, - setq: Some(strategy) + setq: strategy, }, lo..hi), "SELECT" > => state.node(ast::Projection { kind: ast::ProjectionKind::ProjectList(project_items), - setq: Some(strategy), + setq: strategy, }, lo..hi), "SELECT" "VALUE" => state.node(ast::Projection { kind: ast::ProjectionKind::ProjectValue(value), - setq: Some(strategy), + setq: strategy, }, lo..hi), "PIVOT" "AT" => state.node(ast::Projection { kind: ast::ProjectionKind::ProjectPivot(ast::ProjectPivot { key, value }), - setq: None + setq: None, }, lo..hi), } #[inline] -SetQuantifierStrategy: ast::SetQuantifier = { - "ALL" => ast::SetQuantifier::All, - => { - match distinct { - Some(_) => ast::SetQuantifier::Distinct, - None => ast::SetQuantifier::All, - } - } +SetQuantifierStrategy: Option = { + "ALL" => Some(ast::SetQuantifier::All), + => distinct.map(|_|ast::SetQuantifier::Distinct) } #[inline] @@ -436,14 +431,9 @@ GroupClause: Box> = { } } #[inline] -GroupStrategy: ast::GroupingStrategy = { - "ALL" => ast::GroupingStrategy::GroupFull, - => { - match partial { - Some(_) => ast::GroupingStrategy::GroupPartial, - None => ast::GroupingStrategy::GroupFull, - } - } +GroupStrategy: Option = { + "ALL" => Some(ast::GroupingStrategy::GroupFull), + => partial.map(|_| ast::GroupingStrategy::GroupPartial), } #[inline] GroupByKeys: Vec> = { @@ -842,7 +832,6 @@ ExprPrecedence02: Synth = { PathExpr: ast::Path = { => { - let step = ast::PathStep::PathWildCard; ast::Path { root: Box::new(l.data), steps: s @@ -1109,16 +1098,16 @@ PathSteps: Vec = { } PathStep: ast::PathStep = { "." => { - ast::PathStep::PathExpr( ast::PathExpr{ index: Box::new(v) }) + ast::PathStep::PathProject( ast::PathExpr{ index: Box::new(v) }) }, "[" "*" "]" => { - ast::PathStep::PathWildCard + ast::PathStep::PathForEach }, "." "*" => { ast::PathStep::PathUnpivot }, "[" "]" => { - ast::PathStep::PathExpr( ast::PathExpr{ index: Box::new(*expr) }) + ast::PathStep::PathIndex( ast::PathExpr{ index: Box::new(*expr) }) }, }