unwrap anonymous functions with comments in

The indentation of comments when we're done with them might be a bit
off, but it's nothing an autoformat shouldn't be able to fix

Author: treuherz

compiler-core/src/language_server/code_action.rs

@@ -7626,7 +7626,8 @@ impl<'ast> ast::visit::Visit<'ast> for WrapInAnonymousFunction<'ast> {
     }
 }
 
-/// Code action to unwrap trivial one-statement anonymous functions into just a reference to the function called
+/// Code action to unwrap trivial one-statement anonymous functions into just a
+/// reference to the function called
 ///
 /// For example, if the code action was used on the anonymous function here:
 ///
@@ -7651,9 +7652,10 @@ pub struct UnwrapAnonymousFunction<'a> {
 /// Helper struct, a target for [UnwrapAnonymousFunction]
 struct FunctionToUnwrap {
     /// Location of the anonymous function to apply the action to
-    anonymous_function_location: SrcSpan,
-    /// Location of the function being called inside the anonymous function. This will be all that's left after the action.
-    inner_function_location: SrcSpan,
+    outer_function: SrcSpan,
+    /// Location of the function being called inside the anonymous function.
+    /// This will be all that's left after the action, plus any comments.
+    inner_function: SrcSpan,
 }
 
 impl<'a> UnwrapAnonymousFunction<'a> {
@@ -7674,16 +7676,29 @@ impl<'a> UnwrapAnonymousFunction<'a> {
         self.visit_typed_module(&self.module.ast);
 
         let mut actions = Vec::with_capacity(self.functions.len());
-        for target in self.functions {
+        for function in &self.functions {
             let mut edits = TextEdits::new(self.line_numbers);
 
+            // We need to delete the anonymous function's head but preserve
+            // comments between it and the inner function call.
+            edits.delete(self.span_until_comment(SrcSpan {
+                start: function.outer_function.start,
+                end: function.inner_function.start,
+            }));
+
+            // Now we need to delete the inner function call's arguments,
+            // preserving comments before the outer function tail.
+            edits.delete(self.span_until_comment(SrcSpan {
+                start: function.inner_function.end,
+                end: function.outer_function.end - 1,
+            }));
+
+            // To delete the tail we nip one character to make sure we get the
+            // `}`. This could be redundant with the above if there were no
+            // comments, but that's fine.
             edits.delete(SrcSpan {
-                start: target.anonymous_function_location.start,
-                end: target.inner_function_location.start,
-            });
-            edits.delete(SrcSpan {
-                start: target.inner_function_location.end,
-                end: target.anonymous_function_location.end,
+                start: function.outer_function.end - 1,
+                end: function.outer_function.end,
             });
 
             CodeActionBuilder::new("Remove anonymous function wrapper")
@@ -7694,6 +7709,24 @@ impl<'a> UnwrapAnonymousFunction<'a> {
         actions
     }
 
+    // Returns the given span, but with the end point adjusted to the start
+    // of the first comment in the span, if any.
+    fn span_until_comment(&self, span: SrcSpan) -> SrcSpan {
+        let SrcSpan { start, end } = span;
+        let adjusted_end = self
+            .module
+            .extra
+            .first_comment_between(start, end)
+            // The above will return the start of the comment's text, so we need
+            // to step backwards a bit to get the `//`.
+            .map(|comment| comment.start - 2)
+            .unwrap_or(end);
+        SrcSpan {
+            start,
+            end: adjusted_end,
+        }
+    }
+
     /// If an anonymous function can be unwrapped, save it to our list
     ///
     /// We need to ensure our subjects:
@@ -7714,15 +7747,6 @@ impl<'a> UnwrapAnonymousFunction<'a> {
             _ => return,
         }
 
-        // We can't apply to functions with comments in (yet)
-        if self
-            .module
-            .extra
-            .has_comment_between(location.start, location.end)
-        {
-            return;
-        }
-
         // We can only apply to anonymous functions containing a single function call
         let [
             TypedStatement::Expression(TypedExpr::Call {
@@ -7761,8 +7785,8 @@ impl<'a> UnwrapAnonymousFunction<'a> {
         }
 
         self.functions.push(FunctionToUnwrap {
-            anonymous_function_location: *location,
-            inner_function_location: called_function.location(),
+            outer_function: *location,
+            inner_function: called_function.location(),
         })
     }
 }

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

@@ -9412,8 +9412,8 @@ fn op(first a, second b) {
 }
 
 #[test]
-fn dont_unwrap_anonymous_function_with_comment_after() {
-    assert_no_code_actions!(
+fn unwrap_anonymous_function_with_comment_after() {
+    assert_code_action!(
         UNWRAP_ANONYMOUS_FUNCTION,
         "pub fn main() {
   fn(a) {
@@ -9431,12 +9431,54 @@ fn op(a) {
 }
 
 #[test]
-fn dont_unwrap_anonymous_function_with_comment_before() {
-    assert_no_code_actions!(
+fn unwrap_anonymous_function_with_comment_on_line() {
+    assert_code_action!(
         UNWRAP_ANONYMOUS_FUNCTION,
         "pub fn main() {
   fn(a) {
-    // look out!
+    op(a) // look out!
+  }
+}
+
+fn op(a) {
+  todo
+}
+",
+        find_position_of("fn(a)").to_selection()
+    );
+}
+
+#[test]
+fn unwrap_anonymous_function_with_comment_on_head_line() {
+    assert_code_action!(
+        UNWRAP_ANONYMOUS_FUNCTION,
+        "pub fn main() {
+  fn(a) { // look out!
+    op(a)
+  }
+}
+
+fn op(a) {
+  todo
+}
+",
+        find_position_of("fn(a)").to_selection()
+    );
+}
+
+#[test]
+fn unwrap_anonymous_function_with_comments_before() {
+    assert_code_action!(
+        UNWRAP_ANONYMOUS_FUNCTION,
+        "pub fn main() {
+  fn(a) {
+    // look out,
+    // there's a comment!
+
+    // another comment!
+    //here's one without a leading space
+      // here's one indented wrong
+  // here's one indented even wronger
     op(a)
   }
 }

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

@@ -0,0 +1,27 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "pub fn main() {\n  fn(a) {\n    op(a)\n    // look out!\n  }\n}\n\nfn op(a) {\n  todo\n}\n"
+---
+----- BEFORE ACTION
+pub fn main() {
+  fn(a) {
+  ↑      
+    op(a)
+    // look out!
+  }
+}
+
+fn op(a) {
+  todo
+}
+
+
+----- AFTER ACTION
+pub fn main() {
+  op// look out!
+  
+}
+
+fn op(a) {
+  todo
+}

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

@@ -0,0 +1,26 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "pub fn main() {\n  fn(a) { // look out!\n    op(a)\n  }\n}\n\nfn op(a) {\n  todo\n}\n"
+---
+----- BEFORE ACTION
+pub fn main() {
+  fn(a) { // look out!
+  ↑                   
+    op(a)
+  }
+}
+
+fn op(a) {
+  todo
+}
+
+
+----- AFTER ACTION
+pub fn main() {
+  // look out!
+    op
+}
+
+fn op(a) {
+  todo
+}

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

@@ -0,0 +1,26 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "pub fn main() {\n  fn(a) {\n    op(a) // look out!\n  }\n}\n\nfn op(a) {\n  todo\n}\n"
+---
+----- BEFORE ACTION
+pub fn main() {
+  fn(a) {
+  ↑      
+    op(a) // look out!
+  }
+}
+
+fn op(a) {
+  todo
+}
+
+
+----- AFTER ACTION
+pub fn main() {
+  op// look out!
+  
+}
+
+fn op(a) {
+  todo
+}

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

@@ -0,0 +1,39 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "pub fn main() {\n  fn(a) {\n    // look out,\n    // there's a comment!\n\n    // another comment!\n    //here's one without a leading space\n      // here's one indented wrong\n  // here's one indented even wronger\n    op(a)\n  }\n}\n\nfn op(a) {\n  todo\n}\n"
+---
+----- BEFORE ACTION
+pub fn main() {
+  fn(a) {
+  ↑      
+    // look out,
+    // there's a comment!
+
+    // another comment!
+    //here's one without a leading space
+      // here's one indented wrong
+  // here's one indented even wronger
+    op(a)
+  }
+}
+
+fn op(a) {
+  todo
+}
+
+
+----- AFTER ACTION
+pub fn main() {
+  // look out,
+    // there's a comment!
+
+    // another comment!
+    //here's one without a leading space
+      // here's one indented wrong
+  // here's one indented even wronger
+    op
+}
+
+fn op(a) {
+  todo
+}

compiler-core/src/parse/extra.rs

@@ -65,7 +65,7 @@ impl ModuleExtra {
         let mut search_list = &self.comments[..];
         while let Some(index) = inner(search_list, start, end) {
             best = self.comments.get(index).copied();
-            search_list = &search_list[0..index];
+            search_list = search_list.get(0..index).unwrap_or(&[]);
         }
         best
     }