Skip to content

perf(workspace): keep every statement of a document cached and walk once for hover #913

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
psteinroe wants to merge 1 commit into
base: fix/statement-iterator-stack-overflow
Choose a base branch
from
+263 −129
Open

perf(workspace): keep every statement of a document cached and walk once for hover #913

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 1 addition & 22 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion crates/pgls_workspace/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ version = "0.0.0"

[dependencies]
globset = "0.4.16"
lru = "0.16.3"

ignore = { workspace = true }
json-strip-comments = { workspace = true }
Expand Down
186 changes: 113 additions & 73 deletions crates/pgls_workspace/src/workspace/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::{
borrow::Borrow,
collections::HashMap,
fs,
hash::{DefaultHasher, Hash, Hasher},
Expand Down Expand Up @@ -29,8 +30,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::{ColumnSource, Reference, ResolveParams, resolve};
use pgls_catalog::{Catalog, CatalogBase, Origin, Session, TransactionMode};
use std::borrow::Cow;

use pgls_catalog::Snapshot;
Expand Down Expand Up @@ -109,13 +110,22 @@ pub struct WorkspaceServer {
/// The indexed database snapshot the catalog starts from
catalog_base: Mutex<Option<Arc<CatalogBase>>>,

/// The last snapshot built for completions and hover, by the statements it includes
snapshot_before: Mutex<Option<(u64, Arc<Snapshot>)>>,
/// The last snapshot built for completions and hover
snapshot_before: Mutex<Option<SnapshotBefore>>,

#[cfg(feature = "db")]
connection: ConnectionManager,
}

/// A database snapshot with the changes of the statements before a statement.
struct SnapshotBefore {
/// The database snapshot it is built from.
base: Arc<Snapshot>,
/// The text that the statement sees.
key: u64,
snapshot: Arc<Snapshot>,
}

/// The `Workspace` object is long-lived, so we want it to be able to cross
/// unwind boundaries.
/// In return, we have to make sure operations on the workspace either do not
Expand Down Expand Up @@ -463,14 +473,15 @@ impl WorkspaceServer {
/// 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.
/// the catalog and the session as they are before it. Returns the catalog after the walked
/// statements.
fn walk_statements(
&self,
doc: &Document,
snapshot: &Arc<Snapshot>,
end: TextSize,
mut visit: impl FnMut(&pgls_analyser::AnalysableStatement, &Catalog, &Session),
) {
) -> Catalog {
let (search_path, sections) = self
.workspaces()
.settings()
Expand Down Expand Up @@ -541,56 +552,64 @@ impl WorkspaceServer {
session.apply(&stmt.root);
}
}
catalog
}

/// 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(
/// For hover on the statement `id` at `range`: the references whose span contains
/// `position`, and the database snapshot as the statement sees it.
fn hover_references(
&self,
doc: &Document,
snapshot: &Arc<Snapshot>,
id: &StatementId,
range: TextRange,
position: TextSize,
) -> Vec<Reference> {
let mut innermost: Option<(TextRange, Vec<Reference>)> = None;
self.walk_statements(
snapshot: Arc<Snapshot>,
) -> (Vec<Reference>, Arc<Snapshot>) {
let mut found = None;
let catalog = self.walk_statements(
doc,
snapshot,
&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))
{
if stmt.range != range {
return;
}
let references = resolve(ResolveParams {
let position = position - range.start();
let (references, others): (Vec<_>, Vec<_>) = 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, references));
.references
.into_iter()
.partition(|reference| {
reference
.span()
.is_some_and(|span| span.contains_inclusive(position))
});
// Objects that the file didn't change are the same in the database snapshot.
let snapshot = if in_database(&references, &others) {
Arc::clone(&snapshot)
} else {
self.snapshot_before(doc, id, range, Arc::clone(&snapshot), || catalog)
};
found = Some((references, snapshot));
},
);
let Some((range, references)) = innermost else {
return Vec::new();
};
let position = position - range.start();
references
.into_iter()
.filter(|reference| {
reference
.span()
.is_some_and(|span| span.contains_inclusive(position))
})
.collect()
// A statement that doesn't parse is not visited. It changes nothing, so it sees the
// catalog after the walk.
found.unwrap_or_else(|| {
(
Vec::new(),
self.snapshot_before(doc, id, range, snapshot, || catalog),
)
})
}

/// The indexed database snapshot, built once per snapshot.
Expand All @@ -604,53 +623,41 @@ impl WorkspaceServer {
base
}

