perf(format): stop re-parsing and re-tokenizing work the formatter already has (#8)

Statement splitting in sqlclean format now parses the file and checks the
statement count without collecting every statement's facts. The layout
passes now take the printer's own token ranges instead of re-tokenizing the
text it just printed. The final token check and the idempotence re-print
still re-tokenize and compare, so every safety check stays.

Amp-Thread-ID: https://ampcode.com/threads/T-01a18b8a-340f-717d-9ca7-6dcf564a065a

Co-authored-by: Amp <amp@ampcode.com>

Author: kvlonge

crates/sqlclean-cli/src/formatting/main/format_source.rs

@@ -1,6 +1,6 @@
 //! Format one SQL file: each statement in canonical layout, separated by one blank line.
 
-use sqlclean::sql_facts::main::file_facts::file_facts;
+use sqlclean::sql_facts::main::statement_ends::statement_ends;
 use sqlclean::sql_lint::constants::LINT_API_VERSION;
 use sqlclean::sql_lint::main::format_requests::format_requests;
 use sqlclean::sql_lint::models::{FormatRequest, FormatResponse};
@@ -14,13 +14,8 @@ pub fn format_source(
     line_width: usize,
     indent: usize,
 ) -> Result<String, String> {
-    let facts = file_facts(text, dialect)?;
     let characters: Vec<char> = text.chars().collect();
-    let mut boundaries: Vec<usize> = facts
-        .statements
-        .iter()
-        .map(|statement| statement.terminated_end)
-        .collect();
+    let mut boundaries: Vec<usize> = statement_ends(text, dialect)?;
     if let Some(last) = boundaries.last_mut() {
         *last = characters.len();
     }

crates/sqlclean/src/sql_facts/_helpers/collection.rs

@@ -15,18 +15,13 @@ use crate::sql_lint::main::parse_sql::parse_sql;
 use crate::sql_lint::models::ParsedSql;
 
 pub(crate) fn file_facts(sql: &str, dialect: &str) -> Result<FileFacts, String> {
-    let ParsedSql {
-        tokens, statements, ..
-    } = parse_sql(sql, dialect)?;
+    let (
+        ParsedSql {
+            tokens, statements, ..
+        },
+        ranges,
+    ) = parsed_statements(sql, dialect)?;
     let characters: Vec<char> = sql.chars().collect();
-    let ranges: Vec<StatementTokens> = statement_token_ranges(&tokens);
-    if ranges.len() != statements.len() {
-        return Err(format!(
-            "found {} statements but {} statement ranges",
-            statements.len(),
-            ranges.len()
-        ));
-    }
     let facts: Vec<StatementFacts> = statements
         .iter()
         .zip(ranges)
@@ -46,6 +41,28 @@ pub(crate) fn file_facts(sql: &str, dialect: &str) -> Result<FileFacts, String>
     Ok(FileFacts { statements: facts })
 }
 
+pub(crate) fn statement_ends(sql: &str, dialect: &str) -> Result<Vec<usize>, String> {
+    let (_, ranges) = parsed_statements(sql, dialect)?;
+    Ok(ranges.iter().map(|range| range.terminated_end).collect())
+}
+
+/// Parse `sql` and split its tokens into one range per parsed statement.
+fn parsed_statements(
+    sql: &str,
+    dialect: &str,
+) -> Result<(ParsedSql, Vec<StatementTokens>), String> {
+    let parsed: ParsedSql = parse_sql(sql, dialect)?;
+    let ranges: Vec<StatementTokens> = statement_token_ranges(&parsed.tokens);
+    if ranges.len() != parsed.statements.len() {
+        return Err(format!(
+            "found {} statements but {} statement ranges",
+            parsed.statements.len(),
+            ranges.len()
+        ));
+    }
+    Ok((parsed, ranges))
+}
+
 /// One parsed statement, its tokens and the end offset of its terminator.
 struct StatementSource<'a> {
     index: usize,

crates/sqlclean/src/sql_facts/main/mod.rs

@@ -1 +1,2 @@
 pub mod file_facts;
+pub mod statement_ends;

crates/sqlclean/src/sql_facts/main/statement_ends.rs

@@ -0,0 +1,6 @@
+//! Where each statement of one SQL file ends, without collecting its facts.
+
+/// Parse `sql` under `dialect` and return the end of each statement, after its terminator.
+pub fn statement_ends(sql: &str, dialect: &str) -> Result<Vec<usize>, String> {
+    crate::sql_facts::_helpers::collection::statement_ends(sql, dialect)
+}

crates/sqlclean/src/sql_facts/tests/test_file_facts.rs

@@ -1,7 +1,10 @@
 use crate::sql_facts::main::file_facts::file_facts;
+use crate::sql_facts::main::statement_ends::statement_ends;
 use crate::sql_facts::models::{FileFacts, StatementFacts};
 use crate::sql_facts::tests::helpers::text;
-use crate::sql_facts::tests::test_types::{ParseOutcomeTestCase, StatementFactsTestCase};
+use crate::sql_facts::tests::test_types::{
+    ParseOutcomeTestCase, StatementEndsTestCase, StatementFactsTestCase,
+};
 
 const TWO_STATEMENTS: &str = "WITH recent AS (SELECT o.order_id, o.* FROM shop.orders AS o WHERE o.total > 10)\n\
 SELECT r.order_id, UPPER(c.name) AS name, 'é' AS mark\n\
@@ -229,3 +232,50 @@ fn given_sql_text_when_collecting_facts_then_only_unparseable_sql_is_an_error()
         );
     }
 }
