Take into account size of function definition when formatting return type

CHANGELOG.md

@@ -29,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Luau: Fixed incorrect removal of semicolon before compound assignment with parentheses leading to ambiguous syntax error ([#885](https://github.com/JohnnyMorganz/StyLua/issues/885))
 - Luau: Fixed incorrect collapsing of union/intersection type value with comments in a type table leading to a syntax error ([#893](https://github.com/JohnnyMorganz/StyLua/issues/893))
 - Fixed `--verify` panicing due to overflow for very large Hex numbers ([#875](https://github.com/JohnnyMorganz/StyLua/issues/875), [#889](https://github.com/JohnnyMorganz/StyLua/issues/889))
+- Luau: take into account the size of the current function definition when formatting the return type ([#896](https://github.com/JohnnyMorganz/StyLua/issues/896))
 
 ## [0.20.0] - 2024-01-20
 

src/formatters/functions.rs

@@ -750,6 +750,9 @@ pub fn format_function_body(
     function_body: &FunctionBody,
     shape: Shape,
 ) -> FunctionBody {
+    const SINGLE_PAREN_LEN: usize = "(".len();
+    const PARENS_LEN: usize = SINGLE_PAREN_LEN * 2;
+
     // Calculate trivia
     let leading_trivia = vec![create_indent_trivia(ctx, shape)];
 
@@ -814,14 +817,26 @@ pub fn format_function_body(
             shape
         };
 
+        let type_specifiers = function_body
+            .type_specifiers()
+            .map(|x| x.map(|specifier| format_type_specifier(ctx, specifier, parameters_shape)))
+            .collect::<Vec<_>>();
+
+        let return_type_shape = if multiline_params {
+            shape.reset() + SINGLE_PAREN_LEN
+        } else {
+            shape.take_first_line(&formatted_parameters)
+                + PARENS_LEN
+                + type_specifiers.iter().fold(0, |acc, x| {
+                    acc + x.as_ref().map_or(0, |x| x.to_string().len())
+                })
+        };
+
         (
-            function_body
-                .type_specifiers()
-                .map(|x| x.map(|specifier| format_type_specifier(ctx, specifier, parameters_shape)))
-                .collect::<Vec<_>>(),
+            type_specifiers,
             function_body
                 .return_type()
-                .map(|return_type| format_type_specifier(ctx, return_type, shape)),
+                .map(|return_type| format_type_specifier(ctx, return_type, return_type_shape)),
         )
     };
 
@@ -834,7 +849,6 @@ pub fn format_function_body(
         if trivia_util::is_block_empty(function_body.block()) {
             Block::new()
         } else {
-            const PARENS_LEN: usize = "()".len();
             let block_shape = shape.take_first_line(&formatted_parameters) + PARENS_LEN;
 
             #[cfg(feature = "luau")]

tests/inputs-luau/function_types_4.lua

@@ -0,0 +1,2 @@
+-- https://github.com/JohnnyMorganz/StyLua/issues/896
+local function getData(player: Player): (LongReturnValueWithTypeGeneric1<TypeParameter1>, LongReturnValueWithTypeGeneric2<TypeParameter2>) end

tests/snapshots/tests__luau@function_types_4.lua.snap

@@ -0,0 +1,11 @@
+---
+source: tests/tests.rs
+expression: "format(&contents, LuaVersion::Luau)"
+input_file: tests/inputs-luau/function_types_4.lua
+---
+-- https://github.com/JohnnyMorganz/StyLua/issues/896
+local function getData(player: Player): (
+	LongReturnValueWithTypeGeneric1<TypeParameter1>,
+	LongReturnValueWithTypeGeneric2<TypeParameter2>
+) end
+