Add Input mode for collapse_simple_statement

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Luau: Added support for `const` variable assignments (`const x = 1`) and `const function` declarations ([#1102](https://github.com/JohnnyMorganz/StyLua/issues/1102))
+- Added `Input` mode for `collapse_simple_statement`, preserving whether simple statements are single-line or multi-line in the input while still formatting their contents
 
 ### Changed
 

README.md

@@ -319,7 +319,7 @@ StyLua only offers the following options:
 | `call_parentheses`           | `Always`           | Whether parentheses should be applied on function calls with a single string/table argument. Possible options: `Always`, `NoSingleString`, `NoSingleTable`, `None`, `Input`. `Always` applies parentheses in all cases. `NoSingleString` omits parentheses on calls with a single string argument. Similarly, `NoSingleTable` omits parentheses on calls with a single table argument. `None` omits parentheses in both cases. Note: parentheses are still kept in situations where removal can lead to obscurity (e.g. `foo "bar".setup -> foo("bar").setup`, since the index is on the call result, not the string). `Input` removes all automation and preserves parentheses only if they were present in input code: consistency is not enforced. |
 | `space_after_function_names` | `Never`            | Specify whether to add a space between the function name and parentheses. Possible options: `Never`, `Definitions`, `Calls`, or `Always`                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              |
 | `block_newline_gaps`         | `Never`            | Specify whether to preserve leading and trailing newline gaps for blocks. Possible options: `Never`, `Preserve`                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
-| `collapse_simple_statement`  | `Never`            | Specify whether to collapse simple statements. Possible options: `Never`, `FunctionOnly`, `ConditionalOnly`, or `Always`                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              |
+| `collapse_simple_statement`  | `Never`            | Specify whether to collapse simple statements. Possible options: `Never`, `Input`, `FunctionOnly`, `ConditionalOnly`, or `Always`                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     |
 
 Default `stylua.toml`, note you do not need to explicitly specify each option if you want to use the defaults:
 

src/cli/config.rs

@@ -303,7 +303,9 @@ fn load_overrides(config: Config, opt: &Opt) -> Config {
 mod tests {
     use super::*;
     use clap::StructOpt;
-    use stylua_lib::{CallParenType, IndentType, LineEndings, LuaVersion, QuoteStyle};
+    use stylua_lib::{
+        CallParenType, CollapseSimpleStatement, IndentType, LineEndings, LuaVersion, QuoteStyle,
+    };
 
     #[test]
     fn test_override_syntax() {
@@ -360,4 +362,16 @@ mod tests {
         let config = load_overrides(default_config, &override_opt);
         assert_eq!(config.call_parentheses, CallParenType::None);
     }
+
+    #[test]
+    fn test_override_collapse_simple_statement() {
+        let override_opt =
+            Opt::parse_from(vec!["BINARY_NAME", "--collapse-simple-statement", "Input"]);
+        let default_config = Config::new();
+        let config = load_overrides(default_config, &override_opt);
+        assert_eq!(
+            config.collapse_simple_statement,
+            CollapseSimpleStatement::Input
+        );
+    }
 }

src/cli/opt.rs

@@ -278,6 +278,7 @@ convert_enum!(CallParenType, ArgCallParenType, {
 
 convert_enum!(CollapseSimpleStatement, ArgCollapseSimpleStatement, {
     Never,
+    Input,
     FunctionOnly,
     ConditionalOnly,
     Always,

src/context.rs

@@ -155,6 +155,13 @@ impl Context {
         )
     }
 
+    pub fn should_preserve_input_simple_statements(&self) -> bool {
+        matches!(
+            self.config().collapse_simple_statement,
+            CollapseSimpleStatement::Input
+        )
+    }
+
     pub fn should_preserve_leading_block_newline_gaps(&self) -> bool {
         matches!(self.config().block_newline_gaps, BlockNewlineGaps::Preserve)
     }

src/editorconfig.rs

@@ -77,6 +77,7 @@ property_choice! {
 property_choice! {
     CollapseSimpleStatementChoice, "collapse_simple_statement";
     (Never, "never"),
+    (Input, "input"),
     (FunctionOnly, "functiononly"),
     (ConditionalOnly, "conditionalonly"),
     (Always, "always")
@@ -162,6 +163,7 @@ fn load(mut config: Config, properties: &Properties) -> Config {
     if let Ok(collapse_simple_statement) = properties.get::<CollapseSimpleStatementChoice>() {
         config.collapse_simple_statement = match collapse_simple_statement {
             CollapseSimpleStatementChoice::Never => CollapseSimpleStatement::Never,
+            CollapseSimpleStatementChoice::Input => CollapseSimpleStatement::Input,
             CollapseSimpleStatementChoice::FunctionOnly => CollapseSimpleStatement::FunctionOnly,
             CollapseSimpleStatementChoice::ConditionalOnly => {
                 CollapseSimpleStatement::ConditionalOnly
@@ -422,6 +424,17 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_collapse_simple_statement_input() {
+        let mut properties = Properties::new();
+        properties.insert_raw_for_key("collapse_simple_statement", "Input");
+        let config = Config::from(&properties);
+        assert_eq!(
+            config.collapse_simple_statement,
+            CollapseSimpleStatement::Input
+        );
+    }
+
     #[test]
     fn test_collapse_simple_statement_function_only() {
         let mut properties = Properties::new();

src/formatters/functions.rs

@@ -7,6 +7,7 @@ use full_moon::ast::{
     FunctionDeclaration, FunctionName, Index, LastStmt, LocalFunction, MethodCall, Parameter,
     Prefix, Stmt, Suffix, TableConstructor, Var,
 };
+use full_moon::node::Node;
 use full_moon::tokenizer::{Token, TokenKind, TokenReference, TokenType};
 
 #[cfg(feature = "luau")]
@@ -368,6 +369,44 @@ fn function_args_multiline_heuristic(
     singleline_shape.add_width(PAREN_LEN).over_budget()
 }
 
+fn should_preserve_function_arguments_hug(
+    ctx: &Context,
+    parentheses: &ContainedSpan,
+    arguments: &Punctuated<Expression>,
+) -> bool {
+    if !ctx.should_preserve_input_simple_statements()
+        || arguments.is_empty()
+        || !arguments
+            .iter()
+            .any(|argument| matches!(argument, Expression::Function(_)))
+    {
+        return false;
+    }
+
+    let (start_parens, end_parens) = parentheses.tokens();
+    match (
+        start_parens.end_position(),
+        arguments
+            .first()
+            .and_then(|argument| argument.value().start_position()),
+        arguments
+            .last()
+            .and_then(|argument| argument.value().end_position()),
+        end_parens.start_position(),
+    ) {
+        (
+            Some(start_parens_end),
+            Some(first_arg_start),
+            Some(last_arg_end),
+            Some(end_parens_start),
+        ) => {
+            start_parens_end.line() == first_arg_start.line()
+                && last_arg_end.line() == end_parens_start.line()
+        }
+        _ => false,
+    }
+}
+
 /// Formats a singular argument in a [`FunctionArgs`] node, in a multiline fashion
 fn format_argument_multiline(ctx: &Context, argument: &Expression, shape: Shape) -> Expression {
     // First format the argument assuming infinite width
@@ -457,9 +496,12 @@ pub fn format_function_args(
 
             // If there is a comment present anywhere in between the start parentheses and end parentheses, we should keep it multiline
             let force_mutliline = function_args_contains_comments(parentheses, arguments);
+            let preserve_function_arguments_hug =
+                should_preserve_function_arguments_hug(ctx, parentheses, arguments);
 
-            let is_multiline =
-                force_mutliline || function_args_multiline_heuristic(ctx, arguments, shape);
+            let is_multiline = force_mutliline
+                || (!preserve_function_arguments_hug
+                    && function_args_multiline_heuristic(ctx, arguments, shape));
 
             // Handle special case: we want to go multiline, but we have a single argument which is a table constructor
             // In this case, we want to hug the table braces with the parentheses.
@@ -763,6 +805,13 @@ fn block_contains_nested_function(block: &Block) -> bool {
     }
 }
 
+fn node_spans_single_line(node: &impl Node) -> bool {
+    match (node.start_position(), node.end_position()) {
+        (Some(start), Some(end)) => start.line() == end.line(),
+        _ => false,
+    }
+}
+
 pub fn should_collapse_function_body(ctx: &Context, function_body: &FunctionBody) -> bool {
     // Test for presence of any comments
     let require_multiline_function = function_body
@@ -777,10 +826,16 @@ pub fn should_collapse_function_body(ctx: &Context, function_body: &FunctionBody
             .any(trivia_util::trivia_is_comment)
         || trivia_util::contains_comments(function_body.block());
 
+    let input_single_line = node_spans_single_line(function_body);
+    let preserve_input_simple_statements = ctx.should_preserve_input_simple_statements();
+    let should_collapse_empty_function = !preserve_input_simple_statements || input_single_line;
+    let should_collapse_simple_function = ctx.should_collapse_simple_functions()
+        || (preserve_input_simple_statements && input_single_line);
+
     !require_multiline_function
-        && (trivia_util::is_block_empty(function_body.block())
+        && ((trivia_util::is_block_empty(function_body.block()) && should_collapse_empty_function)
             || (trivia_util::is_block_simple(function_body.block())
-                && ctx.should_collapse_simple_functions()
+                && should_collapse_simple_function
                 && !block_contains_nested_function(function_body.block())))
 }
 
@@ -794,6 +849,8 @@ pub fn format_function_body(
     let leading_trivia = vec![create_indent_trivia(ctx, shape)];
 
     let should_collapse = should_collapse_function_body(ctx, function_body);
+    let preserve_input_singleline_function =
+        ctx.should_preserve_input_simple_statements() && node_spans_single_line(function_body);
 
     // Check if the parameters should be placed across multiple lines
     let multiline_params = {
@@ -817,7 +874,8 @@ pub fn format_function_body(
         });
 
         contains_comments
-            || should_parameters_format_multiline(ctx, function_body, shape, should_collapse)
+            || (!preserve_input_singleline_function
+                && should_parameters_format_multiline(ctx, function_body, shape, should_collapse))
     };
 
     // Format the function body block on a single line if its empty, or it is "simple" (and the option has been enabled)
@@ -899,7 +957,8 @@ pub fn format_function_body(
             };
 
             // If the block forces multiline or goes over width, then bail out of singleline formatting and format multiline
-            if block_shape.take_first_line(&block).over_budget()
+            if (!preserve_input_singleline_function
+                && block_shape.take_first_line(&block).over_budget())
                 || trivia_util::spans_multiple_lines(&block)
             {
                 singleline_function = false;

src/formatters/stmt.rs

@@ -40,6 +40,7 @@ use full_moon::{
         punctuated::Punctuated, Block, Call, Do, ElseIf, Expression, FunctionArgs, FunctionCall,
         GenericFor, If, NumericFor, Repeat, Stmt, Suffix, While,
     },
+    node::Node,
     tokenizer::{Token, TokenKind, TokenReference, TokenType},
 };
 
@@ -429,6 +430,13 @@ fn is_if_guard(if_node: &If) -> bool {
         && !trivia_util::contains_comments(if_node.then_token())
 }
 
+fn node_spans_single_line(node: &impl Node) -> bool {
+    match (node.start_position(), node.end_position()) {
+        (Some(start), Some(end)) => start.line() == end.line(),
+        _ => false,
+    }
+}
+
 /// Format an If node
 pub fn format_if(ctx: &Context, if_node: &If, shape: Shape) -> If {
     const IF_LEN: usize = "if ".len();
@@ -444,21 +452,24 @@ pub fn format_if(ctx: &Context, if_node: &If, shape: Shape) -> If {
     let singleline_if_token = fmt_symbol!(ctx, if_node.if_token(), "if ", shape);
     let singleline_condition = format_expression(ctx, &condition, shape + IF_LEN + THEN_LEN);
     let singleline_then_token = fmt_symbol!(ctx, if_node.then_token(), " then", shape);
+    let preserve_input_singleline_conditional =
+        ctx.should_preserve_input_simple_statements() && node_spans_single_line(if_node);
 
     // Determine if we need to hang the condition
     let singleline_shape =
         shape + (IF_LEN + THEN_LEN + strip_trivia(&singleline_condition).to_string().len());
-    let require_multiline_expression = singleline_shape.over_budget()
+    let require_multiline_expression = (!preserve_input_singleline_conditional
+        && singleline_shape.over_budget())
         || if_node.if_token().has_trailing_comments(CommentSearch::All)
         || if_node
             .then_token()
             .has_leading_comments(CommentSearch::All)
         || trivia_util::contains_comments(&condition);
 
-    if !require_multiline_expression
-        && ctx.should_collapse_simple_conditionals()
-        && is_if_guard(if_node)
-    {
+    let should_collapse_simple_conditional =
+        ctx.should_collapse_simple_conditionals() || preserve_input_singleline_conditional;
+
+    if !require_multiline_expression && should_collapse_simple_conditional && is_if_guard(if_node) {
         // Rather than deferring to `format_block()`, since we know that there is only a single Stmt or LastStmt in the block, we can format it immediately
         // We need to modify the formatted LastStmt, since it will have automatically added leading/trailing trivia we don't want
         // We assume that there is only a laststmt present in the block - the callee of this function should have already checked for this
@@ -503,9 +514,10 @@ pub fn format_if(ctx: &Context, if_node: &If, shape: Shape) -> If {
             .with_end_token(end_token);
 
         // See if it fits under the column width. If it does, bail early and return this singleline if
-        if !shape
-            .add_width(strip_trivia(&singleline_if).to_string().len())
-            .over_budget()
+        if preserve_input_singleline_conditional
+            || !shape
+                .add_width(strip_trivia(&singleline_if).to_string().len())
+                .over_budget()
         {
             return singleline_if;
         }

src/lib.rs

@@ -145,6 +145,8 @@ pub enum CollapseSimpleStatement {
     /// Never collapse
     #[default]
     Never,
+    /// Preserve whether simple statements were collapsed in the input
+    Input,
     /// Collapse simple functions onto a single line
     FunctionOnly,
     /// Collapse simple if guards onto a single line
@@ -264,7 +266,8 @@ pub struct Config {
     /// function is called with only one table or string argument (same as no_call_parentheses).
     pub call_parentheses: CallParenType,
     /// Whether we should collapse simple structures like functions or guard statements
-    /// if set to [`CollapseSimpleStatement::None`] structures are never collapsed.
+    /// if set to [`CollapseSimpleStatement::Never`] structures are never collapsed.
+    /// if set to [`CollapseSimpleStatement::Input`] simple structures preserve their input shape.
     /// if set to [`CollapseSimpleStatement::FunctionOnly`] then simple functions (i.e., functions with a single laststmt) can be collapsed
     pub collapse_simple_statement: CollapseSimpleStatement,
     /// Whether we should allow blocks to preserve leading and trailing newline gaps.