Add option to preserve newline gaps for blocks

README.md

@@ -293,6 +293,7 @@ StyLua only offers the following options:
 | `quote_style`                | `AutoPreferDouble` | Quote style for string literals. Possible options: `AutoPreferDouble`, `AutoPreferSingle`, `ForceDouble`, `ForceSingle`. `AutoPrefer` styles will prefer the specified quote style, but fall back to the alternative if it has fewer string escapes. `Force` styles always use the specified style regardless of escapes.                                                                                                                                                                                                                                                                                                                                                                                                                             |
 | `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`                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              |
 
 Default `stylua.toml`, note you do not need to explicitly specify each option if you want to use the defaults:
@@ -307,6 +308,7 @@ quote_style = "AutoPreferDouble"
 call_parentheses = "Always"
 collapse_simple_statement = "Never"
 space_after_function_names = "Never"
+block_newline_gaps = "Never"
 
 [sort_requires]
 enabled = false

src/cli/opt.rs

@@ -1,8 +1,8 @@
 use clap::{ArgEnum, StructOpt};
 use std::path::PathBuf;
 use stylua_lib::{
-    CallParenType, CollapseSimpleStatement, IndentType, LineEndings, LuaVersion, QuoteStyle,
-    SpaceAfterFunctionNames,
+    BlockNewlineGaps, CallParenType, CollapseSimpleStatement, IndentType, LineEndings, LuaVersion,
+    QuoteStyle, SpaceAfterFunctionNames,
 };
 
 lazy_static::lazy_static! {
@@ -186,6 +186,9 @@ pub struct FormatOpts {
     /// Specify whether to collapse simple statements.
     #[structopt(long, arg_enum, ignore_case = true)]
     pub collapse_simple_statement: Option<ArgCollapseSimpleStatement>,
+    /// Specify whether to preserve leading and trailing newline gaps for blocks.
+    #[structopt(long, arg_enum, ignore_case = true)]
+    pub preserve_block_newline_gaps: Option<ArgBlockNewlineGaps>,
     /// Enable requires sorting
     #[structopt(long)]
     pub sort_requires: bool,
@@ -272,6 +275,11 @@ convert_enum!(CollapseSimpleStatement, ArgCollapseSimpleStatement, {
     Always,
 });
 
+convert_enum!(BlockNewlineGaps, ArgBlockNewlineGaps, {
+    Never,
+    Preserve,
+});
+
 convert_enum!(SpaceAfterFunctionNames, ArgSpaceAfterFunctionNames, {
     Never,
     Definitions,

src/context.rs

@@ -1,6 +1,6 @@
 use crate::{
-    shape::Shape, CallParenType, CollapseSimpleStatement, Config, IndentType, LineEndings,
-    Range as FormatRange, SpaceAfterFunctionNames,
+    shape::Shape, BlockNewlineGaps, CallParenType, CollapseSimpleStatement, Config, IndentType,
+    LineEndings, Range as FormatRange, SpaceAfterFunctionNames,
 };
 use full_moon::{
     node::Node,
@@ -154,6 +154,14 @@ impl Context {
             CollapseSimpleStatement::ConditionalOnly | CollapseSimpleStatement::Always
         )
     }
+
+    pub fn should_preserve_leading_block_newline_gaps(&self) -> bool {
+        matches!(self.config().block_newline_gaps, BlockNewlineGaps::Preserve)
+    }
+
+    pub fn should_preserve_trailing_block_newline_gaps(&self) -> bool {
+        matches!(self.config().block_newline_gaps, BlockNewlineGaps::Preserve)
+    }
 }
 
 /// Returns the relevant line ending string from the [`LineEndings`] enum

src/formatters/block.rs

@@ -535,7 +535,7 @@ fn check_stmt_requires_semicolon(
 pub fn format_block(ctx: &Context, block: &Block, shape: Shape) -> Block {
     let mut ctx = *ctx;
     let mut formatted_statements: Vec<(Stmt, Option<TokenReference>)> = Vec::new();
-    let mut found_first_stmt = false;
+    let mut remove_next_stmt_leading_newlines = !ctx.should_preserve_leading_block_newline_gaps();
     let mut stmt_iterator = block.stmts_with_semicolon().peekable();
 
     while let Some((stmt, semi)) = stmt_iterator.next() {
@@ -544,12 +544,12 @@ pub fn format_block(ctx: &Context, block: &Block, shape: Shape) -> Block {
         let shape = shape.reset();
         let mut stmt = format_stmt(&ctx, stmt, shape);
 
-        // If this is the first stmt, then remove any leading newlines
-        if !found_first_stmt {
+        // If this is the first stmt, and leading newlines should be removed, then remove them
+        if remove_next_stmt_leading_newlines {
             if let FormatNode::Normal = ctx.should_format_node(&stmt) {
                 stmt = stmt_remove_leading_newlines(stmt);
             }
-            found_first_stmt = true;
+            remove_next_stmt_leading_newlines = false;
         }
 
         // If we have a semicolon, we need to push all the trailing trivia from the statement
@@ -611,8 +611,9 @@ pub fn format_block(ctx: &Context, block: &Block, shape: Shape) -> Block {
 
             let shape = shape.reset();
             let mut last_stmt = format_last_stmt(&ctx, last_stmt, shape);
-            // If this is the first stmt, then remove any leading newlines
-            if !found_first_stmt && matches!(ctx.should_format_node(&last_stmt), FormatNode::Normal)
+            // If this is the first stmt, and leading newlines should be removed, then remove them
+            if remove_next_stmt_leading_newlines
+                && matches!(ctx.should_format_node(&last_stmt), FormatNode::Normal)
             {
                 last_stmt = last_stmt_remove_leading_newlines(last_stmt);
             }

src/formatters/general.rs

@@ -679,7 +679,7 @@ pub fn format_end_token(
     shape: Shape,
 ) -> TokenReference {
     // Indent any comments leading a token, as these comments are technically part of the function body block
-    let formatted_leading_trivia: Vec<Token> = load_token_trivia(
+    let mut formatted_leading_trivia: Vec<Token> = load_token_trivia(
         ctx,
         current_token.leading_trivia().collect(),
         FormatTokenType::LeadingTrivia,
@@ -698,37 +698,39 @@ pub fn format_end_token(
         shape,
     );
 
-    // Special case for block end tokens:
-    // We will reverse the leading trivia, and keep removing any newlines we find, until we find something else, then we stop.
-    // This is to remove unnecessary newlines at the end of the block.
-    let mut iter = formatted_leading_trivia.iter().rev().peekable();
-
-    let mut formatted_leading_trivia = Vec::new();
-    let mut stop_removal = false;
-    while let Some(x) = iter.next() {
-        match x.token_type() {
-            TokenType::Whitespace { ref characters } => {
-                if !stop_removal
-                    && characters.contains('\n')
-                    && !matches!(
-                        iter.peek().map(|x| x.token_kind()),
-                        Some(TokenKind::SingleLineComment) | Some(TokenKind::MultiLineComment)
-                    )
-                {
-                    continue;
-                } else {
+    if !ctx.should_preserve_trailing_block_newline_gaps() {
+        // Special case for block end tokens:
+        // We will reverse the leading trivia, and keep removing any newlines we find, until we find something else, then we stop.
+        // This is to remove unnecessary newlines at the end of the block.
+        let original_leading_trivia = std::mem::take(&mut formatted_leading_trivia);
+        let mut iter = original_leading_trivia.iter().cloned().rev().peekable();
+
+        let mut stop_removal = false;
+        while let Some(x) = iter.next() {
+            match x.token_type() {
+                TokenType::Whitespace { ref characters } => {
+                    if !stop_removal
+                        && characters.contains('\n')
+                        && !matches!(
+                            iter.peek().map(|x| x.token_kind()),
+                            Some(TokenKind::SingleLineComment) | Some(TokenKind::MultiLineComment)
+                        )
+                    {
+                        continue;
+                    } else {
+                        formatted_leading_trivia.push(x.to_owned());
+                    }
+                }
+                _ => {
                     formatted_leading_trivia.push(x.to_owned());
+                    stop_removal = true; // Stop removing newlines once we have seen some sort of comment
                 }
             }
-            _ => {
-                formatted_leading_trivia.push(x.to_owned());
-                stop_removal = true; // Stop removing newlines once we have seen some sort of comment
-            }
         }
-    }
 
-    // Need to reverse the vector since we reversed the iterator
-    formatted_leading_trivia.reverse();
+        // Need to reverse the vector since we reversed the iterator
+        formatted_leading_trivia.reverse();
+    }
 
     TokenReference::new(
         formatted_leading_trivia,

src/lib.rs

@@ -153,6 +153,19 @@ pub enum CollapseSimpleStatement {
     Always,
 }
 
+/// If blocks should be allowed to have leading and trailing newline gaps.
+#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, Deserialize)]
+#[cfg_attr(all(target_arch = "wasm32", feature = "wasm-bindgen"), wasm_bindgen)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "fromstr", derive(strum::EnumString))]
+pub enum BlockNewlineGaps {
+    /// Never allow leading or trailing newline gaps
+    #[default]
+    Never,
+    /// Preserve both leading and trailing newline gaps if present in input
+    Preserve,
+}
+
 /// An optional formatting range.
 /// If provided, only content within these boundaries (inclusive) will be formatted.
 /// Both boundaries are optional, and are given as byte offsets from the beginning of the file.
@@ -254,6 +267,10 @@ pub struct Config {
     /// if set to [`CollapseSimpleStatement::None`] structures are never collapsed.
     /// 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.
+    /// if set to [`BlockNewlineGaps::Never`] then newline gaps are never allowed at the start or end of blocks.
+    /// if set to [`BlockNewlineGaps::Preserve`] then newline gaps are preserved at the start and end of blocks.
+    pub block_newline_gaps: BlockNewlineGaps,
     /// Configuration for the sort requires codemod
     pub sort_requires: SortRequiresConfig,
     /// Whether we should include a space between the function name and arguments.
@@ -287,6 +304,7 @@ impl Default for Config {
             collapse_simple_statement: CollapseSimpleStatement::default(),
             sort_requires: SortRequiresConfig::default(),
             space_after_function_names: SpaceAfterFunctionNames::default(),
+            block_newline_gaps: BlockNewlineGaps::default(),
         }
     }
 }

tests/inputs-preserve-block-newline-gaps/block-empty-lines.lua

@@ -0,0 +1,28 @@
+function foo()
+
+	local x = 1
+
+
+	return true
+
+end
+
+function bar()
+
+
+	return
+
+
+end
+
+do
+
+	-- comment
+	local x = 1
+
+
+	local foo = bar
+
+	-- comment
+
+end

tests/inputs-preserve-block-newline-gaps/empty-function.lua

@@ -0,0 +1,11 @@
+local function noop() -- comment
+end
+
+function noop()
+	-- comment
+end
+
+call(function()
+	-- comment
+
+end)

tests/inputs-preserve-block-newline-gaps/long-elseif-chain.lua

@@ -0,0 +1,33 @@
+if
+	this_is == very_long_variable_name
+	and to_ensure_that == it_is_broken_into
+	and multiple_lines == in_order_to_see_how_that_looks
+then
+	return false
+elseif
+	this_is == very_long_variable_name
+	and to_ensure_that == it_is_broken_into
+	and multiple_lines == in_order_to_see_how_that_looks
+then
+	local only_a_gap_of_one_newline_is_preserved_below = 1
+	local hurray = true
+
+
+
+
+elseif
+	this_is == very_long_variable_name
+	and to_ensure_that == it_is_broken_into
+	and multiple_lines == in_order_to_see_how_that_looks
+then
+
+
+
+
+
+	local only_a_gap_of_one_newline_is_preserved_above = 1
+	local hurray = true
+else
+
+	return also_preserved_in_else_blocks
+end

tests/snapshots/tests__preserve_block_newline_gaps@block-empty-lines.lua.snap

@@ -0,0 +1,29 @@
+---
+source: tests/tests.rs
+expression: "format_code(&contents,\n        Config {\n            preserve_block_newline_gaps: PreserveBlockNewlineGaps::Always,\n            ..Config::default()\n        }, None, OutputVerification::None).unwrap()"
+input_file: tests/inputs-preserve-block-newline-gaps/block-empty-lines.lua
+---
+function foo()
+
+	local x = 1
+
+	return true
+
+end
+
+function bar()
+
+	return
+
+end
+
+do
+
+	-- comment
+	local x = 1
+
+	local foo = bar
+
+	-- comment
+
+end

tests/snapshots/tests__preserve_block_newline_gaps@empty-function.lua.snap

@@ -0,0 +1,16 @@
+---
+source: tests/tests.rs
+expression: "format_code(&contents,\n        Config {\n            preserve_block_newline_gaps: PreserveBlockNewlineGaps::Always,\n            ..Config::default()\n        }, None, OutputVerification::None).unwrap()"
+input_file: tests/inputs-preserve-block-newline-gaps/empty-function.lua
+---
+local function noop() -- comment
+end
+
+function noop()
+	-- comment
+end
+
+call(function()
+	-- comment
+
+end)

tests/snapshots/tests__preserve_block_newline_gaps@long-elseif-chain.lua.snap

@@ -0,0 +1,31 @@
+---
+source: tests/tests.rs
+expression: "format_code(&contents,\n        Config {\n            preserve_block_newline_gaps: PreserveBlockNewlineGaps::Always,\n            ..Config::default()\n        }, None, OutputVerification::None).unwrap()"
+input_file: tests/inputs-preserve-block-newline-gaps/long-elseif-chain.lua
+---
+if
+	this_is == very_long_variable_name
+	and to_ensure_that == it_is_broken_into
+	and multiple_lines == in_order_to_see_how_that_looks
+then
+	return false
+elseif
+	this_is == very_long_variable_name
+	and to_ensure_that == it_is_broken_into
+	and multiple_lines == in_order_to_see_how_that_looks
+then
+	local only_a_gap_of_one_newline_is_preserved_below = 1
+	local hurray = true
+
+elseif
+	this_is == very_long_variable_name
+	and to_ensure_that == it_is_broken_into
+	and multiple_lines == in_order_to_see_how_that_looks
+then
+
+	local only_a_gap_of_one_newline_is_preserved_above = 1
+	local hurray = true
+else
+
+	return also_preserved_in_else_blocks
+end