collapse_simple_stmt: check if return expressions are "simple"

CHANGELOG.md

@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `.stylua.toml` config resolution now supports looking up config files next to files being formatted, recursively going
   upwards until reaching the current working directory, then stopping (unless `--search-parent-directories` was specified).
   For example, for a file `./src/test.lua`, executing `stylua src/` will look for `./src/stylua.toml` and then `./stylua.toml`.
+- When `collapse_simple_statement` is enabled, if the enclosing block is a return, we will check if the return expression is "simple" (currently, not containing a function definition) ([#898](https://github.com/JohnnyMorganz/StyLua/issues/898))
 
 ### Fixed
 

src/formatters/trivia_util.rs

@@ -134,8 +134,50 @@ pub fn is_block_empty(block: &Block) -> bool {
     block.stmts().next().is_none() && block.last_stmt().is_none()
 }
 
+fn is_expression_simple(expression: &Expression) -> bool {
+    match expression {
+        Expression::Function(_) => false,
+        Expression::FunctionCall(function_call) => {
+            function_call.suffixes().all(|suffix| match suffix {
+                Suffix::Index(_) => true,
+                Suffix::Call(call) => match call {
+                    Call::AnonymousCall(function_args) => match function_args {
+                        FunctionArgs::Parentheses { arguments, .. } => {
+                            arguments.iter().all(is_expression_simple)
+                        }
+                        _ => true,
+                    },
+                    Call::MethodCall(method_call) => match method_call.args() {
+                        FunctionArgs::Parentheses { arguments, .. } => {
+                            arguments.iter().all(is_expression_simple)
+                        }
+                        _ => true,
+                    },
+                    other => unreachable!("unknown node: {:?}", other),
+                },
+                other => unreachable!("unknown node: {:?}", other),
+            })
+        }
+        _ => true,
+    }
+}
+
+fn is_last_stmt_simple(last_stmt: &LastStmt) -> bool {
+    match last_stmt {
+        LastStmt::Break(_) => true,
+        #[cfg(feature = "luau")]
+        LastStmt::Continue(_) => true,
+        LastStmt::Return(r#return) => {
+            r#return.returns().is_empty() || r#return.returns().iter().all(is_expression_simple)
+        }
+        other => unreachable!("unknown node {:?}", other),
+    }
+}
+
 pub fn is_block_simple(block: &Block) -> bool {
-    (block.stmts().next().is_none() && block.last_stmt().is_some())
+    (block.stmts().next().is_none()
+        && block.last_stmt().is_some()
+        && is_last_stmt_simple(block.last_stmt().unwrap()))
         || (block.stmts().count() == 1
             && block.last_stmt().is_none()
             && match block.stmts().next().unwrap() {

tests/inputs-collapse-single-statement/conditional-with-function-1.lua

@@ -0,0 +1,13 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/898
+
+if bar then
+	return function()
+		foo()
+	end
+end
+
+if bar then
+	return Array.filter({}, function()
+		return true
+	end)
+end

tests/snapshots/tests__collapse_single_statement@conditional-with-function-1.lua.snap

@@ -0,0 +1,15 @@
+---
+source: tests/tests.rs
+expression: "format_code(&contents,\n        Config {\n            collapse_simple_statement: CollapseSimpleStatement::Always,\n            ..Config::default()\n        }, None, OutputVerification::None).unwrap()"
+input_file: tests/inputs-collapse-single-statement/conditional-with-function-1.lua
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/898
+
+if bar then
+	return function() foo() end
+end
+
+if bar then
+	return Array.filter({}, function() return true end)
+end
+