squawk_ide/

find_references.rs

use crate::db::parse;
use crate::file::InFile;
use crate::goto_definition;
use crate::location::Location;
use rowan::TextSize;
use salsa::Database as Db;
use squawk_syntax::{
    SyntaxNode,
    ast::{self, AstNode},
};

fn is_reference_node(node: &SyntaxNode) -> bool {
    if ast::AnyNameRef::can_cast(node.kind()) || ast::ConfigValueName::can_cast(node.kind()) {
        return true;
    }

    if let Some(literal) = ast::Literal::cast(node.clone())
        && matches!(literal.kind(), Some(ast::LitKind::PositionalParam(_)))
    {
        return true;
    }

    if let Some(ty) = ast::Type::cast(node.clone()) {
        return match ty {
            ast::Type::BitType(_)
            | ast::Type::CharType(_)
            | ast::Type::DoubleType(_)
            | ast::Type::IntervalType(_)
            | ast::Type::TimeType(_) => true,
            ast::Type::ArrayType(_)
            | ast::Type::ExprType(_)
            | ast::Type::PathType(_)
            | ast::Type::PercentType(_) => false,
        };
    }

    false
}

pub fn find_references(db: &dyn Db, position: InFile<TextSize>) -> Vec<Location> {
    let file = position.file_id;
    let targets = goto_definition::goto_definition(db, position);
    let Some(first) = targets.first() else {
        return vec![];
    };

    let mut refs = targets.to_vec();

    for node in parse(db, file)
        .tree()
        .syntax()
        .descendants()
        .filter(is_reference_node)
    {
        let range = node.text_range();
        let matches = goto_definition::goto_definition(db, InFile::new(file, range.start()))
            .into_iter()
            .any(|location| targets.contains(&location));
        if matches {
            refs.push(Location {
                file,
                range,
                kind: first.kind,
            });
        }
    }
    refs.sort_by_key(|loc| (loc.file != file, loc.range.start()));
    refs
}

#[cfg(test)]
mod test {
    use crate::builtins::builtins_file;
    use crate::db::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;

