diff --git a/crates/squawk_ide/src/find_references.rs b/crates/squawk_ide/src/find_references.rs index 1776d48f..6fa9a767 100644 --- a/crates/squawk_ide/src/find_references.rs +++ b/crates/squawk_ide/src/find_references.rs @@ -1,37 +1,70 @@ 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 { - 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 { + 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, + ¤t_binder, + &builtins_tree, + &builtins_binder, + ) else { return vec![]; }; - let mut refs = vec![]; + let (binder, root) = match target_file { + FileId::Current => (¤t_binder, file.syntax()), + FileId::Builtins => (&builtins_binder, builtins_tree.syntax()), + }; + + let mut refs: Vec = 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(), + }); } }, _ => (), @@ -39,25 +72,37 @@ pub fn find_references(file: &ast::SourceFile, offset: TextSize) -> Vec Option> { + 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 @@ -65,10 +110,13 @@ fn find_targets( #[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 = (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 + "); + } } diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 4b92eb87..50fe5ecc 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/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, diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index a31b6eb0..69a5e906 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/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, diff --git a/crates/squawk_server/src/lib.rs b/crates/squawk_server/src/lib.rs index 048760ad..0c413cb2 100644 --- a/crates/squawk_server/src/lib.rs +++ b/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 = ranges + let locations: Vec = 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(); diff --git a/crates/squawk_wasm/src/lib.rs b/crates/squawk_wasm/src/lib.rs index 7b2d4b2d..25827e64 100644 --- a/crates/squawk_wasm/src/lib.rs +++ b/crates/squawk_wasm/src/lib.rs @@ -245,20 +245,21 @@ pub fn find_references(content: String, line: u32, col: u32) -> Result = 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,