ide: find ref support for builtins

crates/squawk_ide/src/find_references.rs

@@ -1,74 +1,122 @@
 use crate::binder::{self, Binder};
+use crate::builtins::BUILTINS_SQL;
+use crate::goto_definition::{FileId, Location};
 use crate::offsets::token_from_offset;
 use crate::resolve;
-use rowan::{TextRange, TextSize};
+use rowan::TextSize;
 use smallvec::{SmallVec, smallvec};
-use squawk_syntax::SyntaxNode;
 use squawk_syntax::{
     SyntaxNodePtr,
     ast::{self, AstNode},
     match_ast,
 };
 
-pub fn find_references(file: &ast::SourceFile, offset: TextSize) -> Vec<TextRange> {
-    let binder = binder::bind(file);
-    let root = file.syntax();
-    let Some(targets) = find_targets(file, root, offset, &binder) else {
+pub fn find_references(file: &ast::SourceFile, offset: TextSize) -> Vec<Location> {
+    let current_binder = binder::bind(file);
+
+    let builtins_tree = ast::SourceFile::parse(BUILTINS_SQL).tree();
+    let builtins_binder = binder::bind(&builtins_tree);
+
+    let Some((target_file, target_defs)) = find_target_defs(
+        file,
+        offset,
+        &current_binder,
+        &builtins_tree,
+        &builtins_binder,
+    ) else {
         return vec![];
     };
 
-    let mut refs = vec![];
+    let (binder, root) = match target_file {
+        FileId::Current => (&current_binder, file.syntax()),
+        FileId::Builtins => (&builtins_binder, builtins_tree.syntax()),
+    };
+
+    let mut refs: Vec<Location> = vec![];
+
+    if target_file == FileId::Builtins {
+        for ptr in &target_defs {
+            refs.push(Location {
+                file: FileId::Builtins,
+                range: ptr.to_node(builtins_tree.syntax()).text_range(),
+            });
+        }
+    }
+
     for node in file.syntax().descendants() {
         match_ast! {
             match node {
                 ast::NameRef(name_ref) => {
-                    if let Some(found_refs) = resolve::resolve_name_ref_ptrs(&binder, root, &name_ref)
-                        && found_refs.iter().any(|ptr| targets.contains(ptr))
+                    // Check if the ref matches one of the defs
+                    if let Some(found_defs) = resolve::resolve_name_ref_ptrs(binder, root, &name_ref)
+                        && found_defs.iter().any(|def| target_defs.contains(def))
                     {
-                        refs.push(name_ref.syntax().text_range());
+                        refs.push(Location {
+                            file: FileId::Current,
+                            range: name_ref.syntax().text_range(),
+                        });
                     }
                 },
                 ast::Name(name) => {
+                    // Find refs also includes the defs so we have to check.
                     let found = SyntaxNodePtr::new(name.syntax());
-                    if targets.contains(&found) {
-                        refs.push(name.syntax().text_range());
+                    if target_defs.contains(&found) {
+                        refs.push(Location {
+                            file: FileId::Current,
+                            range: name.syntax().text_range(),
+                        });
                     }
                 },
                 _ => (),
             }
         }
     }
 
-    refs.sort_by_key(|range| range.start());
+    refs.sort_by_key(|loc| (loc.file, loc.range.start()));
     refs
 }
 
-fn find_targets(
+fn find_target_defs(
     file: &ast::SourceFile,
-    root: &SyntaxNode,
     offset: TextSize,
-    binder: &Binder,
-) -> Option<SmallVec<[SyntaxNodePtr; 1]>> {
+    current_binder: &Binder,
+    builtins_tree: &ast::SourceFile,
+    builtins_binder: &Binder,
+) -> Option<(FileId, SmallVec<[SyntaxNodePtr; 1]>)> {
     let token = token_from_offset(file, offset)?;
     let parent = token.parent()?;
 
     if let Some(name) = ast::Name::cast(parent.clone()) {
-        return Some(smallvec![SyntaxNodePtr::new(name.syntax())]);
+        return Some((
+            FileId::Current,
+            smallvec![SyntaxNodePtr::new(name.syntax())],
+        ));
     }
 
     if let Some(name_ref) = ast::NameRef::cast(parent.clone()) {
-        return resolve::resolve_name_ref_ptrs(binder, root, &name_ref);
+        for file_id in [FileId::Current, FileId::Builtins] {
+            let (binder, root) = match file_id {
+                FileId::Current => (current_binder, file.syntax()),
+                FileId::Builtins => (builtins_binder, builtins_tree.syntax()),
+            };
+            if let Some(ptrs) = resolve::resolve_name_ref_ptrs(binder, root, &name_ref) {
+                return Some((file_id, ptrs));
+            }
+        }
     }
 
     None
 }
 
 #[cfg(test)]
 mod test {
+    use crate::builtins::BUILTINS_SQL;
     use crate::find_references::find_references;
+    use crate::goto_definition::FileId;
     use crate::test_utils::fixture;
     use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle};
     use insta::assert_snapshot;
+    use rowan::TextRange;
     use squawk_syntax::ast;
 
     #[track_caller]
@@ -83,32 +131,62 @@ mod test {
 
         let offset_usize: usize = offset.into();
 
-        let labels: Vec<String> = (1..=references.len())
-            .map(|i| format!("{}. reference", i))
-            .collect();
+        let mut current_refs = vec![];
+        let mut builtin_refs = vec![];
+        for (i, location) in references.iter().enumerate() {
+            let label_index = i + 1;
+            match location.file {
+                FileId::Current => current_refs.push((label_index, location.range)),
+                FileId::Builtins => builtin_refs.push((label_index, location.range)),
+            }
+        }
 
-        let mut snippet = Snippet::source(&sql).fold(true).annotation(
+        let has_builtins = !builtin_refs.is_empty();
+
+        let mut snippet = Snippet::source(&sql).fold(true);
+        if has_builtins {
+            snippet = snippet.path("current.sql");
+        }
+        snippet = snippet.annotation(
             AnnotationKind::Context
                 .span(offset_usize..offset_usize + 1)
                 .label("0. query"),
         );
+        snippet = annotate_refs(snippet, current_refs);
 
-        for (i, range) in references.iter().enumerate() {
-            snippet = snippet.annotation(
-                AnnotationKind::Context
-                    .span((*range).into())
-                    .label(&labels[i]),
+        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);
+            groups.push(
+                Level::INFO
+                    .primary_title("references")
+                    .element(builtins_snippet),
             );
         }
 
-        let group = Level::INFO.primary_title("references").element(snippet);
         let renderer = Renderer::plain().decor_style(DecorStyle::Unicode);
         renderer
-            .render(&[group])
+            .render(&groups)
             .to_string()
             .replace("info: references", "")
     }
 
+    fn annotate_refs<'a>(
+        mut snippet: Snippet<'a, annotate_snippets::Annotation<'a>>,
+        refs: Vec<(usize, TextRange)>,
+    ) -> Snippet<'a, annotate_snippets::Annotation<'a>> {
+        for (label_index, range) in refs {
+            snippet = snippet.annotation(
+                AnnotationKind::Context
+                    .span(range.into())
+                    .label(format!("{}. reference", label_index)),
+            );
+        }
+        snippet
+    }
+
     #[test]
     fn simple_table_reference() {
         assert_snapshot!(find_refs("
@@ -365,4 +443,28 @@ drop table foo_bar;
           ╰╴           ─── 2. reference
         ");
     }
+
+    #[test]
+    fn builtin_function_references() {
+        assert_snapshot!(find_refs("
+select now$0();
+select now();
+"), @r"
+              ╭▸ current.sql:2:8
+              │
+            2 │ select now();
+              │        ┬─┬
+              │        │ │
+              │        │ 0. query
+              │        1. reference
+            3 │ select now();
+              │        ─── 2. reference
+              ╰╴
+
+              ╭▸ builtin.sql:10798:28
+              │
+        10798 │ create function pg_catalog.now() returns timestamp with time zone
+              ╰╴                           ─── 3. reference
+        ");
+    }
 }

crates/squawk_ide/src/goto_definition.rs

@@ -105,7 +105,7 @@ pub fn goto_definition(file: &ast::SourceFile, offset: TextSize) -> SmallVec<[Lo
     smallvec![]
 }
 
-#[derive(Debug, PartialEq, Clone, Copy, Eq)]
+#[derive(Debug, PartialEq, Clone, Copy, Eq, PartialOrd, Ord)]
 pub enum FileId {
     Current,
     Builtins,

crates/squawk_ide/src/resolve.rs

@@ -12,6 +12,16 @@ use crate::infer::{Type, infer_type_from_expr, infer_type_from_ty};
 pub(crate) use crate::symbols::Schema;
 use crate::symbols::{Name, SymbolKind};
 
+/// Resolves a name reference to its definition(s).
+///
+/// Most of the time returns one result, but can return two definitions
+/// in the case of:
+///
+/// ```sql
+/// select * from t join u using (col);
+/// ```
+///
+/// since `col` is defined in both `t` and `u`.
 pub(crate) fn resolve_name_ref_ptrs(
     binder: &Binder,
     root: &SyntaxNode,

crates/squawk_server/src/lib.rs

@@ -512,15 +512,25 @@ fn handle_references(
     let line_index = LineIndex::new(content);
     let offset = lsp_utils::offset(&line_index, position).unwrap();
 
-    let ranges = find_references(&file, offset);
+    let refs = find_references(&file, offset);
     let include_declaration = params.context.include_declaration;
 
-    let locations: Vec<Location> = ranges
+    let locations: Vec<Location> = refs
         .into_iter()
-        .filter(|range| include_declaration || !range.contains(offset))
-        .map(|range| Location {
-            uri: uri.clone(),
-            range: lsp_utils::range(&line_index, range),
+        .filter(|loc| include_declaration || !loc.range.contains(offset))
+        .filter_map(|loc| {
+            let uri = match loc.file {
+                squawk_ide::goto_definition::FileId::Current => uri.clone(),
+                squawk_ide::goto_definition::FileId::Builtins => builtins_url()?,
+            };
+            let line_index = match loc.file {
+                squawk_ide::goto_definition::FileId::Current => &line_index,
+                squawk_ide::goto_definition::FileId::Builtins => &LineIndex::new(BUILTINS_SQL),
+            };
+            Some(Location {
+                uri,
+                range: lsp_utils::range(line_index, loc.range),
+            })
         })
         .collect();
 

crates/squawk_wasm/src/lib.rs

@@ -245,20 +245,21 @@ pub fn find_references(content: String, line: u32, col: u32) -> Result<JsValue,
     let offset = position_to_offset(&line_index, line, col)?;
     let references = squawk_ide::find_references::find_references(&parse.tree(), offset);
 
+    let builtins_line_index = LineIndex::new(BUILTINS_SQL);
     let locations: Vec<LocationRange> = references
         .iter()
-        .map(|range| {
-            let start = line_index.line_col(range.start());
-            let end = line_index.line_col(range.end());
-            let start_wide = line_index
-                .to_wide(line_index::WideEncoding::Utf16, start)
-                .unwrap();
-            let end_wide = line_index
-                .to_wide(line_index::WideEncoding::Utf16, end)
-                .unwrap();
+        .map(|loc| {
+            let (li, file) = match loc.file {
+                FileId::Current => (&line_index, "current"),
+                FileId::Builtins => (&builtins_line_index, "builtin"),
+            };
+            let start = li.line_col(loc.range.start());
+            let end = li.line_col(loc.range.end());
+            let start_wide = li.to_wide(line_index::WideEncoding::Utf16, start).unwrap();
+            let end_wide = li.to_wide(line_index::WideEncoding::Utf16, end).unwrap();
 
             LocationRange {
-                file: "current".to_string(),
+                file: file.to_string(),
                 start_line: start_wide.line,
                 start_column: start_wide.col,
                 end_line: end_wide.line,