Add 'merge case branches' code action

CHANGELOG.md

@@ -116,6 +116,31 @@
 
 ### Language server
 
+- The language server can now offer a code action to merge consecutive case
+  branches with the same body. For example:
+
+  ```gleam
+  case user {
+    Admin(name:, ..) -> todo
+  //^^^^^^^^^^^^^^^^^^^^^^^^
+    Guest(name:, ..) -> todo
+  //^^^^^^^^^^^^^^^^ Selecting these two branches you can
+  //                 trigger the "Merge case branches" code action
+    _ -> todo
+  }
+  ```
+
+  Triggering the code action would result in the following code:
+
+  ```gleam
+  case user {
+    Admin(name:, ..) | Guest(name:, ..) -> todo
+    _ -> todo
+  }
+  ```
+
+  ([Giacomo Cavalieri](https://github.com/giacomocavalieri))
+
 - The "inline variable" code action can now trigger when used over the let
   keyword of a variable to inline.
   ([Giacomo Cavalieri](https://github.com/giacomocavalieri))

compiler-core/src/language_server/code_action.rs

@@ -9271,3 +9271,162 @@ impl<'ast> ast::visit::Visit<'ast> for ExtractFunction<'ast> {
         }
     }
 }
+
+/// Code action to merge two identical branches together.
+///
+pub struct MergeCaseBranches<'a> {
+    module: &'a Module,
+    params: &'a CodeActionParams,
+    edits: TextEdits<'a>,
+    /// These are the positions of the patterns of all the consecutive branches
+    /// we've determined can be merged, for example if we're mergin the first
+    /// two branches here:
+    ///
+    /// ```gleam
+    ///   case wibble {
+    ///     1 -> todo
+    /// //  ^ this location here
+    ///     20 -> todo
+    /// //  ^^ and this location here
+    ///   _ -> todo
+    /// }
+    /// ```
+    ///
+    /// We need those to delete all the space between each consecutive pattern,
+    /// replacing it with the `|` for alternatives
+    ///
+    patterns_to_merge: Option<Vec<SrcSpan>>,
+}
+
+impl<'a> MergeCaseBranches<'a> {
+    pub fn new(
+        module: &'a Module,
+        line_numbers: &'a LineNumbers,
+        params: &'a CodeActionParams,
+    ) -> Self {
+        Self {
+            module,
+            params,
+            edits: TextEdits::new(line_numbers),
+            patterns_to_merge: None,
+        }
+    }
+
+    pub fn code_actions(mut self) -> Vec<CodeAction> {
+        self.visit_typed_module(&self.module.ast);
+
+        let Some(patterns) = self.patterns_to_merge else {
+            return vec![];
+        };
+
+        for (one, next) in patterns.iter().tuple_windows() {
+            self.edits
+                .replace(SrcSpan::new(one.end, next.start), " | ".into());
+        }
+
+        let mut action = Vec::with_capacity(1);
+        CodeActionBuilder::new("Merge case branches")
+            .kind(CodeActionKind::REFACTOR_REWRITE)
+            .changes(self.params.text_document.uri.clone(), self.edits.edits)
+            .preferred(false)
+            .push_to(&mut action);
+        action
+    }
+
+    fn select_mergeable_branches(&self, clauses: &'a [ast::TypedClause]) -> Option<Vec<SrcSpan>> {
+        let mut clauses = clauses
+            .iter()
+            // We want to skip all the branches at the beginning of the case
+            // expression that the cursor is not hovering over. For example:
+            //
+            // ```gleam
+            // case wibble {
+            //   a -> 1   <- we want to skip this one here that is not selected
+            //   b -> 2
+            //     ^^^^   this is the selection
+            //   _ -> 3
+            //   ^^
+            // }
+            // ```
+            .skip_while(|clause| {
+                let clause_range = self.edits.src_span_to_lsp_range(clause.location);
+                !overlaps(self.params.range, clause_range)
+            })
+            // Then we only want to take the clauses that we're hovering over
+            // with our selection (even partially!)
+            // In the provious example they would be `b -> 2` and `_ -> 3`.
+            .take_while(|clause| {
+                let clause_range = self.edits.src_span_to_lsp_range(clause.location);
+                overlaps(self.params.range, clause_range)
+            });
+
+        let first_hovered_clause = clauses.next()?;
+        let clauses = iter::once(first_hovered_clause)
+            .chain(clauses.take_while(|clause| clauses_can_be_merged(first_hovered_clause, clause)))
+            .map(|clause| clause.pattern_location())
+            .collect_vec();
+
+        if clauses.len() < 2 {
+            None
+        } else {
+            Some(clauses)
+        }
+    }
+}
+
+fn clauses_can_be_merged(one: &ast::TypedClause, other: &ast::TypedClause) -> bool {
+    // Two clauses cannot be merged if any of those has an if guard
+    if one.guard.is_some() || other.guard.is_some() {
+        return false;
+    }
+
+    // Two clauses can only be merged if they define the same variables,
+    // otherwise joining them would result in invalid code.
+    let variables_one = one
+        .bound_variables()
+        .map(|variable| variable.name())
+        .collect::<HashSet<_>>();
+
+    let variables_other = other
+        .bound_variables()
+        .map(|variable| variable.name())
+        .collect::<HashSet<_>>();
+
+    if variables_one != variables_other {
+        return false;
+    }
+
+    // For now we only support joining branches where all the branches have a
+    // `todo` as a body.
+    match (&one.then, &other.then) {
+        (TypedExpr::Todo { message: None, .. }, TypedExpr::Todo { message: None, .. }) => true,
+        _ => false,
+    }
+}
+
+impl<'ast> ast::visit::Visit<'ast> for MergeCaseBranches<'ast> {
+    fn visit_typed_expr_case(
+        &mut self,
+        location: &'ast SrcSpan,
+        type_: &'ast Arc<Type>,
+        subjects: &'ast [TypedExpr],
+        clauses: &'ast [ast::TypedClause],
+        compiled_case: &'ast CompiledCase,
+    ) {
+        // We only trigger the code action if we are within a case expression,
+        // otherwise there's no point in exploring the expression any further.
+        let case_range = self.edits.src_span_to_lsp_range(*location);
+        if !within(self.params.range, case_range) {
+            return;
+        }
+
+        if let result @ Some(_) = self.select_mergeable_branches(clauses) {
+            self.patterns_to_merge = result
+        }
+
+        // We still need to visit the case expression in case we want to apply
+        // the code action to some case expression that is nested in one of its
+        // branches!
+        ast::visit::visit_typed_expr_case(self, location, type_, subjects, clauses, compiled_case);
+    }
+}

