From f1831b31cfb530ce88efee6bfa6e687218a86037 Mon Sep 17 00:00:00 2001
From: Taco de Wolff <tacodewolff@gmail.com>
Date: Fri, 26 Jan 2024 15:53:47 -0300
Subject: [PATCH] JS: keep all bang-comments (including in expressions) and
 move to top of statement list

---
 js/ast.go        |  1 -
 js/ast_test.go   |  4 ++--
 js/parse.go      | 30 ++++++++++++++++++++++++------
 js/parse_test.go |  3 ++-
 4 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/js/ast.go b/js/ast.go
index 63f20d7..a097a01 100644
--- a/js/ast.go
+++ b/js/ast.go
@@ -435,7 +435,6 @@ func (n Comment) JS(w io.Writer) {
 	} else {
 		w.Write(n.Value)
 	}
-	w.Write([]byte("\n"))
 }
 
 // BlockStmt is a block statement.
diff --git a/js/ast_test.go b/js/ast_test.go
index e21f47a..6ed4a3e 100644
--- a/js/ast_test.go
+++ b/js/ast_test.go
@@ -119,10 +119,10 @@ func TestJS(t *testing.T) {
 		{"{}\n;", "{} ;"},
 		{"- - --3", "- - --3;"},
 		{"([,,])=>P", "([,,]) => { return P; };"},
-		{"(t)=>{//!\n}", "(t) => { //!  };"}, // space after //! is newline
+		{"(t)=>{//!\n}", "(t) => { //! };"}, // space after //! is newline
 		{"import();", "import();"},
 		{"0\n.k", "(0).k;"},
-		{"do//!\nwhile(1)", "do //! while (1);"}, // space after //! is newline
+		{"do//!\n; while(1)", "//! do; while (1);"}, // space after //! is newline
 	}
 
 	re := regexp.MustCompile("\n *")
diff --git a/js/parse.go b/js/parse.go
index 6e9174e..9bcd95c 100644
--- a/js/parse.go
+++ b/js/parse.go
@@ -27,6 +27,7 @@ type Parser struct {
 	in, await, yield, deflt, retrn bool
 	assumeArrowFunc                bool
 	allowDirectivePrologue         bool
+	comments                       []IStmt
 
 	stmtLevel int
 	exprLevel int
@@ -91,9 +92,22 @@ func Parse(r *parse.Input, o Options) (*AST, error) {
 func (p *Parser) next() {
 	p.prevLT = false
 	p.tt, p.data = p.l.Next()
-	for p.tt == WhitespaceToken || p.tt == LineTerminatorToken || (p.tt == CommentToken || p.tt == CommentLineTerminatorToken) && (p.exprLevel != 0 || len(p.data) < 3 || p.data[2] != '!') {
-		if p.tt == LineTerminatorToken || p.tt == CommentLineTerminatorToken {
+Loop:
+	for {
+		switch p.tt {
+		case WhitespaceToken:
+			// no-op
+		case LineTerminatorToken:
 			p.prevLT = true
+		case CommentToken, CommentLineTerminatorToken:
+			if 2 < len(p.data) && p.data[2] == '!' {
+				p.comments = append(p.comments, &Comment{p.data})
+			}
+			if p.tt == CommentLineTerminatorToken {
+				p.prevLT = true
+			}
+		default:
+			break Loop
 		}
 		p.tt, p.data = p.l.Next()
 	}
@@ -180,6 +194,10 @@ func (p *Parser) parseModule() (module BlockStmt) {
 	for {
 		switch p.tt {
 		case ErrorToken:
+			if 0 < len(p.comments) {
+				module.List = append(p.comments, module.List...)
+				p.comments = p.comments[:0]
+			}
 			return
 		case ImportToken:
 			p.next()
@@ -567,10 +585,6 @@ func (p *Parser) parseStmt(allowDeclaration bool) (stmt IStmt) {
 	case SemicolonToken:
 		stmt = &EmptyStmt{}
 		p.next()
-	case CommentToken, CommentLineTerminatorToken:
-		// bang comment
-		stmt = &Comment{p.data}
-		p.next()
 	case ErrorToken:
 		stmt = &EmptyStmt{}
 		return
@@ -635,6 +649,10 @@ func (p *Parser) parseStmtList(in string) (list []IStmt) {
 		}
 		list = append(list, p.parseStmt(true))
 	}
+	if 0 < len(p.comments) {
+		list = append(p.comments, list...)
+		p.comments = p.comments[:0]
+	}
 	return
 }
 
diff --git a/js/parse_test.go b/js/parse_test.go
index 8712461..171866e 100644
--- a/js/parse_test.go
+++ b/js/parse_test.go
@@ -24,7 +24,8 @@ func TestParse(t *testing.T) {
 		{"#!/usr/bin/env node", "Stmt(#!/usr/bin/env node)"},
 		{"/* comment */", ""},
 		{"/*! comment */", "Stmt(/*! comment */)"},
-		{"5 + /*! comment */ 6", "Stmt(5+6)"},
+		{"5 + /*! comment */ 6", "Stmt(/*! comment */) Stmt(5+6)"},
+		{"var a = /*! comment */ 6", "Stmt(/*! comment */) Decl(var Binding(a = 6))"},
 		{"{}", "Stmt({ })"},
 		{`"use strict"`, `Stmt("use strict")`},
 		{"var a = b;", "Decl(var Binding(a = b))"},