get fn locations from parser

This increases the blast radius of the change a bit, but means we
don't need to worry about comments any more.

Author: treuherz

compiler-core/src/ast/typed.rs

@@ -76,6 +76,7 @@ pub enum TypedExpr {
         type_: Arc<Type>,
         fun: Box<Self>,
         arguments: Vec<CallArg<Self>>,
+        arguments_start: Option<u32>,
     },
 
     BinOp {

compiler-core/src/ast/untyped.rs

@@ -53,6 +53,7 @@ pub enum UntypedExpr {
         location: SrcSpan,
         fun: Box<Self>,
         arguments: Vec<CallArg<Self>>,
+        arguments_start: u32,
     },
 
     BinOp {
@@ -293,7 +294,7 @@ impl HasLocation for UntypedExpr {
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum FunctionLiteralKind {
     Capture { hole: SrcSpan },
-    Anonymous { head: SrcSpan },
+    Anonymous { head: SrcSpan, body: SrcSpan },
     Use { location: SrcSpan },
 }
 

compiler-core/src/ast/visit.rs

@@ -196,8 +196,9 @@ pub trait Visit<'ast> {
         type_: &'ast Arc<Type>,
         fun: &'ast TypedExpr,
         arguments: &'ast [TypedCallArg],
+        arguments_start: &'ast Option<u32>,
     ) {
-        visit_typed_expr_call(self, location, type_, fun, arguments);
+        visit_typed_expr_call(self, location, type_, fun, arguments, arguments_start);
     }
 
     fn visit_typed_expr_bin_op(
@@ -814,7 +815,8 @@ where
             type_,
             fun,
             arguments,
-        } => v.visit_typed_expr_call(location, type_, fun, arguments),
+            arguments_start,
+        } => v.visit_typed_expr_call(location, type_, fun, arguments, arguments_start),
         TypedExpr::BinOp {
             location,
             type_,
@@ -1049,6 +1051,7 @@ pub fn visit_typed_expr_call<'a, V>(
     _type_: &'a Arc<Type>,
     fun: &'a TypedExpr,
     arguments: &'a [TypedCallArg],
+    _arguments_start: &'a Option<u32>,
 ) where
     V: Visit<'a> + ?Sized,
 {

compiler-core/src/ast_folder.rs

@@ -289,7 +289,8 @@ pub trait UntypedExprFolder: TypeAstFolder + UntypedConstantFolder + PatternFold
                 location,
                 fun,
                 arguments,
-            } => self.fold_call(location, fun, arguments),
+                arguments_start,
+            } => self.fold_call(location, fun, arguments, arguments_start),
 
             UntypedExpr::BinOp {
                 location,
@@ -445,6 +446,7 @@ pub trait UntypedExprFolder: TypeAstFolder + UntypedConstantFolder + PatternFold
                 location,
                 fun,
                 arguments,
+                arguments_start,
             } => {
                 let fun = Box::new(self.fold_expr(*fun));
                 let arguments = arguments
@@ -458,6 +460,7 @@ pub trait UntypedExprFolder: TypeAstFolder + UntypedConstantFolder + PatternFold
                     location,
                     fun,
                     arguments,
+                    arguments_start,
                 }
             }
 
@@ -770,11 +773,13 @@ pub trait UntypedExprFolder: TypeAstFolder + UntypedConstantFolder + PatternFold
         location: SrcSpan,
         fun: Box<UntypedExpr>,
         arguments: Vec<CallArg<UntypedExpr>>,
+        arguments_start: u32,
     ) -> UntypedExpr {
         UntypedExpr::Call {
             location,
             fun,
             arguments,
+            arguments_start,
         }
     }
 

compiler-core/src/format.rs

@@ -1592,6 +1592,7 @@ impl<'comments> Formatter<'comments> {
             fun,
             arguments,
             location,
+            ..
         })) = call.first()
         else {
             // The body of a capture being not a fn shouldn't be possible...

compiler-core/src/inline.rs

@@ -363,7 +363,8 @@ impl Inliner<'_> {
                 type_,
                 fun,
                 arguments,
-            } => self.call(location, type_, fun, arguments),
+                arguments_start,
+            } => self.call(location, type_, fun, arguments, arguments_start),
 
             TypedExpr::BinOp {
                 location,
@@ -581,6 +582,7 @@ impl Inliner<'_> {
         type_: Arc<Type>,
         function: Box<TypedExpr>,
         arguments: Vec<TypedCallArg>,
+        arguments_start: Option<u32>,
     ) -> TypedExpr {
         let arguments = self.arguments(arguments);
 
@@ -699,6 +701,7 @@ impl Inliner<'_> {
             type_,
             fun: Box::new(function),
             arguments,
+            arguments_start,
         }
     }
 
