ide: internal cleanup goto def & find ref tests (#1080)

this should help when we add support for extension stubs

Author: sbdchd

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<Url> {
+pub fn builtins_url(db: &dyn Db) -> Option<Url> {
     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()
 }
 

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<File, Vec<(usize, TextRange)>> = 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(&current_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(&current_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(&current_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

crates/squawk_ide/src/goto_definition.rs

@@ -149,14 +149,15 @@ 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;
     use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle};
     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<File, Vec<(usize, TextRange)>> = 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(&current_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(&current_file).unwrap());
+        }
+        if let Some(current_dests) = dests_by_file.remove(&current_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