From 4e703f3cea154fa6ffeb94735021565327dbc62c Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Fri, 31 Jan 2025 23:19:48 +0100 Subject: [PATCH] fix: Improve inlining of subqueries --- src/formatter.rs | 54 +++++++++++++++++++++------ src/inline_block.rs | 90 +++++++++++++++++++++++++++++++++------------ src/lib.rs | 31 ++++++++++++++++ 3 files changed, 140 insertions(+), 35 deletions(-) diff --git a/src/formatter.rs b/src/formatter.rs index f3fc35d..90fc76a 100644 --- a/src/formatter.rs +++ b/src/formatter.rs @@ -167,7 +167,14 @@ impl<'a> Formatter<'a> { indentation: Indentation::new(options), inline_block: InlineBlock::new( options.max_inline_block, - options.max_inline_arguments.is_none(), + match (options.max_inline_arguments, options.max_inline_top_level) { + (Some(max_inline_args), Some(max_inline_top)) => { + max_inline_args.min(max_inline_top) + } + (Some(max_inline_args), None) => max_inline_args, + (None, Some(max_inline_top)) => max_inline_top, + (None, None) => 0, + }, ), block_level: 0, } @@ -208,17 +215,34 @@ impl<'a> Formatter<'a> { self.add_new_line(query); } + // if we are inside an inline block we decide our behaviour as if we are an inline argument + fn top_level_behavior(&self) -> (bool, bool) { + let span_len = self.top_level_tokens_span(); + let block_len = self.inline_block.cur_len(); + if block_len > 0 { + let limit = self.options.max_inline_arguments.unwrap_or(0); + (limit < block_len, limit < span_len) + } else { + ( + true, + self.options + .max_inline_top_level + .map_or(true, |limit| limit < span_len), + ) + } + } + 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(span_len); + let (newline_before, newline_after) = self.top_level_behavior(); + + if newline_before { + self.indentation.decrease_top_level(); + self.add_new_line(query); + } query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value))); - if self - .options - .max_inline_top_level - .map_or(true, |limit| limit < span_len) - { + if newline_after { + self.indentation.increase_top_level(span_len); self.add_new_line(query); } else { query.push(' '); @@ -226,10 +250,16 @@ impl<'a> Formatter<'a> { } fn format_top_level_reserved_word_no_indent(&mut self, token: &Token<'_>, query: &mut String) { - self.indentation.decrease_top_level(); - self.add_new_line(query); + let (newline_before, newline_after) = self.top_level_behavior(); + + if newline_before { + self.indentation.decrease_top_level(); + self.add_new_line(query); + } query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value))); - self.add_new_line(query); + if newline_after { + self.add_new_line(query); + } } fn format_newline_reserved_word(&mut self, token: &Token<'_>, query: &mut String) { diff --git a/src/inline_block.rs b/src/inline_block.rs index 32f8843..3c9caa2 100644 --- a/src/inline_block.rs +++ b/src/inline_block.rs @@ -1,41 +1,60 @@ use crate::tokenizer::{Token, TokenKind}; +pub(crate) struct BlockInfo { + length: usize, + has_forbidden_tokens: bool, + top_level_token_span: usize, +} + pub(crate) struct InlineBlock { level: usize, inline_max_length: usize, - newline_on_reserved: bool, + newline_on_reserved_limit: usize, + info: Vec, } impl Default for InlineBlock { fn default() -> Self { InlineBlock { + info: Vec::new(), level: 0, inline_max_length: 50, - newline_on_reserved: true, + newline_on_reserved_limit: 0, } } } impl InlineBlock { - pub fn new(inline_max_length: usize, newline_on_reserved: bool) -> Self { + pub fn new(inline_max_length: usize, newline_on_reserved_limit: usize) -> Self { InlineBlock { - level: 0, inline_max_length, - newline_on_reserved, + newline_on_reserved_limit, + ..Default::default() } } + fn is_inline_block(&self, info: &BlockInfo) -> bool { + !info.has_forbidden_tokens + && info.length <= self.inline_max_length + && info.top_level_token_span <= self.newline_on_reserved_limit + } + pub fn begin_if_possible(&mut self, tokens: &[Token<'_>], index: usize) { - if self.level == 0 && self.is_inline_block(tokens, index) { + let info = self.build_info(tokens, index); + if self.level == 0 && self.is_inline_block(&info) { self.level = 1; } else if self.level > 0 { self.level += 1; } else { self.level = 0; } + if self.level > 0 { + self.info.push(info); + } } pub fn end(&mut self) { + self.info.pop(); self.level -= 1; } @@ -43,41 +62,66 @@ impl InlineBlock { self.level > 0 } - fn is_inline_block(&self, tokens: &[Token<'_>], index: usize) -> bool { + /// Get the current inline block length + pub fn cur_len(&self) -> usize { + self.info.last().map_or(0, |info| info.length) + } + + fn build_info(&self, tokens: &[Token<'_>], index: usize) -> BlockInfo { let mut length = 0; let mut level = 0; + let mut top_level_token_span = 0; + let mut start_top_level = -1; + let mut start_span = 0; + let mut has_forbidden_tokens = false; for token in &tokens[index..] { length += token.value.len(); - // Overran max length - if length > self.inline_max_length { - return false; - } - if token.kind == TokenKind::OpenParen { - level += 1; - } else if token.kind == TokenKind::CloseParen { - level -= 1; - if level == 0 { - return true; + match token.kind { + TokenKind::ReservedTopLevel | TokenKind::ReservedTopLevelNoIndent => { + if start_top_level != -1 { + if start_top_level == level { + top_level_token_span = top_level_token_span.max(length - start_span); + start_top_level = -1; + } + } else { + start_top_level = level; + start_span = length; + } } + TokenKind::OpenParen => { + level += 1; + } + TokenKind::CloseParen => { + level -= 1; + if level == 0 { + break; + } + } + _ => {} } if self.is_forbidden_token(token) { - return false; + has_forbidden_tokens = true; } } - false + // broken syntax let's try our best + BlockInfo { + length, + has_forbidden_tokens, + top_level_token_span, + } } fn is_forbidden_token(&self, token: &Token<'_>) -> bool { - token.kind == TokenKind::ReservedTopLevel - || token.kind == TokenKind::LineComment + token.kind == TokenKind::LineComment || token.kind == TokenKind::BlockComment || token.value == ";" - || if self.newline_on_reserved { - token.kind == TokenKind::ReservedNewline + || if self.newline_on_reserved_limit == 0 { + token.kind == TokenKind::ReservedTopLevel + || token.kind == TokenKind::ReservedNewline } else { false } diff --git a/src/lib.rs b/src/lib.rs index 941e8a1..823c2fa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2135,4 +2135,35 @@ from assert_eq!(format(input, &QueryParams::None, &options), expected); } + + #[test] + fn it_formats_blocks_inline_or_not() { + let input = " UPDATE t SET o = ($5 + $6 + $7 + $8),a = CASE WHEN $2 + THEN NULL ELSE COALESCE($3, b) END, b = CASE WHEN $4 THEN NULL ELSE + COALESCE($5, b) END, s = (SELECT true FROM bar WHERE bar.foo = $99), + c = CASE WHEN $6 THEN NULL ELSE COALESCE($7, c) END, + d = CASE WHEN $8 THEN NULL ELSE COALESCE($9, d) END, + e = (SELECT true FROM bar) WHERE id = $1"; + let options = FormatOptions { + max_inline_arguments: Some(50), + max_inline_block: 100, + max_inline_top_level: Some(10), + ..Default::default() + }; + let expected = indoc!( + " + UPDATE t + SET + o = ($5 + $6 + $7 + $8), + a = CASE WHEN $2 THEN NULL ELSE COALESCE($3, b) END, + b = CASE WHEN $4 THEN NULL ELSE COALESCE($5, b) END, + s = (SELECT true FROM bar WHERE bar.foo = $99), + c = CASE WHEN $6 THEN NULL ELSE COALESCE($7, c) END, + d = CASE WHEN $8 THEN NULL ELSE COALESCE($9, d) END, + e = (SELECT true FROM bar) + WHERE id = $1" + ); + + assert_eq!(format(input, &QueryParams::None, &options), expected); + } }