feat: More formatting options (#74)

* chore: Add pretty_assertions

* feat: provide an option to format to a single line

* feat: Make the inline blocks limit tunable

* feat: Add an option to keep arguments inline

* feat: Add an option to inline over the top level tokens

Author: lu-zero

Cargo.toml

@@ -19,6 +19,7 @@ winnow = { version = "0.6.20", features = ["simd"] }
 [dev-dependencies]
 criterion = "0.4"
 indoc = "2.0"
+pretty_assertions = "1.4.1"
 
 [[bench]]
 name = "bench"

src/formatter.rs

@@ -90,10 +90,12 @@ pub(crate) fn format(
             TokenKind::ReservedTopLevel => {
                 formatter.format_top_level_reserved_word(token, &mut formatted_query);
                 formatter.previous_reserved_word = Some(token);
+                formatter.previous_top_level_reserved_word = Some(token);
             }
             TokenKind::ReservedTopLevelNoIndent => {
                 formatter.format_top_level_reserved_word_no_indent(token, &mut formatted_query);
                 formatter.previous_reserved_word = Some(token);
+                formatter.previous_top_level_reserved_word = Some(token);
             }
             TokenKind::ReservedNewline => {
                 formatter.format_newline_reserved_word(token, &mut formatted_query);
@@ -140,27 +142,34 @@ pub(crate) fn format(
 struct Formatter<'a> {
     index: usize,
     previous_reserved_word: Option<&'a Token<'a>>,
+    previous_top_level_reserved_word: Option<&'a Token<'a>>,
     tokens: &'a [Token<'a>],
     params: Params<'a>,
     options: &'a FormatOptions<'a>,
     indentation: Indentation<'a>,
     inline_block: InlineBlock,
+    line_start: usize,
 }
 
 impl<'a> Formatter<'a> {
     fn new(tokens: &'a [Token<'a>], params: &'a QueryParams, options: &'a FormatOptions) -> Self {
         Formatter {
             index: 0,
             previous_reserved_word: None,
+            previous_top_level_reserved_word: None,
             tokens,
             params: Params::new(params),
             options,
             indentation: Indentation::new(options),
-            inline_block: InlineBlock::new(),
+            inline_block: InlineBlock::new(
+                options.max_inline_block,
+                options.max_inline_arguments.is_none(),
+            ),
+            line_start: 0,
         }
     }
 
-    fn format_line_comment(&self, token: &Token<'_>, query: &mut String) {
+    fn format_line_comment(&mut self, token: &Token<'_>, query: &mut String) {
         let is_whitespace_followed_by_special_token =
             self.next_token(1).map_or(false, |current_token| {
                 current_token.kind == TokenKind::Whitespace
@@ -189,7 +198,7 @@ impl<'a> Formatter<'a> {
         self.trim_all_spaces_end(query);
         query.push_str("::");
     }
-    fn format_block_comment(&self, token: &Token<'_>, query: &mut String) {
+    fn format_block_comment(&mut self, token: &Token<'_>, query: &mut String) {
         self.add_new_line(query);
         query.push_str(&self.indent_comment(token.value));
         self.add_new_line(query);
@@ -200,7 +209,16 @@ impl<'a> Formatter<'a> {
         self.add_new_line(query);
         self.indentation.increase_top_level();
         query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
-        self.add_new_line(query);
+        let len = self.top_level_tokens_span();
+        if self
+            .options
+            .max_inline_top_level
+            .map_or(true, |limit| limit < len)
+        {
+            self.add_new_line(query);
+        } else {
+            query.push(' ');
+        }
     }
 
     fn format_top_level_reserved_word_no_indent(&mut self, token: &Token<'_>, query: &mut String) {
@@ -210,8 +228,17 @@ impl<'a> Formatter<'a> {
         self.add_new_line(query);
     }
 
-    fn format_newline_reserved_word(&self, token: &Token<'_>, query: &mut String) {
-        self.add_new_line(query);
+    fn format_newline_reserved_word(&mut self, token: &Token<'_>, query: &mut String) {
+        if self
+            .options
+            .max_inline_arguments
+            .map_or(true, |limit| limit < self.line_len_next(query))
+        {
+            self.add_new_line(query);
+        } else {
+            self.trim_spaces_end(query);
+            query.push(' ');
+        }
         query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
         query.push(' ');
     }
@@ -303,7 +330,13 @@ impl<'a> Formatter<'a> {
 
         if self.inline_block.is_active() {
             self.inline_block.end();
-            self.format_with_space_after(&token, query);
+            if token.value.to_lowercase() == "end" {
+                self.trim_spaces_end(query);
+                query.push(' ');
+                self.format_with_spaces(&token, query);
+            } else {
+                self.format_with_space_after(&token, query);
+            }
         } else {
             self.indentation.decrease_block_level();
             self.add_new_line(query);
@@ -317,7 +350,7 @@ impl<'a> Formatter<'a> {
     }
 
     // Commas start a new line (unless within inline parentheses or SQL "LIMIT" clause)
-    fn format_comma(&self, token: &Token<'_>, query: &mut String) {
+    fn format_comma(&mut self, token: &Token<'_>, query: &mut String) {
         self.trim_spaces_end(query);
         query.push_str(token.value);
         query.push(' ');
@@ -332,6 +365,12 @@ impl<'a> Formatter<'a> {
         {
             return;
         }
+
+        if matches!((self.previous_top_level_reserved_word, self.options.max_inline_arguments),
+            (Some(word), Some(limit)) if word.value.to_lowercase() == "select" && limit > self.line_len_next(query))
+        {
+            return;
+        }
         self.add_new_line(query);
     }
 
@@ -355,11 +394,16 @@ impl<'a> Formatter<'a> {
         }
     }
 
-    fn add_new_line(&self, query: &mut String) {
+    fn add_new_line(&mut self, query: &mut String) {
         self.trim_spaces_end(query);
+        if self.options.inline {
+            query.push(' ');
+            return;
+        }
         if !query.ends_with('\n') {
             query.push('\n');
         }
+        self.line_start = query.len();
         query.push_str(&self.indentation.get_indent());
     }
 
@@ -446,6 +490,20 @@ impl<'a> Formatter<'a> {
         }
     }
 
+    fn line_len_next(&self, query: &str) -> usize {
+        query.len() - self.line_start + self.next_token(1).map_or(0, |t| t.value.len())
+    }
+
+    fn top_level_tokens_span(&self) -> usize {
+        assert_eq!(self.tokens[self.index].kind, TokenKind::ReservedTopLevel);
+
+        self.tokens[self.index..]
+            .iter()
+            .skip(1)
+            .map(|token| token.value.len())
+            .sum()
+    }
+
     fn format_no_change(&self, token: &Token<'_>, query: &mut String) {
         query.push_str(token.value);
     }

src/inline_block.rs

@@ -2,11 +2,27 @@ use crate::tokenizer::{Token, TokenKind};
 
 pub(crate) struct InlineBlock {
     level: usize,
+    inline_max_length: usize,
+    newline_on_reserved: bool,
+}
+
+impl Default for InlineBlock {
+    fn default() -> Self {
+        InlineBlock {
+            level: 0,
+            inline_max_length: 50,
+            newline_on_reserved: true,
+        }
+    }
 }
 
 impl InlineBlock {
-    pub fn new() -> Self {
-        InlineBlock { level: 0 }
+    pub fn new(inline_max_length: usize, newline_on_reserved: bool) -> Self {
+        InlineBlock {
+            level: 0,
+            inline_max_length,
+            newline_on_reserved,
+        }
     }
 
     pub fn begin_if_possible(&mut self, tokens: &[Token<'_>], index: usize) {
@@ -28,16 +44,14 @@ impl InlineBlock {
     }
 
     fn is_inline_block(&self, tokens: &[Token<'_>], index: usize) -> bool {
-        const INLINE_MAX_LENGTH: usize = 50;
-
         let mut length = 0;
         let mut level = 0;
 
         for token in &tokens[index..] {
             length += token.value.len();
 
             // Overran max length
-            if length > INLINE_MAX_LENGTH {
+            if length > self.inline_max_length {
                 return false;
             }
             if token.kind == TokenKind::OpenParen {
@@ -59,9 +73,14 @@ impl InlineBlock {
 
     fn is_forbidden_token(&self, token: &Token<'_>) -> bool {
         token.kind == TokenKind::ReservedTopLevel
-            || token.kind == TokenKind::ReservedNewline
             || token.kind == TokenKind::LineComment
             || token.kind == TokenKind::BlockComment
             || token.value == ";"
+            || if self.newline_on_reserved {
+                token.kind == TokenKind::ReservedNewline
+            } else {
+                false
+            }
+            || token.value.to_lowercase() == "case"
     }
 }