Skip to content

Commit 111e02c

Browse files
authored
ide: goto def on create/drop index (#762)
1 parent 8f58f3d commit 111e02c

File tree

4 files changed

+333
-1
lines changed

4 files changed

+333
-1
lines changed

crates/squawk_ide/src/binder.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ fn bind_file(b: &mut Binder, file: &ast::SourceFile) {
7878
fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
7979
match stmt {
8080
ast::Stmt::CreateTable(create_table) => bind_create_table(b, create_table),
81+
ast::Stmt::CreateIndex(create_index) => bind_create_index(b, create_index),
8182
ast::Stmt::Set(set) => bind_set(b, set),
8283
_ => {}
8384
}
@@ -106,6 +107,28 @@ fn bind_create_table(b: &mut Binder, create_table: ast::CreateTable) {
106107
b.scopes[root].insert(table_name, table_id);
107108
}
108109

110+
fn bind_create_index(b: &mut Binder, create_index: ast::CreateIndex) {
111+
let Some(name) = create_index.name() else {
112+
return;
113+
};
114+
115+
let index_name = Name::new(name.syntax().text().to_string());
116+
let name_ptr = SyntaxNodePtr::new(name.syntax());
117+
118+
let Some(schema) = b.current_search_path().first().cloned() else {
119+
return;
120+
};
121+
122+
let index_id = b.symbols.alloc(Symbol {
123+
kind: SymbolKind::Index,
124+
ptr: name_ptr,
125+
schema,
126+
});
127+
128+
let root = b.root_scope();
129+
b.scopes[root].insert(index_name, index_id);
130+
}
131+
109132
fn item_name(path: &ast::Path) -> Option<Name> {
110133
let segment = path.segment()?;
111134

crates/squawk_ide/src/goto_definition.rs

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,4 +691,221 @@ table t$0;
691691
╰╴ ─ 1. source
692692
");
693693
}
694+
695+
#[test]
696+
fn goto_drop_index() {
697+
assert_snapshot!(goto("
698+
create index idx_name on t(x);
699+
drop index idx_name$0;
700+
"), @r"
701+
╭▸
702+
2 │ create index idx_name on t(x);
703+
│ ──────── 2. destination
704+
3 │ drop index idx_name;
705+
╰╴ ─ 1. source
706+
");
707+
}
708+
709+
#[test]
710+
fn goto_drop_index_with_schema() {
711+
assert_snapshot!(goto(r#"
712+
set search_path to public;
713+
create index idx_name on t(x);
714+
drop index public.idx_name$0;
715+
"#), @r"
716+
╭▸
717+
3 │ create index idx_name on t(x);
718+
│ ──────── 2. destination
719+
4 │ drop index public.idx_name;
720+
╰╴ ─ 1. source
721+
");
722+
}
723+
724+
#[test]
725+
fn goto_drop_index_defined_after() {
726+
assert_snapshot!(goto("
727+
drop index idx_name$0;
728+
create index idx_name on t(x);
729+
"), @r"
730+
╭▸
731+
2 │ drop index idx_name;
732+
│ ─ 1. source
733+
3 │ create index idx_name on t(x);
734+
╰╴ ──────── 2. destination
735+
");
736+
}
737+
738+
#[test]
739+
fn goto_index_definition_returns_self() {
740+
assert_snapshot!(goto("
741+
create index idx_name$0 on t(x);
742+
"), @r"
743+
╭▸
744+
2 │ create index idx_name on t(x);
745+
│ ┬──────┬
746+
│ │ │
747+
│ │ 1. source
748+
╰╴ 2. destination
749+
");
750+
}
751+
752+
#[test]
753+
fn goto_drop_index_with_search_path() {
754+
assert_snapshot!(goto(r#"
755+
create index idx_name on t(x);
756+
set search_path to bar;
757+
create index idx_name on f(x);
758+
set search_path to default;
759+
drop index idx_name$0;
760+
"#), @r"
761+
╭▸
762+
2 │ create index idx_name on t(x);
763+
│ ──────── 2. destination
764+
765+
6 │ drop index idx_name;
766+
╰╴ ─ 1. source
767+
");
768+
}
769+
770+
#[test]
771+
fn goto_drop_index_multiple() {
772+
assert_snapshot!(goto("
773+
create index idx1 on t(x);
774+
create index idx2 on t(y);
775+
drop index idx1, idx2$0;
776+
"), @r"
777+
╭▸
778+
3 │ create index idx2 on t(y);
779+
│ ──── 2. destination
780+
4 │ drop index idx1, idx2;
781+
╰╴ ─ 1. source
782+
");
783+
}
784+
785+
#[test]
786+
fn goto_create_index_table() {
787+
assert_snapshot!(goto("
788+
create table users(id int);
789+
create index idx_users on users$0(id);
790+
"), @r"
791+
╭▸
792+
2 │ create table users(id int);
793+
│ ───── 2. destination
794+
3 │ create index idx_users on users(id);
795+
╰╴ ─ 1. source
796+
");
797+
}
798+
799+
#[test]
800+
fn goto_create_index_table_with_schema() {
801+
assert_snapshot!(goto("
802+
create table public.users(id int);
803+
create index idx_users on public.users$0(id);
804+
"), @r"
805+
╭▸
806+
2 │ create table public.users(id int);
807+
│ ───── 2. destination
808+
3 │ create index idx_users on public.users(id);
809+
╰╴ ─ 1. source
810+
");
811+
}
812+
813+
#[test]
814+
fn goto_create_index_table_with_search_path() {
815+
assert_snapshot!(goto(r#"
816+
set search_path to foo;
817+
create table foo.users(id int);
818+
create index idx_users on users$0(id);
819+
"#), @r"
820+
╭▸
821+
3 │ create table foo.users(id int);
822+
│ ───── 2. destination
823+
4 │ create index idx_users on users(id);
824+
╰╴ ─ 1. source
825+
");
826+
}
827+
828+
#[test]
829+
fn goto_create_index_temp_table() {
830+
assert_snapshot!(goto("
831+
create temp table users(id int);
832+
create index idx_users on users$0(id);
833+
"), @r"
834+
╭▸
835+
2 │ create temp table users(id int);
836+
│ ───── 2. destination
837+
3 │ create index idx_users on users(id);
838+
╰╴ ─ 1. source
839+
");
840+
}
841+
842+
#[test]
843+
fn goto_create_index_column() {
844+
assert_snapshot!(goto("
845+
create table users(id int, email text);
846+
create index idx_email on users(email$0);
847+
"), @r"
848+
╭▸
849+
2 │ create table users(id int, email text);
850+
│ ───── 2. destination
851+
3 │ create index idx_email on users(email);
852+
╰╴ ─ 1. source
853+
");
854+
}
855+
856+
#[test]
857+
fn goto_create_index_first_column() {
858+
assert_snapshot!(goto("
859+
create table users(id int, email text);
860+
create index idx_id on users(id$0);
861+
"), @r"
862+
╭▸
863+
2 │ create table users(id int, email text);
864+
│ ── 2. destination
865+
3 │ create index idx_id on users(id);
866+
╰╴ ─ 1. source
867+
");
868+
}
869+
870+
#[test]
871+
fn goto_create_index_multiple_columns() {
872+
assert_snapshot!(goto("
873+
create table users(id int, email text, name text);
874+
create index idx_users on users(id, email$0, name);
875+
"), @r"
876+
╭▸
877+
2 │ create table users(id int, email text, name text);
878+
│ ───── 2. destination
879+
3 │ create index idx_users on users(id, email, name);
880+
╰╴ ─ 1. source
881+
");
882+
}
883+
884+
#[test]
885+
fn goto_create_index_column_with_schema() {
886+
assert_snapshot!(goto("
887+
create table public.users(id int, email text);
888+
create index idx_email on public.users(email$0);
889+
"), @r"
890+
╭▸
891+
2 │ create table public.users(id int, email text);
892+
│ ───── 2. destination
893+
3 │ create index idx_email on public.users(email);
894+
╰╴ ─ 1. source
895+
");
896+
}
897+
898+
#[test]
899+
fn goto_create_index_column_temp_table() {
900+
assert_snapshot!(goto("
901+
create temp table users(id int, email text);
902+
create index idx_email on users(email$0);
903+
"), @r"
904+
╭▸
905+
2 │ create temp table users(id int, email text);
906+
│ ───── 2. destination
907+
3 │ create index idx_email on users(email);
908+
╰╴ ─ 1. source
909+
");
910+
}
694911
}

crates/squawk_ide/src/resolve.rs

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,35 +6,61 @@ use squawk_syntax::{
66

77
use crate::binder::Binder;
88
use crate::symbols::{Name, Schema, SymbolKind};
9+
use squawk_syntax::SyntaxNode;
910

1011
#[derive(Debug)]
1112
enum NameRefContext {
1213
DropTable,
1314
Table,
15+
DropIndex,
16+
CreateIndex,
17+
CreateIndexColumn,
1418
}
1519

1620
pub(crate) fn resolve_name_ref(binder: &Binder, name_ref: &ast::NameRef) -> Option<SyntaxNodePtr> {
1721
let context = classify_name_ref_context(name_ref)?;
1822

1923
match context {
20-
NameRefContext::DropTable | NameRefContext::Table => {
24+
NameRefContext::DropTable | NameRefContext::Table | NameRefContext::CreateIndex => {
2125
let path = find_containing_path(name_ref)?;
2226
let table_name = extract_table_name(&path)?;
2327
let schema = extract_schema_name(&path);
2428
let position = name_ref.syntax().text_range().start();
2529
resolve_table(binder, &table_name, &schema, position)
2630
}
31+
NameRefContext::DropIndex => {
32+
let path = find_containing_path(name_ref)?;
33+
let index_name = extract_table_name(&path)?;
34+
let schema = extract_schema_name(&path);
35+
let position = name_ref.syntax().text_range().start();
36+
resolve_index(binder, &index_name, &schema, position)
37+
}
38+
NameRefContext::CreateIndexColumn => resolve_create_index_column(binder, name_ref),
2739
}
2840
}
2941

3042
fn classify_name_ref_context(name_ref: &ast::NameRef) -> Option<NameRefContext> {
43+
let mut in_partition_item = false;
44+
3145
for ancestor in name_ref.syntax().ancestors() {
3246
if ast::DropTable::can_cast(ancestor.kind()) {
3347
return Some(NameRefContext::DropTable);
3448
}
3549
if ast::Table::can_cast(ancestor.kind()) {
3650
return Some(NameRefContext::Table);
3751
}
52+
if ast::DropIndex::can_cast(ancestor.kind()) {
53+
return Some(NameRefContext::DropIndex);
54+
}
55+
if ast::PartitionItem::can_cast(ancestor.kind()) {
56+
in_partition_item = true;
57+
}
58+
if ast::CreateIndex::can_cast(ancestor.kind()) {
59+
if in_partition_item {
60+
return Some(NameRefContext::CreateIndexColumn);
61+
}
62+
return Some(NameRefContext::CreateIndex);
63+
}
3864
}
3965

4066
None
@@ -68,6 +94,71 @@ fn resolve_table(
6894
None
6995
}
7096

97+
fn resolve_index(
98+
binder: &Binder,
99+
index_name: &Name,
100+
schema: &Option<Schema>,
101+
position: TextSize,
102+
) -> Option<SyntaxNodePtr> {
103+
let symbols = binder.scopes[binder.root_scope()].get(index_name)?;
104+
105+
if let Some(schema) = schema {
106+
let symbol_id = symbols.iter().copied().find(|id| {
107+
let symbol = &binder.symbols[*id];
108+
symbol.kind == SymbolKind::Index && &symbol.schema == schema
109+
})?;
110+
return Some(binder.symbols[symbol_id].ptr);
111+
} else {
112+
let search_path = binder.search_path_at(position);
113+
for search_schema in search_path {
114+
if let Some(symbol_id) = symbols.iter().copied().find(|id| {
115+
let symbol = &binder.symbols[*id];
116+
symbol.kind == SymbolKind::Index && &symbol.schema == search_schema
117+
}) {
118+
return Some(binder.symbols[symbol_id].ptr);
119+
}
120+
}
121+
}
122+
None
123+
}
124+
125+
fn resolve_create_index_column(binder: &Binder, name_ref: &ast::NameRef) -> Option<SyntaxNodePtr> {
126+
let column_name = Name::new(name_ref.syntax().text().to_string());
127+
128+
let create_index = name_ref
129+
.syntax()
130+
.ancestors()
131+
.find_map(ast::CreateIndex::cast)?;
132+
let relation_name = create_index.relation_name()?;
133+
let path = relation_name.path()?;
134+
135+
let table_name = extract_table_name(&path)?;
136+
let schema = extract_schema_name(&path);
137+
let position = name_ref.syntax().text_range().start();
138+
139+
let table_ptr = resolve_table(binder, &table_name, &schema, position)?;
140+
141+
let root: &SyntaxNode = &name_ref.syntax().ancestors().last()?;
142+
let table_name_node = table_ptr.to_node(root);
143+
144+
let create_table = table_name_node
145+
.ancestors()
146+
.find_map(ast::CreateTable::cast)?;
147+
148+
let table_arg_list = create_table.table_arg_list()?;
149+
150+
for arg in table_arg_list.args() {
151+
if let ast::TableArg::Column(column) = arg
152+
&& let Some(col_name) = column.name()
153+
&& Name::new(col_name.syntax().text().to_string()) == column_name
154+
{
155+
return Some(SyntaxNodePtr::new(col_name.syntax()));
156+
}
157+
}
158+
159+
None
160+
}
161+
71162
fn find_containing_path(name_ref: &ast::NameRef) -> Option<ast::Path> {
72163
for ancestor in name_ref.syntax().ancestors() {
73164
if let Some(path) = ast::Path::cast(ancestor) {

0 commit comments

Comments
 (0)