support if-local and if-const

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Luau: Added support for `if local` / `if const` condition bindings in both `if` statements and `if` expressions (e.g. `if local x = getValue() then ... end`)
+
 ### Fixed
 
 - Fixed range formatting indenting statements one level too shallow when nested inside a function expression within a table constructor ([#1145](https://github.com/JohnnyMorganz/StyLua/issues/1145))

src/formatters/expression.rs

@@ -15,7 +15,10 @@ use std::boxed::Box;
 #[cfg(feature = "luau")]
 use crate::formatters::{
     assignment::calculate_hang_level,
-    luau::{format_type_assertion, format_type_instantiation},
+    luau::{
+        format_if_condition_binding, format_type_assertion, format_type_instantiation,
+        if_condition_binding_width,
+    },
     stmt::remove_condition_parentheses,
     trivia_util::HasInlineComments,
 };
@@ -517,13 +520,17 @@ fn format_else_if_expression_singleline(
     shape: Shape,
 ) -> ElseIfExpression {
     let else_if_token = fmt_symbol!(ctx, else_if_expression.else_if_token(), "elseif ", shape);
+    let binding = else_if_expression
+        .binding()
+        .map(|binding| format_if_condition_binding(ctx, binding, shape));
+    let binding_width = binding.as_ref().map_or(0, if_condition_binding_width);
     let else_if_condition = remove_condition_parentheses(else_if_expression.condition().to_owned());
-    let else_if_condition = format_expression(ctx, &else_if_condition, shape + 7); // 7 = "elseif "
+    let else_if_condition = format_expression(ctx, &else_if_condition, shape + 7 + binding_width); // 7 = "elseif "
     let (then_token, expression) = format_token_expression_sequence(
         ctx,
         else_if_expression.then_token(),
         else_if_expression.expression(),
-        shape.take_first_line(&else_if_condition) + 13, // 13 = "elseif " + " then ",
+        shape.take_first_line(&else_if_condition) + 13 + binding_width, // 13 = "elseif " + " then ",
     );
 
     // Add a space before the then token
@@ -533,6 +540,7 @@ fn format_else_if_expression_singleline(
 
     ElseIfExpression::new(else_if_condition, expression)
         .with_else_if_token(else_if_token)
+        .with_binding(binding)
         .with_then_token(then_token)
 }
 
@@ -602,6 +610,12 @@ fn format_if_expression(ctx: &Context, if_expression: &IfExpression, shape: Shap
     let condition = remove_condition_parentheses(if_expression.condition().to_owned());
     let if_token = fmt_symbol!(ctx, if_expression.if_token(), "if ", shape);
 
+    // Format the `local <name> =` binding (`if local` / `if const`), if present
+    let binding = if_expression
+        .binding()
+        .map(|binding| format_if_condition_binding(ctx, binding, shape));
+    let binding_width = binding.as_ref().map_or(0, if_condition_binding_width);
+
     // Initially format the remainder on a single line
     let singleline_condition = format_expression(ctx, &condition, shape.with_infinite_width());
     let then_token = fmt_symbol!(ctx, if_expression.then_token(), " then ", shape);
@@ -636,7 +650,7 @@ fn format_if_expression(ctx: &Context, if_expression: &IfExpression, shape: Shap
     const ELSE_LENGTH: usize = 6; // " else "
 
     // Determine if we need to hang the expression
-    let singleline_shape = (shape + IF_LENGTH + THEN_LENGTH + ELSE_LENGTH)
+    let singleline_shape = (shape + IF_LENGTH + THEN_LENGTH + ELSE_LENGTH + binding_width)
         .take_first_line(&strip_trivia(&singleline_condition))
         .take_first_line(&strip_trivia(&singleline_expression))
         .take_first_line(&else_ifs.as_ref().map_or(String::new(), |x| {
@@ -667,7 +681,7 @@ fn format_if_expression(ctx: &Context, if_expression: &IfExpression, shape: Shap
         let condition = hang_expression_trailing_newline(
             ctx,
             if_expression.condition(),
-            shape.increment_additional_indent(),
+            (shape + binding_width).increment_additional_indent(),
             Some(1),
         );
         let hanging_shape = shape.reset().increment_additional_indent();
@@ -750,8 +764,25 @@ fn format_if_expression(ctx: &Context, if_expression: &IfExpression, shape: Shap
                                     create_indent_trivia(ctx, hanging_shape),
                                 ]));
 
+                            let binding = else_if_expression.binding().map(|binding| {
+                                let binding =
+                                    format_if_condition_binding(ctx, binding, hanging_shape);
+                                let equal_token = binding
+                                    .equal_token()
+                                    .update_trailing_trivia(FormatTriviaType::Replace(vec![]));
+                                binding.with_equal_token(equal_token)
+                            });
+                            let else_if_token = if binding.is_some() {
+                                else_if_token.update_trailing_trivia(FormatTriviaType::Append(
+                                    vec![Token::new(TokenType::spaces(1))],
+                                ))
+                            } else {
+                                else_if_token
+                            };
+
                             ElseIfExpression::new(else_if_condition, expression)
                                 .with_else_if_token(else_if_token)
+                                .with_binding(binding)
                                 .with_then_token(then_token)
                         } else {
                             singleline_else_if.update_leading_trivia(FormatTriviaType::Append(
@@ -778,6 +809,7 @@ fn format_if_expression(ctx: &Context, if_expression: &IfExpression, shape: Shap
 
         IfExpression::new(condition, expression, else_expression)
             .with_if_token(if_token)
+            .with_binding(binding)
             .with_then_token(then_token)
             .with_else_if(else_ifs)
             .with_else_token(else_token)
@@ -799,6 +831,7 @@ fn format_if_expression(ctx: &Context, if_expression: &IfExpression, shape: Shap
             singleline_else_expression,
         )
         .with_if_token(if_token)
+        .with_binding(binding)
         .with_then_token(then_token)
         .with_else_if(else_ifs)
         .with_else_token(else_token)

src/formatters/luau.rs

@@ -27,9 +27,9 @@ use crate::{
 use full_moon::ast::{
     luau::{
         ExportedTypeDeclaration, ExportedTypeFunction, GenericDeclaration,
-        GenericDeclarationParameter, GenericParameterInfo, IndexedTypeInfo, LuauAttribute,
-        TypeArgument, TypeAssertion, TypeDeclaration, TypeField, TypeFieldKey, TypeFunction,
-        TypeInfo, TypeInstantiation, TypeIntersection, TypeSpecifier, TypeUnion,
+        GenericDeclarationParameter, GenericParameterInfo, IfConditionBinding, IndexedTypeInfo,
+        LuauAttribute, TypeArgument, TypeAssertion, TypeDeclaration, TypeField, TypeFieldKey,
+        TypeFunction, TypeInfo, TypeInstantiation, TypeIntersection, TypeSpecifier, TypeUnion,
     },
     punctuated::Pair,
 };
@@ -1533,6 +1533,57 @@ pub fn format_type_specifier(
         .with_type_info(type_info)
 }
 
+/// The width of the last line of an [`IfConditionBinding`], ignoring comments.
+pub fn if_condition_binding_width(binding: &IfConditionBinding) -> usize {
+    let text = format!(
+        "{} {}{} =",
+        strip_trivia(binding.local_token()),
+        strip_trivia(binding.name()),
+        binding
+            .type_specifier()
+            .map_or(String::new(), |specifier| format!(
+                ": {}",
+                strip_trivia(specifier.type_info())
+            )),
+    );
+    text.lines().last().map_or(0, |line| line.chars().count())
+}
+
+/// Formats an [`IfConditionBinding`] node - the `local <name> =` / `const <name> =` part of an
+/// `if local` / `if const` binding, as in `if local player = getPlayer() then ... end`.
+/// The formatted binding is `<local|const> <name>[: <type>] = ` (with a trailing space), so it
+/// slots directly between the `if `/`elseif ` token and the condition expression.
+pub fn format_if_condition_binding(
+    ctx: &Context,
+    binding: &IfConditionBinding,
+    shape: Shape,
+) -> IfConditionBinding {
+    // Preserve the keyword (`local` or `const`), normalising to a single trailing space
+    let local_token = format_token_reference(ctx, binding.local_token(), shape)
+        .update_trailing_trivia(FormatTriviaType::Append(vec![Token::new(
+            TokenType::spaces(1),
+        )]));
+    let name_shape = shape
+        + (strip_trivia(binding.local_token())
+            .to_string()
+            .chars()
+            .count()
+            + 1);
+    let name = format_token_reference(ctx, binding.name(), name_shape);
+    let specifier_shape = name_shape + strip_trivia(binding.name()).to_string().chars().count();
+    let type_specifier = binding
+        .type_specifier()
+        .map(|type_specifier| format_type_specifier(ctx, type_specifier, specifier_shape));
+    let equal_token = fmt_symbol!(ctx, binding.equal_token(), " = ", shape);
+
+    binding
+        .to_owned()
+        .with_local_token(local_token)
+        .with_name(name)
+        .with_type_specifier(type_specifier)
+        .with_equal_token(equal_token)
+}
+
 pub fn format_type_instantiation(
     ctx: &Context,
     type_instantiation: &TypeInstantiation,