From da686774dc72f191b81603684541631e8a2007a1 Mon Sep 17 00:00:00 2001
From: Russell Davis <551404+russelldavis@users.noreply.github.com>
Date: Thu, 25 May 2023 13:13:53 -0700
Subject: [PATCH] Formatter: fix indentation of comments at end of range
 (#54250)

---
 src/services/formatting/formatting.ts         | 19 ++++++++++++-------
 .../fourslash/formatSelectionWithTrivia3.ts   | 11 +++++++++++
 .../fourslash/formatSelectionWithTrivia4.ts   | 12 ++++++++++++
 .../fourslash/formatSelectionWithTrivia5.ts   | 12 ++++++++++++
 .../fourslash/formatSelectionWithTrivia6.ts   | 10 ++++++++++
 .../fourslash/formatSelectionWithTrivia7.ts   | 11 +++++++++++
 .../fourslash/formatSelectionWithTrivia8.ts   | 11 +++++++++++
 7 files changed, 79 insertions(+), 7 deletions(-)
 create mode 100644 tests/cases/fourslash/formatSelectionWithTrivia3.ts
 create mode 100644 tests/cases/fourslash/formatSelectionWithTrivia4.ts
 create mode 100644 tests/cases/fourslash/formatSelectionWithTrivia5.ts
 create mode 100644 tests/cases/fourslash/formatSelectionWithTrivia6.ts
 create mode 100644 tests/cases/fourslash/formatSelectionWithTrivia7.ts
 create mode 100644 tests/cases/fourslash/formatSelectionWithTrivia8.ts

diff --git a/src/services/formatting/formatting.ts b/src/services/formatting/formatting.ts
index cce023cc0e5..0f066e4807e 100644
--- a/src/services/formatting/formatting.ts
+++ b/src/services/formatting/formatting.ts
@@ -515,17 +515,22 @@ function formatSpanWorker(
         processNode(enclosingNode, enclosingNode, startLine, undecoratedStartLine, initialIndentation, delta);
     }
 
-    if (!formattingScanner.isOnToken()) {
+    // Leading trivia items get attached to and processed with the token that proceeds them. If the
+    // range ends in the middle of some leading trivia, the token that proceeds them won't be in the
+    // range and thus won't get processed. So we process those remaining trivia items here.
+    const remainingTrivia = formattingScanner.getCurrentLeadingTrivia();
+    if (remainingTrivia) {
         const indentation = SmartIndenter.nodeWillIndentChild(options, enclosingNode, /*child*/ undefined, sourceFile, /*indentByDefault*/ false)
             ? initialIndentation + options.indentSize!
             : initialIndentation;
-        const leadingTrivia = formattingScanner.getCurrentLeadingTrivia();
-        if (leadingTrivia) {
-            indentTriviaItems(leadingTrivia, indentation, /*indentNextTokenOrTrivia*/ false,
-                item => processRange(item, sourceFile.getLineAndCharacterOfPosition(item.pos), enclosingNode, enclosingNode, /*dynamicIndentation*/ undefined!));
-            if (options.trimTrailingWhitespace !== false) {
-                trimTrailingWhitespacesForRemainingRange(leadingTrivia);
+        indentTriviaItems(remainingTrivia, indentation, /*indentNextTokenOrTrivia*/ true,
+            item => {
+                processRange(item, sourceFile.getLineAndCharacterOfPosition(item.pos), enclosingNode, enclosingNode, /*dynamicIndentation*/ undefined!);
+                insertIndentation(item.pos, indentation, /*lineAdded*/ false);
             }
+        );
+        if (options.trimTrailingWhitespace !== false) {
+            trimTrailingWhitespacesForRemainingRange(remainingTrivia);
         }
     }
 
diff --git a/tests/cases/fourslash/formatSelectionWithTrivia3.ts b/tests/cases/fourslash/formatSelectionWithTrivia3.ts
new file mode 100644
index 00000000000..6aa0332e9b4
--- /dev/null
+++ b/tests/cases/fourslash/formatSelectionWithTrivia3.ts
@@ -0,0 +1,11 @@
+/// <reference path="fourslash.ts" />
+
+// Tests comment indentation with range ending before next token (end block)
+
+////if (true) {
+/////*begin*/// test comment
+/////*end*/}
+
+format.selection('begin', 'end');
+
+verify.currentFileContentIs("if (true) {\n    // test comment\n}");
diff --git a/tests/cases/fourslash/formatSelectionWithTrivia4.ts b/tests/cases/fourslash/formatSelectionWithTrivia4.ts
new file mode 100644
index 00000000000..523db34ed6d
--- /dev/null
+++ b/tests/cases/fourslash/formatSelectionWithTrivia4.ts
@@ -0,0 +1,12 @@
+/// <reference path="fourslash.ts" />
+
+// Tests comment indentation with range ending before next token (statement)
+
+////if (true) {
+/////*begin*/// test comment
+/////*end*/console.log();
+////}
+
+format.selection('begin', 'end');
+
+verify.currentFileContentIs("if (true) {\n    // test comment\nconsole.log();\n}");
diff --git a/tests/cases/fourslash/formatSelectionWithTrivia5.ts b/tests/cases/fourslash/formatSelectionWithTrivia5.ts
new file mode 100644
index 00000000000..8b39383c5cd
--- /dev/null
+++ b/tests/cases/fourslash/formatSelectionWithTrivia5.ts
@@ -0,0 +1,12 @@
+/// <reference path="fourslash.ts" />
+
+// Tests comment indentation with range ending before next token (statement w/ preceding whitespace)
+
+////if (true) {
+/////*begin*/// test comment
+/////*end*/    console.log();
+////}
+
+format.selection('begin', 'end');
+
+verify.currentFileContentIs("if (true) {\n    // test comment\n    console.log();\n}");
diff --git a/tests/cases/fourslash/formatSelectionWithTrivia6.ts b/tests/cases/fourslash/formatSelectionWithTrivia6.ts
new file mode 100644
index 00000000000..475052b9620
--- /dev/null
+++ b/tests/cases/fourslash/formatSelectionWithTrivia6.ts
@@ -0,0 +1,10 @@
+/// <reference path="fourslash.ts" />
+
+// Tests comment indentation with range ending before next token (end of file)
+
+/////*begin*/    // test comment
+/////*end*/
+
+format.selection('begin', 'end');
+
+verify.currentFileContentIs("// test comment\n");
diff --git a/tests/cases/fourslash/formatSelectionWithTrivia7.ts b/tests/cases/fourslash/formatSelectionWithTrivia7.ts
new file mode 100644
index 00000000000..91c04b6051e
--- /dev/null
+++ b/tests/cases/fourslash/formatSelectionWithTrivia7.ts
@@ -0,0 +1,11 @@
+/// <reference path="fourslash.ts" />
+
+// Tests comment indentation with range ending before next token (same line)
+
+////if (true) {
+/////*begin*/// test comment/*end*/
+////}
+
+format.selection('begin', 'end');
+
+verify.currentFileContentIs("if (true) {\n    // test comment\n}");
diff --git a/tests/cases/fourslash/formatSelectionWithTrivia8.ts b/tests/cases/fourslash/formatSelectionWithTrivia8.ts
new file mode 100644
index 00000000000..1107bd655fc
--- /dev/null
+++ b/tests/cases/fourslash/formatSelectionWithTrivia8.ts
@@ -0,0 +1,11 @@
+/// <reference path="fourslash.ts" />
+
+// Same as formatSelectionWithTrivia2, but range is immediately proceeded by a token
+
+/////*begin*/;
+////    
+/////*end*/console.log();
+
+format.selection('begin', 'end');
+
+verify.currentFileContentIs(";\n\nconsole.log();");