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