Handle else blocks and comments in unbuffered code

src/printer.ts

@@ -1823,8 +1823,7 @@ export class PugPrinter {
     return val.slice(0, -1);
   }
 
-  // Since every line is parsed independently, babel will throw a SyntaxError if the line of code is only valid when there is another statement after it. This is a hack to get babel to properly parse what would otherwise be an invalid standalone JS line (e.g., `if (foo)`)
-  private async formatRawCodeWithFallback(
+  private async formatRawCodeWithFallbackNoElse(
     val: string,
     useSemi: boolean,
   ): Promise<string> {
@@ -1844,14 +1843,38 @@ export class PugPrinter {
       // At this point, we know the SyntaxError is from the fact that there's no statement after our val's statement, implying we likely need a block after it. Using an empty block to get babel to parse it without affecting the code semantics.
       // Example: `if (foo)` is not valid JS on its own, but `if (foo) {}` is.
       try {
-        val = await this.formatRawCode(val + '{}', useSemi);
+        // Look for comments at the end of the line, since we have to insert the block in between the statement and the comment or else the block will potentially be commented out
+        const commentIndex: number = val.search(/\/(?:\/|\*).*$/);
+        if (commentIndex === -1) {
+          val += '{}';
+        } else {
+          val = `${val.slice(0, commentIndex)}{}${val.slice(commentIndex)}`;
+        }
+
+        val = await this.formatRawCode(val, useSemi);
+
+        /*
+        Strip out the empty block, which prettier has now formatted to split across two lines, and if there was a comment, it's now at the end of the second line. The first \s is to account for the space prettier inserted between the statement and the empty block.
+        Input:
+        `if (foo)   // comment`
 
-        // Dynamically find the index of the last `{` in the code, which is the start of the empty block, since it's been reformatted at this point and a newline has likely been added between (shouldn't rely on that behavior though)
-        // Also account for any newly-introduced whitespace right before the empty block
-        const cutoffIndex: number = val.search(/\s?{\s?}\s?$/);
-        if (cutoffIndex === -1) throw new Error('No empty block found');
+        Transformed:
+        `if (foo)   {}// comment`
 
-        return val.slice(0, cutoffIndex);
+        Initial formatted:
+        `if (foo) {
+         } // comment`
+
+        Final formatted:
+        `if (foo) // comment`
+        */
+        const ma: RegExpExecArray | null = /\s{\n?}(.*)$/.exec(val);
+        if (!ma) throw new Error('No follow-up block found');
+
+        const comment: string | undefined = ma[1];
+
+        val = val.slice(0, ma.index) + (comment ?? '');
+        return val.trim();
       } catch (secondError: unknown) {
         logger.debug('[PugPrinter] fallback format error', secondError);
         // throw original error since our fallback didn't work
@@ -1860,6 +1883,25 @@ export class PugPrinter {
     }
   }
 
+  // Since every line is parsed independently, babel will throw a SyntaxError if the line of code is only valid when there is another statement after it, or if the line starts with `else if` or `else`. This is a hack to get babel to properly parse what would otherwise be an invalid standalone JS line (e.g., `if (foo)`, `else if (bar)`, `else`)
+  private async formatRawCodeWithFallback(
+    val: string,
+    useSemi: boolean,
+  ): Promise<string> {
+    if (val.startsWith('else')) {
+      // If the code starts with `else`, then we can format the code without the `else` keyword, and then add it back onto the start.
+      // We can call the same helper function in each case, just with different inputs, so we can easily handle all `if`, `else if`, and `else` cases without having to write out each one.
+      const noElse: string = await this.formatRawCodeWithFallbackNoElse(
+        val.slice(4),
+        useSemi,
+      );
+      // `noElse` will either be an empty string or it will contain a comment. Now we just prepend `else` onto the start and trim in case `noElse` is empty
+      return `else ${noElse}`.trim();
+    } else {
+      return await this.formatRawCodeWithFallbackNoElse(val, useSemi);
+    }
+  }
+
   private async code(token: CodeToken): Promise<string> {
     let result: string = this.computedIndent;
     if (!token.mustEscape && token.buffer) {

tests/unbuffered-code/formatted.pug

@@ -1,5 +1,22 @@
-- if (obj?.name)
+- if (obj?.name) /* Block comment */
   p #{ obj.name } is the name, if it exists
+- if (obj?.name2) // Double-slash comment
+  p #{ obj.name2 } is the second name, if it exists
+- if (obj?.name3)
+  p #{ obj.name3 } is the third name, if it exists
+- else if (obj?.color) /* Comment */
+  p #{ obj.color } is the color, if it exists
+- else // Comment, plus the else has whitespace before it
+  p No name or color
+- if (obj?.color2)
+  p #{ obj.color2 } is the second color, if it exists
+- else
+  p Else without a comment
+- for (const k of obj?.items || []) // Comment
+  p #{ k }
 - for (const k of obj?.items || [])
   p #{ k }
+- let n = 2;
+- while (n--) // Comment!
+  p #{ n }
 p This should always appear

tests/unbuffered-code/unbuffered-code.test.ts

@@ -2,7 +2,7 @@ import { compareFiles } from 'tests/common';
 import { describe, expect, it } from 'vitest';
 
 describe('Unbuffered code', () => {
-  it('should handle JS statements where the isolated line is only parsable when there is another statement afterwards', async () => {
+  it('should handle JS statements where the isolated line is only parsable when there is another statement before/afterwards', async () => {
     const { expected, actual } = await compareFiles(import.meta.url);
     expect(actual).toBe(expected);
   });

tests/unbuffered-code/unformatted.pug

@@ -1,5 +1,22 @@
-- if (         obj?.name      )
+- if (         obj?.name      )      /* Block comment */
   p #{ obj.name } is the name, if it exists
+- if (         obj?.name2      )    // Double-slash comment
+  p #{ obj.name2 } is the second name, if it exists
+- if (         obj?.name3      )
+  p #{ obj.name3 } is the third name, if it exists
+- else if (   obj?.color   )   /* Comment */
+  p #{ obj.color } is the color, if it exists
+-       else     // Comment, plus the else has whitespace before it
+  p No name or color
+-  if (         obj?.color2      )
+  p #{ obj.color2 } is the second color, if it exists
+-    else
+  p Else without a comment
+- for (const k of        obj?.items         || [])    // Comment
+  p #{ k }
 - for (const k of        obj?.items         || [])
   p #{ k }
+- let n = 2;
+- while (   n--    )    // Comment!
+  p #{ n }
 p This should always appear