fix: Keep the previous tokens per-block

Author: lu-zero

src/formatter.rs

@@ -89,21 +89,19 @@ 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);
+                formatter.indentation.set_previous_top_level(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);
+                formatter.indentation.set_previous_top_level(token);
             }
             TokenKind::ReservedNewline => {
                 formatter.format_newline_reserved_word(token, &mut formatted_query);
-                formatter.previous_reserved_word = Some(token);
+                formatter.indentation.set_previous_reserved(token);
             }
             TokenKind::Reserved => {
                 formatter.format_with_spaces(token, &mut formatted_query);
-                formatter.previous_reserved_word = Some(token);
+                formatter.indentation.set_previous_reserved(token);
             }
             TokenKind::OpenParen => {
                 formatter.format_opening_parentheses(token, &mut formatted_query);
@@ -141,8 +139,6 @@ 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>,
@@ -155,8 +151,6 @@ 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,
@@ -389,14 +383,15 @@ impl<'a> Formatter<'a> {
             return;
         }
         if self
-            .previous_reserved_word
+            .indentation
+            .previous_reserved()
             .map(|word| word.value.to_lowercase() == "limit")
             .unwrap_or(false)
         {
             return;
         }
 
-        if matches!((self.previous_top_level_reserved_word, self.options.max_inline_arguments),
+        if matches!((self.indentation.previous_top_level_reserved(), self.options.max_inline_arguments),
             (Some(word), Some(limit)) if ["select", "from"].contains(&word.value.to_lowercase().as_str()) &&
                 limit > self.indentation.span())
         {

src/indentation.rs

@@ -1,9 +1,16 @@
-use crate::{FormatOptions, Indent, SpanInfo};
+use crate::{tokenizer::Token, FormatOptions, Indent, SpanInfo};
+
+#[derive(Debug, Default)]
+struct PreviousTokens<'a> {
+    top_level_reserved: Option<&'a Token<'a>>,
+    reserved: Option<&'a Token<'a>>,
+}
 
 pub(crate) struct Indentation<'a> {
     options: &'a FormatOptions<'a>,
     indent_types: Vec<IndentType>,
     top_level_span: Vec<SpanInfo>,
+    previous: Vec<PreviousTokens<'a>>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -18,6 +25,7 @@ impl<'a> Indentation<'a> {
             options,
             indent_types: Vec::new(),
             top_level_span: Vec::new(),
+            previous: Vec::new(),
         }
     }
 
@@ -37,18 +45,21 @@ impl<'a> Indentation<'a> {
 
     pub fn increase_block_level(&mut self) {
         self.indent_types.push(IndentType::BlockLevel);
+        self.previous.push(Default::default());
     }
 
     pub fn decrease_top_level(&mut self) {
         if self.indent_types.last() == Some(&IndentType::TopLevel) {
             self.indent_types.pop();
             self.top_level_span.pop();
+            self.previous.pop();
         }
     }
 
     pub fn decrease_block_level(&mut self) {
         while !self.indent_types.is_empty() {
             let kind = self.indent_types.pop();
+            self.previous.pop();
             if kind != Some(IndentType::TopLevel) {
                 break;
             }
@@ -58,6 +69,53 @@ impl<'a> Indentation<'a> {
     pub fn reset_indentation(&mut self) {
         self.indent_types.clear();
         self.top_level_span.clear();
+        self.previous.clear();
+    }
+
+    pub fn set_previous_reserved(&mut self, token: &'a Token<'a>) {
+        if let Some(previous) = self.previous.last_mut() {
+            previous.reserved = Some(token);
+        } else {
+            self.previous.push(PreviousTokens {
+                top_level_reserved: None,
+                reserved: Some(token),
+            });
+        }
+    }
+
+    pub fn set_previous_top_level(&mut self, token: &'a Token<'a>) {
+        if let Some(previous) = self.previous.last_mut() {
+            previous.top_level_reserved = Some(token);
+        } else {
+            self.previous.push(PreviousTokens {
+                top_level_reserved: Some(token),
+                reserved: Some(token),
+            });
+        }
+    }
+
+    pub fn previous_reserved(&'a self) -> Option<&'a Token<'a>> {
+        if let Some(PreviousTokens {
+            reserved,
+            top_level_reserved: _,
+        }) = self.previous.last()
+        {
+            reserved.as_deref()
+        } else {
+            None
+        }
+    }
+
+    pub fn previous_top_level_reserved(&'a self) -> Option<&'a Token<'a>> {
+        if let Some(PreviousTokens {
+            top_level_reserved,
+            reserved: _,
+        }) = self.previous.last()
+        {
+            top_level_reserved.as_deref()
+        } else {
+            None
+        }
     }
 
     /// The full span between two top level tokens

src/lib.rs

@@ -2053,6 +2053,49 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 
+    #[test]
+    fn format_nested_select_nested_blocks() {
+        let input =
+            "WITH a AS ( SELECT a, b, c FROM t WHERE a > 100 ), aa AS ( SELECT field FROM table ),
+            bb AS ( SELECT count(*) as c FROM d ), cc AS ( INSERT INTO C (a, b, c, d) VALUES (1 ,2 ,3 ,4) )
+        SELECT b, field FROM a, aa;";
+        let max_line = 20;
+        let options = FormatOptions {
+            max_inline_block: max_line,
+            max_inline_arguments: Some(max_line),
+            max_inline_top_level: Some(max_line / 2),
+            joins_as_top_level: true,
+            ..Default::default()
+        };
+        let expected = indoc! {
+            "
+            WITH
+            a AS (
+              SELECT a, b, c
+              FROM t
+              WHERE a > 100
+            ),
+            aa AS (
+              SELECT field
+              FROM table
+            ),
+            bb AS (
+              SELECT
+                count(*) as c
+              FROM d
+            ),
+            cc AS (
+              INSERT INTO
+                C (a, b, c, d)
+              VALUES
+                (1, 2, 3, 4)
+            )
+            SELECT b, field
+            FROM a, aa;"
+        };
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
+
     #[test]
     fn it_converts_keywords_nothing_when_no_option_passed_in() {
         let input = "select distinct * frOM foo left join bar WHERe cola > 1 and colb = 3";