From 98fcc6faf40fcc77d757fe1f8f0ecb56f6b2170b Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 24 Oct 2026 07:20:54 +0000 Subject: [PATCH] perf(workspace): cache the walk of a document for hover, inlay hints and semantic tokens Hover, inlay hints and semantic tokens walked the file from the start and resolved every statement again on each request. The walk is now cached per document: a checkpoint of the catalog and the session every few statements, and the references of each resolved statement, keyed by a hash of the statements before it. A request resolves at most the statements since the last valid checkpoint before its range, and a repeated request only looks the references up. Inlay hints and semantic tokens are computed from the cached references. --- crates/pgls_inlay_hints/src/lib.rs | 28 +- crates/pgls_semantic_tokens/src/lib.rs | 12 +- .../src/workspace/server.perf_tests.rs | 398 ++++++++++++++ crates/pgls_workspace/src/workspace/server.rs | 207 ++------ .../src/workspace/server.tests.rs | 123 ++++- .../src/workspace/server/document.rs | 27 + .../src/workspace/server/editor_walk.rs | 489 ++++++++++++++++++ 7 files changed, 1112 insertions(+), 172 deletions(-) create mode 100644 crates/pgls_workspace/src/workspace/server.perf_tests.rs create mode 100644 crates/pgls_workspace/src/workspace/server/editor_walk.rs diff --git a/crates/pgls_inlay_hints/src/lib.rs b/crates/pgls_inlay_hints/src/lib.rs index 5125c13d8..bdec7656b 100644 --- a/crates/pgls_inlay_hints/src/lib.rs +++ b/crates/pgls_inlay_hints/src/lib.rs @@ -47,22 +47,39 @@ pub fn inlay_hints(params: InlayHintsParams<'_>) -> Vec { collect_references: true, }) .references; + inlay_hints_of_references(&references, params.sql) +} + +/// The inlay hints of the `references` of the statement `sql`, which +/// [`pgls_catalog::resolve::resolve`] collected. Ordered by offset. +pub fn inlay_hints_of_references(references: &[Reference], sql: &str) -> Vec { let calls: Vec<_> = references - .into_iter() + .iter() .filter_map(|reference| match reference { Reference::Function { location, function, arguments, .. - } => Some((location, function, arguments)), + } => Some((*location, function, arguments)), _ => None, }) + // Only functions with named parameters get hints, which most built-in functions don't. + .filter(|(_, function, _)| { + function.signature.as_ref().is_some_and(|signature| { + signature.arguments.iter().any(|argument| { + argument + .name + .as_deref() + .is_some_and(|name| !name.is_empty()) + }) + }) + }) .collect(); if calls.is_empty() { return Vec::new(); } - let Ok(scan) = pgls_query::scan(params.sql) else { + let Ok(scan) = pgls_query::scan(sql) else { return Vec::new(); }; let tokens: Vec = scan @@ -98,7 +115,7 @@ pub fn inlay_hints(params: InlayHintsParams<'_>) -> Vec { else { continue; }; - if is_reference_to(params.sql, tokens, name) { + if is_reference_to(sql, tokens, name) { continue; } let Ok(offset) = u32::try_from(tokens[0].start) else { @@ -118,7 +135,8 @@ pub fn inlay_hints(params: InlayHintsParams<'_>) -> Vec { /// The tokens of each argument of the call whose name starts at `location`, without a leading /// `DISTINCT` or `ALL` of an aggregate. fn argument_tokens(tokens: &[ScanToken], location: i32) -> Option> { - let name = tokens.iter().position(|token| token.start == location)?; + let name = tokens.partition_point(|token| token.start < location); + tokens.get(name).filter(|token| token.start == location)?; // The name is at most `schema.name`. let open = name + tokens[name..] diff --git a/crates/pgls_semantic_tokens/src/lib.rs b/crates/pgls_semantic_tokens/src/lib.rs index 917503880..6a02b8156 100644 --- a/crates/pgls_semantic_tokens/src/lib.rs +++ b/crates/pgls_semantic_tokens/src/lib.rs @@ -61,13 +61,18 @@ pub fn semantic_tokens(params: SemanticTokensParams<'_>) -> Vec { collect_references: true, }) .references; + semantic_tokens_of_references(&references, params.sql) +} +/// The semantic tokens of the `references` of the statement `sql`, which +/// [`pgls_catalog::resolve::resolve`] collected. Ordered by their start, and they don't overlap. +pub fn semantic_tokens_of_references(references: &[Reference], sql: &str) -> Vec { let mut tokens = Vec::new(); - for reference in &references { + for reference in references { let Some(span) = reference.span() else { continue; }; - let Some(text) = params.sql.get(std::ops::Range::::from(span)) else { + let Some(text) = sql.get(std::ops::Range::::from(span)) else { continue; }; let parts: Vec = name_parts(text) @@ -127,8 +132,7 @@ pub fn semantic_tokens(params: SemanticTokensParams<'_>) -> Vec { } // `table.column`, but not `alias.column`. [table, _] - if relation - .is_some_and(|(_, name)| identifier(¶ms.sql[*table]) == *name) => + if relation.is_some_and(|(_, name)| identifier(&sql[*table]) == *name) => { push(Some(table), SemanticTokenKind::Class, schema); } diff --git a/crates/pgls_workspace/src/workspace/server.perf_tests.rs b/crates/pgls_workspace/src/workspace/server.perf_tests.rs new file mode 100644 index 000000000..678a94833 --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server.perf_tests.rs @@ -0,0 +1,398 @@ +//! Timings of the resolver-based editor features on large files. +//! +//! Run with +//! `cargo test --release -p pgls_workspace --lib perf_tests -- --ignored --nocapture --test-threads=1`. +//! `PERF_GROUPS` sets the sizes of the migration files (4 top-level statements and a function +//! body per group), `PERF_CALLS` the calls of the huge statement, and `PERF_SEED_ROWS` the rows +//! of the seed `INSERT`s. + +use std::{sync::Arc, time::Instant}; + +use pgls_configuration::{Merge, PartialConfiguration, database::PartialDatabaseConfiguration}; +use pgls_fs::PgLSPath; +use pgls_text_size::{TextRange, TextSize}; +use sqlx::PgPool; + +use super::document::AnalyserDiagnosticsMapper; +use crate::{ + Workspace, + features::code_actions::CodeActionsParams, + features::inlay_hints::PullInlayHintsParams, + features::on_hover::OnHoverParams, + features::semantic_tokens::PullSemanticTokensParams, + workspace::{ + OpenFileParams, RegisterProjectFolderParams, UpdateSettingsParams, server::WorkspaceServer, + }, +}; + +fn workspace(test_db: &PgPool) -> WorkspaceServer { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + ..Default::default() + }); + let workspace = WorkspaceServer::new(); + workspace + .register_project_folder(RegisterProjectFolderParams { + path: None, + set_as_current_workspace: true, + }) + .unwrap(); + workspace + .update_settings(UpdateSettingsParams { + configuration: conf, + gitignore_matches: vec![], + vcs_base_path: None, + workspace_directory: None, + }) + .unwrap(); + workspace +} + +/// A migration-style file: per group a table, an index, a SQL function and a query. +fn migration_file(groups: usize) -> String { + let mut sql = String::new(); + for i in 0..groups { + sql.push_str(&format!( + "create table t{i} (id bigint primary key, name text not null, status text, created_at timestamptz default now());\n\ + create index t{i}_name_idx on t{i} (name);\n\ + create function f{i}(a int, b text) returns int language sql stable as $$ select count(*)::int from t{i} where id = a and name = b $$;\n\ + select t.id, t.name, f{i}(1, 'x') from t{i} t where t.status = 'active' order by t.created_at;\n" + )); + } + sql +} + +/// One huge statement: a select with many function calls. +fn huge_statement(calls: usize) -> String { + let mut sql = String::from("create table big (id int, a int, b text);\nselect\n"); + for i in 0..calls { + if i > 0 { + sql.push_str(",\n"); + } + sql.push_str(&format!(" coalesce(a, {i}) + length(b) + abs(id) as c{i}")); + } + sql.push_str("\nfrom big;\n"); + sql +} + +/// A seed file: one INSERT with many rows, each with an enum literal and a function call. +fn seed_insert(rows: usize) -> String { + let mut sql = String::from( + "create type seed_status as enum ('active', 'inactive');\n\ + create table seed (id int, name text, status seed_status, created_at timestamptz);\n\ + insert into seed (id, name, status, created_at) values\n", + ); + for i in 0..rows { + if i > 0 { + sql.push_str(",\n"); + } + sql.push_str(&format!(" ({i}, 'name {i}', 'active', now())")); + } + sql.push_str(";\nselect name from seed;\n"); + sql +} + +fn ms(start: Instant) -> f64 { + start.elapsed().as_secs_f64() * 1000.0 +} + +/// The fastest of three runs. +fn time(f: impl Fn() -> T) -> (f64, T) { + let mut best = f64::MAX; + let mut result = None; + for _ in 0..3 { + let start = Instant::now(); + result = Some(f()); + best = best.min(ms(start)); + } + (best, result.unwrap()) +} + +/// The fastest of three runs that each start without a cached walk, and the fastest of three +/// repeated runs. +fn time_cold_warm(workspace: &WorkspaceServer, f: impl Fn() -> T) -> (f64, f64, T) { + let (cold, _) = time(|| { + workspace.editor_walks.lock().unwrap().clear(); + f() + }); + f(); + let (warm, result) = time(&f); + (cold, warm, result) +} + +fn time_once(f: impl FnOnce() -> T) -> (f64, T) { + let start = Instant::now(); + let result = f(); + (ms(start), result) +} + +fn measure(workspace: &WorkspaceServer, name: &str, content: &str) { + let path = PgLSPath::new(format!("{name}.sql")); + let (open, ()) = time_once(|| { + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: content.into(), + version: 1, + }) + .unwrap() + }); + let statements = { + let documents = workspace.documents.read().unwrap(); + documents + .get(&path) + .unwrap() + .iter(AnalyserDiagnosticsMapper) + .count() + }; + println!( + "\n== {name}: {} bytes, {} lines, {statements} statements (incl. function bodies); open {open:.1} ms", + content.len(), + content.lines().count() + ); + + // The parse/mapper cost alone: iterate the document twice (the second is "warm"). + let (t, _) = time(|| { + let documents = workspace.documents.read().unwrap(); + documents + .get(&path) + .unwrap() + .iter(AnalyserDiagnosticsMapper) + .count() + }); + println!( + " {:<40} {t:>9.1} ms", + "iterate document (AnalyserDiagnosticsMapper)" + ); + + let snapshot = { + let settings = workspace.workspaces(); + let settings = settings.settings().unwrap(); + workspace + .connection + .with_pool(&settings.db, |pool| workspace.snapshot.load(pool)) + .unwrap() + .unwrap() + }; + // Positions: the first and the last query of the file. + let first = TextSize::try_from(content.find("select").unwrap()).unwrap(); + let last = TextSize::try_from(content.rfind("select").unwrap()).unwrap(); + // A visible range of ~60 lines around a position. + let visible = |at: TextSize| { + let start = usize::from(at); + let end = content[start..] + .char_indices() + .filter(|(_, c)| *c == '\n') + .nth(60) + .map_or(content.len(), |(i, _)| start + i); + TextRange::new(at, TextSize::try_from(end).unwrap()) + }; + + for (label, at) in [("start", first), ("end", last)] { + { + let (cold, warm, hints) = time_cold_warm(workspace, || { + workspace + .pull_inlay_hints(PullInlayHintsParams { + path: path.clone(), + range: visible(at), + }) + .unwrap() + .hints + .len() + }); + println!( + " {:<40} {cold:>9.1} ms cold {warm:>9.1} ms warm ({hints} hints)", + format!("inlay hints (60 lines) @{label}") + ); + } + let (cold, warm, blocks) = time_cold_warm(workspace, || { + workspace + .on_hover(OnHoverParams { + path: path.clone(), + position: at + TextSize::from(8), + }) + .unwrap() + .markdown_blocks + .len() + }); + println!( + " {:<40} {cold:>9.1} ms cold {warm:>9.1} ms warm ({blocks} blocks)", + format!("hover @{label}") + ); + // The parts of hover: the resolved references, and the materialized snapshot. + let (cold, warm, _) = time_cold_warm(workspace, || { + let documents = workspace.documents.read().unwrap(); + let doc = documents.get(&path).unwrap(); + workspace + .references_at(&path, doc, &snapshot, at + TextSize::from(8)) + .len() + }); + println!( + " {:<40} {cold:>9.1} ms cold {warm:>9.1} ms warm", + format!(" hover: references_at @{label}") + ); + let (t, _) = time(|| { + let documents = workspace.documents.read().unwrap(); + let doc = documents.get(&path).unwrap(); + // Clear the one-entry cache, as hovering another statement would. + *workspace.snapshot_before.lock().unwrap() = None; + workspace.snapshot_before(doc, at, Arc::clone(&snapshot)); + }); + println!( + " {:<40} {t:>9.1} ms", + format!(" hover: snapshot_before (miss) @{label}") + ); + let (cold, warm, tokens) = time_cold_warm(workspace, || { + workspace + .pull_semantic_tokens(PullSemanticTokensParams { + path: path.clone(), + range: Some(visible(at)), + }) + .unwrap() + .tokens + .len() + }); + println!( + " {:<40} {cold:>9.1} ms cold {warm:>9.1} ms warm ({tokens} tokens)", + format!("semantic tokens range (60 lines) @{label}") + ); + let (t, actions) = time(|| { + workspace + .pull_code_actions(CodeActionsParams { + path: path.clone(), + cursor_position: at + TextSize::from(8), + only: vec![], + skip: vec![], + }) + .unwrap() + .actions + .len() + }); + println!( + " {:<40} {t:>9.1} ms ({actions} actions)", + format!("code actions @{label}") + ); + } + { + let (cold, warm, tokens) = time_cold_warm(workspace, || { + workspace + .pull_semantic_tokens(PullSemanticTokensParams { + path: path.clone(), + range: None, + }) + .unwrap() + .tokens + .len() + }); + println!( + " {:<40} {cold:>9.1} ms cold {warm:>9.1} ms warm ({tokens} tokens)", + "semantic tokens full" + ); + } + + // Typing in the last statement: each edit changes it, and the editor asks for the tokens and + // hints of the visible range. + let mut edits = Vec::new(); + for version in 2..5 { + let mut edited = content.to_string(); + edited.insert_str( + usize::from(last) + 6, + &" ".repeat(usize::try_from(version).unwrap()), + ); + workspace + .change_file(crate::workspace::ChangeFileParams { + path: path.clone(), + version, + content: edited, + }) + .unwrap(); + let (t, _) = time_once(|| { + workspace + .pull_semantic_tokens(PullSemanticTokensParams { + path: path.clone(), + range: Some(visible(last)), + }) + .unwrap(); + workspace + .pull_inlay_hints(PullInlayHintsParams { + path: path.clone(), + range: visible(last), + }) + .unwrap(); + }); + edits.push(t); + } + let best = edits.iter().copied().fold(f64::MAX, f64::min); + println!( + " {:<40} {best:>9.1} ms", + "edit @end, then tokens + hints @end" + ); +} + +#[ignore = "timing benchmark"] +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn perf_resolver_features(test_db: PgPool) { + let workspace = workspace(&test_db); + // Load the snapshot once, so that it doesn't count in the first request. + let path = PgLSPath::new("warmup.sql"); + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: "select 1;".into(), + version: 1, + }) + .unwrap(); + workspace + .pull_inlay_hints(PullInlayHintsParams { + path: path.clone(), + range: TextRange::up_to(9.into()), + }) + .unwrap(); + workspace + .on_hover(OnHoverParams { + path, + position: 2.into(), + }) + .unwrap(); + + let groups: Vec = std::env::var("PERF_GROUPS") + .ok() + .map(|groups| groups.split(',').map(|n| n.parse().unwrap()).collect()) + .unwrap_or_else(|| vec![250, 1000, 4000]); + let calls: usize = std::env::var("PERF_CALLS").map_or(3000, |n| n.parse().unwrap()); + // Hover iterates the document with `ParseIterator`, which recurses once per statement, so + // measure on a thread with a large stack. + std::thread::scope(|scope| { + std::thread::Builder::new() + .stack_size(1 << 30) + .spawn_scoped(scope, || { + for groups in groups { + measure( + &workspace, + &format!("migration_{groups}"), + &migration_file(groups), + ); + } + measure(&workspace, &format!("huge_{calls}"), &huge_statement(calls)); + if let Ok(rows) = std::env::var("PERF_SEED_ROWS") { + for rows in rows.split(',') { + let rows: usize = rows.parse().unwrap(); + measure(&workspace, &format!("seed_{rows}"), &seed_insert(rows)); + } + } + }) + .unwrap() + .join() + .unwrap(); + }); +} diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index e0447e456..ebdf44a87 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -31,8 +31,8 @@ use pgls_analyser::{ Analyser, AnalyserConfig, AnalyserParams, AnalysisResult, Direction, FileKind, LinterDiagnostic, LinterOptions, MigrationSection, RuleFix, }; -use pgls_catalog::resolve::{Reference, ResolveParams, resolve}; -use pgls_catalog::{Catalog, CatalogBase, Session, TransactionMode}; +use pgls_catalog::resolve::Reference; +use pgls_catalog::{Catalog, CatalogBase, Session}; use std::borrow::Cow; use pgls_catalog::Snapshot; @@ -91,6 +91,7 @@ mod connection_key; #[cfg(feature = "db")] mod connection_manager; pub(crate) mod document; +mod editor_walk; mod migration; mod pg_query; #[cfg(feature = "db")] @@ -115,6 +116,9 @@ pub struct WorkspaceServer { /// The last snapshot built for completions and hover, by the statements it includes snapshot_before: Mutex)>>, + /// The cached walks of the documents for hover, inlay hints, and semantic tokens + editor_walks: editor_walk::EditorWalks, + #[cfg(feature = "db")] connection: ConnectionManager, } @@ -143,6 +147,7 @@ impl WorkspaceServer { snapshot: SnapshotManager::new(), catalog_base: Mutex::default(), snapshot_before: Mutex::default(), + editor_walks: Mutex::default(), connection: ConnectionManager::new(), } } @@ -156,6 +161,7 @@ impl WorkspaceServer { snapshot: SnapshotManager::new(), catalog_base: Mutex::default(), snapshot_before: Mutex::default(), + editor_walks: Mutex::default(), } } @@ -462,138 +468,44 @@ impl WorkspaceServer { .collect() } - /// Walks the statements of `doc` that start before `end` like the analyser does: the catalog - /// and the session follow the statements, so that each statement sees the objects created - /// earlier in the file. Each migration section gets a session of its own, and the changes of - /// a down section don't carry into the next section. `visit` gets every parsed statement with - /// the catalog and the session as they are before it. - fn walk_statements( - &self, - doc: &Document, - snapshot: &Arc, - end: TextSize, - mut visit: impl FnMut(&pgls_analyser::AnalysableStatement, &Catalog, &Session), - ) { - let (search_path, sections) = self - .workspaces() - .settings() - .map(|settings| { - ( - settings.typecheck.search_path.clone(), - document_sections(settings, doc), - ) - }) - .unwrap_or_default(); - let search_path = snapshot.expand_search_path(&search_path); - let new_session = |section: Option<&MigrationSection>| { - Session::new(search_path.clone()).with_transaction_mode( - section.map_or(TransactionMode::Unknown, MigrationSection::transaction_mode), - ) - }; - - // Editor features show what names most likely refer to, so a `DO` block or a statement - // that doesn't parse doesn't turn them off for the rest of the file. - let mut catalog = Catalog::new(Some(self.catalog_base(snapshot))).best_effort(); - let mut session = new_session(None); - let mut section: Option<&MigrationSection> = None; - // The catalog after the up sections so far, while walking a down section. - let mut up_catalog: Option = None; - for (stmt, _, unparsable) in doc.iter(AnalyserDiagnosticsMapper) { - let Some(start) = stmt.as_ref().map(|stmt| stmt.range.start()).or_else(|| { - unparsable - .as_ref() - .map(|unparsable| unparsable.range.start()) - }) else { - continue; - }; - if start >= end { - break; - } - let current = sections - .iter() - .find(|section| section.range.contains(start)); - if current != section { - if let Some(up_catalog) = up_catalog.take() { - catalog = up_catalog; - } - if current.is_some_and(|section| section.direction == Direction::Down) { - up_catalog = Some(catalog.clone()); - } - session = new_session(current); - section = current; - } - - if let Some(unparsable) = unparsable { - if pgls_analyser::may_change_catalog(&unparsable.sql) { - catalog.taint(); - } - continue; - } - let Some(stmt) = stmt else { - continue; - }; - // The names of statements with identifier parameters like `:name` are placeholders. - if stmt.has_identifier_parameters { - if !pgls_analyser::is_plain_query(&stmt.root) { - catalog.taint(); - } - continue; - } - visit(&stmt, &catalog, &session); - // Statements in the bodies of SQL functions are not executed. - if stmt.function.is_none() { - catalog.apply(&stmt.root, session.search_path()); - session.apply(&stmt.root); - } - } - } - /// The references whose span contains `position`, in the innermost statement at `position` /// (the body of a SQL function rather than its `CREATE FUNCTION`). fn references_at( &self, + path: &PgLSPath, doc: &Document, snapshot: &Arc, position: TextSize, ) -> Vec { - let mut innermost: Option<(TextRange, TextSize, Vec)> = None; - self.walk_statements( + let statements = self.walk_references( + path, doc, snapshot, - position + TextSize::from(1), - |stmt, catalog, session| { - let Some(sql) = stmt.sql.as_deref() else { - return; - }; - if !stmt.range.contains(position) - || innermost - .as_ref() - .is_some_and(|(range, _, _)| !range.contains_range(stmt.range)) - { - return; - } - let references = resolve(ResolveParams { - stmt: &stmt.root, - catalog, - search_path: session.search_path(), - function: stmt.function.as_ref(), - sql: Some(sql), - collect_references: true, - }) - .references; - innermost = Some((stmt.range, stmt.statement_offset(position), references)); - }, + TextRange::at(position, TextSize::from(1)), ); - let Some((_, position, references)) = innermost else { + let mut innermost: Option<&editor_walk::WalkedStatement> = None; + for statement in &statements { + if statement.range.contains(position) + && innermost.is_none_or(|innermost| innermost.range.contains_range(statement.range)) + { + innermost = Some(statement); + } + } + let Some(statement) = innermost else { return Vec::new(); }; - references - .into_iter() + let position = statement + .offsets + .to_statement(position - statement.range.start()); + statement + .references + .iter() .filter(|reference| { reference .span() .is_some_and(|span| span.contains_inclusive(position)) }) + .cloned() .collect() } @@ -944,6 +856,7 @@ impl Workspace for WorkspaceServer { documents .remove(¶ms.path) .ok_or_else(WorkspaceError::not_found)?; + self.editor_walks.lock().unwrap().remove(¶ms.path); Ok(()) } @@ -1196,6 +1109,7 @@ impl Workspace for WorkspaceServer { #[cfg(feature = "db")] fn invalidate_schema_cache(&self, all: bool) -> Result<(), WorkspaceError> { + self.editor_walks.lock().unwrap().clear(); if all { // Clear all schemas - both db-loaded and json-loaded // DB completions always take precedence when a connection is available, @@ -1217,6 +1131,7 @@ impl Workspace for WorkspaceServer { #[cfg(not(feature = "db"))] fn invalidate_schema_cache(&self, _all: bool) -> Result<(), WorkspaceError> { + self.editor_walks.lock().unwrap().clear(); self.snapshot.clear_all(); Ok(()) } @@ -1774,7 +1689,7 @@ impl Workspace for WorkspaceServer { let position_in_stmt = doc .statement_offsets(&stmt_id) .to_statement(params.position - range.start()); - let references = self.references_at(doc, &snapshot, params.position); + let references = self.references_at(¶ms.path, doc, &snapshot, params.position); let snapshot = self.snapshot_before(doc, range.start(), snapshot); let markdown_blocks = pgls_hover::on_hover(pgls_hover::OnHoverParams { @@ -1836,34 +1751,18 @@ impl Workspace for WorkspaceServer { }; let mut hints = Vec::new(); - self.walk_statements( - doc, - &snapshot, - params.range.end(), - |stmt, catalog, session| { - let Some(sql) = stmt.sql.as_deref() else { - return; - }; - if stmt.range.intersect(params.range).is_none() { - return; - } - hints.extend( - pgls_inlay_hints::inlay_hints(pgls_inlay_hints::InlayHintsParams { - stmt: &stmt.root, - sql, - catalog, - search_path: session.search_path(), - function: stmt.function.as_ref(), - }) + for statement in self.walk_references(¶ms.path, doc, &snapshot, params.range) { + hints.extend( + pgls_inlay_hints::inlay_hints_of_references(&statement.references, statement.sql()) .into_iter() .map(|hint| InlayHint { - offset: stmt.document_offset(hint.offset), + offset: statement.offsets.to_document(hint.offset) + + statement.range.start(), label: hint.label, kind: hint.kind.into(), }), - ); - }, - ); + ); + } hints.sort_by_key(|hint| hint.offset); Ok(PullInlayHintsResult { hints }) @@ -1915,29 +1814,21 @@ impl Workspace for WorkspaceServer { .range .unwrap_or_else(|| TextRange::up_to(TextSize::of(doc.get_document_content()))); let mut tokens = Vec::new(); - self.walk_statements(doc, &snapshot, range.end(), |stmt, catalog, session| { - let Some(sql) = stmt.sql.as_deref() else { - return; - }; - if stmt.range.intersect(range).is_none() { - return; - } + for statement in self.walk_references(¶ms.path, doc, &snapshot, range) { tokens.extend( - pgls_semantic_tokens::semantic_tokens(pgls_semantic_tokens::SemanticTokensParams { - stmt: &stmt.root, - sql, - catalog, - search_path: session.search_path(), - function: stmt.function.as_ref(), - }) + pgls_semantic_tokens::semantic_tokens_of_references( + &statement.references, + statement.sql(), + ) .into_iter() .map(|token| SemanticToken { - range: stmt.document_range(token.range), + range: statement.offsets.range_to_document(token.range) + + statement.range.start(), kind: token.kind.into(), default_library: token.default_library, }), ); - }); + } // The body of a SQL function is a statement of its own, inside its `CREATE FUNCTION`. tokens.sort_by_key(|token| (token.range.start(), token.range.end())); @@ -2059,3 +1950,7 @@ fn restore_named_parameters( #[cfg(all(test, feature = "db"))] #[path = "server.tests.rs"] mod tests; + +#[cfg(all(test, feature = "db"))] +#[path = "server.perf_tests.rs"] +mod perf_tests; diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index a1cd0aabf..05f057246 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -16,7 +16,7 @@ use pgls_configuration::{ use pgls_configuration::plpgsql_check::PartialPlPgSqlCheckConfiguration; use pgls_diagnostics::Diagnostic; use pgls_fs::PgLSPath; -use pgls_text_size::TextRange; +use pgls_text_size::{TextRange, TextSize}; use sqlx::{Executor, PgPool}; use crate::{ @@ -1721,7 +1721,7 @@ select * from auth.users; let position = content .find('|') - .map(|idx| pgls_text_size::TextSize::new(idx as u32)) + .map(|idx| TextSize::new(idx as u32)) .expect("Unable to find cursor position in test content"); let sanitized_content = content.replace('|', ""); @@ -1781,7 +1781,7 @@ async fn completions_see_the_ddl_before_the_statement(test_db: PgPool) { let labels = |content: &str| { let position = content .find('|') - .map(|idx| pgls_text_size::TextSize::new(idx as u32)) + .map(|idx| TextSize::new(idx as u32)) .expect("Unable to find cursor position in test content"); workspace .open_file(OpenFileParams { @@ -2260,10 +2260,7 @@ async fn test_editor_features_after_unmodelled_statements(test_db: PgPool) { .expect("Unable to open test file"); let call = content.find("call").unwrap(); - let range = TextRange::new( - u32::try_from(call).unwrap().into(), - pgls_text_size::TextSize::of(content), - ); + let range = TextRange::new(u32::try_from(call).unwrap().into(), TextSize::of(content)); let hints = workspace .pull_inlay_hints(PullInlayHintsParams { path: path.clone(), @@ -2476,3 +2473,115 @@ async fn test_inlay_hints_follow_migration_sections(test_db: PgPool) { .collect(); assert_eq!(labels, ["x:1", "y:2", "a:3", "b:4", "x:5", "y:6"]); } + +/// The cached walk of a document gives the same tokens and hints as walking it from the start, +/// after edits before, inside, and after the parts of the file it has cached. +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_editor_walk_cache_follows_edits(test_db: PgPool) { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + ..Default::default() + }); + let cached = get_test_workspace(Some(conf.clone())).expect("Unable to create test workspace"); + let path = PgLSPath::new("test.sql"); + + // 300 top-level statements, so that the walk keeps checkpoints. + let file = |renamed: Option| { + let mut content = String::from( + "create function point(x int, y int) returns int language sql as 'select x + y';\n", + ); + for i in 0..100 { + let column = if renamed == Some(i) { "other" } else { "id" }; + content.push_str(&format!( + "create table t{i} ({column} int);\nselect id, point(1, 2) from t{i};\nselect id from t{} ;\n", + i / 2 + )); + } + content + }; + let results = |workspace: &WorkspaceServer, content: &str, range: Option| { + let tokens = workspace + .pull_semantic_tokens(PullSemanticTokensParams { + path: path.clone(), + range, + }) + .expect("Unable to pull semantic tokens") + .tokens + .into_iter() + .map(|token| (token.range, token.kind)) + .collect::>(); + let hints = workspace + .pull_inlay_hints(PullInlayHintsParams { + path: path.clone(), + range: range.unwrap_or(TextRange::up_to(TextSize::of(content))), + }) + .expect("Unable to pull inlay hints") + .hints; + (tokens, hints) + }; + let fresh = |content: &str, range: Option| { + let workspace = + get_test_workspace(Some(conf.clone())).expect("Unable to create test workspace"); + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: content.into(), + version: 1, + }) + .expect("Unable to open test file"); + results(&workspace, content, range) + }; + let tail = |content: &str| { + let start = content.find("create table t90 ").unwrap(); + Some(TextRange::new( + u32::try_from(start).unwrap().into(), + TextSize::of(content), + )) + }; + + let content = file(None); + cached + .open_file(OpenFileParams { + path: path.clone(), + content: content.clone(), + version: 1, + }) + .expect("Unable to open test file"); + assert_eq!( + results(&cached, &content, tail(&content)), + fresh(&content, tail(&content)) + ); + assert_eq!(results(&cached, &content, None), fresh(&content, None)); + + // Renaming a column changes what the statements after it resolve to. + for (version, renamed) in [(2, 60), (3, 95), (4, 3), (5, 60)] { + let content = file(Some(renamed)); + cached + .change_file(crate::workspace::ChangeFileParams { + path: path.clone(), + version, + content: content.clone(), + }) + .expect("Unable to change test file"); + assert_eq!( + results(&cached, &content, tail(&content)), + fresh(&content, tail(&content)), + "tail after renaming in t{renamed}" + ); + assert_eq!( + results(&cached, &content, None), + fresh(&content, None), + "everything after renaming in t{renamed}" + ); + } +} diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 5ec25226f..639018518 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -84,6 +84,23 @@ impl Document { self.iter_with_filter(mapper, NoFilter) } + /// The statements from the top-level statement at index `root` of [`Document::statement_ranges`] + /// on. + pub fn iter_from<'a, M>(&'a self, mapper: M, root: usize) -> ParseIterator<'a, M, NoFilter> + where + M: StatementMapper<'a>, + { + let mut iter = ParseIterator::new(self, mapper, NoFilter); + iter.ranges = self.ranges.get(root..).unwrap_or_default().iter(); + iter + } + + /// The ranges of the top-level statements, in order. The statements in the bodies of SQL + /// functions are not included. + pub fn statement_ranges(&self) -> &[TextRange] { + &self.ranges + } + pub fn iter_with_filter<'a, M, F>(&'a self, mapper: M, filter: F) -> ParseIterator<'a, M, F> where M: StatementMapper<'a>, @@ -197,6 +214,16 @@ where } } +/// The id and range of each statement, without parsing it. +pub struct IdMapper; +impl<'a> StatementMapper<'a> for IdMapper { + type Output = (StatementId, TextRange); + + fn map(&self, _parser: &'a Document, id: StatementId, range: TextRange) -> Self::Output { + (id, range) + } +} + pub struct DefaultMapper; impl<'a> StatementMapper<'a> for DefaultMapper { type Output = (StatementId, TextRange, String); diff --git a/crates/pgls_workspace/src/workspace/server/editor_walk.rs b/crates/pgls_workspace/src/workspace/server/editor_walk.rs new file mode 100644 index 000000000..33b921561 --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/editor_walk.rs @@ -0,0 +1,489 @@ +//! The walk of a document for hover, inlay hints, and semantic tokens: the references of its +//! statements, resolved against the catalog as the statements before them leave it. +//! +//! Walking a long file from the start for every request is slow, so the walk is cached per +//! document. Every few statements it keeps a checkpoint of the catalog and the session, and it +//! keeps the references of every statement it resolved. Both are keyed by a hash of the +//! statements before them, so that an edit only invalidates what comes after it. A request for a +//! range of the file then resolves at most the statements since the last checkpoint before it, +//! and a repeated request only looks the references up. + +use std::{ + hash::{DefaultHasher, Hash, Hasher}, + sync::{Arc, Mutex}, +}; + +use pgls_analyser::{AnalysableStatement, Direction, MigrationSection, StatementOffsets}; +use pgls_catalog::{ + Catalog, Session, Snapshot, TransactionMode, + resolve::{Reference, ResolveParams, resolve}, +}; +use pgls_fs::PgLSPath; +use pgls_text_size::TextRange; +use rustc_hash::{FxHashMap, FxHashSet}; + +use super::{ + StatementId, WorkspaceServer, + document::{AnalyserDiagnosticsMapper, Document, IdMapper, StatementMapper}, + document_sections, + sql_function::get_sql_fn_bodies, +}; + +/// A statement of the walk with its references. +pub(super) struct WalkedStatement { + pub id: StatementId, + /// The range of the statement in the document. + pub range: TextRange, + /// Where the text of the statement is in the document. + pub offsets: StatementOffsets, + pub references: Arc<[Reference]>, +} + +impl WalkedStatement { + /// The statement text, which the references point into. + pub fn sql(&self) -> &str { + self.id.content() + } +} + +/// The cached walk of one document. +#[derive(Default)] +pub(super) struct WalkCache { + /// The snapshot the walk starts from. Holding it also keeps another snapshot from being + /// allocated where it is, so comparing the pointers is enough. + snapshot: Option>, + /// A hash of the settings the walk depends on. + key: u64, + /// The state before some of the top-level statements, ordered by statement. + checkpoints: Vec, + /// The statements walked so far, by [`statement_key`]. + statements: FxHashMap, +} + +#[derive(Clone)] +struct CachedStatement { + offsets: StatementOffsets, + /// `None` for statements that are skipped: those that don't parse, and those with + /// identifier parameters. + references: Option>, + /// The statements of the function body of a top-level statement, with their ranges relative + /// to it. + children: Arc<[(StatementId, TextRange)]>, +} + +/// The state of the walk before a top-level statement. +#[derive(Clone)] +struct Checkpoint { + /// The index of the statement in [`Document::statement_ranges`]. + root: usize, + /// The hash of the statements before it, [`WalkHashes::before`]. + hash: u64, + state: WalkState, +} + +#[derive(Clone)] +struct WalkState { + catalog: Catalog, + session: Session, + /// The index of the current migration section. + section: Option, + /// The catalog after the up sections so far, while walking a down section. + up_catalog: Option, +} + +/// The hashes of the top-level statements: of everything before each one, and of everything +/// before it together with the migration section it is in, which decide what it resolves to. +struct WalkHashes { + before: Vec, + with_section: Vec, +} + +impl WalkHashes { + fn new(doc: &Document, sections: &[MigrationSection], key: u64, roots: usize) -> Self { + let content = doc.get_document_content(); + let mut before = Vec::with_capacity(roots + 1); + let mut with_section = Vec::with_capacity(roots); + let mut hash = key; + for range in doc.statement_ranges().iter().take(roots) { + before.push(hash); + let mut hasher = DefaultHasher::new(); + hash.hash(&mut hasher); + section_identity(sections, section_of(sections, range.start())).hash(&mut hasher); + let with = hasher.finish(); + with_section.push(with); + let mut hasher = DefaultHasher::new(); + with.hash(&mut hasher); + content[*range].hash(&mut hasher); + hash = hasher.finish(); + } + before.push(hash); + Self { + before, + with_section, + } + } +} + +/// Saves a checkpoint every this many top-level statements, but at most about 32 per document: +/// each holds a copy of the catalog. +fn checkpoint_interval(roots: usize) -> usize { + (roots / 32).max(64) +} + +/// The index of the migration section that contains `offset`. +fn section_of(sections: &[MigrationSection], offset: pgls_text_size::TextSize) -> Option { + sections + .iter() + .position(|section| section.range.contains(offset)) +} + +/// What a statement's resolution depends on about its section: not where it is, which changes +/// with every edit before it. +fn section_identity( + sections: &[MigrationSection], + section: Option, +) -> Option<( + usize, + std::mem::Discriminant, + std::mem::Discriminant, +)> { + let section = section?; + let found = sections.get(section)?; + Some(( + section, + std::mem::discriminant(&found.direction), + std::mem::discriminant(&found.transaction_mode()), + )) +} + +/// The key of the references of a statement: the hash of its top-level statement with its +/// section, and the statement itself, whose id includes the function it is the body of. +fn statement_key(root_hash: u64, id: &StatementId) -> u64 { + let mut hasher = DefaultHasher::new(); + root_hash.hash(&mut hasher); + id.hash(&mut hasher); + hasher.finish() +} + +impl WorkspaceServer { + /// The statements of `doc` whose range intersects `range`, with their references, in the + /// order of the walk: each top-level statement before the statements of its function body. + /// + /// The walk is like the analyser's: the catalog and the session follow the statements, so + /// that each statement sees the objects created earlier in the file. Each migration section + /// gets a session of its own, and the changes of a down section don't carry into the next + /// section. + pub(super) fn walk_references( + &self, + path: &PgLSPath, + doc: &Document, + snapshot: &Arc, + range: TextRange, + ) -> Vec { + let (search_path, sections) = self + .workspaces() + .settings() + .map(|settings| { + ( + settings.typecheck.search_path.clone(), + document_sections(settings, doc), + ) + }) + .unwrap_or_default(); + let search_path = snapshot.expand_search_path(&search_path); + let key = { + let mut hasher = DefaultHasher::new(); + search_path.hash(&mut hasher); + hasher.finish() + }; + + let cache = Arc::clone( + self.editor_walks + .lock() + .unwrap() + .entry(path.clone()) + .or_default(), + ); + let mut cache = cache.lock().unwrap(); + if cache.key != key + || !cache + .snapshot + .as_ref() + .is_some_and(|cached| Arc::ptr_eq(cached, snapshot)) + { + *cache = WalkCache { + snapshot: Some(Arc::clone(snapshot)), + key, + ..WalkCache::default() + }; + } + + let ranges = doc.statement_ranges(); + // The top-level statements that may intersect the range, and those before them. + let first = ranges.partition_point(|root| root.end() < range.start()); + let end = ranges.partition_point(|root| !is_after(*root, range)); + let hashes = WalkHashes::new(doc, §ions, key, end); + + if let Some(statements) = cached_statements(&cache, doc, &hashes, first, end, range) { + return statements; + } + + // Start from the last valid checkpoint before the range. + let checkpoint = cache + .checkpoints + .iter() + .rev() + .find(|checkpoint| { + checkpoint.root <= first + && hashes.before.get(checkpoint.root) == Some(&checkpoint.hash) + }) + .cloned(); + let (start, mut state) = match checkpoint { + Some(checkpoint) => (checkpoint.root, checkpoint.state), + None => ( + 0, + WalkState { + // Editor features show what names most likely refer to, so a `DO` block + // or a statement that doesn't parse doesn't turn them off for the rest of + // the file. + catalog: Catalog::new(Some(self.catalog_base(snapshot))).best_effort(), + session: new_session(&search_path, §ions, None), + section: None, + up_catalog: None, + }, + ), + }; + + let interval = checkpoint_interval(ranges.len()); + let mut statements = Vec::new(); + let mut root = start.wrapping_sub(1); + for (id, stmt_range) in doc.iter_from(IdMapper, start) { + if is_after(stmt_range, range) { + break; + } + if id.is_root() { + root = root.wrapping_add(1); + if root > start && root % interval == 0 { + save_checkpoint(&mut cache, root, hashes.before[root], &state); + } + } + let in_range = stmt_range.intersect(range).is_some(); + // The statements of function bodies before the range change nothing. + if id.is_child() && !in_range { + continue; + } + + let current = section_of(§ions, stmt_range.start()); + if current != state.section { + if let Some(up_catalog) = state.up_catalog.take() { + state.catalog = up_catalog; + } + if current + .and_then(|current| sections.get(current)) + .is_some_and(|section| section.direction == Direction::Down) + { + state.up_catalog = Some(state.catalog.clone()); + } + state.session = new_session(&search_path, §ions, current); + state.section = current; + } + + let key = statement_key(hashes.with_section[root], &id); + let (stmt, _, unparsable) = AnalyserDiagnosticsMapper.map(doc, id.clone(), stmt_range); + let skipped = match (&stmt, &unparsable) { + (_, Some(unparsable)) => { + if pgls_analyser::may_change_catalog(&unparsable.sql) { + state.catalog.taint(); + } + true + } + (None, None) => true, + // The names of statements with identifier parameters like `:name` are + // placeholders. + (Some(stmt), None) if stmt.has_identifier_parameters => { + if !pgls_analyser::is_plain_query(&stmt.root) { + state.catalog.taint(); + } + true + } + (Some(_), None) => false, + }; + if in_range { + let cached = match cache.statements.get(&key) { + Some(cached) => cached.clone(), + None => { + let children = match (&stmt, id.is_root()) { + (Some(stmt), true) => get_sql_fn_bodies(&stmt.root, id.content()) + .into_iter() + .map(|body| (id.create_child(&body.body), body.range)) + .collect(), + _ => Arc::default(), + }; + let cached = CachedStatement { + offsets: stmt + .as_ref() + .map(|stmt| stmt.offsets.clone()) + .unwrap_or_default(), + references: stmt + .as_ref() + .filter(|_| !skipped) + .map(|stmt| resolve_statement(stmt, &state).into()), + children, + }; + cache.statements.insert(key, cached.clone()); + cached + } + }; + if let Some(references) = cached.references { + statements.push(WalkedStatement { + id, + range: stmt_range, + offsets: cached.offsets, + references, + }); + } + } + let Some(stmt) = stmt.filter(|_| !skipped) else { + continue; + }; + // Statements in the bodies of SQL functions are not executed. + if stmt.function.is_none() { + state.catalog.apply(&stmt.root, state.session.search_path()); + state.session.apply(&stmt.root); + } + } + + prune(&mut cache, doc, §ions, key); + statements + } +} + +/// Whether the statement starts after the range. A statement that starts where a non-empty range +/// ends is after it. +fn is_after(statement: TextRange, range: TextRange) -> bool { + statement.start() > range.end() || (statement.start() == range.end() && !range.is_empty()) +} + +fn new_session( + search_path: &[String], + sections: &[MigrationSection], + section: Option, +) -> Session { + Session::new(search_path.to_vec()).with_transaction_mode( + section + .and_then(|section| sections.get(section)) + .map_or(TransactionMode::Unknown, MigrationSection::transaction_mode), + ) +} + +fn resolve_statement(stmt: &AnalysableStatement, state: &WalkState) -> Vec { + let Some(sql) = stmt.sql.as_deref() else { + return Vec::new(); + }; + resolve(ResolveParams { + stmt: &stmt.root, + catalog: &state.catalog, + search_path: state.session.search_path(), + function: stmt.function.as_ref(), + sql: Some(sql), + collect_references: true, + }) + .references +} + +/// The statements in `range` from the cache, if all of them are cached. +fn cached_statements( + cache: &WalkCache, + doc: &Document, + hashes: &WalkHashes, + first: usize, + end: usize, + range: TextRange, +) -> Option> { + let content = doc.get_document_content(); + let mut statements = Vec::new(); + for (root, root_range) in doc + .statement_ranges() + .iter() + .enumerate() + .take(end) + .skip(first) + { + let id = StatementId::new(&content[*root_range]); + let cached = cache + .statements + .get(&statement_key(hashes.with_section[root], &id))?; + let children = Arc::clone(&cached.children); + if let Some(references) = &cached.references { + statements.push(WalkedStatement { + id, + range: *root_range, + offsets: cached.offsets.clone(), + references: Arc::clone(references), + }); + } + for (child, child_range) in children.iter() { + let child_range = *child_range + root_range.start(); + if is_after(child_range, range) || child_range.intersect(range).is_none() { + continue; + } + let cached = cache + .statements + .get(&statement_key(hashes.with_section[root], child))?; + if let Some(references) = &cached.references { + statements.push(WalkedStatement { + id: child.clone(), + range: child_range, + offsets: cached.offsets.clone(), + references: Arc::clone(references), + }); + } + } + } + Some(statements) +} + +fn save_checkpoint(cache: &mut WalkCache, root: usize, hash: u64, state: &WalkState) { + let index = cache + .checkpoints + .partition_point(|checkpoint| checkpoint.root < root); + let checkpoint = Checkpoint { + root, + hash, + state: state.clone(), + }; + match cache.checkpoints.get_mut(index) { + Some(existing) if existing.root == root => { + if existing.hash != hash { + *existing = checkpoint; + } + } + _ => cache.checkpoints.insert(index, checkpoint), + } +} + +/// Drops the references of statements that are no longer in the document, once there are many +/// more than the document has statements. +fn prune(cache: &mut WalkCache, doc: &Document, sections: &[MigrationSection], key: u64) { + let roots = doc.statement_ranges().len(); + if cache.statements.len() <= 2 * roots + 1024 { + return; + } + let hashes = WalkHashes::new(doc, sections, key, roots); + let current: FxHashSet = doc + .iter_from(IdMapper, 0) + .scan(usize::MAX, |root, (id, _)| { + if id.is_root() { + *root = root.wrapping_add(1); + } + Some(statement_key(hashes.with_section[*root], &id)) + }) + .collect(); + cache.statements.retain(|key, _| current.contains(key)); + cache + .checkpoints + .retain(|checkpoint| hashes.before.get(checkpoint.root) == Some(&checkpoint.hash)); +} + +/// The cached walks of the open documents. +pub(super) type EditorWalks = Mutex>>>;