add action to unwrap a trivial anonymous function

Author: treuherz

compiler-core/src/language_server/code_action.rs

@@ -3518,7 +3518,7 @@ impl VariablesNames {
         }
         variables
     }
-    
+
     fn from_expr(expr: &TypedExpr) -> Self {
         let mut variables = Self {
             names: HashSet::new(),
@@ -8288,3 +8288,107 @@ impl<'ast> ast::visit::Visit<'ast> for WrapInAnonymousFunction<'ast> {
         }
     }
 }
+
+pub struct UnwrapAnonymousFunction<'a> {
+    module: &'a Module,
+    line_numbers: &'a LineNumbers,
+    params: &'a CodeActionParams,
+    targets: Vec<FnToUnwrap>,
+}
+
+struct FnToUnwrap {
+    fn_location: SrcSpan,
+    inner_function_location: SrcSpan,
+}
+
+impl<'a> UnwrapAnonymousFunction<'a> {
+    pub fn new(
+        module: &'a Module,
+        line_numbers: &'a LineNumbers,
+        params: &'a CodeActionParams,
+    ) -> Self {
+        Self {
+            module,
+            line_numbers,
+            params,
+            targets: vec![],
+        }
+    }
+
+    pub fn code_actions(mut self) -> Vec<CodeAction> {
+        self.visit_typed_module(&self.module.ast);
+
+        let mut actions = Vec::with_capacity(self.targets.len());
+        for target in self.targets {
+            let mut edits = TextEdits::new(self.line_numbers);
+
+            edits.delete(SrcSpan {
+                start: target.fn_location.start,
+                end: target.inner_function_location.start,
+            });
+            edits.delete(SrcSpan {
+                start: target.inner_function_location.end,
+                end: target.fn_location.end,
+            });
+
+            CodeActionBuilder::new("Unwrap anonymous function")
+                .kind(CodeActionKind::REFACTOR_REWRITE)
+                .changes(self.params.text_document.uri.clone(), edits.edits)
+                .push_to(&mut actions);
+        }
+        actions
+    }
+}
+
+impl<'ast> ast::visit::Visit<'ast> for UnwrapAnonymousFunction<'ast> {
+    fn visit_typed_expr_fn(
+        &mut self,
+        location: &'ast SrcSpan,
+        _type_: &'ast Arc<Type>,
+        kind: &'ast FunctionLiteralKind,
+        arguments: &'ast [TypedArg],
+        body: &'ast Vec1<TypedStatement>,
+        _return_annotation: &'ast Option<ast::TypeAst>,
+    ) {
+        match kind {
+            FunctionLiteralKind::Anonymous { .. } => (),
+            _ => return,
+        }
+
+        // We need the existing argument list for the fn to be a 1:1 match for the args we pass
+        // to the called function. We figure out what the call-arg list needs to look like here,
+        // and bail out if any incoming args are discarded.
+        let mut expected_arguments = Vec::with_capacity(arguments.len());
+        for a in arguments {
+            match &a.names {
+                ArgNames::Named { name, .. } => expected_arguments.push(name),
+                ArgNames::NamedLabelled { name, .. } => expected_arguments.push(name),
+                ArgNames::Discard { .. } => return,
+                ArgNames::LabelledDiscard { .. } => return,
+            }
+        }
+
+        // match fn bodies with only a single function call
+        if let [stmt] = body.as_slice()
+            && let TypedStatement::Expression(expr) = stmt
+            && let TypedExpr::Call { fun, arguments, .. } = expr
+        {
+            let mut call_arguments = Vec::with_capacity(arguments.len());
+            for a in arguments {
+                match &a.value {
+                    TypedExpr::Var { name, .. } => call_arguments.push(name),
+                    _ => return,
+                }
+            }
+
+            if call_arguments != expected_arguments {
+                return;
+            }
+
+            self.targets.push(FnToUnwrap {
+                fn_location: *location,
+                inner_function_location: fun.location(),
+            })
+        }
+    }
+}

compiler-core/src/language_server/engine.rs