compiler-core/src/language_server/engine.rs

@@ -13,7 +13,7 @@ use crate::{
     io::{BeamCompiler, CommandExecutor, FileSystemReader, FileSystemWriter},
     language_server::{
         code_action::{
-            AddOmittedLabels, CollapseNestedCase, ExtractFunction, RemoveBlock,
+            AddOmittedLabels, CollapseNestedCase, ExtractFunction, MergeCaseBranches, RemoveBlock,
             RemovePrivateOpaque, RemoveUnreachableCaseClauses,
         },
         compiler::LspProjectCompiler,
@@ -420,6 +420,7 @@ where
                 &this.error,
                 &mut actions,
             );
+            actions.extend(MergeCaseBranches::new(module, &lines, &params).code_actions());
             actions.extend(FixBinaryOperation::new(module, &lines, &params).code_actions());
             actions
                 .extend(FixTruncatedBitArraySegment::new(module, &lines, &params).code_actions());

compiler-core/src/language_server/tests/action.rs

@@ -136,6 +136,7 @@ const COLLAPSE_NESTED_CASE: &str = "Collapse nested case";
 const REMOVE_UNREACHABLE_CLAUSES: &str = "Remove unreachable clauses";
 const ADD_OMITTED_LABELS: &str = "Add omitted labels";
 const EXTRACT_FUNCTION: &str = "Extract function";
+const MERGE_CASE_BRANCHES: &str = "Merge case branches";
 
 macro_rules! assert_code_action {
     ($title:expr, $code:literal, $range:expr $(,)?) => {
@@ -10800,3 +10801,110 @@ fn wibble() -> Nil
         find_position_of("wibble").to_selection()
     );
 }
+
+#[test]
+fn merge_case_branch() {
+    assert_code_action!(
+        MERGE_CASE_BRANCHES,
+        r#"pub fn go(n: Int) {
+  case n {
+    1 -> todo
+    2 -> todo
+    _ -> todo
+  }
+  }"#,
+        find_position_of("1").select_until(find_position_of("2"))
+    );
+}
+
+#[test]
+fn merge_case_branch_will_not_merge_branches_with_guards() {
+    assert_no_code_actions!(
+        MERGE_CASE_BRANCHES,
+        r#"pub fn go(n: Int) {
+  case n {
+    1 if True -> todo
+    2 -> todo
+    _ -> todo
+  }
+  }"#,
+        find_position_of("1").select_until(find_position_of("2"))
+    );
+}
+
+#[test]
+fn merge_case_branch_will_not_merge_branches_defining_different_variables() {
+    assert_no_code_actions!(
+        MERGE_CASE_BRANCHES,
+        r#"pub fn go(result) {
+  case result {
+    Ok(value) -> todo
+    Error(error) -> todo
+    _ -> todo
+  }
+  }"#,
+        find_position_of("Ok").select_until(find_position_of("error"))
+    );
+}
+
+#[test]
+fn merge_case_branch_can_merge_branches_defining_the_same_variables() {
+    assert_code_action!(
+        MERGE_CASE_BRANCHES,
+        r#"pub fn go(result) {
+  case result {
+    [Ok(value), ..] -> todo
+    [_, Error(value)] -> todo
+    _ -> todo
+  }
+}"#,
+        find_position_of("Ok").select_until(find_position_of("todo").nth_occurrence(2))
+    );
+}
+
+#[test]
+fn merge_case_branch_can_merge_multiple_branches() {
+    assert_code_action!(
+        MERGE_CASE_BRANCHES,
+        r#"pub fn go(result) {
+  case result {
+    [_] -> 1
+    [Ok(value), ..] -> todo
+    [_, Error(value)] -> todo
+    [_, _, Error(value)] -> todo
+    [_, _] -> 1
+    _ -> 2
+  }
+}"#,
+        find_position_of("todo").select_until(find_position_of("todo").nth_occurrence(3))
+    );
+}
+
+#[test]
+fn merge_case_branch_does_not_pop_up_with_a_single_selected_branch() {
+    assert_no_code_actions!(
+        MERGE_CASE_BRANCHES,
+        r#"pub fn go(result) {
+  case result {
+    [] -> todo
+    _ -> 2
+  }
+}"#,
+        find_position_of("[]").to_selection()
+    );
+}
+
+#[test]
+fn merge_case_branch_works_with_existing_alternative_patterns() {
+    assert_code_action!(
+        MERGE_CASE_BRANCHES,
+        r#"pub fn go(result) {
+  case result {
+    [] | [_, _, ..]-> todo
+    [_] -> todo
+    _ -> 2
+  }
+}"#,
+        find_position_of("[]").select_until(find_position_of("[_]"))
+    );
+}

