Detect which variables are used outside the extracted function

Author: GearsDatapacks

compiler-core/src/language_server/code_action.rs

@@ -8417,15 +8417,40 @@ pub struct ExtractFunction<'a> {
     module: &'a Module,
     params: &'a CodeActionParams,
     edits: TextEdits<'a>,
-    extract: Extract<'a>,
+    extract: Option<ExtractedFunction<'a>>,
     function_end_position: Option<u32>,
 }
 
+struct ExtractedFunction<'a> {
+    parameters: Vec<(EcoString, Arc<Type>)>,
+    returned_variables: Vec<(EcoString, Arc<Type>)>,
+    value: ExtractedValue<'a>,
+}
+
+impl<'a> ExtractedFunction<'a> {
+    fn new(value: ExtractedValue<'a>) -> Self {
+        Self {
+            value,
+            parameters: Vec::new(),
+            returned_variables: Vec::new(),
+        }
+    }
+
+    fn location(&self) -> SrcSpan {
+        match &self.value {
+            ExtractedValue::Expression(expression) => expression.location(),
+            ExtractedValue::Statements(statements) => SrcSpan::new(
+                statements.first().location().start,
+                statements.last().location().end,
+            ),
+        }
+    }
+}
+
 #[derive(Debug)]
-enum Extract<'a> {
-    None,
+enum ExtractedValue<'a> {
     Expression(&'a TypedExpr),
-    Statements(Vec<&'a TypedStatement>),
+    Statements(Vec1<&'a TypedStatement>),
 }
 
 impl<'a> ExtractFunction<'a> {
@@ -8438,7 +8463,7 @@ impl<'a> ExtractFunction<'a> {
             module,
             params,
             edits: TextEdits::new(line_numbers),
-            extract: Extract::None,
+            extract: None,
             function_end_position: None,
         }
     }
@@ -8454,10 +8479,19 @@ impl<'a> ExtractFunction<'a> {
             return Vec::new();
         };
 
-        match std::mem::replace(&mut self.extract, Extract::None) {
-            Extract::None => return Vec::new(),
-            Extract::Expression(expression) => self.extract_expression(expression, end),
-            Extract::Statements(statements) => self.extract_statements(statements, end),
+        let Some(extracted) = self.extract.take() else {
+            return Vec::new();
+        };
+        match extracted.value {
+            ExtractedValue::Expression(expression) => {
+                self.extract_expression(expression, extracted.parameters, end)
+            }
+            ExtractedValue::Statements(statements) => self.extract_statements(
+                statements,
+                extracted.parameters,
+                extracted.returned_variables,
+                end,
+            ),
         }
 
         let mut action = Vec::with_capacity(1);
@@ -8469,18 +8503,21 @@ impl<'a> ExtractFunction<'a> {
         action
     }
 
-    fn extract_expression(&mut self, expression: &TypedExpr, function_end: u32) {
-        let referenced_variables = referenced_variables(expression);
-
+    fn extract_expression(
+        &mut self,
+        expression: &TypedExpr,
+        parameters: Vec<(EcoString, Arc<Type>)>,
+        function_end: u32,
+    ) {
         let expression_code = code_at(self.module, expression.location());
 
-        let arguments = referenced_variables.iter().map(|(name, _)| name).join(", ");
+        let arguments = parameters.iter().map(|(name, _)| name).join(", ");
         let call = format!("function({arguments})");
         self.edits.replace(expression.location(), call);
 
         let mut printer = Printer::new(&self.module.ast.names);
 
-        let parameters = referenced_variables
+        let parameters = parameters
             .iter()
             .map(|(name, type_)| eco_format!("{name}: {}", printer.print_type(type_)))
             .join(", ");
@@ -8495,40 +8532,97 @@ impl<'a> ExtractFunction<'a> {
         self.edits.insert(function_end, function);
     }
 
-    fn extract_statements(&mut self, statements: Vec<&TypedStatement>, function_end: u32) {
-        let Some(first) = statements.first() else {
-            return;
-        };
-        let Some(last) = statements.last() else {
-            return;
-        };
+    fn extract_statements(
+        &mut self,
+        statements: Vec1<&TypedStatement>,
+        parameters: Vec<(EcoString, Arc<Type>)>,
+        returned_variables: Vec<(EcoString, Arc<Type>)>,
+        function_end: u32,
+    ) {
+        let first = statements.first();
+        let last = statements.last();
 
         let location = SrcSpan::new(first.location().start, last.location().end);
 
-        let referenced_variables = referenced_variables_for_statements(&statements, location);
-
         let code = code_at(self.module, location);
 
-        let arguments = referenced_variables.iter().map(|(name, _)| name).join(", ");
-        let call = format!("function({arguments})");
+        let returns_anything = !returned_variables.is_empty();
+
+        let (return_type, return_value) = match returned_variables.as_slice() {
+            [] => (type_::nil(), "Nil".into()),
+            [(name, type_)] => (type_.clone(), name.clone()),
+            _ => {
+                let values = returned_variables.iter().map(|(name, _)| name).join(", ");
+                let type_ = type_::tuple(
+                    returned_variables
+                        .into_iter()
+                        .map(|(_, type_)| type_)
+                        .collect(),
+                );
+
+                (type_, eco_format!("#({values})"))
+            }
+        };
+
+        let arguments = parameters.iter().map(|(name, _)| name).join(", ");
+
+        let call = if returns_anything {
+            format!("let {return_value} = function({arguments})")
+        } else {
+            format!("function({arguments})")
+        };
         self.edits.replace(location, call);
 
         let mut printer = Printer::new(&self.module.ast.names);
 
-        let parameters = referenced_variables
+        let parameters = parameters
             .iter()
             .map(|(name, type_)| eco_format!("{name}: {}", printer.print_type(type_)))
             .join(", ");
-        let return_type = printer.print_type(&last.type_());
+
+        let return_type = printer.print_type(&return_type);
 
         let function = format!(
             "\n\nfn function({parameters}) -> {return_type} {{
   {code}
+  {return_value}
 }}"
         );
 
         self.edits.insert(function_end, function);
     }
+
+    fn register_referenced_variable(
+        &mut self,
+        name: &EcoString,
+        type_: &Arc<Type>,
+        location: SrcSpan,
+        definition_location: SrcSpan,
+    ) {
+        let Some(extracted) = &mut self.extract else {
+            return;
+        };
+
+        let extracted_location = extracted.location();
+
+        let variables = if extracted_location.contains_span(location)
+            && !extracted_location.contains_span(definition_location)
+        {
+            &mut extracted.parameters
+        } else if extracted_location.contains_span(definition_location)
+            && !extracted_location.contains_span(location)
+        {
+            &mut extracted.returned_variables
+        } else {
+            return;
+        };
+
+        if variables.iter().any(|(variable, _)| variable == name) {
+            return;
+        }
+
+        variables.push((name.clone(), type_.clone()));
+    }
 }
 
 impl<'ast> ast::visit::Visit<'ast> for ExtractFunction<'ast> {
@@ -8543,16 +8637,14 @@ impl<'ast> ast::visit::Visit<'ast> for ExtractFunction<'ast> {
     }
 
     fn visit_typed_expr(&mut self, expression: &'ast TypedExpr) {
-        match &self.extract {
-            Extract::None => {
-                let range = self.edits.src_span_to_lsp_range(expression.location());
+        if self.extract.is_none() {
+            let range = self.edits.src_span_to_lsp_range(expression.location());
 
-                if within(range, self.params.range) {
-                    self.extract = Extract::Expression(expression);
-                    return;
-                }
+            if within(range, self.params.range) {
+                self.extract = Some(ExtractedFunction::new(ExtractedValue::Expression(
+                    expression,
+                )));
             }
-            Extract::Expression(_) | Extract::Statements(_) => {}
         }
         ast::visit::visit_typed_expr(self, expression);
     }
@@ -8561,81 +8653,43 @@ impl<'ast> ast::visit::Visit<'ast> for ExtractFunction<'ast> {
         let range = self.edits.src_span_to_lsp_range(statement.location());
         if within(range, self.params.range) {
             match &mut self.extract {
-                Extract::None => {
-                    self.extract = Extract::Statements(vec![statement]);
+                None => {
+                    self.extract = Some(ExtractedFunction::new(ExtractedValue::Statements(vec1![
+                        statement,
+                    ])));
                 }
-                Extract::Expression(expression) => {
-                    if expression.location().contains_span(statement.location()) {
-                        return;
-                    }
-
-                    self.extract = Extract::Statements(vec![statement]);
-                }
-                Extract::Statements(statements) => {
+                Some(ExtractedFunction {
+                    value: ExtractedValue::Expression(_),
+                    ..
+                }) => {}
+                Some(ExtractedFunction {
+                    value: ExtractedValue::Statements(statements),
+                    ..
+                }) => {
                     statements.push(statement);
                 }
             }
-        } else {
-            ast::visit::visit_typed_statement(self, statement);
-        }
-    }
-}
-
-fn referenced_variables(expression: &TypedExpr) -> Vec<(EcoString, Arc<Type>)> {
-    let mut references = ReferencedVariables::new(expression.location());
-    references.visit_typed_expr(expression);
-    references.variables
-}
-
-fn referenced_variables_for_statements(
-    statements: &[&TypedStatement],
-    location: SrcSpan,
-) -> Vec<(EcoString, Arc<Type>)> {
-    let mut references = ReferencedVariables::new(location);
-    for statement in statements {
-        references.visit_typed_statement(*statement);
-    }
-    references.variables
-}
-
-struct ReferencedVariables {
-    variables: Vec<(EcoString, Arc<Type>)>,
-    location: SrcSpan,
-}
-
-impl ReferencedVariables {
-    fn new(location: SrcSpan) -> Self {
-        Self {
-            variables: Vec::new(),
-            location,
-        }
-    }
-
-    fn register(&mut self, name: &EcoString, type_: &Arc<Type>, definition_location: SrcSpan) {
-        if self.location.contains_span(definition_location) {
-            return;
-        }
-
-        if !self
-            .variables
-            .iter()
-            .any(|(variable_name, _)| variable_name == name)
-        {
-            self.variables.push((name.clone(), type_.clone()))
         }
+        ast::visit::visit_typed_statement(self, statement);
     }
-}
 
-impl<'ast> ast::visit::Visit<'ast> for ReferencedVariables {
     fn visit_typed_expr_var(
         &mut self,
-        _location: &'ast SrcSpan,
+        location: &'ast SrcSpan,
         constructor: &'ast ValueConstructor,
         name: &'ast EcoString,
     ) {
         match &constructor.variant {
-            type_::ValueConstructorVariant::LocalVariable { location, .. } => {
-                self.register(name, &constructor.type_, *location);
+            type_::ValueConstructorVariant::LocalVariable {
+                location: definition_location,
+                ..
+            } => {
+                self.register_referenced_variable(
+                    name,
+                    &constructor.type_,
+                    *location,
+                    *definition_location,
+                );
             }
             type_::ValueConstructorVariant::ModuleConstant { .. }
             | type_::ValueConstructorVariant::LocalConstant { .. }

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

@@ -10299,3 +10299,33 @@ pub fn do_things(a, b) {
         find_position_of("let").select_until(find_position_of("+ a * b\n").under_char('\n'))
     );
 }
+
+#[test]
+fn extract_function_which_uses_multiple_extracted_variables() {
+    assert_code_action!(
+        EXTRACT_FUNCTION,
+        "
+pub fn do_things(a, b) {
+  let wibble = a + b
+  let wobble = a * b
+  wobble / wibble
+}
+",
+        find_position_of("let").select_until(find_position_of("* b\n").under_char('\n'))
+    );
+}
+
+#[test]
+fn extract_function_which_uses_no_extracted_variables() {
+    assert_code_action!(
+        EXTRACT_FUNCTION,
+        "
+pub fn do_things(a, b) {
+  let x = a + b
+  echo x
+  a
+}
+",
+        find_position_of("let").select_until(find_position_of("echo x\n").under_char('\n'))
+    );
+}

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

@@ -18,12 +18,13 @@ pub fn do_things(a, b) {
 ----- AFTER ACTION
 
 pub fn do_things(a, b) {
-  function(a, b)
+  let result = function(a, b)
   result + 3
 }
 
 fn function(a: Int, b: Int) -> Int {
   let a = 10 + a
   let b = 10 + b
   let result = a * b
+  result
 }