Add option to preserve newline gaps for blocks

src/cli/opt.rs

@@ -1,6 +1,9 @@
 use clap::{ArgEnum, StructOpt};
 use std::path::PathBuf;
-use stylua_lib::{CallParenType, CollapseSimpleStatement, IndentType, LineEndings, QuoteStyle};
+use stylua_lib::{
+    CallParenType, CollapseSimpleStatement, IndentType, LineEndings, PreserveBlockNewlineGaps,
+    QuoteStyle,
+};
 
 lazy_static::lazy_static! {
     static ref NUM_CPUS: String = num_cpus::get().to_string();
@@ -180,6 +183,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<ArgPreserveBlockNewlineGaps>,
     /// Enable requires sorting
     #[structopt(long)]
     pub sort_requires: bool,
@@ -250,6 +256,13 @@ convert_enum!(CollapseSimpleStatement, ArgCollapseSimpleStatement, {
     Always,
 });
 
+convert_enum!(PreserveBlockNewlineGaps, ArgPreserveBlockNewlineGaps, {
+    Never,
+    AlwaysLeading,
+    AlwaysTrailing,
+    Always,
+});
+
 #[cfg(test)]
 mod tests {
     use super::Opt;

src/context.rs

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

src/formatters/block.rs

@@ -452,7 +452,7 @@ fn last_stmt_remove_leading_newlines(last_stmt: LastStmt) -> LastStmt {
 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() {
@@ -461,12 +461,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;
         }
 
         // Need to check next statement if it is a function call, with a parameters expression as the prefix
@@ -563,8 +563,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

@@ -643,7 +643,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,
@@ -662,37 +662,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

@@ -94,6 +94,23 @@ 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 PreserveBlockNewlineGaps {
+    /// Never allow leading or trailing newline gaps
+    #[default]
+    Never,
+    /// Always preserve leading newline gaps if present in input
+    AlwaysLeading,
+    /// Always preserve trailing newline gaps if present in input
+    AlwaysTrailing,
+    /// Always preserve both leading and trailing newline gaps if present in input
+    Always,
+}
+
 /// 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.
@@ -176,6 +193,12 @@ 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 [`PreserveBlockNewlineGaps::Never`] then newline gaps are never allowed at the start or end of blocks.
+    /// if set to [`PreserveBlockNewlineGaps::AlwaysLeading`] then newline gaps are preserved at the start blocks.
+    /// if set to [`PreserveBlockNewlineGaps::AlwaysTrailing`] then newline gaps are preserved at the end of blocks.
+    /// if set to [`PreserveBlockNewlineGaps::Always`] then newline gaps are preserved at the start and end of blocks.
+    pub preserve_block_newline_gaps: PreserveBlockNewlineGaps,
     /// Configuration for the sort requires codemod
     pub sort_requires: SortRequiresConfig,
 }
@@ -346,6 +369,7 @@ impl Default for Config {
             call_parentheses: CallParenType::default(),
             collapse_simple_statement: CollapseSimpleStatement::default(),
             sort_requires: SortRequiresConfig::default(),
+            preserve_block_newline_gaps: PreserveBlockNewlineGaps::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