+
+#[test]
+fn given_sql_text_when_finding_statement_ends_then_each_ends_after_its_terminator() {
+    let test_cases = [
+        StatementEndsTestCase {
+            description: "a comment after a terminator starts the next statement's text",
+            sql: "SELECT 'é;' AS a; -- next; note\nSELECT 2 /* end; */",
+            expected_statements: Some(&["SELECT 'é;' AS a;", " -- next; note\nSELECT 2"]),
+        },
+        StatementEndsTestCase {
+            description: "empty statements between semicolons are skipped",
+            sql: ";; SELECT 1;;SELECT 2;",
+            expected_statements: Some(&[";; SELECT 1;", ";SELECT 2;"]),
+        },
+        StatementEndsTestCase {
+            description: "a file of only comments has no statements",
+            sql: "-- only a note; nothing else\n/* still; nothing */\n",
+            expected_statements: Some(&[]),
+        },
+        StatementEndsTestCase {
+            description: "an unparseable statement is an error, not a split",
+            sql: "SELECT 1; SELECT FROM WHERE (",
+            expected_statements: None,
+        },
+    ];
+    for test_case in &test_cases {
+        let characters: Vec<char> = test_case.sql.chars().collect();
+        let statements: Option<Vec<String>> = statement_ends(test_case.sql, "duckdb")
+            .map(|ends| {
+                let starts = std::iter::once(0).chain(ends.iter().copied());
+                starts
+                    .zip(ends.iter().copied())
+                    .map(|(start, end)| characters[start..end].iter().collect())
+                    .collect()
+            })
+            .ok();
+        let texts: Option<Vec<&str>> = statements
+            .as_ref()
+            .map(|statements| statements.iter().map(String::as_str).collect());
+        assert_eq!(
+            texts.as_deref(),
+            test_case.expected_statements,
+            "{}",
+            test_case.description
+        );
+    }
+}

crates/sqlclean/src/sql_facts/tests/test_types.rs

@@ -19,3 +19,9 @@ pub(crate) struct ParseOutcomeTestCase {
     pub(crate) sql: &'static str,
     pub(crate) expected_statements: Option<usize>,
 }
+
+pub(crate) struct StatementEndsTestCase {
+    pub(crate) description: &'static str,
+    pub(crate) sql: &'static str,
+    pub(crate) expected_statements: Option<&'static [&'static str]>,
+}

