feat: Correctly format array type specifier

Author: lu-zero

src/formatter.rs

@@ -76,7 +76,6 @@ pub(crate) fn format(
             formatter.format_no_change(token, &mut formatted_query);
             continue;
         }
-
         match token.kind {
             TokenKind::Whitespace => {
                 // ignore (we do our own whitespace formatting)
@@ -112,8 +111,8 @@ pub(crate) fn format(
             TokenKind::Placeholder => {
                 formatter.format_placeholder(token, &mut formatted_query);
             }
-            TokenKind::DoubleColon => {
-                formatter.format_double_colon(token, &mut formatted_query);
+            TokenKind::TypeSpecifier => {
+                formatter.format_type_specifier(token, &mut formatted_query);
             }
             _ => match token.value {
                 "," => {
@@ -189,9 +188,9 @@ impl<'a> Formatter<'a> {
         self.add_new_line(query);
     }
 
-    fn format_double_colon(&self, _token: &Token<'_>, query: &mut String) {
+    fn format_type_specifier(&self, token: &Token<'_>, query: &mut String) {
         self.trim_all_spaces_end(query);
-        query.push_str("::");
+        query.push_str(token.value);
     }
     fn format_block_comment(&mut self, token: &Token<'_>, query: &mut String) {
         self.add_new_line(query);

src/lib.rs

@@ -397,6 +397,22 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 
+    #[test]
+    fn it_formats_type_specifiers() {
+        let input = "SELECT id,  ARRAY [] :: UUID [] FROM UNNEST($1  ::  UUID   []);";
+        let options = FormatOptions::default();
+        let expected = indoc!(
+            "
+            SELECT
+              id,
+              ARRAY[]::UUID[]
+            FROM
+              UNNEST($1::UUID[]);"
+        );
+
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
+
     #[test]
     fn it_formats_limit_of_single_value_and_offset() {
         let input = "LIMIT 5 OFFSET 8;";

src/tokenizer.rs

@@ -17,6 +17,7 @@ pub(crate) fn tokenize<'a>(
 ) -> Vec<Token<'a>> {
     let mut tokens: Vec<Token> = Vec::new();
 
+    let mut last_non_whitespace_token = None;
     let mut last_reserved_token = None;
     let mut last_reserved_top_level_token = None;
 
@@ -27,7 +28,7 @@ pub(crate) fn tokenize<'a>(
     // Keep processing the string until it is empty
     while let Ok(mut result) = get_next_token(
         &mut input,
-        tokens.last().cloned(),
+        last_non_whitespace_token.clone(),
         last_reserved_token.clone(),
         last_reserved_top_level_token.clone(),
         named_placeholders,
@@ -49,6 +50,10 @@ pub(crate) fn tokenize<'a>(
             _ => {}
         }
 
+        if result.kind != TokenKind::Whitespace {
+            last_non_whitespace_token = Some(result.clone());
+        }
+
         tokens.push(result);
 
         if let Ok(Some(result)) = opt(get_whitespace_token).parse_next(&mut input) {
@@ -68,7 +73,7 @@ pub(crate) struct Token<'a> {
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) enum TokenKind {
-    DoubleColon,
+    TypeSpecifier,
     Whitespace,
     String,
     Reserved,
@@ -119,6 +124,7 @@ fn get_next_token<'a>(
 ) -> Result<Token<'a>> {
     alt((
         get_comment_token,
+        |input: &mut _| get_type_specifier_token(input, previous_token.clone()),
         get_string_token,
         get_open_paren_token,
         get_close_paren_token,
@@ -131,20 +137,36 @@ fn get_next_token<'a>(
                 last_reserved_top_level_token.clone(),
             )
         },
-        get_double_colon_token,
         get_operator_token,
         |input: &mut _| get_placeholder_token(input, named_placeholders),
         get_word_token,
         get_any_other_char,
     ))
     .parse_next(input)
 }
-fn get_double_colon_token<'i>(input: &mut &'i str) -> Result<Token<'i>> {
-    "::".parse_next(input).map(|token| Token {
-        kind: TokenKind::DoubleColon,
-        value: token,
-        key: None,
-    })
+fn get_type_specifier_token<'i>(
+    input: &mut &'i str,
+    previous_token: Option<Token<'i>>,
+) -> Result<Token<'i>> {
+    if previous_token.is_some_and(|token| {
+        ![
+            TokenKind::CloseParen,
+            TokenKind::Placeholder,
+            TokenKind::Reserved,
+            TokenKind::String,
+            TokenKind::TypeSpecifier,
+            TokenKind::Word,
+        ]
+        .contains(&token.kind)
+    }) {
+        fail.parse_next(input)
+    } else {
+        alt(("::", "[]")).parse_next(input).map(|token| Token {
+            kind: TokenKind::TypeSpecifier,
+            value: token,
+            key: None,
+        })
+    }
 }
 fn get_whitespace_token<'i>(input: &mut &'i str) -> Result<Token<'i>> {
     take_while(1.., char::is_whitespace)