lsp: add code action for case else clause

crates/squawk_ide/Cargo.toml

@@ -13,6 +13,7 @@ rust-version.workspace = true
 
 [dependencies]
 squawk-syntax.workspace = true
+squawk-linter.workspace = true
 rowan.workspace = true
 line-index.workspace = true
 annotate-snippets.workspace = true

crates/squawk_ide/src/code_actions.rs

@@ -0,0 +1,154 @@
+use rowan::TextSize;
+use squawk_linter::Edit;
+use squawk_syntax::{
+    SyntaxKind,
+    ast::{self, AstNode},
+};
+
+#[derive(Debug, Clone)]
+pub struct CodeAction {
+    pub title: String,
+    pub edits: Vec<Edit>,
+}
+
+pub fn code_actions(file: ast::SourceFile, offset: TextSize) -> Option<Vec<CodeAction>> {
+    let mut actions = vec![];
+    remove_else_clause(&mut actions, &file, offset);
+    Some(actions)
+}
+
+fn remove_else_clause(
+    actions: &mut Vec<CodeAction>,
+    file: &ast::SourceFile,
+    offset: TextSize,
+) -> Option<()> {
+    let else_token = file
+        .syntax()
+        .token_at_offset(offset)
+        .find(|x| x.kind() == SyntaxKind::ELSE_KW)?;
+    let parent = else_token.parent()?;
+    let else_clause = ast::ElseClause::cast(parent)?;
+
+    let mut edits = vec![];
+    edits.push(Edit::delete(else_clause.syntax().text_range()));
+    if let Some(token) = else_token.prev_token() {
+        if token.kind() == SyntaxKind::WHITESPACE {
+            edits.push(Edit::delete(token.text_range()));
+        }
+    }
+
+    actions.push(CodeAction {
+        title: "Remove `else` clause".to_owned(),
+        edits,
+    });
+    Some(())
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::test_utils::fixture;
+    use insta::assert_snapshot;
+    use rowan::TextSize;
+    use squawk_syntax::ast;
+
+    fn apply_code_action(
+        f: impl Fn(&mut Vec<CodeAction>, &ast::SourceFile, TextSize) -> Option<()>,
+        sql: &str,
+    ) -> String {
+        let (offset, sql) = fixture(sql);
+        let parse = ast::SourceFile::parse(&sql);
+        assert_eq!(parse.errors(), vec![]);
+        let file: ast::SourceFile = parse.tree();
+
+        let mut actions = vec![];
+        f(&mut actions, &file, offset);
+
+        assert!(
+            !actions.is_empty(),
+            "We should always have actions for `apply_code_action`. If you want to ensure there are no actions, use `code_action_not_applicable` instead."
+        );
+
+        let action = &actions[0];
+        let mut result = sql.clone();
+
+        let mut edits = action.edits.clone();
+        edits.sort_by_key(|e| e.text_range.start());
+        check_overlap(&edits);
+        edits.reverse();
+
+        for edit in edits {
+            let start: usize = edit.text_range.start().into();
+            let end: usize = edit.text_range.end().into();
+            let replacement = edit.text.as_deref().unwrap_or("");
+            result.replace_range(start..end, replacement);
+        }
+
+        let reparse = ast::SourceFile::parse(&result);
+        assert_eq!(
+            reparse.errors(),
+            vec![],
+            "Code actions shouldn't cause syntax errors"
+        );
+
+        result
+    }
+
+    // There's an invariant where the edits can't overlap.
+    // For example, if we have an edit that deletes the full `else clause` and
+    // another edit that deletes the `else` keyword and they overlap, then
+    // vscode doesn't surface the code action.
+    fn check_overlap(edits: &[Edit]) {
+        for (edit_i, edit_j) in edits.iter().zip(edits.iter().skip(1)) {
+            if let Some(intersection) = edit_i.text_range.intersect(edit_j.text_range) {
+                assert!(
+                    intersection.is_empty(),
+                    "Edit ranges must not overlap: {:?} and {:?} intersect at {:?}",
+                    edit_i.text_range,
+                    edit_j.text_range,
+                    intersection
+                );
+            }
+        }
+    }
+
+    fn code_action_not_applicable(
+        f: impl Fn(&mut Vec<CodeAction>, &ast::SourceFile, TextSize) -> Option<()>,
+        sql: &str,
+    ) -> bool {
+        let (offset, sql) = fixture(sql);
+        let parse = ast::SourceFile::parse(&sql);
+        assert_eq!(parse.errors(), vec![]);
+        let file: ast::SourceFile = parse.tree();
+
+        let mut actions = vec![];
+        f(&mut actions, &file, offset);
+        actions.is_empty()
+    }
+
+    #[test]
+    fn remove_else_clause_() {
+        assert_snapshot!(apply_code_action(
+            remove_else_clause,
+            "select case x when true then 1 else$0 2 end;"),
+            @"select case x when true then 1 end;"
+        );
+    }
+
+    #[test]
+    fn remove_else_clause_before_token() {
+        assert_snapshot!(apply_code_action(
+            remove_else_clause,
+            "select case x when true then 1 $0else 2 end;"),
+            @"select case x when true then 1 end;"
+        );
+    }
+
+    #[test]
+    fn remove_else_clause_not_applicable() {
+        assert!(code_action_not_applicable(
+            remove_else_clause,
+            "select case x when true then 1 else 2 end$0;"
+        ));
+    }
+}