/// The database snapshot as the statement at `statement_start` sees it: with the objects
/// that the statements before it create, change, or drop.
fn snapshot_before(
/// The database snapshot as the statement `id` at `range` sees it: with the objects that the
/// statements before it create, change, or drop. `catalog` returns the catalog that the
/// statement sees, and is only called when the snapshot is not cached.
fn snapshot_before<C: Borrow<Catalog>>(
&self,
doc: &Document,
statement_start: TextSize,
id: &StatementId,
range: TextRange,
snapshot: Arc<Snapshot>,
catalog: impl FnOnce() -> C,
) -> Arc<Snapshot> {
let statements: Vec<_> = doc
.iter(ExecuteStatementMapper)
.take_while(|(_, range, _, _)| range.end() <= statement_start)
.filter(|(id, _, _, _)| !id.is_child())
.filter_map(|(_, _, sql, ast)| Some((sql, ast?)))
.collect();
if statements.is_empty() {
return snapshot;
}

let search_path = self
.workspaces()
.settings()
.map(|settings| settings.typecheck.search_path.clone())
.unwrap_or_default();

// The statement sees the text before it, and the body of a SQL function also sees its
// `CREATE FUNCTION`.
let mut hasher = DefaultHasher::new();
Arc::as_ptr(&snapshot).hash(&mut hasher);
search_path.hash(&mut hasher);
for (sql, _) in &statements {
sql.hash(&mut hasher);
}
doc.get_document_content()[..usize::from(range.start())].hash(&mut hasher);
id.parent_content().hash(&mut hasher);
let key = hasher.finish();

let mut cached = self.snapshot_before.lock().unwrap();
if let Some((cached_key, cached_snapshot)) = cached.as_ref()
&& *cached_key == key
if let Some(cached) = cached.as_ref()
&& cached.key == key
&& Arc::ptr_eq(&cached.base, &snapshot)
{
return Arc::clone(cached_snapshot);
return Arc::clone(&cached.snapshot);
}

let mut catalog = Catalog::new(Some(self.catalog_base(&snapshot)));
let mut session = Session::new(snapshot.expand_search_path(&search_path));
for (_, ast) in &statements {
catalog.apply(ast, session.search_path());
session.apply(ast);
}
let result = catalog.snapshot().unwrap_or(snapshot);
*cached = Some((key, Arc::clone(&result)));
let result = catalog()
.borrow()
.snapshot()
.unwrap_or_else(|| Arc::clone(&snapshot));
*cached = Some(SnapshotBefore {
base: snapshot,
key,
snapshot: Arc::clone(&result),
});
result
}

Expand Down Expand Up @@ -913,6 +920,8 @@ impl Workspace for WorkspaceServer {
params.vcs_base_path,
params.gitignore_matches.as_slice(),
)?;
// The search path and the migration sections change what a statement sees.
*self.snapshot_before.lock().unwrap() = None;

Ok(())
}
Expand Down Expand Up @@ -1707,7 +1716,10 @@ impl Workspace for WorkspaceServer {
}
Some((id, range, cst)) => {
let position = params.position - range.start();
let snapshot = self.snapshot_before(parsed_doc, range.start(), snapshot);
let snapshot =
self.snapshot_before(parsed_doc, &id, range, Arc::clone(&snapshot), || {
self.walk_statements(parsed_doc, &snapshot, range.start(), |_, _, _| {})
});

let items = pgls_completions::complete(pgls_completions::CompletionParams {
position,
Expand Down Expand Up @@ -1756,8 +1768,8 @@ impl Workspace for WorkspaceServer {
match doc.iter_at(WithCSTandASTMapper, params.position).last() {
Some((stmt_id, range, ts_tree, maybe_ast)) => {
let position_in_stmt = params.position - range.start();
let references = self.references_at(doc, &snapshot, params.position);
let snapshot = self.snapshot_before(doc, range.start(), snapshot);
let (references, snapshot) =
self.hover_references(doc, &stmt_id, range, params.position, snapshot);

let markdown_blocks = pgls_hover::on_hover(pgls_hover::OnHoverParams {
ts_tree: &ts_tree,
Expand Down Expand Up @@ -1844,6 +1856,34 @@ impl Workspace for WorkspaceServer {
}
}

/// Whether hover finds what `references` resolve to in the database snapshot: objects that the
/// file didn't change, and columns of CTEs and subqueries. `others` are the other references of
/// the statement.
fn in_database(references: &[Reference], others: &[Reference]) -> bool {
// The statement references each relation that it takes columns from.
let unchanged_relation = |schema: &str, name: &str| {
references.iter().chain(others).any(|reference| {
matches!(
reference,
Reference::Relation { relation, .. }
if relation.schema == schema
&& relation.name == name
&& relation.origin == Origin::Database
)
})
};
!references.is_empty()
&& references.iter().all(|reference| match reference {
Reference::Relation { relation, .. } => relation.origin == Origin::Database,
Reference::Function { function, .. } => function.origin == Origin::Database,
Reference::Column {
source: ColumnSource::Relation { schema, name },
..
} => unchanged_relation(schema, name),
Reference::Column { .. } => true,
})
}

/// The migration sections of `doc`, like the up and down migrations of dbmate and goose.
fn document_sections(settings: &Settings, doc: &Document) -> Vec<MigrationSection> {
settings
Expand Down
Loading
Loading