crates/sqlclean/src/sql_lint/_helpers/formatter.rs

@@ -187,11 +187,11 @@ impl Layout<'_> {
         authored: &AuthoredSql<'_>,
         oracle: &OracleLayout,
     ) -> Result<(String, BTreeMap<usize, KeywordRecase>, TokenTexts), String> {
-        let (printed, recases, texts) = print_authored_tokens(authored, oracle, self.dialect)?;
+        let print = print_authored_tokens(authored, oracle, self.dialect)?;
         let no_widths: HashMap<String, usize> = HashMap::new();
         let token_widths = self.wrap.map_or(&no_widths, |wrap| wrap.token_widths);
-        let Some(mut sql) = PrintedSql::parse(&printed, self.dialect, token_widths)? else {
-            return Ok((printed, recases, texts));
+        let Some(mut sql) = PrintedSql::new(&print.sql, &print.tokens, token_widths) else {
+            return Ok((print.sql, print.recases, print.texts));
         };
         sql.set_indent(self.indent);
         sql.attach_commas_before_comments();
@@ -203,7 +203,7 @@ impl Layout<'_> {
             wrap_lines(&mut sql, wrap.line_width);
         }
         sql.lead_operators_of_multiline_operands();
-        Ok((sql.render(), recases, texts))
+        Ok((sql.render(), print.recases, print.texts))
     }
 }
 

crates/sqlclean/src/sql_lint/_helpers/printed_layout.rs

@@ -1,8 +1,8 @@
 //! Printed SQL as tokens and gaps, with clause layout rules and line-width wrapping.
 
 use std::collections::HashMap;
+use std::ops::Range;
 
-use polyglot_sql::Dialect;
 use polyglot_sql::tokens::TokenType;
 
 use crate::sql_lint::constants::OPEN_PARENTHESIS;