    #[must_use]
    #[track_caller]
    fn find_refs(sql: &str) -> String {
        let fixture = Fixture::new(sql);
        let marker = fixture.marker();
        let offset = marker.offset_before();
        let query_span = marker.range();
        let db = fixture.db();
        let current_file = offset.file_id;

        let references = find_references(db, offset);

        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() {
            refs_by_file
                .entry(location.file)
                .or_default()
                .push((i + 1, location.range));
        }

        let multi_file = refs_by_file.len() > 1 || !refs_by_file.contains_key(&current_file);

        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());
        }
        snippet = snippet.annotation(AnnotationKind::Context.span(query_span).label("0. query"));
        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)];

        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(other_snippet),
            );
        }

        let renderer = Renderer::plain().decor_style(DecorStyle::Unicode);
        renderer
            .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!("{label_index}. reference")),
            );
        }
        snippet
    }

    #[test]
    fn simple_table_reference() {
        assert_snapshot!(find_refs("
create table t();
drop table t$0;
"), @"
          ╭▸ 
        2 │ create table t();
          │              ─ 1. reference
        3 │ drop table t;
          │            ┬
          │            │
          │            0. query
          ╰╴           2. reference
        ");
    }

    #[test]
    fn multiple_references() {
        assert_snapshot!(find_refs("
create table users();
drop table users$0;
table users;
"), @r"
          ╭▸ 
        2 │ create table users();
          │              ───── 1. reference
        3 │ drop table users;
          │            ┬───┬
          │            │   │
          │            │   0. query
          │            2. reference
        4 │ table users;
          ╰╴      ───── 3. reference
        ");
    }

    #[test]
    fn join_using_column() {
        assert_snapshot!(find_refs("
create table t(id int);
create table u(id int);
select * from t join u using (id$0);
"), @r"
          ╭▸ 
        2 │ create table t(id int);
          │                ── 1. reference
        3 │ create table u(id int);
          │                ── 2. reference
        4 │ select * from t join u using (id);
          │                               ┬┬
          │                               ││
          │                               │0. query
          ╰╴                              3. reference
        ");
    }

    #[test]
    fn find_from_definition() {
        assert_snapshot!(find_refs("
create table t$0();
drop table t;
"), @r"
          ╭▸ 
        2 │ create table t();
          │              ┬
          │              │
          │              0. query
          │              1. reference
        3 │ drop table t;
          ╰╴           ─ 2. reference
        ");
    }

    #[test]
    fn with_schema_qualified() {
        assert_snapshot!(find_refs("
create table public.users();
drop table public.users$0;
table users;
"), @r"
          ╭▸ 
        2 │ create table public.users();
          │                     ───── 1. reference
        3 │ drop table public.users;
          │                   ┬───┬
          │                   │   │
          │                   │   0. query
          │                   2. reference
        4 │ table users;
          ╰╴      ───── 3. reference
        ");
    }

    #[test]
    fn temp_table_shadows_public() {
        assert_snapshot!(find_refs("
create table t();
create temp table t$0();
drop table t;
"), @"
          ╭▸ 
        3 │ create temp table t();
          │                   ┬
          │                   │
          │                   0. query
          │                   1. reference
        4 │ drop table t;
          ╰╴           ─ 2. reference
        ");
    }

    #[test]
    fn different_schema_no_match() {
        assert_snapshot!(find_refs("
create table foo.t();
create table bar.t$0();
"), @r"
          ╭▸ 
        3 │ create table bar.t();
          │                  ┬
          │                  │
          │                  0. query
          ╰╴                 1. reference
        ");
    }

    #[test]
    fn with_search_path() {
        assert_snapshot!(find_refs("
set search_path to myschema;
create table myschema.users$0();
drop table users;
"), @r"
          ╭▸ 
        3 │ create table myschema.users();
          │                       ┬───┬
          │                       │   │
          │                       │   0. query
          │                       1. reference
        4 │ drop table users;
          ╰╴           ───── 2. reference
        ");
    }

    #[test]
    fn temp_table_with_pg_temp_schema() {
        assert_snapshot!(find_refs("
create temp table t();
drop table pg_temp.t$0;
"), @r"
          ╭▸ 
        2 │ create temp table t();
          │                   ─ 1. reference
        3 │ drop table pg_temp.t;
          │                    ┬
          │                    │
          │                    0. query
          ╰╴                   2. reference
        ");
    }

    #[test]
    fn case_insensitive() {
        assert_snapshot!(find_refs("
create table Users();
drop table USERS$0;
table users;
"), @r"
          ╭▸ 
        2 │ create table Users();
          │              ───── 1. reference
        3 │ drop table USERS;
          │            ┬───┬
          │            │   │
          │            │   0. query
          │            2. reference
        4 │ table users;
          ╰╴      ───── 3. reference
        ");
    }
    #[test]
    fn case_insensitive_part_2() {
        // we should see refs for `drop table` and `table`
        assert_snapshot!(find_refs(r#"
create table actors();
create table "Actors"();
drop table ACTORS$0;
table actors;
"#), @r#"
          ╭▸ 
        2 │ create table actors();
          │              ────── 1. reference
        3 │ create table "Actors"();
        4 │ drop table ACTORS;
          │            ┬────┬
          │            │    │
          │            │    0. query
          │            2. reference
        5 │ table actors;
          ╰╴      ────── 3. reference
        "#);
    }

    #[test]
    fn case_insensitive_with_schema() {
        assert_snapshot!(find_refs("
create table Public.Users();
drop table PUBLIC.USERS$0;
table public.users;
"), @r"
          ╭▸ 
        2 │ create table Public.Users();
          │                     ───── 1. reference
        3 │ drop table PUBLIC.USERS;
          │                   ┬───┬
          │                   │   │
          │                   │   0. query
          │                   2. reference
        4 │ table public.users;
          ╰╴             ───── 3. reference
        ");
    }

    #[test]
    fn no_partial_match() {
        assert_snapshot!(find_refs("
create table t$0();
create table temp_t();
"), @r"
          ╭▸ 
        2 │ create table t();
          │              ┬
          │              │
          │              0. query
          ╰╴             1. reference
        ");
    }

    #[test]
    fn identifier_boundaries() {
        assert_snapshot!(find_refs("
create table foo$0();
drop table foo;
drop table foo1;
drop table barfoo;
drop table foo_bar;
"), @r"
          ╭▸ 
        2 │ create table foo();
          │              ┬─┬
          │              │ │
          │              │ 0. query
          │              1. reference
        3 │ drop table foo;
          ╰╴           ─── 2. reference
        ");
    }

    #[test]
    fn builtin_function_references() {
        assert_snapshot!(find_refs("
-- include-builtins
select now$0();
select now();
"), @"
              ╭▸ current.sql:3:8
              │
            3 │ select now();
              │        ┬─┬
              │        │ │
              │        │ 0. query
              │        1. reference
            4 │ select now();
              │        ─── 2. reference
              ╰╴

              ╭▸ builtins.sql:11089:28
              │
        11089 │ create function pg_catalog.now() returns timestamp with time zone
              ╰╴                           ─── 3. reference
        ");
    }

    #[test]
    fn bit() {
        assert_snapshot!(find_refs("
create type pg_catalog.bit$0;

create function pg_catalog.bit(bigint, integer) returns bit
  language internal;

create function pg_catalog.bit(bit, integer, boolean) returns bit
  language internal;

create function pg_catalog.bit(integer, integer) returns bit
  language internal;
"), @"
           ╭▸ 
         2 │ create type pg_catalog.bit;
           │                        ┬─┬
           │                        │ │
           │                        │ 0. query
           │                        1. reference
         3 │
         4 │ create function pg_catalog.bit(bigint, integer) returns bit
           │                                                         ─── 2. reference
           ‡
         7 │ create function pg_catalog.bit(bit, integer, boolean) returns bit
           │                                ─── 3. reference               ─── 4. reference
           ‡
        10 │ create function pg_catalog.bit(integer, integer) returns bit
           ╰╴                                                         ─── 5. reference
        ");
    }

    #[test]
    fn positional_param_and_named_param() {
        assert_snapshot!(find_refs("
create function f(x$0 int) returns int language sql return x + $1;
"), @"
          ╭▸ 
        2 │ create function f(x int) returns int language sql return x + $1;
          │                   ┬                                      ┬   ── 3. reference
          │                   │                                      │
          │                   0. query                               2. reference
          ╰╴                  1. reference
        ");
    }

    #[test]
    fn char() {
        assert_snapshot!(find_refs("
create type pg_catalog.bpchar$0;

select '1'::char;
select '1'::bpchar;
"), @"
          ╭▸ 
        2 │ create type pg_catalog.bpchar;
          │                        ┬────┬
          │                        │    │
          │                        │    0. query
          │                        1. reference
        3 │
        4 │ select '1'::char;
          │             ──── 2. reference
        5 │ select '1'::bpchar;
          ╰╴            ────── 3. reference
        ");
    }

    #[test]
    fn access_method_ref() {
        assert_snapshot!(find_refs("
create access method heap2$0 type table handler heap_tableam_handler;
drop access method heap2;
"), @"
          ╭▸ 
        2 │ create access method heap2 type table handler heap_tableam_handler;
          │                      ┬───┬
          │                      │   │
          │                      │   0. query
          │                      1. reference
        3 │ drop access method heap2;
          ╰╴                   ───── 2. reference
        ");
    }

    #[test]
    fn channel_ref() {
        assert_snapshot!(find_refs("
listen updates$0;
notify updates;
"), @"
          ╭▸ 
        2 │ listen updates;
          │        ┬─────┬
          │        │     │
          │        │     0. query
          │        1. reference
        3 │ notify updates;
          ╰╴       ─────── 2. reference
        ");
    }

    #[test]
    fn column_name_ref() {
        assert_snapshot!(find_refs("
create table t(id$0 int);
create table u(t_id int references t(id));
"), @"
          ╭▸ 
        2 │ create table t(id int);
          │                ┬┬
          │                ││
          │                │0. query
          │                1. reference
        3 │ create table u(t_id int references t(id));
          ╰╴                                     ── 2. reference
        ");
    }

    #[test]
    fn cursor_ref() {
        assert_snapshot!(find_refs("
declare c$0 scroll cursor for select * from t;
fetch forward 5 from c;
"), @"
          ╭▸ 
        2 │ declare c scroll cursor for select * from t;
          │         ┬
          │         │
          │         0. query
          │         1. reference
        3 │ fetch forward 5 from c;
          ╰╴                     ─ 2. reference
        ");
    }

    #[test]
    fn database_ref() {
        assert_snapshot!(find_refs("
create database mydb$0;
drop database mydb;
"), @"
          ╭▸ 
        2 │ create database mydb;
          │                 ┬──┬
          │                 │  │
          │                 │  0. query
          │                 1. reference
        3 │ drop database mydb;
          ╰╴              ──── 2. reference
        ");
    }

    #[test]
    fn event_trigger_ref() {
        assert_snapshot!(find_refs("
create event trigger et$0 on ddl_command_start execute function f();
drop event trigger et;
"), @"
          ╭▸ 
        2 │ create event trigger et on ddl_command_start execute function f();
          │                      ┬┬
          │                      ││
          │                      │0. query
          │                      1. reference
        3 │ drop event trigger et;
          ╰╴                   ── 2. reference
        ");
    }

    #[test]
    fn extension_ref() {
        assert_snapshot!(find_refs("
create extension myext$0;
drop extension myext;
"), @"
          ╭▸ 
        2 │ create extension myext;
          │                  ┬───┬
          │                  │   │
          │                  │   0. query
          │                  1. reference
        3 │ drop extension myext;
          ╰╴               ───── 2. reference
        ");
    }

    #[test]
    fn foreign_data_wrapper_ref() {
        assert_snapshot!(find_refs("
create foreign data wrapper fdw$0;
create server srv foreign data wrapper fdw;
"), @"
          ╭▸ 
        2 │ create foreign data wrapper fdw;
          │                             ┬─┬
          │                             │ │
          │                             │ 0. query
          │                             1. reference
        3 │ create server srv foreign data wrapper fdw;
          ╰╴                                       ─── 2. reference
        ");
    }

    #[test]
    fn json_path_name_ref() {
        assert_snapshot!(find_refs("
select * from json_table(
  '{}'::jsonb, '$' as root$0
  columns (value text path '$')
  plan (root)
);
"), @"
          ╭▸ 
        3 │   '{}'::jsonb, '$' as root
          │                       ┬──┬
          │                       │  │
          │                       │  0. query
          │                       1. reference
        4 │   columns (value text path '$')
        5 │   plan (root)
          ╰╴        ──── 2. reference
        ");
    }

    #[test]
    fn language_ref() {
        assert_snapshot!(find_refs("
create language mylang$0;
create function f() returns int language mylang as $$x$$;
"), @"
          ╭▸ 
        2 │ create language mylang;
          │                 ┬────┬
          │                 │    │
          │                 │    0. query
          │                 1. reference
        3 │ create function f() returns int language mylang as $$x$$;
          ╰╴                                         ────── 2. reference
        ");
    }

    #[test]
    fn policy_ref() {
        assert_snapshot!(find_refs("
create table t(c int);
create policy p$0 on t;
drop policy p on t;
"), @"
          ╭▸ 
        3 │ create policy p on t;
          │               ┬
          │               │
          │               0. query
          │               1. reference
        4 │ drop policy p on t;
          ╰╴            ─ 2. reference
        ");
    }

    #[test]
    fn prepared_statement_ref() {
        assert_snapshot!(find_refs("
prepare stmt$0 as select 1;
execute stmt;
"), @"
          ╭▸ 
        2 │ prepare stmt as select 1;
          │         ┬──┬
          │         │  │
          │         │  0. query
          │         1. reference
        3 │ execute stmt;
          ╰╴        ──── 2. reference
        ");
    }

    #[test]
    fn publication_ref() {
        assert_snapshot!(find_refs("
create table t(id int);
create publication pub$0 for table t;
alter publication pub add table t;
"), @"
          ╭▸ 
        3 │ create publication pub for table t;
          │                    ┬─┬
          │                    │ │
          │                    │ 0. query
          │                    1. reference
        4 │ alter publication pub add table t;
          ╰╴                  ─── 2. reference
        ");
    }

    #[test]
    fn role_ref() {
        assert_snapshot!(find_refs("
create role reader$0;
drop role reader;
"), @"
          ╭▸ 
        2 │ create role reader;
          │             ┬────┬
          │             │    │
          │             │    0. query
          │             1. reference
        3 │ drop role reader;
          ╰╴          ────── 2. reference
        ");
    }

    #[test]
    fn rule_ref() {
        assert_snapshot!(find_refs("
create table t(a int);
create rule r$0 as on select to t do instead nothing;
drop rule r on t;
"), @"
          ╭▸ 
        3 │ create rule r as on select to t do instead nothing;
          │             ┬
          │             │
          │             0. query
          │             1. reference
        4 │ drop rule r on t;
          ╰╴          ─ 2. reference
        ");
    }

    #[test]
    fn savepoint_ref() {
        assert_snapshot!(find_refs("
begin;
savepoint sp$0;
release savepoint sp;
"), @"
          ╭▸ 
        3 │ savepoint sp;
          │           ┬┬
          │           ││
          │           │0. query
          │           1. reference
        4 │ release savepoint sp;
          ╰╴                  ── 2. reference
        ");
    }

    #[test]
    fn schema_ref() {
        assert_snapshot!(find_refs("
create schema app$0;
drop schema app;
"), @"
          ╭▸ 
        2 │ create schema app;
          │               ┬─┬
          │               │ │
          │               │ 0. query
          │               1. reference
        3 │ drop schema app;
          ╰╴            ─── 2. reference
        ");
    }

    #[test]
    fn server_ref() {
        assert_snapshot!(find_refs("
create server myserver$0 foreign data wrapper fdw;
drop server myserver;
"), @"
          ╭▸ 
        2 │ create server myserver foreign data wrapper fdw;
          │               ┬──────┬
          │               │      │
          │               │      0. query
          │               1. reference
        3 │ drop server myserver;
          ╰╴            ──────── 2. reference
        ");
    }

    #[test]
    fn subscription_ref() {
        assert_snapshot!(find_refs("
create subscription sub$0 connection $$host=localhost$$ publication pub;
alter subscription sub refresh publication;
"), @"
          ╭▸ 
        2 │ create subscription sub connection $$host=localhost$$ publication pub;
          │                     ┬─┬
          │                     │ │
          │                     │ 0. query
          │                     1. reference
        3 │ alter subscription sub refresh publication;
          ╰╴                   ─── 2. reference
        ");
    }

    #[test]
    fn tablespace_ref() {
        assert_snapshot!(find_refs("
create tablespace ts$0 location '/tmp/ts';
drop tablespace ts;
"), @"
          ╭▸ 
        2 │ create tablespace ts location '/tmp/ts';
          │                   ┬┬
          │                   ││
          │                   │0. query
          │                   1. reference
        3 │ drop tablespace ts;
          ╰╴                ── 2. reference
        ");
    }

    #[test]
    fn trigger_ref() {
        assert_snapshot!(find_refs("
create trigger tr$0 before insert on t for each row execute function f();
drop trigger tr on t;
"), @"
          ╭▸ 
        2 │ create trigger tr before insert on t for each row execute function f();
          │                ┬┬
          │                ││
          │                │0. query
          │                1. reference
        3 │ drop trigger tr on t;
          ╰╴             ── 2. reference
        ");
    }

    #[test]
    fn vertex_table_ref() {
        assert_snapshot!(find_refs("
create table v1(id int);
create table e1(source_id int references v1);
create property graph g
  vertex tables (v1 as source_vertex)
  edge tables (e1 source source_vertex$0 destination source_vertex);
"), @"
          ╭▸ 
        5 │   vertex tables (v1 as source_vertex)
          │                        ───────────── 1. reference
        6 │   edge tables (e1 source source_vertex destination source_vertex);
          │                          ┬───────────┬             ───────────── 3. reference
          │                          │           │
          │                          │           0. query
          ╰╴                         2. reference
        ");
    }

    #[test]
    fn window_ref() {
        assert_snapshot!(find_refs("
select row_number() over w
window w$0 as (partition by 1);
"), @"
          ╭▸ 
        2 │ select row_number() over w
          │                          ─ 1. reference
        3 │ window w as (partition by 1);
          │        ┬
          │        │
          │        0. query
          ╰╴       2. reference
        ");
    }
}