ide: goto def on create/drop index

crates/squawk_ide/src/binder.rs

@@ -78,6 +78,7 @@ fn bind_file(b: &mut Binder, file: &ast::SourceFile) {
 fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) {
     match stmt {
         ast::Stmt::CreateTable(create_table) => bind_create_table(b, create_table),
+        ast::Stmt::CreateIndex(create_index) => bind_create_index(b, create_index),
         ast::Stmt::Set(set) => bind_set(b, set),
         _ => {}
     }
@@ -106,6 +107,28 @@ fn bind_create_table(b: &mut Binder, create_table: ast::CreateTable) {
     b.scopes[root].insert(table_name, table_id);
 }
 
+fn bind_create_index(b: &mut Binder, create_index: ast::CreateIndex) {
+    let Some(name) = create_index.name() else {
+        return;
+    };
+
+    let index_name = Name::new(name.syntax().text().to_string());
+    let name_ptr = SyntaxNodePtr::new(name.syntax());
+
+    let Some(schema) = b.current_search_path().first().cloned() else {
+        return;
+    };
+
+    let index_id = b.symbols.alloc(Symbol {
+        kind: SymbolKind::Index,
+        ptr: name_ptr,
+        schema,
+    });
+
+    let root = b.root_scope();
+    b.scopes[root].insert(index_name, index_id);
+}
+
 fn item_name(path: &ast::Path) -> Option<Name> {
     let segment = path.segment()?;
 

crates/squawk_ide/src/goto_definition.rs

@@ -691,4 +691,221 @@ table t$0;
           ╰╴      ─ 1. source
         ");
     }
+
+    #[test]
+    fn goto_drop_index() {
+        assert_snapshot!(goto("
+create index idx_name on t(x);
+drop index idx_name$0;
+"), @r"
+          ╭▸ 
+        2 │ create index idx_name on t(x);
+          │              ──────── 2. destination
+        3 │ drop index idx_name;
+          ╰╴                  ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_drop_index_with_schema() {
+        assert_snapshot!(goto(r#"
+set search_path to public;
+create index idx_name on t(x);
+drop index public.idx_name$0;
+"#), @r"
+          ╭▸ 
+        3 │ create index idx_name on t(x);
+          │              ──────── 2. destination
+        4 │ drop index public.idx_name;
+          ╰╴                         ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_drop_index_defined_after() {
+        assert_snapshot!(goto("
+drop index idx_name$0;
+create index idx_name on t(x);
+"), @r"
+          ╭▸ 
+        2 │ drop index idx_name;
+          │                   ─ 1. source
+        3 │ create index idx_name on t(x);
+          ╰╴             ──────── 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_index_definition_returns_self() {
+        assert_snapshot!(goto("
+create index idx_name$0 on t(x);
+"), @r"
+          ╭▸ 
+        2 │ create index idx_name on t(x);
+          │              ┬──────┬
+          │              │      │
+          │              │      1. source
+          ╰╴             2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_drop_index_with_search_path() {
+        assert_snapshot!(goto(r#"
+create index idx_name on t(x);
+set search_path to bar;
+create index idx_name on f(x);
+set search_path to default;
+drop index idx_name$0;
+"#), @r"
+          ╭▸ 
+        2 │ create index idx_name on t(x);
+          │              ──────── 2. destination
+          ‡
+        6 │ drop index idx_name;
+          ╰╴                  ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_drop_index_multiple() {
+        assert_snapshot!(goto("
+create index idx1 on t(x);
+create index idx2 on t(y);
+drop index idx1, idx2$0;
+"), @r"
+          ╭▸ 
+        3 │ create index idx2 on t(y);
+          │              ──── 2. destination
+        4 │ drop index idx1, idx2;
+          ╰╴                    ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_table() {
+        assert_snapshot!(goto("
+create table users(id int);
+create index idx_users on users$0(id);
+"), @r"
+          ╭▸ 
+        2 │ create table users(id int);
+          │              ───── 2. destination
+        3 │ create index idx_users on users(id);
+          ╰╴                              ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_table_with_schema() {
+        assert_snapshot!(goto("
+create table public.users(id int);
+create index idx_users on public.users$0(id);
+"), @r"
+          ╭▸ 
+        2 │ create table public.users(id int);
+          │                     ───── 2. destination
+        3 │ create index idx_users on public.users(id);
+          ╰╴                                     ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_table_with_search_path() {
+        assert_snapshot!(goto(r#"
+set search_path to foo;
+create table foo.users(id int);
+create index idx_users on users$0(id);
+"#), @r"
+          ╭▸ 
+        3 │ create table foo.users(id int);
+          │                  ───── 2. destination
+        4 │ create index idx_users on users(id);
+          ╰╴                              ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_temp_table() {
+        assert_snapshot!(goto("
+create temp table users(id int);
+create index idx_users on users$0(id);
+"), @r"
+          ╭▸ 
+        2 │ create temp table users(id int);
+          │                   ───── 2. destination
+        3 │ create index idx_users on users(id);
+          ╰╴                              ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_column() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+create index idx_email on users(email$0);
+"), @r"
+          ╭▸ 
+        2 │ create table users(id int, email text);
+          │                            ───── 2. destination
+        3 │ create index idx_email on users(email);
+          ╰╴                                    ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_first_column() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+create index idx_id on users(id$0);
+"), @r"
+          ╭▸ 
+        2 │ create table users(id int, email text);
+          │                    ── 2. destination
+        3 │ create index idx_id on users(id);
+          ╰╴                              ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_multiple_columns() {
+        assert_snapshot!(goto("
+create table users(id int, email text, name text);
+create index idx_users on users(id, email$0, name);
+"), @r"
+          ╭▸ 
+        2 │ create table users(id int, email text, name text);
+          │                            ───── 2. destination
+        3 │ create index idx_users on users(id, email, name);
+          ╰╴                                        ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_column_with_schema() {
+        assert_snapshot!(goto("
+create table public.users(id int, email text);
+create index idx_email on public.users(email$0);
+"), @r"
+          ╭▸ 
+        2 │ create table public.users(id int, email text);
+          │                                   ───── 2. destination
+        3 │ create index idx_email on public.users(email);
+          ╰╴                                           ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_index_column_temp_table() {
+        assert_snapshot!(goto("
+create temp table users(id int, email text);
+create index idx_email on users(email$0);
+"), @r"
+          ╭▸ 
+        2 │ create temp table users(id int, email text);
+          │                                 ───── 2. destination
+        3 │ create index idx_email on users(email);
+          ╰╴                                    ─ 1. source
+        ");
+    }
 }

crates/squawk_ide/src/resolve.rs

@@ -6,35 +6,61 @@ use squawk_syntax::{
 
 use crate::binder::Binder;
 use crate::symbols::{Name, Schema, SymbolKind};
+use squawk_syntax::SyntaxNode;
 
 #[derive(Debug)]
 enum NameRefContext {
     DropTable,
     Table,
+    DropIndex,
+    CreateIndex,
+    CreateIndexColumn,
 }
 
 pub(crate) fn resolve_name_ref(binder: &Binder, name_ref: &ast::NameRef) -> Option<SyntaxNodePtr> {
     let context = classify_name_ref_context(name_ref)?;
 
     match context {
-        NameRefContext::DropTable | NameRefContext::Table => {
+        NameRefContext::DropTable | NameRefContext::Table | NameRefContext::CreateIndex => {
             let path = find_containing_path(name_ref)?;
             let table_name = extract_table_name(&path)?;
             let schema = extract_schema_name(&path);
             let position = name_ref.syntax().text_range().start();
             resolve_table(binder, &table_name, &schema, position)
         }
+        NameRefContext::DropIndex => {
+            let path = find_containing_path(name_ref)?;
+            let index_name = extract_table_name(&path)?;
+            let schema = extract_schema_name(&path);
+            let position = name_ref.syntax().text_range().start();
+            resolve_index(binder, &index_name, &schema, position)
+        }
+        NameRefContext::CreateIndexColumn => resolve_create_index_column(binder, name_ref),
     }
 }
 
 fn classify_name_ref_context(name_ref: &ast::NameRef) -> Option<NameRefContext> {
+    let mut in_partition_item = false;
+
     for ancestor in name_ref.syntax().ancestors() {
         if ast::DropTable::can_cast(ancestor.kind()) {
             return Some(NameRefContext::DropTable);
         }
         if ast::Table::can_cast(ancestor.kind()) {
             return Some(NameRefContext::Table);
         }
+        if ast::DropIndex::can_cast(ancestor.kind()) {
+            return Some(NameRefContext::DropIndex);
+        }
+        if ast::PartitionItem::can_cast(ancestor.kind()) {
+            in_partition_item = true;
+        }
+        if ast::CreateIndex::can_cast(ancestor.kind()) {
+            if in_partition_item {
+                return Some(NameRefContext::CreateIndexColumn);
+            }
+            return Some(NameRefContext::CreateIndex);
+        }
     }
 
     None
@@ -68,6 +94,71 @@ fn resolve_table(
     None
 }
 
+fn resolve_index(
+    binder: &Binder,
+    index_name: &Name,
+    schema: &Option<Schema>,
+    position: TextSize,
+) -> Option<SyntaxNodePtr> {
+    let symbols = binder.scopes[binder.root_scope()].get(index_name)?;
+
+    if let Some(schema) = schema {
+        let symbol_id = symbols.iter().copied().find(|id| {
+            let symbol = &binder.symbols[*id];
+            symbol.kind == SymbolKind::Index && &symbol.schema == schema
+        })?;
+        return Some(binder.symbols[symbol_id].ptr);
+    } else {
+        let search_path = binder.search_path_at(position);
+        for search_schema in search_path {
+            if let Some(symbol_id) = symbols.iter().copied().find(|id| {
+                let symbol = &binder.symbols[*id];
+                symbol.kind == SymbolKind::Index && &symbol.schema == search_schema
+            }) {
+                return Some(binder.symbols[symbol_id].ptr);
+            }
+        }
+    }
+    None
+}
+
+fn resolve_create_index_column(binder: &Binder, name_ref: &ast::NameRef) -> Option<SyntaxNodePtr> {
+    let column_name = Name::new(name_ref.syntax().text().to_string());
+
+    let create_index = name_ref
+        .syntax()
+        .ancestors()
+        .find_map(ast::CreateIndex::cast)?;
+    let relation_name = create_index.relation_name()?;
+    let path = relation_name.path()?;
+
+    let table_name = extract_table_name(&path)?;
+    let schema = extract_schema_name(&path);
+    let position = name_ref.syntax().text_range().start();
+
+    let table_ptr = resolve_table(binder, &table_name, &schema, position)?;
+
+    let root: &SyntaxNode = &name_ref.syntax().ancestors().last()?;
+    let table_name_node = table_ptr.to_node(root);
+
+    let create_table = table_name_node
+        .ancestors()
+        .find_map(ast::CreateTable::cast)?;
+
+    let table_arg_list = create_table.table_arg_list()?;
+
+    for arg in table_arg_list.args() {
+        if let ast::TableArg::Column(column) = arg
+            && let Some(col_name) = column.name()
+            && Name::new(col_name.syntax().text().to_string()) == column_name
+        {
+            return Some(SyntaxNodePtr::new(col_name.syntax()));
+        }
+    }
+
+    None
+}
+
 fn find_containing_path(name_ref: &ast::NameRef) -> Option<ast::Path> {
     for ancestor in name_ref.syntax().ancestors() {
         if let Some(path) = ast::Path::cast(ancestor) {