From d8f820e14dd230492533b5037cde0bb920e492b4 Mon Sep 17 00:00:00 2001
From: JohnnyMorganz <johnnymorganz@outlook.com>
Date: Sun, 17 Nov 2024 17:47:44 +0100
Subject: [PATCH] collapse_simple_stmt: check if return expressions are
 "simple" (#923)

* Add test case

* Check if return expression is simple under collapse_simple_stmt

* Update changelog and snapshots

* Update changelog link
---
 CHANGELOG.md                                  |  1 +
 src/formatters/trivia_util.rs                 | 44 ++++++++++++++++++-
 .../conditional-with-function-1.lua           | 13 ++++++
 ...ement@conditional-with-function-1.lua.snap | 15 +++++++
 4 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 tests/inputs-collapse-single-statement/conditional-with-function-1.lua
 create mode 100644 tests/snapshots/tests__collapse_single_statement@conditional-with-function-1.lua.snap

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f9208118..4c5a7e88 100644
--- a/CHANGELOG.md
+++ b/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
 
diff --git a/src/formatters/trivia_util.rs b/src/formatters/trivia_util.rs
index 31981706..0b76dcb6 100644
--- a/src/formatters/trivia_util.rs
+++ b/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() {
diff --git a/tests/inputs-collapse-single-statement/conditional-with-function-1.lua b/tests/inputs-collapse-single-statement/conditional-with-function-1.lua
new file mode 100644
index 00000000..9961929d
--- /dev/null
+++ b/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
diff --git a/tests/snapshots/tests__collapse_single_statement@conditional-with-function-1.lua.snap b/tests/snapshots/tests__collapse_single_statement@conditional-with-function-1.lua.snap
new file mode 100644
index 00000000..0d27dcf7
--- /dev/null
+++ b/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
+