@@ -1534,6 +1537,7 @@ impl InlinableExpression {
                     .iter()
                     .map(|argument| argument.to_call_arg(Self::to_expression))
                     .collect(),
+                arguments_start: None,
             },
         }
     }

compiler-core/src/language_server/code_action.rs

@@ -102,6 +102,28 @@ fn count_indentation(code: &str, line_numbers: &LineNumbers, line: u32) -> usize
     indent_size
 }
 
+// Given a string and a position in it, if the position points to whitespace,
+// this function returns the next position which doesn't.
+fn next_nonwhitespace(string: &EcoString, position: u32) -> u32 {
+    let mut n = position;
+    let mut chars = string[position as usize..].chars();
+    while chars.next().is_some_and(char::is_whitespace) {
+        n += 1;
+    }
+    n
+}
+
+// Given a string and a position in it, if the position points after whitespace,
+// this function returns the previous position which doesn't.
+fn previous_nonwhitespace(string: &EcoString, position: u32) -> u32 {
+    let mut n = position;
+    let mut chars = string[..position as usize].chars();
+    while chars.next_back().is_some_and(char::is_whitespace) {
+        n -= 1;
+    }
+    n
+}
+
 /// Code action to remove literal tuples in case subjects, essentially making
 /// the elements of the tuples into the case's subjects.
 ///
@@ -842,6 +864,7 @@ impl<'ast> ast::visit::Visit<'ast> for FillInMissingLabelledArgs<'ast> {
         type_: &'ast Arc<Type>,
         fun: &'ast TypedExpr,
         arguments: &'ast [TypedCallArg],
+        arguments_start: &'ast Option<u32>,
     ) {
         let call_range = self.edits.src_span_to_lsp_range(*location);
         if !within(self.params.range, call_range) {
@@ -864,7 +887,7 @@ impl<'ast> ast::visit::Visit<'ast> for FillInMissingLabelledArgs<'ast> {
         // we're inside a nested call.
         let previous = self.use_right_hand_side_location;
         self.use_right_hand_side_location = None;
-        ast::visit::visit_typed_expr_call(self, location, type_, fun, arguments);
+        ast::visit::visit_typed_expr_call(self, location, type_, fun, arguments, arguments_start);
         self.use_right_hand_side_location = previous;
     }
 
@@ -1263,7 +1286,7 @@ impl<'ast> ast::visit::Visit<'ast> for AddAnnotations<'_> {
         let location = match kind {
             // Function captures don't need any type annotations
             FunctionLiteralKind::Capture { .. } => return,
-            FunctionLiteralKind::Anonymous { head } => head,
+            FunctionLiteralKind::Anonymous { head, .. } => head,
             FunctionLiteralKind::Use { location } => location,
         };
 
@@ -5269,6 +5292,7 @@ impl<'ast> ast::visit::Visit<'ast> for GenerateFunction<'ast> {
         type_: &'ast Arc<Type>,
         fun: &'ast TypedExpr,
         arguments: &'ast [TypedCallArg],
+        arguments_start: &'ast Option<u32>,
     ) {
         // If the function being called is invalid we need to generate a
         // function that has the proper labels.
@@ -5302,8 +5326,7 @@ impl<'ast> ast::visit::Visit<'ast> for GenerateFunction<'ast> {
                 _ => {}
             }
         }
-
-        ast::visit::visit_typed_expr_call(self, location, type_, fun, arguments);
+        ast::visit::visit_typed_expr_call(self, location, type_, fun, arguments, arguments_start);
     }
 }
 
@@ -5626,6 +5649,7 @@ where
         type_: &'ast Arc<Type>,
         fun: &'ast TypedExpr,
         arguments: &'ast [TypedCallArg],
+        arguments_start: &'ast Option<u32>,
     ) {
         // If the function being called is invalid we need to generate a
         // function that has the proper labels.
@@ -5639,7 +5663,14 @@ where
                 );
             }
         } else {
-            ast::visit::visit_typed_expr_call(self, location, type_, fun, arguments);
+            ast::visit::visit_typed_expr_call(
+                self,
+                location,
+                type_,
+                fun,
+                arguments,
+                arguments_start,
+            );
         }
     }
 
