From 3bc2d4860b4c68266dbf3ecdc5000ab3c1613229 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Mon, 20 Apr 2026 23:52:37 -0400 Subject: [PATCH] ide: internal cleanup goto def & find ref tests this should help when we add support for extension stubs --- crates/squawk_ide/src/builtins.rs | 8 +- crates/squawk_ide/src/find_references.rs | 51 ++++++----- crates/squawk_ide/src/goto_definition.rs | 109 ++++++++++++----------- 3 files changed, 90 insertions(+), 78 deletions(-) diff --git a/crates/squawk_ide/src/builtins.rs b/crates/squawk_ide/src/builtins.rs index 531abb07..dac7a3f9 100644 --- a/crates/squawk_ide/src/builtins.rs +++ b/crates/squawk_ide/src/builtins.rs @@ -6,7 +6,7 @@ use url::Url; use crate::db::File; -pub(crate) const BUILTINS_SQL: &str = include_str!("generated/builtins.sql"); +const BUILTINS_SQL: &str = include_str!("generated/builtins.sql"); #[salsa::tracked] pub fn builtins_file(db: &dyn Db) -> File { @@ -15,13 +15,15 @@ pub fn builtins_file(db: &dyn Db) -> File { #[cfg(not(target_arch = "wasm32"))] #[salsa::tracked] -pub fn builtins_url(_db: &dyn Db) -> Option { +pub fn builtins_url(db: &dyn Db) -> Option { let strategy = etcetera::base_strategy::choose_base_strategy().ok()?; let config_dir = strategy.config_dir(); let cache_dir = config_dir.join("squawk/stubs"); let path = cache_dir.join("builtins.sql"); std::fs::create_dir_all(&cache_dir).ok()?; - std::fs::write(&path, BUILTINS_SQL).ok()?; + let builtins = builtins_file(db); + let builtins_sql = builtins.content(db); + std::fs::write(&path, builtins_sql.as_ref()).ok()?; Url::from_file_path(path).ok() } diff --git a/crates/squawk_ide/src/find_references.rs b/crates/squawk_ide/src/find_references.rs index 2be7fecb..c99322d7 100644 --- a/crates/squawk_ide/src/find_references.rs +++ b/crates/squawk_ide/src/find_references.rs @@ -77,60 +77,65 @@ fn find_target_defs( #[cfg(test)] mod test { - use crate::builtins::{BUILTINS_SQL, builtins_file}; + use crate::builtins::builtins_file; use crate::db::{Database, File}; use crate::find_references::find_references; use crate::test_utils::fixture; use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle}; use insta::assert_snapshot; use rowan::TextRange; + use rustc_hash::FxHashMap; #[track_caller] fn find_refs(sql: &str) -> String { let (mut offset, sql) = fixture(sql); offset = offset.checked_sub(1.into()).unwrap_or_default(); let db = Database::default(); - let file = File::new(&db, sql.clone().into()); - assert_eq!(crate::db::parse(&db, file).errors(), vec![]); - - let references = find_references(&db, file, offset); + let current_file = File::new(&db, sql.clone().into()); + assert_eq!(crate::db::parse(&db, current_file).errors(), vec![]); + let references = find_references(&db, current_file, offset); let offset_usize: usize = offset.into(); - let mut current_refs = vec![]; - let mut builtin_refs = vec![]; - let builtins_file = builtins_file(&db); + let mut file_paths = FxHashMap::default(); + file_paths.insert(current_file, "current.sql"); + file_paths.insert(builtins_file(&db), "builtins.sql"); + + let mut refs_by_file: FxHashMap> = FxHashMap::default(); for (i, location) in references.iter().enumerate() { - let label_index = i + 1; - if location.file == file { - current_refs.push((label_index, location.range)); - } else if location.file == builtins_file { - builtin_refs.push((label_index, location.range)); - } + refs_by_file + .entry(location.file) + .or_default() + .push((i + 1, location.range)); } - let has_builtins = !builtin_refs.is_empty(); + let multi_file = refs_by_file.len() > 1 || !refs_by_file.contains_key(¤t_file); let mut snippet = Snippet::source(&sql).fold(true); - if has_builtins { - snippet = snippet.path("current.sql"); + if multi_file { + snippet = snippet.path(*file_paths.get(¤t_file).unwrap()); } snippet = snippet.annotation( AnnotationKind::Context .span(offset_usize..offset_usize + 1) .label("0. query"), ); - snippet = annotate_refs(snippet, current_refs); + if let Some(current_refs) = refs_by_file.remove(¤t_file) { + snippet = annotate_refs(snippet, current_refs); + } let mut groups = vec![Level::INFO.primary_title("references").element(snippet)]; - if has_builtins { - let builtins_snippet = Snippet::source(BUILTINS_SQL).path("builtin.sql").fold(true); - let builtins_snippet = annotate_refs(builtins_snippet, builtin_refs); + for (ref_file, refs) in refs_by_file { + let path = file_paths.get(&ref_file).unwrap(); + let other_snippet = Snippet::source(ref_file.content(&db).as_ref()) + .path(*path) + .fold(true); + let other_snippet = annotate_refs(other_snippet, refs); groups.push( Level::INFO .primary_title("references") - .element(builtins_snippet), + .element(other_snippet), ); } @@ -429,7 +434,7 @@ select now(); │ ─── 2. reference ╰╴ - ╭▸ builtin.sql:11089:28 + ╭▸ builtins.sql:11089:28 │ 11089 │ create function pg_catalog.now() returns timestamp with time zone ╰╴ ─── 3. reference diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 6d48a2b8..0737ce93 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -149,7 +149,7 @@ fn find_following_commit_or_rollback( #[cfg(test)] mod test { - use crate::builtins::{BUILTINS_SQL, builtins_file}; + use crate::builtins::builtins_file; use crate::db::{Database, File}; use crate::goto_definition::goto_definition; use crate::test_utils::fixture; @@ -157,6 +157,7 @@ mod test { use insta::assert_snapshot; use log::info; use rowan::TextRange; + use rustc_hash::FxHashMap; #[track_caller] fn goto(sql: &str) -> String { @@ -171,60 +172,64 @@ mod test { // marker after the item we're trying to go to def on. offset = offset.checked_sub(1.into()).unwrap_or_default(); let db = Database::default(); - let file = File::new(&db, sql.clone().into()); - assert_eq!(crate::db::parse(&db, file).errors(), vec![]); - let results = goto_definition(&db, file, offset); - if !results.is_empty() { - let offset: usize = offset.into(); - let mut current_dests = vec![]; - let mut builtin_dests = vec![]; - let builtins_file = builtins_file(&db); - for (i, location) in results.iter().enumerate() { - let label_index = i + 2; - if location.file == file { - current_dests.push((label_index, location.range)); - } else if location.file == builtins_file { - builtin_dests.push((label_index, location.range)); - } - } - - let has_builtins = !builtin_dests.is_empty(); + let current_file = File::new(&db, sql.into()); + assert_eq!(crate::db::parse(&db, current_file).errors(), vec![]); + let results = goto_definition(&db, current_file, offset); + if results.is_empty() { + return None; + } - let mut snippet = Snippet::source(&sql).fold(true); - if has_builtins { - // only show the current file when we have two file types, aka current and builtins - snippet = snippet.path("current.sql"); - } else { - snippet = annotate_destinations(snippet, current_dests); - } - snippet = snippet.annotation( - AnnotationKind::Context - .span(offset..offset + 1) - .label("1. source"), - ); + let mut file_paths = FxHashMap::default(); + file_paths.insert(current_file, "current.sql"); + file_paths.insert(builtins_file(&db), "builtins.sql"); + + let offset: usize = offset.into(); + let mut dests_by_file: FxHashMap> = FxHashMap::default(); + for (i, location) in results.iter().enumerate() { + dests_by_file + .entry(location.file) + .or_default() + .push((i + 2, location.range)); + } - let mut groups = vec![Level::INFO.primary_title("definition").element(snippet)]; + let multi_file = dests_by_file.len() > 1 || !dests_by_file.contains_key(¤t_file); - if has_builtins { - let builtins_snippet = Snippet::source(BUILTINS_SQL).path("builtin.sql").fold(true); - let builtins_snippet = annotate_destinations(builtins_snippet, builtin_dests); - groups.push( - Level::INFO - .primary_title("definition") - .element(builtins_snippet), - ); - } + let mut snippet = Snippet::source(current_file.content(&db).as_ref()).fold(true); + if multi_file { + snippet = snippet.path(*file_paths.get(¤t_file).unwrap()); + } + if let Some(current_dests) = dests_by_file.remove(¤t_file) { + snippet = annotate_destinations(snippet, current_dests); + } + snippet = snippet.annotation( + AnnotationKind::Context + .span(offset..offset + 1) + .label("1. source"), + ); - let renderer = Renderer::plain().decor_style(DecorStyle::Unicode); - return Some( - renderer - .render(&groups) - .to_string() - // hacky cleanup to make the text shorter - .replace("info: definition", ""), + let mut groups = vec![Level::INFO.primary_title("definition").element(snippet)]; + + for (dest_file, dests) in dests_by_file { + let path = file_paths.get(&dest_file).unwrap(); + let other_snippet = Snippet::source(dest_file.content(&db).as_ref()) + .path(*path) + .fold(true); + let other_snippet = annotate_destinations(other_snippet, dests); + groups.push( + Level::INFO + .primary_title("definition") + .element(other_snippet), ); } - None + + let renderer = Renderer::plain().decor_style(DecorStyle::Unicode); + Some( + renderer + .render(&groups) + .to_string() + // hacky cleanup to make the text shorter + .replace("info: definition", ""), + ) } fn goto_not_found(sql: &str) { @@ -870,7 +875,7 @@ select now$0(); │ ─ 1. source ╰╴ - ╭▸ builtin.sql:11089:28 + ╭▸ builtins.sql:11089:28 │ 11089 │ create function pg_catalog.now() returns timestamp with time zone ╰╴ ─── 2. destination @@ -888,7 +893,7 @@ select current_timestamp$0; │ ─ 1. source ╰╴ - ╭▸ builtin.sql:11089:28 + ╭▸ builtins.sql:11089:28 │ 11089 │ create function pg_catalog.now() returns timestamp with time zone ╰╴ ─── 2. destination @@ -989,7 +994,7 @@ select * from t where current_timestamp$0 > t.created_at; │ ─ 1. source ╰╴ - ╭▸ builtin.sql:11089:28 + ╭▸ builtins.sql:11089:28 │ 11089 │ create function pg_catalog.now() returns timestamp with time zone ╰╴ ─── 2. destination