fmt: literals & binary operators

crates/squawk_fmt/src/fmt.rs

@@ -1,6 +1,6 @@
 use itertools::Itertools;
 use rowan::Direction;
-use squawk_syntax::ast::{self, AstNode};
+use squawk_syntax::ast::{self, AstNode, LitKind};
 use squawk_syntax::quote::quote_column_alias;
 use squawk_syntax::{SyntaxKind, SyntaxNode, SyntaxToken};
 use tiny_pretty::Doc;
@@ -79,7 +79,7 @@ fn build_create_table<'a>(create_table: &ast::CreateTable) -> Doc<'a> {
 
 fn build_table_arg<'a>(create_table: ast::TableArg) -> Doc<'a> {
     match create_table {
-        ast::TableArg::Column(column) => Doc::text(column.name().unwrap().syntax().to_string())
+        ast::TableArg::Column(column) => build_name(column.name().unwrap())
             .append(Doc::space())
             .append(Doc::text(column.ty().unwrap().syntax().to_string())),
         ast::TableArg::LikeClause(_like_clause) => todo!(),
@@ -250,7 +250,7 @@ fn build_expr<'a>(expr: ast::Expr) -> Doc<'a> {
         }
         // ast::Expr::FieldExpr(field_expr) => todo!(),
         // ast::Expr::IndexExpr(index_expr) => todo!(),
-        // ast::Expr::Literal(literal) => todo!(),
+        ast::Expr::Literal(literal) => build_literal(literal),
         // ast::Expr::NameRef(name_ref) => todo!(),
         // ast::Expr::ParenExpr(paren_expr) => todo!(),
         ast::Expr::PostfixExpr(postfix_expr) => {
@@ -402,48 +402,124 @@ fn build_unicode_normal_form<'a>(form: ast::UnicodeNormalForm) -> Doc<'a> {
     }
 }
 
+fn build_keyword_node<'a>(node: &SyntaxNode) -> Doc<'a> {
+    let mut docs: Vec<Doc<'a>> = vec![];
+    for el in node.children_with_tokens() {
+        match el {
+            rowan::NodeOrToken::Token(token) => match token.kind() {
+                SyntaxKind::WHITESPACE => continue,
+                SyntaxKind::COMMENT => {
+                    if !docs.is_empty() {
+                        docs.push(Doc::space());
+                    }
+                    docs.push(Doc::text(token.text().to_string()));
+                }
+                _ => {
+                    if !docs.is_empty() {
+                        docs.push(Doc::space());
+                    }
+                    docs.push(Doc::text(token.text().to_ascii_lowercase()));
+                }
+            },
+            rowan::NodeOrToken::Node(_) => (),
+        }
+    }
+    Doc::list(docs)
+}
+
 fn build_op<'a>(op: ast::BinOp) -> Doc<'a> {
     match op {
-        ast::BinOp::And(_) => todo!(),
-        ast::BinOp::AtTimeZone(_) => todo!(),
-        ast::BinOp::Caret(_) => todo!(),
-        ast::BinOp::Collate(_) => todo!(),
-        ast::BinOp::ColonColon(_) => todo!(),
-        ast::BinOp::ColonEq(_) => todo!(),
+        ast::BinOp::And(_) => Doc::text("and"),
+        ast::BinOp::AtTimeZone(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::Caret(_) => Doc::text("^"),
+        ast::BinOp::Collate(_) => Doc::text("collate"),
+        ast::BinOp::ColonColon(_) => Doc::text("::"),
+        ast::BinOp::ColonEq(_) => Doc::text(":="),
         ast::BinOp::CustomOp(custom_op) => Doc::text(custom_op.syntax().to_string()),
-        ast::BinOp::Eq(_) => todo!(),
-        ast::BinOp::FatArrow(_) => todo!(),
-        ast::BinOp::Gteq(_) => todo!(),
-        ast::BinOp::Ilike(_) => todo!(),
-        ast::BinOp::In(_) => todo!(),
-        ast::BinOp::Is(_) => todo!(),
-        ast::BinOp::IsDistinctFrom(_) => todo!(),
-        ast::BinOp::IsNot(_) => todo!(),
-        ast::BinOp::IsNotDistinctFrom(_) => todo!(),
-        ast::BinOp::LAngle(_) => todo!(),
-        ast::BinOp::Like(_) => todo!(),
-        ast::BinOp::Lteq(_) => todo!(),
-        ast::BinOp::Minus(_) => todo!(),
-        ast::BinOp::Neq(_) => todo!(),
-        ast::BinOp::Neqb(_) => todo!(),
-        ast::BinOp::NotIlike(_) => todo!(),
-        ast::BinOp::NotIn(_) => todo!(),
-        ast::BinOp::NotLike(_) => todo!(),
-        ast::BinOp::NotSimilarTo(_) => todo!(),
-        ast::BinOp::OperatorCall(_) => todo!(),
-        ast::BinOp::Or(_) => todo!(),
-        ast::BinOp::Overlaps(_) => todo!(),
-        ast::BinOp::Percent(_) => todo!(),
+        ast::BinOp::Eq(_) => Doc::text("="),
+        ast::BinOp::FatArrow(_) => Doc::text("=>"),
+        ast::BinOp::Gteq(_) => Doc::text(">="),
+        ast::BinOp::Ilike(_) => Doc::text("ilike"),
+        ast::BinOp::In(_) => Doc::text("in"),
+        ast::BinOp::Is(_) => Doc::text("is"),
+        ast::BinOp::IsDistinctFrom(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::IsNot(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::IsNotDistinctFrom(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::LAngle(_) => Doc::text("<"),
+        ast::BinOp::Like(_) => Doc::text("like"),
+        ast::BinOp::Lteq(_) => Doc::text("<="),
+        ast::BinOp::Minus(_) => Doc::text("-"),
+        ast::BinOp::Neq(_) => Doc::text("!="),
+        ast::BinOp::Neqb(_) => Doc::text("<>"),
+        ast::BinOp::NotIlike(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::NotIn(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::NotLike(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::NotSimilarTo(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::OperatorCall(op) => Doc::text(op.syntax().to_string()),
+        ast::BinOp::Or(_) => Doc::text("or"),
+        ast::BinOp::Overlaps(_) => Doc::text("overlaps"),
+        ast::BinOp::Percent(_) => Doc::text("%"),
         ast::BinOp::Plus(_) => Doc::text("+"),
-        ast::BinOp::RAngle(_) => todo!(),
-        ast::BinOp::SimilarTo(_) => todo!(),
-        ast::BinOp::Slash(_) => todo!(),
-        ast::BinOp::Star(_) => todo!(),
+        ast::BinOp::RAngle(_) => Doc::text(">"),
+        ast::BinOp::SimilarTo(n) => build_keyword_node(n.syntax()),
+        ast::BinOp::Slash(_) => Doc::text("/"),
+        ast::BinOp::Star(_) => Doc::text("*"),
     }
 }
 
 fn build_literal<'a>(lit: ast::Literal) -> Doc<'a> {
-    Doc::text(lit.syntax().to_string())
+    let Some(kind) = lit.kind() else {
+        return Doc::nil();
+    };
+    match kind {
+        LitKind::Default(_) => Doc::text("default"),
+        LitKind::False(_) => Doc::text("false"),
+        LitKind::IntNumber(t) => Doc::text(t.text().to_string()),
+        LitKind::Null(_) => Doc::text("null"),
+        LitKind::NumericNumber(t) => Doc::text(t.text().to_string()),
+        LitKind::PositionalParam(t) => Doc::text(t.text().to_string()),
+        LitKind::True(_) => Doc::text("true"),
+        LitKind::BitString(_)
+        | LitKind::ByteString(_)
+        | LitKind::DollarQuotedString(_)
+        | LitKind::EscString(_)
+        | LitKind::String(_)
+        | LitKind::UnicodeEscString(_) => build_string_literal(&lit),
+    }
+}
+
+fn build_string_literal<'a>(lit: &ast::Literal) -> Doc<'a> {
+    let parts: Vec<Doc<'a>> = lit
+        .syntax()
+        .children_with_tokens()
+        .filter_map(|el| match el {
+            rowan::NodeOrToken::Token(t) if t.kind() != SyntaxKind::WHITESPACE => {
+                Some(Doc::text(format_string_token(&t)))
+            }
+            _ => None,
+        })
+        .collect();
+    Doc::list(Itertools::intersperse(parts.into_iter(), Doc::hard_line()).collect())
+}
+
+fn format_string_token(t: &SyntaxToken) -> String {
+    let text = t.text();
+    if matches!(
+        t.kind(),
+        SyntaxKind::STRING | SyntaxKind::DOLLAR_QUOTED_STRING
+    ) {
+        return text.to_string();
+    }
+    match text.find('\'') {
+        Some(idx) => {
+            let (prefix, rest) = text.split_at(idx);
+            let mut s = String::with_capacity(text.len());
+            s.push_str(&prefix.to_ascii_lowercase());
+            s.push_str(rest);
+            s
+        }
+        None => text.to_string(),
+    }
 }
 
 fn build_name<'a>(name: ast::Name) -> Doc<'a> {

crates/squawk_fmt/tests/after/select_expr.snap

@@ -14,6 +14,38 @@ select
   -- bin expr
   1 + 1,
   2 @@@ 2,
+  true and false,
+  ts at time zone 'UTC',
+  2 ^ 3,
+  name collate "C",
+  1 = 1,
+  1 >= 1,
+  'foo' ilike 'F%',
+  1 in (1, 2),
+  x is true,
+  1 is distinct from 2,
+  1 is /* comment */ distinct from 2,
+  x is not true,
+  1 is not distinct from 2,
+  1 < 2,
+  'foo' like 'f%',
+  1 <= 2,
+  1 - 1,
+  1 != 2,
+  1 <> 2,
+  'foo' not ilike 'F%',
+  1 not in (1, 2),
+  'foo' not like 'f%',
+  'foo' not similar to 'f%',
+  1 operator(+) 1,
+  1 operator(public.+) 1,
+  true or false,
+  (1, 2) overlaps (3, 4),
+  10 % 3,
+  2 > 1,
+  'foo' similar to 'f%',
+  6 / 2,
+  2 * 3,
   -- call expr
   date_trunc('month', now()),
   -- case expr

crates/squawk_fmt/tests/after/select_literals.snap

@@ -0,0 +1,60 @@
+---
+source: crates/squawk_fmt/tests/tests.rs
+input_file: crates/squawk_fmt/tests/before/select_literals.sql
+---
+-- null
+select null;
+-- true
+select true;
+-- false
+select false;
+-- int number
+select
+  42;
+-- numeric number
+select 3.14;
+-- string
+select 'hello';
+-- esc string
+select
+  e'esc\n';
+-- unicode esc string
+select u&'unicode';
+-- dollar quoted string
+select
+  $$dollar$$;
+-- dollar quoted string with tag
+select
+  $tag$body$tag$;
+-- bit string
+select b'1010';
+-- byte string
+select
+  x'AF';
+-- positional param
+select $1;
+-- string continuation
+select
+  'foo'
+  'bar';
+-- string continuation, multiple
+select
+  'one'
+  'two'
+  'three';
+-- esc string continuation
+select
+  e'esc\n'
+  'tail';
+-- unicode esc string continuation
+select
+  u&'uni'
+  'tail';
+-- bit string continuation
+select
+  b'1010'
+  '0101';
+-- byte string continuation
+select
+  x'AF'
+  'BE';

crates/squawk_fmt/tests/before/select_expr.sql

@@ -10,6 +10,38 @@ select
   -- bin expr
   1 + 1,
   2@@@2,
+  true  and  false,
+  ts  at  time  zone  'UTC',
+  2  ^  3,
+  name  collate  "C",
+  1  =  1,
+  1  >=  1,
+  'foo'  ilike  'F%',
+  1  in  (1, 2),
+  x  is  true,
+  1  is  distinct  from  2,
+  1  is  /* comment */  distinct  from  2,
+  x  is  not  true,
+  1  is  not  distinct  from  2,
+  1  <  2,
+  'foo'  like  'f%',
+  1  <=  2,
+  1  -  1,
+  1  !=  2,
+  1  <>  2,
+  'foo'  not  ilike  'F%',
+  1  not  in  (1, 2),
+  'foo'  not  like  'f%',
+  'foo'  not  similar  to  'f%',
+  1  operator(+)  1,
+  1  operator(public.+)  1,
+  true  or  false,
+  (1, 2)  overlaps  (3, 4),
+  10  %  3,
+  2  >  1,
+  'foo'  similar  to  'f%',
+  6  /  2,
+  2  *  3,
   -- call expr
   date_trunc('month', now()),
   -- case expr

crates/squawk_fmt/tests/before/select_literals.sql

@@ -0,0 +1,45 @@
+-- null
+select null;
+-- true
+select true;
+-- false
+select false;
+-- int number
+select 42;
+-- numeric number
+select 3.14;
+-- string
+select 'hello';
+-- esc string
+select E'esc\n';
+-- unicode esc string
+select U&'unicode';
+-- dollar quoted string
+select $$dollar$$;
+-- dollar quoted string with tag
+select $tag$body$tag$;
+-- bit string
+select B'1010';
+-- byte string
+select X'AF';
+-- positional param
+select $1;
+-- string continuation
+select 'foo'
+  'bar';
+-- string continuation, multiple
+select 'one'
+  'two'
+  'three';
+-- esc string continuation
+select E'esc\n'
+  'tail';
+-- unicode esc string continuation
+select U&'uni'
+  'tail';
+-- bit string continuation
+select B'1010'
+  '0101';
+-- byte string continuation
+select X'AF'
+  'BE';

crates/squawk_fmt/tests/tests.rs

@@ -52,17 +52,18 @@ fn tokens_equivalent(before: (TokenKind, &str), after: (TokenKind, &str)) -> boo
 
     // We convert `select 1 "foo"` to `select 1 foo` so we need to do some quote
     // munging
-    let unquote = |kind: TokenKind, text: &str| match kind {
-        TokenKind::QuotedIdent { .. } => text
-            .strip_prefix('"')
-            .and_then(|t| t.strip_suffix('"'))
-            .map(|t| t.to_string()),
-        TokenKind::Ident => Some(text.to_string()),
-        _ => None,
-    };
+    fn unquote<'a>(kind: &TokenKind, text: &'a str) -> Option<&'a str> {
+        match kind {
+            TokenKind::QuotedIdent { .. } => {
+                text.strip_prefix('"').and_then(|t| t.strip_suffix('"'))
+            }
+            TokenKind::Ident => Some(text),
+            _ => None,
+        }
+    }
 
-    match (unquote(bkind, btext), unquote(akind, atext)) {
-        (Some(b), Some(a)) => b.eq_ignore_ascii_case(&a),
+    match (unquote(&bkind, btext), unquote(&akind, atext)) {
+        (Some(b), Some(a)) => b.eq_ignore_ascii_case(a),
         _ => false,
     }
 }