Fix span computation

src/formatter.rs

@@ -148,7 +148,6 @@ struct Formatter<'a> {
     options: &'a FormatOptions<'a>,
     indentation: Indentation<'a>,
     inline_block: InlineBlock,
-    span_len: usize,
     block_level: usize,
 }
 
@@ -166,7 +165,6 @@ impl<'a> Formatter<'a> {
                 options.max_inline_block,
                 options.max_inline_arguments.is_none(),
             ),
-            span_len: 0,
             block_level: 0,
         }
     }
@@ -207,15 +205,15 @@ impl<'a> Formatter<'a> {
     }
 
     fn format_top_level_reserved_word(&mut self, token: &Token<'_>, query: &mut String) {
+        let span_len = self.top_level_tokens_span();
         self.indentation.decrease_top_level();
         self.add_new_line(query);
-        self.indentation.increase_top_level();
+        self.indentation.increase_top_level(span_len);
         query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
-        self.span_len = self.top_level_tokens_span();
         if self
             .options
             .max_inline_top_level
-            .map_or(true, |limit| limit < self.span_len)
+            .map_or(true, |limit| limit < span_len)
         {
             self.add_new_line(query);
         } else {
@@ -234,7 +232,7 @@ impl<'a> Formatter<'a> {
         if self
             .options
             .max_inline_arguments
-            .map_or(true, |limit| limit < self.span_len)
+            .map_or(true, |limit| limit < self.indentation.top_level_span())
         {
             self.add_new_line(query);
         } else {
@@ -371,7 +369,8 @@ impl<'a> Formatter<'a> {
         }
 
         if matches!((self.previous_top_level_reserved_word, self.options.max_inline_arguments),
-            (Some(word), Some(limit)) if ["select", "from"].contains(&word.value.to_lowercase().as_str()) && limit > self.span_len)
+            (Some(word), Some(limit)) if ["select", "from"].contains(&word.value.to_lowercase().as_str()) &&
+                limit > self.indentation.top_level_span())
         {
             return;
         }
@@ -509,7 +508,7 @@ impl<'a> Formatter<'a> {
                     block_level = block_level.saturating_sub(1);
                     block_level > self.block_level
                 }
-                TokenKind::ReservedTopLevel => false,
+                TokenKind::ReservedTopLevel => block_level != self.block_level,
                 _ => true,
             })
             .map(|token| token.value.len())

src/indentation.rs

@@ -3,6 +3,7 @@ use crate::{FormatOptions, Indent};
 pub(crate) struct Indentation<'a> {
     options: &'a FormatOptions<'a>,
     indent_types: Vec<IndentType>,
+    top_level_span: Vec<usize>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -16,6 +17,7 @@ impl<'a> Indentation<'a> {
         Indentation {
             options,
             indent_types: Vec::new(),
+            top_level_span: Vec::new(),
         }
     }
 
@@ -28,8 +30,9 @@ impl<'a> Indentation<'a> {
         }
     }
 
-    pub fn increase_top_level(&mut self) {
+    pub fn increase_top_level(&mut self, span: usize) {
         self.indent_types.push(IndentType::TopLevel);
+        self.top_level_span.push(span);
     }
 
     pub fn increase_block_level(&mut self) {
@@ -39,6 +42,7 @@ impl<'a> Indentation<'a> {
     pub fn decrease_top_level(&mut self) {
         if self.indent_types.last() == Some(&IndentType::TopLevel) {
             self.indent_types.pop();
+            self.top_level_span.pop();
         }
     }
 
@@ -53,5 +57,10 @@ impl<'a> Indentation<'a> {
 
     pub fn reset_indentation(&mut self) {
         self.indent_types.clear();
+        self.top_level_span.clear();
+    }
+
+    pub fn top_level_span(&self) -> usize {
+        self.top_level_span.last().map_or(0, |span| *span)
     }
 }

src/lib.rs

@@ -1,4 +1,4 @@
 //! This crate is a port of https://github.com/kufii/sql-formatter-plus
 //! written in Rust. It is intended to be usable as a pure-Rust library
 //! for formatting SQL queries.
 
@@ -1991,7 +1991,8 @@
         };
         let expected = indoc! {
             "
-            WITH a AS (
+            WITH
+            a AS (
               SELECT a, b, c
               FROM t
               WHERE a > 100

src/tokenizer.rs

@@ -594,6 +594,7 @@ fn get_top_level_reserved_token_no_indent<'i>(input: &mut &'i str) -> PResult<To
         terminated("MINUS", end_of_word),
         terminated("UNION", end_of_word),
         terminated("UNION ALL", end_of_word),
+        terminated("WITH", end_of_word),
         terminated("$$", end_of_word),
     ))
     .parse_next(&mut uc_input);
@@ -987,7 +988,6 @@ fn get_plain_reserved_one_token<'i>(input: &mut &'i str) -> PResult<Token<'i>> {
 
         'W' => alt((
             terminated("WHEN", end_of_word),
-            terminated("WITH", end_of_word),
             terminated("WORK", end_of_word),
             terminated("WRITE", end_of_word),
         ))