Merge PR #242: Move AliasAS processing out of formatting phase

nene · web-flow · commit 1bb5d230efb8 · 2022-06-20T16:23:07.000+03:00
diff --git a/src/core/AliasAs.ts b/src/core/AliasAs.ts
@@ -1,20 +1,66 @@
-import type { AliasMode } from 'src/types';
+import type { AliasMode, FormatOptions } from 'src/types';
+import AsTokenFactory from './AsTokenFactory';
 
-import { isCommand, isToken, type Token, TokenType } from './token';
+import { isCommand, isToken, type Token, TokenType, EOF_TOKEN, isReserved } from './token';
 
-export interface TokenStream {
-  isWithinSelect(): boolean;
-  getPreviousReservedToken(): Token;
-  tokenLookBehind(n?: number): Token;
-  tokenLookAhead(n?: number): Token;
-}
-
-/** Decides addition and removal of AS tokens */
+/** Adds and removes AS tokens as configured by aliasAs option */
 export default class AliasAs {
-  constructor(private aliasAs: AliasMode, private formatter: TokenStream) {}
+  private index = 0;
+  private tokens: Token[] = [];
+  private previousReservedToken: Token = EOF_TOKEN;
+  private previousCommandToken: Token = EOF_TOKEN;
+  private asTokenFactory: AsTokenFactory;
+  private aliasAs: AliasMode;
+
+  constructor(cfg: FormatOptions, tokens: Token[]) {
+    this.aliasAs = cfg.aliasAs;
+    this.asTokenFactory = new AsTokenFactory(cfg.keywordCase, tokens);
+    this.tokens = tokens;
+  }
+
+  /** Returns tokens with AS tokens added/removed as needed */
+  public process(): Token[] {
+    const processedTokens: Token[] = [];
+
+    for (this.index = 0; this.index < this.tokens.length; this.index++) {
+      const token = this.tokens[this.index];
+
+      if (isReserved(token)) {
+        this.previousReservedToken = token;
+        if (token.type === TokenType.RESERVED_COMMAND) {
+          this.previousCommandToken = token;
+        }
+      }
+
+      if (isToken.AS(token)) {
+        if (!this.shouldRemove()) {
+          processedTokens.push(token);
+        }
+      } else if (
+        token.type === TokenType.IDENT ||
+        token.type === TokenType.NUMBER ||
+        token.type === TokenType.STRING ||
+        token.type === TokenType.VARIABLE
+      ) {
+        if (this.shouldAddBefore(token)) {
+          processedTokens.push(this.asTokenFactory.token());
+        }
+
+        processedTokens.push(token);
+
+        if (this.shouldAddAfter()) {
+          processedTokens.push(this.asTokenFactory.token());
+        }
+      } else {
+        processedTokens.push(token);
+      }
+    }
+
+    return processedTokens;
+  }
 
   /** True when AS keyword should be added *before* current token */
-  public shouldAddBefore(token: Token): boolean {
+  private shouldAddBefore(token: Token): boolean {
     return this.isMissingTableAlias(token) || this.isMissingSelectColumnAlias(token);
   }
 
@@ -31,7 +77,7 @@ export default class AliasAs {
     const nextToken = this.lookAhead();
     return (
       (this.aliasAs === 'always' || this.aliasAs === 'select') &&
-      this.formatter.isWithinSelect() &&
+      this.isWithinSelect() &&
       token.type === TokenType.IDENT &&
       (isToken.END(prevToken) ||
         ((prevToken.type === TokenType.IDENT || prevToken.type === TokenType.NUMBER) &&
@@ -40,16 +86,16 @@ export default class AliasAs {
   }
 
   /** True when AS keyword should be added *after* current token */
-  public shouldAddAfter(): boolean {
+  private shouldAddAfter(): boolean {
     return this.isEdgeCaseCTE() || this.isEdgeCaseCreateTable() || this.isMissingTypeCastAs();
   }
 
   // checks for CAST(«expression» [AS] type)
   private isMissingTypeCastAs(): boolean {
     return (
       this.aliasAs === 'never' &&
-      this.formatter.isWithinSelect() &&
-      isToken.CAST(this.formatter.getPreviousReservedToken()) &&
+      this.isWithinSelect() &&
+      isToken.CAST(this.getPreviousReservedToken()) &&
       isToken.AS(this.lookAhead()) &&
       (this.lookAhead(2).type === TokenType.IDENT ||
         this.lookAhead(2).type === TokenType.RESERVED_KEYWORD) &&
@@ -79,23 +125,31 @@ export default class AliasAs {
   }
 
   /* True when the current AS token should be discarded */
-  public shouldRemove(): boolean {
+  private shouldRemove(): boolean {
     return this.aliasAs === 'never' || (this.aliasAs === 'select' && this.isRemovableNonSelectAs());
   }
 
   private isRemovableNonSelectAs(): boolean {
     return (
       this.lookBehind().value === ')' && // ) [AS] alias but not SELECT (a) [AS] alpha
-      !this.formatter.isWithinSelect() &&
+      !this.isWithinSelect() &&
       this.lookAhead().value !== '(' // skip WITH foo [AS] ( ...
     );
   }
 
-  private lookBehind(n?: number): Token {
-    return this.formatter.tokenLookBehind(n);
+  public getPreviousReservedToken(): Token {
+    return this.previousReservedToken;
+  }
+
+  public isWithinSelect(): boolean {
+    return isToken.SELECT(this.previousCommandToken);
+  }
+
+  private lookBehind(n = 1): Token {
+    return this.lookAhead(-n);
   }
 
-  private lookAhead(n?: number): Token {
-    return this.formatter.tokenLookAhead(n);
+  private lookAhead(n = 1): Token {
+    return this.tokens[this.index + n] || EOF_TOKEN;
   }
 }
diff --git a/src/core/Formatter.ts b/src/core/Formatter.ts
@@ -3,11 +3,10 @@ import Params from './Params';
 import Tokenizer from './Tokenizer';
 import formatCommaPositions from './formatCommaPositions';
 import formatAliasPositions from './formatAliasPositions';
-import AsTokenFactory from './AsTokenFactory';
 import Parser, { type Statement } from './Parser';
 import StatementFormatter from './StatementFormatter';
-import { type Token } from './token';
 import { indentString } from './config';
+import AliasAs from './AliasAs';
 
 /** Main formatter class that produces a final output string from list of tokens */
 export default class Formatter {
@@ -44,18 +43,17 @@ export default class Formatter {
    */
   public format(query: string): string {
     const tokens = this.cachedTokenizer().tokenize(query);
-    const ast = new Parser(tokens).parse();
-    const formattedQuery = this.formatAst(ast, tokens);
+    const processedTokens = new AliasAs(this.cfg, tokens).process();
+    const ast = new Parser(processedTokens).parse();
+    const formattedQuery = this.formatAst(ast);
     const finalQuery = this.postFormat(formattedQuery);
 
     return finalQuery.trimEnd();
   }
 
-  private formatAst(statements: Statement[], tokens: Token[]): string {
-    const asTokenFactory = new AsTokenFactory(this.cfg.keywordCase, tokens);
-
+  private formatAst(statements: Statement[]): string {
     return statements
-      .map(stat => new StatementFormatter(this.cfg, this.params, asTokenFactory).format(stat))
+      .map(stat => new StatementFormatter(this.cfg, this.params).format(stat))
       .join('\n'.repeat(this.cfg.linesBetweenQueries + 1));
   }
 
diff --git a/src/core/StatementFormatter.ts b/src/core/StatementFormatter.ts
@@ -6,8 +6,6 @@ import InlineBlock from './InlineBlock';
 import Params from './Params';
 import { isReserved, isCommand, isToken, type Token, TokenType, EOF_TOKEN } from './token';
 import toTabularFormat from './tabularStyle';
-import AliasAs from './AliasAs';
-import AsTokenFactory from './AsTokenFactory';
 import { type Statement } from './Parser';
 import { indentString, isTabularStyle } from './config';
 import WhitespaceBuilder, { WS } from './WhitespaceBuilder';
@@ -17,9 +15,7 @@ export default class StatementFormatter {
   private cfg: FormatOptions;
   private indentation: Indentation;
   private inlineBlock: InlineBlock;
-  private aliasAs: AliasAs;
   private params: Params;
-  private asTokenFactory: AsTokenFactory;
   private query: WhitespaceBuilder;
 
   private currentNewline = true;
@@ -28,13 +24,11 @@ export default class StatementFormatter {
   private tokens: Token[] = [];
   private index = -1;
 
-  constructor(cfg: FormatOptions, params: Params, asTokenFactory: AsTokenFactory) {
+  constructor(cfg: FormatOptions, params: Params) {
     this.cfg = cfg;
     this.indentation = new Indentation(indentString(cfg));
     this.inlineBlock = new InlineBlock(this.cfg.expressionWidth);
-    this.aliasAs = new AliasAs(this.cfg.aliasAs, this);
     this.params = params;
-    this.asTokenFactory = asTokenFactory;
     this.query = new WhitespaceBuilder(this.indentation);
   }
 
@@ -99,18 +93,10 @@ export default class StatementFormatter {
   }
 
   /**
-   * Formats word tokens + any potential AS tokens for aliases
+   * Formats ident/string/number/variable tokens
    */
   private formatWord(token: Token) {
-    if (this.aliasAs.shouldAddBefore(token)) {
-      this.query.add(this.show(this.asTokenFactory.token()), WS.SPACE);
-    }
-
     this.query.add(this.show(token), WS.SPACE);
-
-    if (this.aliasAs.shouldAddAfter()) {
-      this.query.add(this.show(this.asTokenFactory.token()), WS.SPACE);
-    }
   }
 
   /**
@@ -230,13 +216,9 @@ export default class StatementFormatter {
   }
 
   /**
-   * Formats a Reserved Keyword onto query, skipping AS if disabled
+   * Formats a Reserved Keyword onto query
    */
   private formatKeyword(token: Token) {
-    if (isToken.AS(token) && this.aliasAs.shouldRemove()) {
-      return;
-    }
-
     this.query.add(this.show(token), WS.SPACE);
   }
 
@@ -439,22 +421,22 @@ export default class StatementFormatter {
   }
 
   /** Returns the latest encountered reserved keyword token */
-  public getPreviousReservedToken(): Token {
+  private getPreviousReservedToken(): Token {
     return this.previousReservedToken;
   }
 
   /** True when currently within SELECT command */
-  public isWithinSelect(): boolean {
+  private isWithinSelect(): boolean {
     return isToken.SELECT(this.previousCommandToken);
   }
 
   /** Fetches nth previous token from the token stream */
-  public tokenLookBehind(n = 1): Token {
+  private tokenLookBehind(n = 1): Token {
     return this.tokens[this.index - n] || EOF_TOKEN;
   }
 
   /** Fetches nth next token from the token stream */
-  public tokenLookAhead(n = 1): Token {
+  private tokenLookAhead(n = 1): Token {
     return this.tokens[this.index + n] || EOF_TOKEN;
   }
 }