@@ -6102,12 +6133,7 @@ impl<'a> InlineVariable<'a> {
         }
 
         let mut location = assignment.location;
-
-        let mut chars = self.module.code[location.end as usize..].chars();
-        // Delete any whitespace after the removed statement
-        while chars.next().is_some_and(char::is_whitespace) {
-            location.end += 1;
-        }
+        location.end = next_nonwhitespace(&self.module.code, location.end);
 
         self.edits.delete(location);
 
@@ -6336,6 +6362,7 @@ impl<'ast> ast::visit::Visit<'ast> for ConvertToPipe<'ast> {
         _type_: &'ast Arc<Type>,
         fun: &'ast TypedExpr,
         arguments: &'ast [TypedCallArg],
+        _arguments_start: &'ast Option<u32>,
     ) {
         if arguments.iter().any(|arg| arg.is_capture_hole()) {
             return;
@@ -8313,6 +8340,7 @@ impl<'ast> ast::visit::Visit<'ast> for WrapInAnonymousFunction<'ast> {
         _type: &'ast Arc<Type>,
         fun: &'ast TypedExpr,
         arguments: &'ast [TypedCallArg],
+        _arguments_start: &'ast Option<u32>,
     ) {
         // We only need to do this interception for explicit calls, so if any
         // of our arguments are explicit we re-enter the visitor as usual.
@@ -8355,11 +8383,15 @@ pub struct UnwrapAnonymousFunction<'a> {
 
 /// Helper struct, a target for [UnwrapAnonymousFunction]
 struct FunctionToUnwrap {
-    /// Location of the anonymous function to apply the action to
+    /// Location of the anonymous function to apply the action to.
     outer_function: SrcSpan,
+    /// Location of the opening brace of the anonymous function.
+    outer_function_body_start: u32,
     /// 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,
+    // Location of the opening parenthesis of the inner function's argument list.
+    inner_function_arguments_start: u32,
 }
 
 impl<'a> UnwrapAnonymousFunction<'a> {
@@ -8383,25 +8415,29 @@ impl<'a> UnwrapAnonymousFunction<'a> {
         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 {
+            // We need to delete the anonymous function's head and the opening
+            // brace but preserve comments between it and the inner function call.
+            // We set our endpoint at the start of the function body, and move
+            // it on through any whitespace.
+            let head_deletion_end =
+                next_nonwhitespace(&self.module.code, function.outer_function_body_start + 1);
+            edits.delete(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.
+                end: head_deletion_end,
+            });
+
+            // Delete the inner function call's arguments.
+            edits.delete(SrcSpan {
+                start: function.inner_function_arguments_start,
+                end: function.inner_function.end,
+            });
+
+            // To delete the tail we remove the function end (the '}') and any
+            // whitespace before it.
+            let tail_deletion_start =
+                previous_nonwhitespace(&self.module.code, function.outer_function.end - 1);
             edits.delete(SrcSpan {
-                start: function.outer_function.end - 1,
+                start: tail_deletion_start,
                 end: function.outer_function.end,
             });
 
@@ -8413,24 +8449,6 @@ 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:
@@ -8446,16 +8464,17 @@ impl<'a> UnwrapAnonymousFunction<'a> {
         arguments: &'a [TypedArg],
         body: &'a Vec1<TypedStatement>,
     ) {
-        match kind {
-            FunctionLiteralKind::Anonymous { .. } => (),
+        let outer_body = match kind {
+            FunctionLiteralKind::Anonymous { body, .. } => body,
             _ => return,
-        }
+        };
 
         // We can only apply to anonymous functions containing a single function call
         let [
             TypedStatement::Expression(TypedExpr::Call {
-                fun: called_function,
+                location: call_location,
                 arguments: call_arguments,
+                arguments_start: Some(arguments_start),
                 ..
             }),
         ] = body.as_slice()
@@ -8490,7 +8509,9 @@ impl<'a> UnwrapAnonymousFunction<'a> {
 
         self.functions.push(FunctionToUnwrap {
             outer_function: *location,
-            inner_function: called_function.location(),
+            outer_function_body_start: outer_body.start,
+            inner_function: *call_location,
+            inner_function_arguments_start: *arguments_start,
         })
     }
 }

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

@@ -9976,7 +9976,6 @@ fn wobble(i) {
     );
 }
 
-
 #[test]
 fn wrap_final_pipeline_step_in_anonymous_function() {
     assert_code_action!(