lsp: go to def for case exprs

PLAN.md

@@ -1056,8 +1056,6 @@ select x from "t";
 
 gives:
 
-Note: we have to be mindful of casing here,
-
 ```sql
 select "x" from "t";
 ```

crates/squawk_ide/src/goto_definition.rs

@@ -0,0 +1,162 @@
+use rowan::{TextRange, TextSize};
+use squawk_syntax::{
+    SyntaxKind, SyntaxToken,
+    ast::{self, AstNode},
+};
+
+pub fn goto_definition(file: ast::SourceFile, offset: TextSize) -> Option<TextRange> {
+    let token = token_from_offset(&file, offset)?;
+    let parent = token.parent()?;
+
+    // goto def on case exprs
+    if (token.kind() == SyntaxKind::WHEN_KW && parent.kind() == SyntaxKind::WHEN_CLAUSE)
+        || (token.kind() == SyntaxKind::ELSE_KW && parent.kind() == SyntaxKind::ELSE_CLAUSE)
+        || (token.kind() == SyntaxKind::END_KW && parent.kind() == SyntaxKind::CASE_EXPR)
+    {
+        for parent in token.parent_ancestors() {
+            if let Some(case_expr) = ast::CaseExpr::cast(parent) {
+                if let Some(case_token) = case_expr.case_token() {
+                    return Some(case_token.text_range());
+                }
+            }
+        }
+    }
+
+    return None;
+}
+
+fn token_from_offset(file: &ast::SourceFile, offset: TextSize) -> Option<SyntaxToken> {
+    let mut token = file.syntax().token_at_offset(offset).right_biased()?;
+    // want to be lenient in case someone clicks the trailing `;` of a line
+    // instead of an identifier
+    if token.kind() == SyntaxKind::SEMICOLON {
+        token = token.prev_token()?;
+    }
+    return Some(token);
+}
+
+#[cfg(test)]
+mod test {
+    use crate::goto_definition::goto_definition;
+    use annotate_snippets::{AnnotationKind, Level, Renderer, Snippet, renderer::DecorStyle};
+    use insta::assert_snapshot;
+    use log::info;
+    use rowan::TextSize;
+    use squawk_syntax::ast;
+
+    // TODO: we should probably use something else since `$0` is valid syntax, maybe `%0`?
+    const MARKER: &str = "$0";
+
+    fn fixture(sql: &str) -> (Option<TextSize>, String) {
+        if let Some(pos) = sql.find(MARKER) {
+            return (
+                Some(TextSize::new((pos - 1) as u32)),
+                sql.replace(MARKER, ""),
+            );
+        }
+        (None, sql.to_string())
+    }
+
+    #[track_caller]
+    fn goto(sql: &str) -> String {
+        goto_(sql).expect("should always find a definition")
+    }
+
+    #[track_caller]
+    fn goto_(sql: &str) -> Option<String> {
+        info!("starting");
+        let (offset, sql) = fixture(sql);
+        let parse = ast::SourceFile::parse(&sql);
+        assert_eq!(parse.errors(), vec![]);
+        let file: ast::SourceFile = parse.tree();
+        let Some(offset) = offset else {
+            info!("offset not found, did you put a marker `$0` in the sql?");
+            return None;
+        };
+        if let Some(result) = goto_definition(file, offset) {
+            let offset: usize = offset.into();
+            let group = Level::INFO.primary_title("definition").element(
+                Snippet::source(&sql)
+                    .fold(true)
+                    .annotation(
+                        AnnotationKind::Context
+                            .span(result.into())
+                            .label("2. destination"),
+                    )
+                    .annotation(
+                        AnnotationKind::Context
+                            .span(offset..offset + 1)
+                            .label("1. source"),
+                    ),
+            );
+            let renderer = Renderer::plain().decor_style(DecorStyle::Unicode);
+            return Some(
+                renderer
+                    .render(&[group])
+                    .to_string()
+                    // hacky cleanup to make the text shorter
+                    .replace("info: definition", ""),
+            );
+        }
+        None
+    }
+
+    fn goto_not_found(sql: &str) {
+        assert!(goto_(sql).is_none(), "Should not find a definition");
+    }
+
+    #[test]
+    fn goto_case_when() {
+        assert_snapshot!(goto("
+select case when$0 x > 1 then 1 else 2 end;
+"), @r"
+          ╭▸ 
+        2 │ select case when x > 1 then 1 else 2 end;
+          │        ┬───    ─ 1. source
+          │        │
+          ╰╴       2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_case_else() {
+        assert_snapshot!(goto("
+select case when x > 1 then 1 else$0 2 end;
+"), @r"
+          ╭▸ 
+        2 │ select case when x > 1 then 1 else 2 end;
+          ╰╴       ──── 2. destination       ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_case_end() {
+        assert_snapshot!(goto("
+select case when x > 1 then 1 else 2 end$0;
+"), @r"
+          ╭▸ 
+        2 │ select case when x > 1 then 1 else 2 end;
+          ╰╴       ──── 2. destination             ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_case_end_trailing_semi() {
+        assert_snapshot!(goto("
+select case when x > 1 then 1 else 2 end;$0
+"), @r"
+          ╭▸ 
+        2 │ select case when x > 1 then 1 else 2 end;
+          ╰╴       ──── 2. destination              ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_case_then_not_found() {
+        goto_not_found(
+            "
+select case when x > 1 then$0 1 else 2 end;
+",
+        )
+    }
+}

crates/squawk_ide/src/lib.rs

@@ -1 +1,2 @@
 pub mod expand_selection;
+pub mod goto_definition;

crates/squawk_parser/src/grammar.rs

@@ -4993,10 +4993,14 @@ fn partition_option(p: &mut Parser<'_>) {
 }
 
 fn opt_inherits_tables(p: &mut Parser<'_>) {
+    let m = p.start();
     if p.eat(INHERITS_KW) {
         p.expect(L_PAREN);
         path_name_ref_list(p);
         p.expect(R_PAREN);
+        m.complete(p, INHERITS);
+    } else {
+        m.abandon(p);
     }
 }
 
@@ -12295,16 +12299,20 @@ fn update(p: &mut Parser<'_>, m: Option<Marker>) -> CompletedMarker {
     // [ FROM from_item [, ...] ]
     opt_from_clause(p);
     // [ WHERE condition | WHERE CURRENT OF cursor_name ]
+    opt_where_or_current_of(p);
+    // [ RETURNING { * | output_expression [ [ AS ] output_name ] } [, ...] ]
+    opt_returning_clause(p);
+    m.complete(p, UPDATE)
+}
+
+fn opt_where_or_current_of(p: &mut Parser<'_>) {
     if p.at(WHERE_KW) {
         if p.nth_at(1, CURRENT_KW) {
             opt_where_current_of(p);
         } else {
             opt_where_clause(p);
         }
     }
-    // [ RETURNING { * | output_expression [ [ AS ] output_name ] } [, ...] ]
-    opt_returning_clause(p);
-    m.complete(p, UPDATE)
 }
 
 fn with(p: &mut Parser<'_>, m: Option<Marker>) -> Option<CompletedMarker> {
@@ -12355,24 +12363,22 @@ fn delete(p: &mut Parser<'_>, m: Option<Marker>) -> CompletedMarker {
         }
     }
     // [ WHERE condition | WHERE CURRENT OF cursor_name ]
-    if p.at(WHERE_KW) {
-        if p.nth_at(1, CURRENT_KW) {
-            opt_where_current_of(p);
-        } else {
-            opt_where_clause(p);
-        }
-    }
+    opt_where_or_current_of(p);
     opt_returning_clause(p);
     m.complete(p, DELETE)
 }
 
 // WHERE CURRENT OF cursor_name
 fn opt_where_current_of(p: &mut Parser<'_>) {
+    let m = p.start();
     if p.eat(WHERE_KW) {
         if p.eat(CURRENT_KW) {
             p.expect(OF_KW);
             name_ref(p);
         }
+        m.complete(p, WHERE_CURRENT_OF);
+    } else {
+        m.abandon(p);
     }
 }
 

crates/squawk_parser/tests/snapshots/tests__create_foreign_table_ok.snap

@@ -189,25 +189,26 @@ SOURCE_FILE
       WHITESPACE "\n"
       R_PAREN ")"
     WHITESPACE "\n  "
-    INHERITS_KW "inherits"
-    WHITESPACE " "
-    L_PAREN "("
-    PATH
+    INHERITS
+      INHERITS_KW "inherits"
+      WHITESPACE " "
+      L_PAREN "("
       PATH
+        PATH
+          PATH_SEGMENT
+            NAME_REF
+              IDENT "foo"
+        DOT "."
         PATH_SEGMENT
           NAME_REF
-            IDENT "foo"
-      DOT "."
-      PATH_SEGMENT
-        NAME_REF
-          IDENT "bar"
-    COMMA ","
-    WHITESPACE " "
-    PATH
-      PATH_SEGMENT
-        NAME_REF
-          IDENT "bar"
-    R_PAREN ")"
+            IDENT "bar"
+      COMMA ","
+      WHITESPACE " "
+      PATH
+        PATH_SEGMENT
+          NAME_REF
+            IDENT "bar"
+      R_PAREN ")"
     WHITESPACE "\n  "
     SERVER_KW "server"
     WHITESPACE " "

crates/squawk_parser/tests/snapshots/tests__create_table_ok.snap

@@ -394,31 +394,32 @@ SOURCE_FILE
                 INT_KW "int"
       R_PAREN ")"
     WHITESPACE "\n"
-    INHERITS_KW "inherits"
-    WHITESPACE " "
-    L_PAREN "("
-    PATH
+    INHERITS
+      INHERITS_KW "inherits"
+      WHITESPACE " "
+      L_PAREN "("
       PATH
+        PATH
+          PATH_SEGMENT
+            NAME_REF
+              IDENT "foo"
+        DOT "."
         PATH_SEGMENT
           NAME_REF
-            IDENT "foo"
-      DOT "."
-      PATH_SEGMENT
-        NAME_REF
-          IDENT "bar"
-    COMMA ","
-    WHITESPACE " "
-    PATH
-      PATH_SEGMENT
-        NAME_REF
-          IDENT "bar"
-    COMMA ","
-    WHITESPACE " "
-    PATH
-      PATH_SEGMENT
-        NAME_REF
-          IDENT "buzz"
-    R_PAREN ")"
+            IDENT "bar"
+      COMMA ","
+      WHITESPACE " "
+      PATH
+        PATH_SEGMENT
+          NAME_REF
+            IDENT "bar"
+      COMMA ","
+      WHITESPACE " "
+      PATH
+        PATH_SEGMENT
+          NAME_REF
+            IDENT "buzz"
+      R_PAREN ")"
   SEMICOLON ";"
   WHITESPACE "\n\n"
   CREATE_TABLE

crates/squawk_parser/tests/snapshots/tests__delete_ok.snap

@@ -477,14 +477,15 @@ SOURCE_FILE
           NAME_REF
             IDENT "invoices"
     WHITESPACE " \n"
-    WHERE_KW "where"
-    WHITESPACE " "
-    CURRENT_KW "current"
-    WHITESPACE " "
-    OF_KW "of"
-    WHITESPACE " "
-    NAME_REF
-      IDENT "invoice_cursor"
+    WHERE_CURRENT_OF
+      WHERE_KW "where"
+      WHITESPACE " "
+      CURRENT_KW "current"
+      WHITESPACE " "
+      OF_KW "of"
+      WHITESPACE " "
+      NAME_REF
+        IDENT "invoice_cursor"
   SEMICOLON ";"
   WHITESPACE "\n\n"
   COMMENT "-- returning"
@@ -931,14 +932,15 @@ SOURCE_FILE
           NAME_REF
             IDENT "tasks"
     WHITESPACE " "
-    WHERE_KW "WHERE"
-    WHITESPACE " "
-    CURRENT_KW "CURRENT"
-    WHITESPACE " "
-    OF_KW "OF"
-    WHITESPACE " "
-    NAME_REF
-      IDENT "c_tasks"
+    WHERE_CURRENT_OF
+      WHERE_KW "WHERE"
+      WHITESPACE " "
+      CURRENT_KW "CURRENT"
+      WHITESPACE " "
+      OF_KW "OF"
+      WHITESPACE " "
+      NAME_REF
+        IDENT "c_tasks"
   SEMICOLON ";"
   WHITESPACE "\n\n"
   DELETE