Skip to content

Commit bc0c221

Browse files
sbdchdsbdchd
authored
ide: goto def with create & drop function (#766)
1 parent 7b71eb7 commit bc0c221

File tree

4 files changed

+156
-0
lines changed

4 files changed

+156
-0
lines changed

‎crates/squawk_ide/src/binder.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
7979
match stmt {
8080
ast::Stmt::CreateTable(create_table) => bind_create_table(b, create_table),
8181
ast::Stmt::CreateIndex(create_index) => bind_create_index(b, create_index),
82+
ast::Stmt::CreateFunction(create_function) => bind_create_function(b, create_function),
8283
ast::Stmt::Set(set) => bind_set(b, set),
8384
_ => {}
8485
}
@@ -129,6 +130,31 @@ fn bind_create_index(b: &mut Binder, create_index: ast::CreateIndex) {
129130
b.scopes[root].insert(index_name, index_id);
130131
}
131132

133+
fn bind_create_function(b: &mut Binder, create_function: ast::CreateFunction) {
134+
let Some(path) = create_function.path() else {
135+
return;
136+
};
137+
138+
let Some(function_name) = item_name(&path) else {
139+
return;
140+
};
141+
142+
let name_ptr = path_to_ptr(&path);
143+
144+
let Some(schema) = b.current_search_path().first().cloned() else {
145+
return;
146+
};
147+
148+
let function_id = b.symbols.alloc(Symbol {
149+
kind: SymbolKind::Function,
150+
ptr: name_ptr,
151+
schema,
152+
});
153+
154+
let root = b.root_scope();
155+
b.scopes[root].insert(function_name, function_id);
156+
}
157+
132158
fn item_name(path: &ast::Path) -> Option<Name> {
133159
let segment = path.segment()?;
134160

‎crates/squawk_ide/src/goto_definition.rs‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -908,4 +908,94 @@ create index idx_email on users(email$0);
908908
╰╴ ─ 1. source
909909
");
910910
}
911+
912+
#[test]
913+
fn goto_drop_function() {
914+
assert_snapshot!(goto("
915+
create function foo() returns int as $$ select 1 $$ language sql;
916+
drop function foo$0();
917+
"), @r"
918+
╭▸
919+
2 │ create function foo() returns int as $$ select 1 $$ language sql;
920+
│ ─── 2. destination
921+
3 │ drop function foo();
922+
╰╴ ─ 1. source
923+
");
924+
}
925+
926+
#[test]
927+
fn goto_drop_function_with_schema() {
928+
assert_snapshot!(goto("
929+
set search_path to public;
930+
create function foo() returns int as $$ select 1 $$ language sql;
931+
drop function public.foo$0();
932+
"), @r"
933+
╭▸
934+
3 │ create function foo() returns int as $$ select 1 $$ language sql;
935+
│ ─── 2. destination
936+
4 │ drop function public.foo();
937+
╰╴ ─ 1. source
938+
");
939+
}
940+
941+
#[test]
942+
fn goto_drop_function_defined_after() {
943+
assert_snapshot!(goto("
944+
drop function foo$0();
945+
create function foo() returns int as $$ select 1 $$ language sql;
946+
"), @r"
947+
╭▸
948+
2 │ drop function foo();
949+
│ ─ 1. source
950+
3 │ create function foo() returns int as $$ select 1 $$ language sql;
951+
╰╴ ─── 2. destination
952+
");
953+
}
954+
955+
#[test]
956+
fn goto_function_definition_returns_self() {
957+
assert_snapshot!(goto("
958+
create function foo$0() returns int as $$ select 1 $$ language sql;
959+
"), @r"
960+
╭▸
961+
2 │ create function foo() returns int as $$ select 1 $$ language sql;
962+
│ ┬─┬
963+
│ │ │
964+
│ │ 1. source
965+
╰╴ 2. destination
966+
");
967+
}
968+
969+
#[test]
970+
fn goto_drop_function_with_search_path() {
971+
assert_snapshot!(goto("
972+
create function foo() returns int as $$ select 1 $$ language sql;
973+
set search_path to bar;
974+
create function foo() returns int as $$ select 1 $$ language sql;
975+
set search_path to default;
976+
drop function foo$0();
977+
"), @r"
978+
╭▸
979+
2 │ create function foo() returns int as $$ select 1 $$ language sql;
980+
│ ─── 2. destination
981+
‡
982+
6 │ drop function foo();
983+
╰╴ ─ 1. source
984+
");
985+
}
986+
987+
#[test]
988+
fn goto_drop_function_multiple() {
989+
assert_snapshot!(goto("
990+
create function foo() returns int as $$ select 1 $$ language sql;
991+
create function bar() returns int as $$ select 1 $$ language sql;
992+
drop function foo(), bar$0();
993+
"), @r"
994+
╭▸
995+
3 │ create function bar() returns int as $$ select 1 $$ language sql;
996+
│ ─── 2. destination
997+
4 │ drop function foo(), bar();
998+
╰╴ ─ 1. source
999+
");
1000+
}
9111001
}

‎crates/squawk_ide/src/resolve.rs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ enum NameRefContext {
1414
DropTable,
1515
Table,
1616
DropIndex,
17+
DropFunction,
1718
CreateIndex,
1819
CreateIndexColumn,
1920
}
@@ -36,6 +37,13 @@ pub(crate) fn resolve_name_ref(binder: &Binder, name_ref: &ast::NameRef) -> Opti
3637
let position = name_ref.syntax().text_range().start();
3738
resolve_index(binder, &index_name, &schema, position)
3839
}
40+
NameRefContext::DropFunction => {
41+
let path = find_containing_path(name_ref)?;
42+
let function_name = extract_table_name(&path)?;
43+
let schema = extract_schema_name(&path);
44+
let position = name_ref.syntax().text_range().start();
45+
resolve_function(binder, &function_name, &schema, position)
46+
}
3947
NameRefContext::CreateIndexColumn => resolve_create_index_column(binder, name_ref),
4048
}
4149
}
@@ -53,6 +61,9 @@ fn classify_name_ref_context(name_ref: &ast::NameRef) -> Option<NameRefContext>
5361
if ast::DropIndex::can_cast(ancestor.kind()) {
5462
return Some(NameRefContext::DropIndex);
5563
}
64+
if ast::DropFunction::can_cast(ancestor.kind()) {
65+
return Some(NameRefContext::DropFunction);
66+
}
5667
if ast::PartitionItem::can_cast(ancestor.kind()) {
5768
in_partition_item = true;
5869
}
@@ -123,6 +134,34 @@ fn resolve_index(
123134
None
124135
}
125136

137+
fn resolve_function(
138+
binder: &Binder,
139+
function_name: &Name,
140+
schema: &Option<Schema>,
141+
position: TextSize,
142+
) -> Option<SyntaxNodePtr> {
143+
let symbols = binder.scopes[binder.root_scope()].get(function_name)?;
144+
145+
if let Some(schema) = schema {
146+
let symbol_id = symbols.iter().copied().find(|id| {
147+
let symbol = &binder.symbols[*id];
148+
symbol.kind == SymbolKind::Function && &symbol.schema == schema
149+
})?;
150+
return Some(binder.symbols[symbol_id].ptr);
151+
} else {
152+
let search_path = binder.search_path_at(position);
153+
for search_schema in search_path {
154+
if let Some(symbol_id) = symbols.iter().copied().find(|id| {
155+
let symbol = &binder.symbols[*id];
156+
symbol.kind == SymbolKind::Function && &symbol.schema == search_schema
157+
}) {
158+
return Some(binder.symbols[symbol_id].ptr);
159+
}
160+
}
161+
}
162+
None
163+
}
164+
126165
fn resolve_create_index_column(binder: &Binder, name_ref: &ast::NameRef) -> Option<SyntaxNodePtr> {
127166
let column_name = Name::new(name_ref.syntax().text().to_string());
128167

‎crates/squawk_ide/src/symbols.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ fn normalize_identifier(text: &str) -> SmolStr {
4141
pub(crate) enum SymbolKind {
4242
Table,
4343
Index,
44+
Function,
4445
}
4546

4647
#[derive(Clone, Debug)]

0 commit comments

Comments
 (0)