fix: ignore comments when post-processing tokens

src/languages/bigquery/bigquery.formatter.ts

@@ -1,6 +1,6 @@
 import { DialectOptions } from '../../dialect.js';
 import { expandPhrases } from '../../expandPhrases.js';
-import { EOF_TOKEN, isToken, Token, TokenType } from '../../lexer/token.js';
+import { isToken, prevNonCommentToken, Token, TokenType } from '../../lexer/token.js';
 import { functions } from './bigquery.functions.js';
 import { dataTypes, keywords } from './bigquery.keywords.js';
 
@@ -204,16 +204,14 @@ function postProcess(tokens: Token[]): Token[] {
 
 // Converts OFFSET token inside array from RESERVED_CLAUSE to RESERVED_FUNCTION_NAME
 // See: https://cloud.google.com/bigquery/docs/reference/standard-sql/functions-and-operators#array_subscript_operator
+// Comments between the "[" and OFFSET are skipped, so "arr[/* c */ OFFSET(0)]"
+// is recognized the same way as "arr[OFFSET(0)]".
 function detectArraySubscripts(tokens: Token[]) {
-  let prevToken = EOF_TOKEN;
-  return tokens.map(token => {
-    if (token.text === 'OFFSET' && prevToken.text === '[') {
-      prevToken = token;
+  return tokens.map((token, i) => {
+    if (token.text === 'OFFSET' && prevNonCommentToken(tokens, i).text === '[') {
       return { ...token, type: TokenType.RESERVED_FUNCTION_NAME };
-    } else {
-      prevToken = token;
-      return token;
     }
+    return token;
   });
 }
 

src/languages/clickhouse/clickhouse.formatter.ts

@@ -1,6 +1,12 @@
 import { DialectOptions } from '../../dialect.js';
 import { expandPhrases } from '../../expandPhrases.js';
-import { EOF_TOKEN, isToken, Token, TokenType } from '../../lexer/token.js';
+import {
+  isToken,
+  nextNonCommentToken,
+  prevNonCommentToken,
+  Token,
+  TokenType,
+} from '../../lexer/token.js';
 import { functions } from './clickhouse.functions.js';
 import { dataTypes, keywords } from './clickhouse.keywords.js';
 
@@ -303,8 +309,8 @@ export const clickhouse: DialectOptions = {
  */
 function postProcess(tokens: Token[]): Token[] {
   return tokens.map((token, i) => {
-    const nextToken = tokens[i + 1] || EOF_TOKEN;
-    const prevToken = tokens[i - 1] || EOF_TOKEN;
+    const nextToken = nextNonCommentToken(tokens, i);
+    const prevToken = prevNonCommentToken(tokens, i);
 
     // If we have queries like
     // > GRANT SELECT, INSERT ON db.table TO john

src/languages/mariadb/likeMariaDb.ts

@@ -1,15 +1,21 @@
-import { EOF_TOKEN, isToken, Token, TokenType } from '../../lexer/token.js';
+import {
+  isToken,
+  nextNonCommentToken,
+  prevNonCommentToken,
+  Token,
+  TokenType,
+} from '../../lexer/token.js';
 
 // Shared functionality used by all MariaDB-like SQL dialects.
 
 export function postProcess(tokens: Token[]) {
   return tokens.map((token, i) => {
-    const nextToken = tokens[i + 1] || EOF_TOKEN;
+    const nextToken = nextNonCommentToken(tokens, i);
     if (isToken.SET(token) && nextToken.text === '(') {
       // This is SET datatype, not SET statement
       return { ...token, type: TokenType.RESERVED_FUNCTION_NAME };
     }
-    const prevToken = tokens[i - 1] || EOF_TOKEN;
+    const prevToken = prevNonCommentToken(tokens, i);
     if (isToken.VALUES(token) && prevToken.text === '=') {
       // This is VALUES() function, not VALUES clause
       return { ...token, type: TokenType.RESERVED_FUNCTION_NAME };

src/languages/spark/spark.formatter.ts

@@ -1,6 +1,12 @@
 import { DialectOptions } from '../../dialect.js';
 import { expandPhrases } from '../../expandPhrases.js';
-import { EOF_TOKEN, isToken, Token, TokenType } from '../../lexer/token.js';
+import {
+  isToken,
+  nextNonCommentToken,
+  prevNonCommentToken,
+  Token,
+  TokenType,
+} from '../../lexer/token.js';
 import { dataTypes, keywords } from './spark.keywords.js';
 import { functions } from './spark.functions.js';
 
@@ -155,8 +161,8 @@ export const spark: DialectOptions = {
 
 function postProcess(tokens: Token[]) {
   return tokens.map((token, i) => {
-    const prevToken = tokens[i - 1] || EOF_TOKEN;
-    const nextToken = tokens[i + 1] || EOF_TOKEN;
+    const prevToken = prevNonCommentToken(tokens, i);
+    const nextToken = nextNonCommentToken(tokens, i);
 
     // [WINDOW](...)
     if (isToken.WINDOW(token) && nextToken.type === TokenType.OPEN_PAREN) {

src/lexer/token.ts

@@ -111,3 +111,29 @@ export const isReserved = (type: TokenType): boolean =>
 
 export const isLogicalOperator = (type: TokenType): boolean =>
   type === TokenType.AND || type === TokenType.OR || type === TokenType.XOR;
+
+/** Checks if token is a comment */
+export const isComment = (type: TokenType): boolean =>
+  type === TokenType.LINE_COMMENT ||
+  type === TokenType.BLOCK_COMMENT ||
+  type === TokenType.DISABLE_COMMENT;
+
+/** Returns the closest non-comment token preceding the given index, or EOF_TOKEN. */
+export const prevNonCommentToken = (tokens: Token[], index: number): Token => {
+  for (let i = index - 1; i >= 0; i--) {
+    if (!isComment(tokens[i].type)) {
+      return tokens[i];
+    }
+  }
+  return EOF_TOKEN;
+};
+
+/** Returns the closest non-comment token following the given index, or EOF_TOKEN. */
+export const nextNonCommentToken = (tokens: Token[], index: number): Token => {
+  for (let i = index + 1; i < tokens.length; i++) {
+    if (!isComment(tokens[i].type)) {
+      return tokens[i];
+    }
+  }
+  return EOF_TOKEN;
+};

test/behavesLikeMariaDbFormatter.ts

@@ -191,4 +191,22 @@ export default function behavesLikeMariaDbFormatter(format: FormatFn) {
       GRANT ALL ON *.* TO user2;
     `);
   });
+  // Comments between a function name and its parenthesis must not stop the
+  // dialect post-processing from recognizing the function.
+  it('formats VALUES() as a function with a comment before it', () => {
+    expect(format('UPDATE t SET a = /*x*/ VALUES(b);')).toBe(dedent`
+      UPDATE t
+      SET
+        a = /*x*/ VALUES(b);
+    `);
+  });
+
+  it('formats VALUES() as a function with a line comment before it', () => {
+    expect(format('UPDATE t SET a = --x\nVALUES(b);')).toBe(dedent`
+      UPDATE t
+      SET
+        a = --x
+        VALUES(b);
+    `);
+  });
 }

test/bigquery.test.ts

@@ -636,4 +636,27 @@ describe('BigQueryFormatter', () => {
       expect(format(input, { linesBetweenQueries: 0 })).toBe(input);
     });
   });
+  describe('BigQuery array subscripts with comments', () => {
+    // A comment between "[" and OFFSET used to leave OFFSET as a clause token,
+    // which made the parser throw a Parse error.
+    it('formats a block comment before OFFSET', () => {
+      expect(format('SELECT arr[/* c */ OFFSET(0)] FROM t;')).toBe(dedent`
+        SELECT
+          arr[/* c */ OFFSET(0)]
+        FROM
+          t;
+      `);
+    });
+
+    it('formats a line comment before OFFSET', () => {
+      expect(format('SELECT arr[-- c\nOFFSET(0)] FROM t;')).toBe(dedent`
+        SELECT
+          arr[ -- c
+            OFFSET(0)
+          ]
+        FROM
+          t;
+      `);
+    });
+  });
 });

test/clickhouse.test.ts

@@ -1737,4 +1737,47 @@ describe('ClickhouseFormatter', () => {
         tuple();
     `);
   });
+  describe('comments in post-processed statements', () => {
+    // Comments between tokens must not affect the post-processing that decides
+    // whether SELECT is a clause or a privilege keyword.
+    it('formats CHECK GRANT with a comment before SELECT', () => {
+      expect(format('CHECK GRANT /* c */ SELECT ON db.table')).toBe(dedent`
+        CHECK GRANT
+          /* c */
+          SELECT ON db.table
+      `);
+    });
+
+    // Comments must not stop SET( from being recognized as a function.
+    it('formats SET( as a function with a comment before the parens', () => {
+      expect(format('SELECT SET /* c */ (100) FROM t;')).toBe(dedent`
+        SELECT
+          SET/* c */ (100)
+        FROM
+          t;
+      `);
+    });
+
+    it('formats EXPLAIN AST with a comment before SELECT', () => {
+      expect(format('EXPLAIN AST /*x*/ SELECT 1;')).toBe(dedent`
+        EXPLAIN AST /*x*/ SELECT 1;
+      `);
+    });
+
+    it('formats EXPLAIN AST with a line comment before SELECT', () => {
+      expect(format('EXPLAIN AST --x\nSELECT 1;')).toBe(dedent`
+        EXPLAIN AST --x
+        SELECT 1;
+      `);
+    });
+
+    it('formats GRANT with a comment before SELECT', () => {
+      expect(format('GRANT /*x*/ SELECT ON db.table TO john;')).toBe(dedent`
+        GRANT
+          /*x*/
+          SELECT ON db.table
+        TO john;
+      `);
+    });
+  });
 });

test/mysql.test.ts

@@ -114,4 +114,12 @@ describe('MySqlFormatter', () => {
       DROP DEFAULT;
     `);
   });
+  describe('comments in post-processed statements', () => {
+    // Comments must not stop SET( from being recognized as a function.
+    it('formats SET( as a function with a comment before the parens', () => {
+      expect(format('CREATE TABLE t (a SET /* c */ (1,2));')).toBe(dedent`
+        CREATE TABLE t (a SET/* c */ (1, 2));
+      `);
+    });
+  });
 });

test/spark.test.ts

@@ -152,4 +152,26 @@ describe('SparkFormatter', () => {
       ALTER COLUMN FirstName COMMENT "new comment";
     `);
   });
+  describe('Spark WINDOW() with comments', () => {
+    // A comment between WINDOW and its parenthesis used to leave WINDOW a clause
+    // token, so it was formatted as a clause instead of a function call.
+    it('formats WINDOW() as a function with a block comment before the parens', () => {
+      expect(format('SELECT window /*x*/ (time) FROM tbl;')).toBe(dedent`
+        SELECT
+          window/*x*/ (time)
+        FROM
+          tbl;
+      `);
+    });
+
+    it('formats WINDOW() as a function with a line comment before the parens', () => {
+      expect(format('SELECT window --x\n(time) FROM tbl;')).toBe(dedent`
+        SELECT
+          window --x
+          (time)
+        FROM
+          tbl;
+      `);
+    });
+  });
 });