Skip to content

Commit

Permalink
fix: Improve inlining of subqueries
Browse files Browse the repository at this point in the history
  • Loading branch information
@lu-zero
lu-zero committed Jan 31, 2025
1 parent c71c64c commit 4e703f3
Show file tree
Hide file tree
Showing 3 changed files with 140 additions and 35 deletions.
54 changes: 42 additions & 12 deletions src/formatter.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand Down Expand Up @@ -208,28 +215,51 @@ 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(' ');
}
}

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) {
Expand Down
90 changes: 67 additions & 23 deletions src/inline_block.rs
View file
Original file line number Diff line number Diff line change
@@ -1,83 +1,127 @@
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<BlockInfo>,
}

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;
}

pub fn is_active(&self) -> bool {
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
}
Expand Down
31 changes: 31 additions & 0 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

0 comments on commit 4e703f3

Please sign in to comment.