crates/squawk_ide/src/expand_selection.rs

@@ -281,8 +281,8 @@ fn extend_list_item(node: &SyntaxNode) -> Option<TextRange> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::test_utils::fixture;
     use insta::assert_debug_snapshot;
-    use rowan::TextSize;
     use squawk_syntax::{SourceFile, ast::AstNode};
 
     fn expand(sql: &str) -> Vec<String> {
@@ -306,14 +306,6 @@ mod tests {
         results
     }
 
-    fn fixture(sql: &str) -> (TextSize, String) {
-        const MARKER: &str = "$0";
-        if let Some(pos) = sql.find(MARKER) {
-            return (TextSize::new(pos as u32), sql.replace(MARKER, ""));
-        }
-        panic!("No marker found in test SQL");
-    }
-
     #[test]
     fn simple() {
         assert_debug_snapshot!(expand(r#"select $01 + 1"#), @r#"

crates/squawk_ide/src/goto_definition.rs

@@ -38,25 +38,12 @@ fn token_from_offset(file: &ast::SourceFile, offset: TextSize) -> Option<SyntaxT
 #[cfg(test)]
 mod test {
     use crate::goto_definition::goto_definition;
+    use crate::test_utils::fixture;
     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")
@@ -65,14 +52,13 @@ mod test {
     #[track_caller]
     fn goto_(sql: &str) -> Option<String> {
         info!("starting");
-        let (offset, sql) = fixture(sql);
+        let (mut offset, sql) = fixture(sql);
+        // For go to def we want the previous character since we usually put the
+        // marker after the item we're trying to go to def on.
+        offset = offset.checked_sub(1.into()).unwrap_or_default();
         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(

crates/squawk_ide/src/lib.rs

@@ -1,2 +1,5 @@
+pub mod code_actions;
 pub mod expand_selection;
 pub mod goto_definition;
+#[cfg(test)]
+pub mod test_utils;

crates/squawk_ide/src/test_utils.rs

@@ -0,0 +1,12 @@
+use rowan::TextSize;
+
+// TODO: we should probably use something else since `$0` is valid syntax, maybe `%0`?
+const MARKER: &str = "$0";
+
+#[track_caller]
+pub(crate) fn fixture(sql: &str) -> (TextSize, String) {
+    if let Some(pos) = sql.find(MARKER) {
+        return (TextSize::new(pos as u32), sql.replace(MARKER, ""));
+    }
+    panic!("No marker found in SQL. Did you forget to add a marker `$0`?");
+}

crates/squawk_linter/src/lib.rs

@@ -228,6 +228,7 @@ impl Fix {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Edit {
     pub text_range: TextRange,
+    // TODO: does this need to be an Option?
     pub text: Option<String>,
 }
 impl Edit {
@@ -243,6 +244,12 @@ impl Edit {
             text: Some(text.into()),
         }
     }
+    pub fn delete(text_range: TextRange) -> Self {
+        Self {
+            text_range,
+            text: None,
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]

crates/squawk_server/src/lib.rs

@@ -17,6 +17,7 @@ use lsp_types::{
     request::{CodeActionRequest, GotoDefinition, Request, SelectionRangeRequest},
 };
 use rowan::TextRange;
+use squawk_ide::code_actions::code_actions;
 use squawk_ide::goto_definition::goto_definition;
 use squawk_syntax::{Parse, SourceFile};
 use std::collections::HashMap;
@@ -235,12 +236,25 @@ fn handle_selection_range(
 fn handle_code_action(
     connection: &Connection,
     req: lsp_server::Request,
-    _documents: &HashMap<Url, DocumentState>,
+    documents: &HashMap<Url, DocumentState>,
 ) -> Result<()> {
     let params: CodeActionParams = serde_json::from_value(req.params)?;
     let uri = params.text_document.uri;
 
-    let mut actions = Vec::new();
+    let mut actions: CodeActionResponse = Vec::new();
+
+    let content = documents.get(&uri).map_or("", |doc| &doc.content);
+    let parse: Parse<SourceFile> = SourceFile::parse(content);
+    let file = parse.tree();
+    let line_index = LineIndex::new(content);
+    let offset = lsp_utils::offset(&line_index, params.range.start).unwrap();
+
+    let ide_actions = code_actions(file, offset).unwrap_or_default();
+
+    for action in ide_actions {
+        let lsp_action = lsp_utils::code_action(&line_index, uri.clone(), action);
+        actions.push(CodeActionOrCommand::CodeAction(lsp_action));
+    }
 
     for mut diagnostic in params
         .context

crates/squawk_server/src/lsp_utils.rs

@@ -1,7 +1,8 @@
-use std::ops::Range;
+use std::{collections::HashMap, ops::Range};
 
 use line_index::{LineIndex, TextRange, TextSize};
 use log::warn;
+use lsp_types::{CodeAction, CodeActionKind, Url, WorkspaceEdit};
 
 fn text_range(index: &LineIndex, range: lsp_types::Range) -> Option<TextRange> {
     let start = offset(index, range.start)?;
@@ -36,6 +37,38 @@ pub(crate) fn offset(index: &LineIndex, position: lsp_types::Position) -> Option
     Some(line_range.start() + clamped_len)
 }
 
+pub(crate) fn code_action(
+    line_index: &LineIndex,
+    uri: Url,
+    action: squawk_ide::code_actions::CodeAction,
+) -> lsp_types::CodeAction {
+    CodeAction {
+        title: action.title,
+        kind: Some(CodeActionKind::QUICKFIX),
+        diagnostics: None,
+        edit: Some(WorkspaceEdit {
+            changes: Some({
+                let mut changes = HashMap::new();
+                let edits = action
+                    .edits
+                    .into_iter()
+                    .map(|edit| lsp_types::TextEdit {
+                        range: range(line_index, edit.text_range),
+                        new_text: edit.text.unwrap_or_default(),
+                    })
+                    .collect();
+                changes.insert(uri, edits);
+                changes
+            }),
+            ..Default::default()
+        }),
+        command: None,
+        is_preferred: Some(true),
+        disabled: None,
+        data: None,
+    }
+}
+
 pub(crate) fn range(line_index: &LineIndex, range: TextRange) -> lsp_types::Range {
     let start = line_index.line_col(range.start());
     let end = line_index.line_col(range.end());