From 3239e351606014460e2bbc99dfb7dc72a5d44f3d Mon Sep 17 00:00:00 2001
From: Luca Barbato <lu_zero@gentoo.org>
Date: Wed, 15 Jan 2025 18:53:00 +0100
Subject: [PATCH] feat: Add an option to inline over the top level tokens

---
 src/formatter.rs | 21 ++++++++++++++++++++-
 src/lib.rs       | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/src/formatter.rs b/src/formatter.rs
index f61bdfa..5957e96 100644
--- a/src/formatter.rs
+++ b/src/formatter.rs
@@ -209,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) {
@@ -485,6 +494,16 @@ impl<'a> Formatter<'a> {
         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);
     }
diff --git a/src/lib.rs b/src/lib.rs
index 28ac600..bfb167b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -57,6 +57,10 @@ pub struct FormatOptions<'a> {
     ///
     /// Default: None
     pub max_inline_arguments: Option<usize>,
+    /// Inline the argument at the top level if they would fit a line
+    ///
+    /// Default: None
+    pub max_inline_top_level: Option<usize>,
 }
 
 impl<'a> Default for FormatOptions<'a> {
@@ -69,6 +73,7 @@ impl<'a> Default for FormatOptions<'a> {
             inline: false,
             max_inline_block: 50,
             max_inline_arguments: None,
+            max_inline_top_level: None,
         }
     }
 }
@@ -221,6 +226,35 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 
+    #[test]
+    fn inline_arguments_when_possible() {
+        let input = indoc! {
+            "
+            SELECT
+              a,
+              b,
+              c,
+              d,
+              e,
+              f,
+              g,
+              h
+            FROM foo;"
+        };
+        let options = FormatOptions {
+            max_inline_arguments: Some(50),
+            max_inline_top_level: Some(20),
+            ..Default::default()
+        };
+        let expected = indoc! {
+            "
+            SELECT
+              a, b, c, d, e, f, g, h
+            FROM foo;"
+        };
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
+
     #[test]
     fn it_formats_select_with_complex_where() {
         let input = indoc!(