From 27d90d174c78fdbee7841fb68a7d3f52851d3bcd Mon Sep 17 00:00:00 2001
From: Josh Pschorr <josh@pschorr.dev>
Date: Thu, 5 Dec 2024 15:59:05 -0800
Subject: [PATCH] Add type to `ast::Lit::EmbeddedDocLit` (#525)

---
 .../src/ast_to_dot.rs                         |  2 +-
 partiql-ast/src/ast.rs                        |  2 +-
 partiql-ast/src/pretty.rs                     | 31 ++++++++++++++++---
 partiql-logical-planner/src/lower.rs          |  2 +-
 partiql-parser/src/parse/partiql.lalrpop      | 10 +++++-
 5 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/extension/partiql-extension-visualize/src/ast_to_dot.rs b/extension/partiql-extension-visualize/src/ast_to_dot.rs
index 74dd1283..82d477dd 100644
--- a/extension/partiql-extension-visualize/src/ast_to_dot.rs
+++ b/extension/partiql-extension-visualize/src/ast_to_dot.rs
@@ -185,7 +185,7 @@ fn lit_to_str(ast: &ast::Lit) -> String {
         Lit::FloatLit(l) => l.to_string(),
         Lit::DoubleLit(l) => l.to_string(),
         Lit::BoolLit(l) => (if *l { "TRUE" } else { "FALSE" }).to_string(),
-        Lit::EmbeddedDocLit(l) => format!("`{}`", l),
+        Lit::EmbeddedDocLit(l, typ) => format!("`{}`::{}", l, type_to_str(typ)),
         Lit::CharStringLit(l) => format!("'{}'", l),
         Lit::NationalCharStringLit(l) => format!("'{}'", l),
         Lit::BitStringLit(l) => format!("b'{}'", l),
diff --git a/partiql-ast/src/ast.rs b/partiql-ast/src/ast.rs
index 25475279..3b641072 100644
--- a/partiql-ast/src/ast.rs
+++ b/partiql-ast/src/ast.rs
@@ -444,7 +444,7 @@ pub enum Lit {
     #[visit(skip)]
     BoolLit(bool),
     #[visit(skip)]
-    EmbeddedDocLit(String),
+    EmbeddedDocLit(String, Type),
     #[visit(skip)]
     CharStringLit(String),
     #[visit(skip)]
diff --git a/partiql-ast/src/pretty.rs b/partiql-ast/src/pretty.rs
index 3f1c818f..53b83be8 100644
--- a/partiql-ast/src/pretty.rs
+++ b/partiql-ast/src/pretty.rs
@@ -394,7 +394,7 @@ impl PrettyDoc for Lit {
             Lit::FloatLit(inner) => arena.text(inner.to_string()),
             Lit::DoubleLit(inner) => arena.text(inner.to_string()),
             Lit::BoolLit(inner) => arena.text(inner.to_string()),
-            Lit::EmbeddedDocLit(inner) => inner.pretty_doc(arena), // TODO better pretty for embedded doc: https://github.com/partiql/partiql-lang-rust/issues/508
+            Lit::EmbeddedDocLit(inner, _typ) => inner.pretty_doc(arena), // TODO better pretty for embedded doc: https://github.com/partiql/partiql-lang-rust/issues/508
             Lit::CharStringLit(inner) => inner.pretty_doc(arena),
             Lit::NationalCharStringLit(inner) => inner.pretty_doc(arena),
             Lit::BitStringLit(inner) => inner.pretty_doc(arena),
@@ -420,9 +420,32 @@ impl PrettyDoc for Type {
     {
         match self {
             Type::CustomType(cty) => cty.pretty_doc(arena),
-            _ => {
-                todo!("Non-custom type type")
-            }
+            Type::NullType => arena.text("NULL"),
+            Type::BooleanType => arena.text("BOOL"),
+            Type::Integer2Type => arena.text("INT2"),
+            Type::Integer4Type => arena.text("INT4"),
+            Type::Integer8Type => arena.text("INT8"),
+            Type::DecimalType => arena.text("DECIMAL"),
+            Type::NumericType => arena.text("NUMERIC"),
+            Type::RealType => arena.text("REAL"),
+            Type::DoublePrecisionType => arena.text("DOUBLE PRECISION"),
+            Type::TimestampType => arena.text("TIMESTAMP"),
+            Type::CharacterType => arena.text("CHAR"),
+            Type::CharacterVaryingType => arena.text("VARCHAR"),
+            Type::MissingType => arena.text("MISSING"),
+            Type::StringType => arena.text("STRING"),
+            Type::SymbolType => arena.text("SYMBOL"),
+            Type::BlobType => arena.text("BLOB"),
+            Type::ClobType => arena.text("CLOB"),
+            Type::DateType => arena.text("DATE"),
+            Type::TimeType => arena.text("TIME"),
+            Type::ZonedTimestampType => arena.text("TIMESTAMPTZ"),
+            Type::StructType => arena.text("STRUCT"),
+            Type::TupleType => arena.text("TUPLE"),
+            Type::ListType => arena.text("LIST"),
+            Type::SexpType => arena.text("SEXP"),
+            Type::BagType => arena.text("BAG"),
+            Type::AnyType => arena.text("ANY"),
         }
     }
 }
diff --git a/partiql-logical-planner/src/lower.rs b/partiql-logical-planner/src/lower.rs
index cfbd19ff..d6da1caa 100644
--- a/partiql-logical-planner/src/lower.rs
+++ b/partiql-logical-planner/src/lower.rs
@@ -1922,7 +1922,7 @@ fn lit_to_lit(lit: &Lit) -> Result<logical::Lit, AstTransformError> {
         Lit::FloatLit(f) => logical::Lit::Double(OrderedFloat::from(*f as f64)),
         Lit::DoubleLit(f) => logical::Lit::Double(OrderedFloat::from(*f)),
         Lit::BoolLit(b) => logical::Lit::Bool(*b),
-        Lit::EmbeddedDocLit(s) => {
+        Lit::EmbeddedDocLit(s, _) => {
             logical::Lit::BoxDocument(s.clone().into_bytes(), "Ion".to_string())
         }
         Lit::CharStringLit(s) => logical::Lit::String(s.clone()),
diff --git a/partiql-parser/src/parse/partiql.lalrpop b/partiql-parser/src/parse/partiql.lalrpop
index 7bfc0175..11057b75 100644
--- a/partiql-parser/src/parse/partiql.lalrpop
+++ b/partiql-parser/src/parse/partiql.lalrpop
@@ -1301,7 +1301,15 @@ LiteralNumber: ast::Lit = {
 
 #[inline]
 LiteralEmbeddedDoc: ast::Lit = {
-    <ion:"EmbeddedDoc"> => ast::Lit::EmbeddedDocLit(ion.to_owned()),
+    <ion:"EmbeddedDoc"> => {
+        let ion_typ = ast::Type::CustomType(ast::CustomType {
+            parts: vec![ast::CustomTypePart::Name(ast::SymbolPrimitive {
+                value: "Ion".to_string(),
+                case: ast::CaseSensitivity::CaseInsensitive,
+            })],
+        });
+        ast::Lit::EmbeddedDocLit(ion.to_owned(), ion_typ)
+    }
 }