ide: goto def for begin atomic function defs

crates/squawk_ide/src/ast_nav.rs

@@ -40,6 +40,17 @@ pub(crate) fn find_cte_with_table(
     None
 }
 
+pub(crate) fn ancestors_outside_own_with_clause(
+    node: &SyntaxNode,
+) -> impl Iterator<Item = SyntaxNode> {
+    let mut prev_was_with_clause = false;
+    node.ancestors().filter(move |ancestor| {
+        let skip = prev_was_with_clause;
+        prev_was_with_clause = ast::WithClause::can_cast(ancestor.kind());
+        !skip
+    })
+}
+
 pub(crate) fn iter_values_columns(values: &ast::Values) -> impl Iterator<Item = (Name, ast::Expr)> {
     values
         .row_list()

crates/squawk_ide/src/classify.rs

@@ -1004,6 +1004,9 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
             // (anything in SELECT except FROM clause)
             return Some(NameRefClass::SelectColumn);
         }
+        if ast::Values::can_cast(ancestor.kind()) {
+            return Some(NameRefClass::SelectColumn);
+        }
         if ast::CompoundSelect::can_cast(ancestor.kind())
             && in_order_by_clause
             && let Some(parent) = node.parent()

crates/squawk_ide/src/goto_definition.rs

@@ -1975,6 +1975,162 @@ end;
         ");
     }
 
+    #[test]
+    fn goto_function_param_in_begin_atomic_insert_values() {
+        assert_snapshot!(goto("
+create table t (a int);
+create function f(x int) returns void
+  begin atomic
+    insert into t values (x$0);
+  end;
+"), @"
+          ╭▸ 
+        3 │ create function f(x int) returns void
+          │                   ─ 2. destination
+        4 │   begin atomic
+        5 │     insert into t values (x);
+          ╰╴                          ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_function_param_in_begin_atomic_update() {
+        assert_snapshot!(goto("
+create table t (a int);
+create function f(x int) returns void
+  begin atomic
+    update t set a = x$0;
+  end;
+"), @"
+          ╭▸ 
+        3 │ create function f(x int) returns void
+          │                   ─ 2. destination
+        4 │   begin atomic
+        5 │     update t set a = x;
+          ╰╴                     ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_function_param_in_begin_atomic_delete_where() {
+        assert_snapshot!(goto("
+create table t (a int);
+create function f(x int) returns void
+  begin atomic
+    delete from t where a = x$0;
+  end;
+"), @"
+          ╭▸ 
+        3 │ create function f(x int) returns void
+          │                   ─ 2. destination
+        4 │   begin atomic
+        5 │     delete from t where a = x;
+          ╰╴                            ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_function_param_in_begin_atomic_merge() {
+        assert_snapshot!(goto("
+create table t (a int);
+create function f(x int) returns void
+  begin atomic
+    merge into t
+    using t u on t.a = x$0
+    when matched then
+      do nothing;
+  end;
+"), @"
+          ╭▸ 
+        3 │ create function f(x int) returns void
+          │                   ─ 2. destination
+          ‡
+        6 │     using t u on t.a = x
+          ╰╴                       ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_function_param_in_begin_atomic_values() {
+        assert_snapshot!(goto("
+create function f(x int) returns int
+  begin atomic
+    values (x$0);
+  end;
+"), @"
+          ╭▸ 
+        2 │ create function f(x int) returns int
+          │                   ─ 2. destination
+        3 │   begin atomic
+        4 │     values (x);
+          ╰╴            ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_begin_atomic_column_shadows_function_param() {
+        assert_snapshot!(goto("
+create table t (x int);
+create function f(x int) returns void
+  begin atomic
+    delete from t where x$0 = 1;
+  end;
+"), @"
+          ╭▸ 
+        2 │ create table t (x int);
+          │                 ─ 2. destination
+          ‡
+        5 │     delete from t where x = 1;
+          ╰╴                        ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_function_param_in_cte_insert_values_not_outer_select_column() {
+        assert_snapshot!(goto("
+create table t (a int);
+create table u (x int);
+create function f(x int) returns int
+  begin atomic
+    with c as (
+      insert into t values (x$0) returning 1
+    )
+    select 1
+    from u;
+  end;
+"), @"
+          ╭▸ 
+        4 │ create function f(x int) returns int
+          │                   ─ 2. destination
+          ‡
+        7 │       insert into t values (x) returning 1
+          ╰╴                            ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_function_param_in_cte_insert_values_not_outer_update_column() {
+        assert_snapshot!(goto("
+create table t (a int);
+create table u (x int);
+create function f(x int) returns void
+  begin atomic
+    with c as (
+      insert into t values (x$0)
+    )
+    update u
+    set x = 1;
+  end;
+"), @"
+          ╭▸ 
+        4 │ create function f(x int) returns void
+          │                   ─ 2. destination
+          ‡
+        7 │       insert into t values (x)
+          ╰╴                            ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_positional_param_unnamed_param() {
         assert_snapshot!(goto("

crates/squawk_ide/src/resolve.rs

@@ -2813,7 +2813,7 @@ fn resolve_select_column_ptr(
     // Walk up through enclosing selects so a column in a correlated subquery
     // resolves against an outer query's from clause:
     // `select (select (select a)) from foo`
-    for ancestor in column_name_ref.syntax().ancestors() {
+    for ancestor in ast_nav::ancestors_outside_own_with_clause(column_name_ref.syntax()) {
         let Some(from_clause) = select_like_from_clause(&ancestor) else {
             continue;
         };
@@ -2835,7 +2835,7 @@ fn resolve_select_column_ptr(
     // A correlated subquery can reference the target relation of an enclosing
     // DML statement, e.g. `update foo set a = (select b)` where `b` is `foo.b`
     let in_file = InFile::new(file, column_name_ref);
-    for ancestor in column_name_ref.syntax().ancestors() {
+    for ancestor in ast_nav::ancestors_outside_own_with_clause(column_name_ref.syntax()) {
         match ancestor.kind() {
             SyntaxKind::UPDATE => return resolve_update_column_ptr(db, in_file),
             SyntaxKind::DELETE => return resolve_delete_column_ptr(db, in_file),
@@ -5497,7 +5497,11 @@ fn resolve_update_column_ptr(
         )?;
     }
 
-    resolve_column_for_path(db, InFile::new(file, &path), column_name)
+    resolve_column_for_path(db, InFile::new(file, &path), column_name).or_else(|| {
+        (!is_set_target)
+            .then(|| resolve_enclosing_function_param(InFile::new(file, column_name_ref)))
+            .flatten()
+    })
 }
 
 fn resolve_delete_column_ptr(
@@ -5535,6 +5539,7 @@ fn resolve_delete_column_ptr(
     }
 
     resolve_column_for_path(db, InFile::new(file, &path), column_name)
+        .or_else(|| resolve_enclosing_function_param(InFile::new(file, column_name_ref)))
 }
 
 fn resolve_delete_table_name_ptr(
@@ -5622,7 +5627,11 @@ fn resolve_merge_column_ptr(
     }
 
     let path = merge.table_relation_name()?.table_name_ref()?.path_ref()?;
-    resolve_column_for_path(db, InFile::new(file, &path), column_name)
+    resolve_column_for_path(db, InFile::new(file, &path), column_name).or_else(|| {
+        (!is_set_target && !in_insert_column_list)
+            .then(|| resolve_enclosing_function_param(InFile::new(file, column_name_ref)))
+            .flatten()
+    })
 }
 
 // TODO: I think we could use trait(s) here to simplify this and have the