@@ -87,43 +87,36 @@ struct Line {
 }
 
 impl<'a> PrintedSql<'a> {
-    pub(crate) fn parse(
+    /// Split printed `sql` at its token byte ranges; `None` when it has no tokens.
+    pub(crate) fn new(
         sql: &str,
-        dialect: &Dialect,
+        printed_tokens: &[(Range<usize>, TokenType)],
         token_widths: &'a HashMap<String, usize>,
-    ) -> Result<Option<Self>, String> {
-        let tokens = dialect
-            .tokenize(sql)
-            .map_err(|error| format!("native formatter tokenization failed: {error}"))?;
-        if tokens.is_empty() {
-            return Ok(None);
-        }
-        let characters: Vec<char> = sql.chars().collect();
-        let mut texts: Vec<String> = Vec::with_capacity(tokens.len());
-        let mut gaps: Vec<String> = Vec::with_capacity(tokens.len() + 1);
+    ) -> Option<Self> {
+        if printed_tokens.is_empty() {
+            return None;
+        }
+        let mut texts: Vec<String> = Vec::with_capacity(printed_tokens.len());
+        let mut gaps: Vec<String> = Vec::with_capacity(printed_tokens.len() + 1);
         let mut previous_end = 0_usize;
-        for token in &tokens {
-            gaps.push(characters[previous_end..token.span.start].iter().collect());
-            texts.push(
-                characters[token.span.start..token.span.end]
-                    .iter()
-                    .collect(),
-            );
-            previous_end = token.span.end;
-        }
-        gaps.push(characters[previous_end..].iter().collect());
+        for (range, _) in printed_tokens {
+            gaps.push(sql[previous_end..range.start].to_owned());
+            texts.push(sql[range.clone()].to_owned());
+            previous_end = range.end;
+        }
+        gaps.push(sql[previous_end..].to_owned());
         let (depths, partners) = bracket_structure(&texts);
         let locked = vec![false; texts.len() + 1];
-        Ok(Some(Self {
-            kinds: tokens.iter().map(|token| token.token_type).collect(),
+        Some(Self {
+            kinds: printed_tokens.iter().map(|(_, kind)| *kind).collect(),
             tokens: texts,
             gaps,
             depths,
             partners,
             locked,
             token_widths,
             indent: PARSED_INDENT,
-        }))
+        })
     }
 
     fn len(&self) -> usize {

crates/sqlclean/src/sql_lint/_helpers/token_layout.rs

@@ -1,6 +1,7 @@
 //! Print authored SQL tokens with whitespace taken from a parse-tree layout oracle.
 
 use std::collections::{BTreeMap, BTreeSet};
+use std::ops::Range;
 
 use polyglot_sql::Dialect;
 use polyglot_sql::tokens::{Token, TokenType};
@@ -67,6 +68,15 @@ impl Lexeme {
     }
 }
 
+/// Authored tokens printed against an oracle, with what the token check needs about them.
+pub(crate) struct AuthoredPrint {
+    pub(crate) sql: String,
+    /// The byte range and type of each printed token; comments and whitespace lie between them.
+    pub(crate) tokens: Vec<(Range<usize>, TokenType)>,
+    pub(crate) recases: BTreeMap<usize, KeywordRecase>,
+    pub(crate) texts: TokenTexts,
+}
+
 /// The source text and case-fold flag of every token, read once for the token check.
 pub(crate) struct TokenTexts {
     texts: Vec<String>,
@@ -227,7 +237,7 @@ pub(crate) fn print_authored_tokens(
     source: &AuthoredSql<'_>,
     oracle: &OracleLayout,
     dialect: &Dialect,
-) -> Result<(String, BTreeMap<usize, KeywordRecase>, TokenTexts), String> {
+) -> Result<AuthoredPrint, String> {
     let AuthoredSql {
         sql,
         tokens,
@@ -241,6 +251,7 @@ pub(crate) fn print_authored_tokens(
     let mut output = String::with_capacity(sql.len() + sql.len() / 4);
     let string_breaks = string_breaks(sql, tokens);
     let mut recases: BTreeMap<usize, KeywordRecase> = BTreeMap::new();
+    let mut printed_tokens: Vec<(Range<usize>, TokenType)> = Vec::with_capacity(tokens.len());
     for (index, lexeme) in authored.iter().enumerate() {
         let forced_break = string_breaks.contains(&index) && !separators[index].contains('\n');
         let forced: String;
@@ -281,7 +292,9 @@ pub(crate) fn print_authored_tokens(
                 },
             );
         }
+        let start = output.len();
         output.push_str(&printed_text(lexeme, guide));
+        printed_tokens.push((start..output.len(), tokens[index].token_type));
         let trailing = &attached[index].trailing;
         for (position, comment) in trailing.iter().enumerate() {
             output.push(' ');
@@ -312,7 +325,12 @@ pub(crate) fn print_authored_tokens(
         texts.texts.push(lexeme.raw);
         texts.foldable.push(lexeme.foldable);
     }
-    Ok((output, recases, texts))
+    Ok(AuthoredPrint {
+        sql: output,
+        tokens: printed_tokens,
+        recases,
+        texts,
+    })
 }
 
 /// Refuse output that changes anything but layout, fold-set case, and recorded keyword recases.

crates/sqlclean/src/sql_lint/tests/test_comment_formatter.rs

@@ -74,6 +74,12 @@ fn given_trailing_comments_after_the_last_token_when_formatting_then_no_newline_
             expected_sql: "SELECT 1 AS order_id\nFROM orders /* latest orders */",
             expected_changed: false,
         },
+        test_types::FormatTestCase {
+            description: "multibyte comments and literals keep every later token in place",
+            sql: "select 'café' as naïve, -- prix en €\n ü.x from t as ü /* 日本 */ where ü.y = 'ß'",
+            expected_sql: "SELECT\n  'café' AS naïve, -- prix en €\n  ü.x\nFROM t AS ü /* 日本 */\nWHERE ü.y = 'ß'",
+            expected_changed: true,
+        },
     ];
     for test_case in &test_cases {
         let response = format_json(