properly deal with label shorthand

Author: giacomocavalieri

compiler-core/src/ast.rs

@@ -1530,7 +1530,18 @@ where
 {
     #[must_use]
     pub fn uses_label_shorthand(&self) -> bool {
-        self.label.is_some() && self.location == self.value.location()
+        self.label_shorthand_name().is_some()
+    }
+
+    /// If the call arg is defined using a label shorthand, this will return the
+    /// label name.
+    ///
+    fn label_shorthand_name(&self) -> Option<EcoString> {
+        if !self.is_implicit() && self.location == self.value.location() {
+            self.label.clone()
+        } else {
+            None
+        }
     }
 }
 
@@ -2533,6 +2544,32 @@ impl<A> Pattern<A> {
     }
 }
 
+/// A variable bound inside a pattern.
+#[derive(Debug, Clone)]
+pub enum BoundVariable {
+    /// A record's labelled field introduced with the shorthand syntax.
+    ShorthandLabel { name: EcoString, location: SrcSpan },
+    /// Any other variable.
+    Regular { name: EcoString, location: SrcSpan },
+}
+
+impl BoundVariable {
+    pub fn name(&self) -> EcoString {
+        match self {
+            BoundVariable::ShorthandLabel { name, .. } | BoundVariable::Regular { name, .. } => {
+                name.clone()
+            }
+        }
+    }
+
+    pub fn location(&self) -> SrcSpan {
+        match self {
+            BoundVariable::ShorthandLabel { location, .. }
+            | BoundVariable::Regular { location, .. } => *location,
+        }
+    }
+}
+
 impl TypedPattern {
     pub fn definition_location(&self) -> Option<DefinitionLocation> {
         match self {
@@ -2707,28 +2744,34 @@ impl TypedPattern {
         }
     }
 
-    pub fn bound_variables(&self) -> Vec<(EcoString, SrcSpan)> {
+    pub fn bound_variables(&self) -> Vec<BoundVariable> {
         let mut variables = Vec::new();
         self.collect_bound_variables(&mut variables);
         variables
     }
 
-    fn collect_bound_variables(&self, variables: &mut Vec<(EcoString, SrcSpan)>) {
+    fn collect_bound_variables(&self, variables: &mut Vec<BoundVariable>) {
         match self {
             Pattern::Int { .. }
             | Pattern::Float { .. }
             | Pattern::String { .. }
             | Pattern::Discard { .. }
             | Pattern::Invalid { .. } => {}
 
-            Pattern::Variable { name, location, .. } => variables.push((name.clone(), *location)),
+            Pattern::Variable { name, location, .. } => variables.push(BoundVariable::Regular {
+                name: name.clone(),
+                location: *location,
+            }),
             Pattern::BitArraySize { .. } => {}
             Pattern::Assign {
                 name,
                 pattern,
                 location,
             } => {
-                variables.push((name.clone(), *location));
+                variables.push(BoundVariable::Regular {
+                    name: name.clone(),
+                    location: *location,
+                });
                 pattern.collect_bound_variables(variables);
             }
             Pattern::List { elements, tail, .. } => {
@@ -2741,7 +2784,14 @@ impl TypedPattern {
             }
             Pattern::Constructor { arguments, .. } => {
                 for argument in arguments {
-                    argument.value.collect_bound_variables(variables);
+                    if let Some(name) = argument.label_shorthand_name() {
+                        variables.push(BoundVariable::ShorthandLabel {
+                            name,
+                            location: argument.location,
+                        })
+                    } else {
+                        argument.value.collect_bound_variables(variables);
+                    }
                 }
             }
             Pattern::Tuple { elements, .. } => {
@@ -2760,11 +2810,17 @@ impl TypedPattern {
                 right_location,
                 ..
             } => {
-                if let Some(assignment) = left_side_assignment {
-                    variables.push(assignment.clone());
+                if let Some((name, location)) = left_side_assignment {
+                    variables.push(BoundVariable::Regular {
+                        name: name.clone(),
+                        location: *location,
+                    });
                 }
                 match right_side_assignment {
-                    AssignName::Variable(name) => variables.push((name.clone(), *right_location)),
+                    AssignName::Variable(name) => variables.push(BoundVariable::Regular {
+                        name: name.clone(),
+                        location: *right_location,
+                    }),
                     AssignName::Discard(_) => {}
                 }
             }

compiler-core/src/javascript/decision.rs

@@ -702,10 +702,10 @@ pub fn let_<'a>(
     // We must generate this after we generate the code for the decision tree
     // itself as we might be re-binding variables which are used in the checks
     // to determine whether the pattern matches or not.
-    let beginning_assignments = pattern.bound_variables().into_iter().map(|(variable, _)| {
+    let beginning_assignments = pattern.bound_variables().into_iter().map(|bound_variable| {
         docvec![
             "let ",
-            expression_generator.local_var(&variable),
+            expression_generator.local_var(&bound_variable.name()),
             ";",
             line()
         ]

compiler-core/src/language_server/code_action.rs

@@ -3,11 +3,11 @@ use std::{collections::HashSet, iter, sync::Arc};
 use crate::{
     Error, STDLIB_PACKAGE_NAME, analyse,
     ast::{
-        self, ArgNames, AssignName, AssignmentKind, BitArraySegmentTruncation, CallArg, CustomType,
-        FunctionLiteralKind, ImplicitCallArgOrigin, Import, PIPE_PRECEDENCE, Pattern,
-        PatternUnusedArguments, PipelineAssignmentKind, RecordConstructor, SrcSpan, TodoKind,
-        TypedArg, TypedAssignment, TypedClauseGuard, TypedExpr, TypedModuleConstant, TypedPattern,
-        TypedPipelineAssignment, TypedRecordConstructor, TypedStatement, TypedUse,
+        self, ArgNames, AssignName, AssignmentKind, BitArraySegmentTruncation, BoundVariable,
+        CallArg, CustomType, FunctionLiteralKind, ImplicitCallArgOrigin, Import, PIPE_PRECEDENCE,
+        Pattern, PatternUnusedArguments, PipelineAssignmentKind, RecordConstructor, SrcSpan,
+        TodoKind, TypedArg, TypedAssignment, TypedClauseGuard, TypedExpr, TypedModuleConstant,
+        TypedPattern, TypedPipelineAssignment, TypedRecordConstructor, TypedStatement, TypedUse,
         visit::Visit as _,
     },
     build::{Located, Module},
@@ -7687,27 +7687,22 @@ struct Collapsed<'a> {
     ///
     outer_guard: &'a Option<TypedClauseGuard>,
 
-    /// In this example it's `username`.
-    ///
-    matched_variable_name: EcoString,
-
-    /// The span covering the definition of the pattern variable being matched
-    /// on:
+    /// The pattern variable being matched on:
     ///
     /// ```gleam
     /// case something {
     ///   User(username: _, NotAdmin) -> "Stranger!!"
     ///   User(username:, Admin) if wibble ->
     ///        ┬───────
-    ///        ╰─ `matched_variable_span`
+    ///        ╰─ `matched_variable`
     ///     case username {
     ///       "Joe" -> "Hello, Joe!"
     ///       _ -> "I don't know you, " <> username
     ///     }
     /// }
     /// ```
     ///
-    matched_variable_span: SrcSpan,
+    matched_variable: BoundVariable,
 
     /// The span covering the entire pattern that is bringing the matched
     /// variable in scope:
@@ -7756,8 +7751,7 @@ impl<'a> CollapseNestedCase<'a> {
         let Some(Collapsed {
             outer_clause_span,
             outer_guard,
-            ref matched_variable_name,
-            matched_variable_span,
+            ref matched_variable,
             matched_pattern_span,
             inner_clauses,
         }) = self.collapsed
@@ -7794,15 +7788,37 @@ impl<'a> CollapseNestedCase<'a> {
         // can't simply write `Ok(2 | 3)` but we have to write `Ok(2) | Ok(3)`!
 
         let pattern_text: String = self.code_at(matched_pattern_span).into();
-
-        let variable_start_in_pattern = matched_variable_span.start - matched_pattern_span.start;
-        let variable_length = matched_variable_span.end - matched_variable_span.start;
-        let variable_end_in_pattern = variable_start_in_pattern + variable_length;
-        let variable_range = variable_start_in_pattern as usize..variable_end_in_pattern as usize;
+        let matched_variable_span = matched_variable.location();
 
         let pattern_with_variable = |new_content: String| {
             let mut new_pattern = pattern_text.clone();
-            new_pattern.replace_range(variable_range.clone(), &new_content);
+
+            match matched_variable {
+                BoundVariable::Regular { .. } => {
+                    // If the variable is a regular variable we'll have to replace
+                    // it entirely with the new pattern taking its place.
+                    let variable_start_in_pattern =
+                        matched_variable_span.start - matched_pattern_span.start;
+                    let variable_length = matched_variable_span.end - matched_variable_span.start;
+                    let variable_end_in_pattern = variable_start_in_pattern + variable_length;
+                    let replaced_range =
+                        variable_start_in_pattern as usize..variable_end_in_pattern as usize;
+
+                    new_pattern.replace_range(replaced_range, &new_content);
+                }
+
+                BoundVariable::ShorthandLabel { .. } => {
+                    // But if it's introduced using the shorthand syntax we can't
+                    // just replace it's location with the new pattern: we would be
+                    // removing the label!!
+                    // So we instead insert the pattern right after the label.
+                    new_pattern.insert_str(
+                        (matched_variable_span.end - matched_pattern_span.start) as usize,
+                        &format!(" {new_content}"),
+                    );
+                }
+            }
+
             new_pattern
         };
 
@@ -7813,7 +7829,7 @@ impl<'a> CollapseNestedCase<'a> {
             // everything together with ` | `.
 
             let references_to_matched_variable =
-                FindVariableReferences::new(matched_variable_span, matched_variable_name.clone())
+                FindVariableReferences::new(matched_variable_span, matched_variable.name())
                     .find(&clause.then);
 
             let new_patterns = iter::once(&clause.pattern)
@@ -7827,7 +7843,7 @@ impl<'a> CollapseNestedCase<'a> {
 
                     let mut pattern_code = self.code_at(pattern_location).to_string();
                     if !references_to_matched_variable.is_empty() {
-                        pattern_code = format!("{pattern_code} as {matched_variable_name}");
+                        pattern_code = format!("{pattern_code} as {}", matched_variable.name());
                     };
                     pattern_with_variable(pattern_code)
                 })
@@ -7927,10 +7943,10 @@ impl<'a> CollapseNestedCase<'a> {
 
         // That variable must be one the variables we brought into scope in this
         // branch.
-        let (variable_name, variable_location) = pattern
+        let variable = pattern
             .iter()
             .flat_map(|pattern| pattern.bound_variables())
-            .find(|(variable, _)| variable == name)?;
+            .find(|variable| variable.name() == *name)?;
 
         // There's one last condition to trigger the code action: we must
         // actually be with the cursor over the pattern or the nested case
@@ -7956,8 +7972,7 @@ impl<'a> CollapseNestedCase<'a> {
             Some(Collapsed {
                 outer_clause_span: *location,
                 outer_guard: guard,
-                matched_variable_name: variable_name,
-                matched_variable_span: variable_location,
+                matched_variable: variable,
                 matched_pattern_span: pattern_location,
                 inner_clauses: clauses,
             })

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

@@ -9290,6 +9290,58 @@ pub type Wibble {
     );
 }
 
+#[test]
+fn collapse_nested_case_does_not_remove_labels() {
+    assert_code_action!(
+        COLLAPSE_NESTED_CASE,
+        "pub fn main(x) {
+  case x {
+    Wibble(field2:, field: wibble) ->
+      case wibble {
+        1 -> 2
+        2 -> 4
+        _ -> -1
+      }
+
+    Wobble -> todo
+  }
+}
+
+pub type Wibble {
+  Wibble(field: Int, field2: String)
+  Wobble
+}
+",
+        find_position_of("field").to_selection()
+    );
+}
+
+#[test]
+fn collapse_nested_case_does_not_remove_labels_with_shorthand_syntax() {
+    assert_code_action!(
+        COLLAPSE_NESTED_CASE,
+        "pub fn main(x) {
+  case x {
+    Wibble(field2:, field:) ->
+      case field {
+        1 -> 2
+        2 -> 4
+        _ -> -1
+      }
+
+    Wobble -> todo
+  }
+}
+
+pub type Wibble {
+  Wibble(field: Int, field2: String)
+  Wobble
+}
+",
+        find_position_of("field").to_selection()
+    );
+}
+
 #[test]
 fn collapse_nested_case_works_with_alternative_patterns() {
     assert_code_action!(

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

@@ -0,0 +1,40 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "pub fn main(x) {\n  case x {\n    Wibble(field2:, field: wibble) ->\n      case wibble {\n        1 -> 2\n        2 -> 4\n        _ -> -1\n      }\n\n    Wobble -> todo\n  }\n}\n\npub type Wibble {\n  Wibble(field: Int, field2: String)\n  Wobble\n}\n"
+---
+----- BEFORE ACTION
+pub fn main(x) {
+  case x {
+    Wibble(field2:, field: wibble) ->
+           ↑                         
+      case wibble {
+        1 -> 2
+        2 -> 4
+        _ -> -1
+      }
+
+    Wobble -> todo
+  }
+}
+
+pub type Wibble {
+  Wibble(field: Int, field2: String)
+  Wobble
+}
+
+
+----- AFTER ACTION
+pub fn main(x) {
+  case x {
+    Wibble(field2:, field: 1) -> 2
+    Wibble(field2:, field: 2) -> 4
+    Wibble(field2:, field: _) -> -1
+
+    Wobble -> todo
+  }
+}
+
+pub type Wibble {
+  Wibble(field: Int, field2: String)
+  Wobble
+}