@@ -48,9 +48,9 @@ use super::{
         FillInMissingLabelledArgs, FillUnusedFields, FixBinaryOperation,
         FixTruncatedBitArraySegment, GenerateDynamicDecoder, GenerateFunction, GenerateJsonEncoder,
         GenerateVariant, InlineVariable, InterpolateString, LetAssertToCase, PatternMatchOnValue,
-        RedundantTupleInCaseSubject, RemoveEchos, RemoveUnusedImports, UseLabelShorthandSyntax,
-        WrapInAnonymousFunction, WrapInBlock, code_action_add_missing_patterns,
-        code_action_convert_qualified_constructor_to_unqualified,
+        RedundantTupleInCaseSubject, RemoveEchos, RemoveUnusedImports, UnwrapAnonymousFunction,
+        UseLabelShorthandSyntax, WrapInAnonymousFunction, WrapInBlock,
+        code_action_add_missing_patterns, code_action_convert_qualified_constructor_to_unqualified,
         code_action_convert_unqualified_constructor_to_qualified, code_action_import_module,
         code_action_inexhaustive_let_to_case,
     },
@@ -448,6 +448,7 @@ where
             actions.extend(RemoveBlock::new(module, &lines, &params).code_actions());
             actions.extend(RemovePrivateOpaque::new(module, &lines, &params).code_actions());
             actions.extend(WrapInAnonymousFunction::new(module, &lines, &params).code_actions());
+            actions.extend(UnwrapAnonymousFunction::new(module, &lines, &params).code_actions());
             GenerateDynamicDecoder::new(module, &lines, &params, &mut actions).code_actions();
             GenerateJsonEncoder::new(
                 module,

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

@@ -135,6 +135,7 @@ const REMOVE_OPAQUE_FROM_PRIVATE_TYPE: &str = "Remove opaque from private type";
 const COLLAPSE_NESTED_CASE: &str = "Collapse nested case";
 const REMOVE_UNREACHABLE_BRANCHES: &str = "Remove unreachable branches";
 const WRAP_IN_ANONYMOUS_FUNCTION: &str = "Wrap in anonymous function";
+const UNWRAP_ANONYMOUS_FUNCTION: &str = "Unwrap anonymous function";
 
 macro_rules! assert_code_action {
     ($title:expr, $code:literal, $range:expr $(,)?) => {
@@ -9830,4 +9831,60 @@ fn op_factory(a: Int, b: Int, c: Int) -> fn(Int) -> Int {
 ",
         find_position_of("op_factory").to_selection()
     );
-}
+}
+
+#[test]
+fn unwrap_trivial_anonymous_function() {
+    assert_code_action!(
+        UNWRAP_ANONYMOUS_FUNCTION,
+        "import gleam/list
+
+pub fn main() {
+    list.map([1, 2, 3], fn(int) { op(int) })
+}
+
+fn op(i: Int) -> Int {
+    todo
+}
+",
+        find_position_of("fn(int)").to_selection()
+    );
+}
+
+#[test]
+fn unwrap_anonymous_function_unavailable_when_args_discarded() {
+    assert_no_code_actions!(
+        UNWRAP_ANONYMOUS_FUNCTION,
+        "import gleam/list
+
+pub fn main() {
+    list.index_map([1, 2, 3], fn(_, int) { op(int) })
+}
+
+fn op(i: Int) -> Int {
+    todo
+}
+",
+        find_position_of("fn(_, int)").to_selection()
+    );
+}
+
+#[test]
+fn unwrap_anonymous_function_unavailable_with_different_args() {
+    assert_no_code_actions!(
+        UNWRAP_ANONYMOUS_FUNCTION,
+        "import gleam/list
+
+const another_int = 7
+        
+pub fn main() {
+    list.map([1, 2, 3], fn(int) { op(another_int) })
+}
+
+fn op(i: Int) -> Int {
+    todo
+}
+",
+        find_position_of("fn(int)").to_selection()
+    );
+}

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

@@ -0,0 +1,27 @@
+---
+source: compiler-core/src/language_server/tests/action.rs
+expression: "import gleam/list\n\npub fn main() {\n    list.map([1, 2, 3], fn(int) { op(int) })\n}\n\nfn op(i: Int) -> Int {\n    todo\n}\n"
+---
+----- BEFORE ACTION
+import gleam/list
+
+pub fn main() {
+    list.map([1, 2, 3], fn(int) { op(int) })
+                        ↑                   
+}
+
+fn op(i: Int) -> Int {
+    todo
+}
+
+
+----- AFTER ACTION
+import gleam/list
+
+pub fn main() {
+    list.map([1, 2, 3], op)
+}
+
+fn op(i: Int) -> Int {
+    todo
+}