compiler-core/src/language_server/tests/snapshots/gleam_core__language_server__tests__action__merge_case_branch.snap

@@ -0,0 +1,23 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "pub fn go(n: Int) {\n  case n {\n    1 -> todo\n    2 -> todo\n    _ -> todo\n  }\n  }"
+---
+----- BEFORE ACTION
+pub fn go(n: Int) {
+  case n {
+    1 -> todo
+    ▔▔▔▔▔▔▔▔▔
+    2 -> todo
+▔▔▔▔↑        
+    _ -> todo
+  }
+  }
+
+
+----- AFTER ACTION
+pub fn go(n: Int) {
+  case n {
+    1 | 2 -> todo
+    _ -> todo
+  }
+  }

compiler-core/src/language_server/tests/snapshots/gleam_core__language_server__tests__action__merge_case_branch_can_merge_branches_defining_the_same_variables.snap

@@ -0,0 +1,23 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "pub fn go(result) {\n  case result {\n    [Ok(value), ..] -> todo\n    [_, Error(value)] -> todo\n    _ -> todo\n  }\n}"
+---
+----- BEFORE ACTION
+pub fn go(result) {
+  case result {
+    [Ok(value), ..] -> todo
+     ▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔
+    [_, Error(value)] -> todo
+▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔↑   
+    _ -> todo
+  }
+}
+
+
+----- AFTER ACTION
+pub fn go(result) {
+  case result {
+    [Ok(value), ..] | [_, Error(value)] -> todo
+    _ -> todo
+  }
+}