Cache block objects when used for minor performance boost

Signed-off-by: worksofliam <[email protected]>

Author: worksofliam

src/language/sql/statement.ts

@@ -1,3 +1,4 @@
+import Tokenizer from "sql-formatter/lib/src/lexer/Tokenizer";
 import SQLTokeniser, { NameTypes } from "./tokens";
 import { CTEReference, CallableReference, ClauseType, ClauseTypeWord, IRange, ObjectRef, QualifiedObject, StatementType, StatementTypeWord, Token } from "./types";
 
@@ -8,6 +9,15 @@ const tokenIs = (token: Token|undefined, type: string, value?: string) => {
 export default class Statement {
 	public type: StatementType = StatementType.Unknown;
 	private label: string|undefined;
+	private _blockTokens: Token[]|undefined;
+
+	private get blockTokens() {
+		if (!this._blockTokens) {
+			this._blockTokens = SQLTokeniser.createBlocks(this.tokens.slice(0));
+		}
+
+		return this._blockTokens;
+	}
 
   constructor(public tokens: Token[], public range: IRange) {
 		this.tokens = this.tokens.filter(newToken => newToken.type !== `newline`);
@@ -101,62 +111,34 @@ export default class Statement {
 		return currentClause;
 	}
 
-	getBlockRangeAt(offset: number) {
-		let start = -1;
-		let end = -1;
-
-		// Get the current token for the provided offset
-		let i = this.tokens.findIndex((token, i) => (offset >= token.range.start && offset <= token.range.end) || (offset > token.range.end && this.tokens[i+1] && offset < this.tokens[i+1].range.start));
-
-		let depth = 0;
-
-		if (tokenIs(this.tokens[i], `closebracket`)) {
-			i--;
-		}
-
-		if (tokenIs(this.tokens[i], `openbracket`)) {
-			start = i+1;
-			i++;
-		} else {
-			for (let x = i; x >= 0; x--) {
-				if (tokenIs(this.tokens[x], `openbracket`)) {
-					if (depth === 0) {
-						start = x+1;
-						break;
+	getBlockRangeAt(offset: number): IRange|undefined {
+		const blockContainsOffset = (cOffset: number, block: Token[]): IRange|undefined => {
+			const tokenInOffset = block.find(token => cOffset >= token.range.start && cOffset <= token.range.end);
+			if (tokenInOffset) {
+				if (tokenInOffset.type === `block`) {
+					if (tokenInOffset.block!.length > 0) {
+						return blockContainsOffset(cOffset, tokenInOffset.block!);
 					} else {
-						depth--;
+						const rawEnd = this.tokens.findIndex(token => token.range.end === tokenInOffset.range.end);
+						return {
+							start: rawEnd,
+							end: rawEnd
+						}
 					}
-				} else
-				if (tokenIs(this.tokens[x], `closebracket`)) {
-					depth++;
-				}
-			}
-		}
-
-		depth = 0;
-
-		for (let x = i; x <= this.tokens.length; x++) {
-			if (tokenIs(this.tokens[x], `openbracket`)) {
-				depth++;
-			} else
-			if (tokenIs(this.tokens[x], `closebracket`)) {
-				if (depth === 0) {
-					end = x;
-					break;
 				} else {
-					depth--;
+					const rawStart = this.tokens.findIndex(token => token.range.start === block[0].range.start);
+					const rawEnd = this.tokens.findIndex(token => token.range.end === block[block.length-1].range.end);
+					return {
+						start: rawStart,
+						end: rawEnd+1
+					};
 				}
 			}
-		}
 
-		if (start === -1 || end === -1) {
 			return undefined;
-		} else {
-			return {
-				start,
-				end
-			}
 		}
+
+		return blockContainsOffset(offset, this.blockTokens);
 	}
 
 	getCallableDetail(offset: number, withBlocks = false): CallableReference {
@@ -176,13 +158,22 @@ export default class Statement {
 	}
 
 	getBlockAt(offset: number): Token[] {
-		const range = this.getBlockRangeAt(offset);
+		const expandBlock = (tokens: Token[]): Token[] => {
+			const block = tokens.filter(token => token.type === `block`);
+			if (block.length > 0) {
+				return block.reduce((acc, token) => acc.concat(expandBlock(token.block!)), tokens);
+			}
 
-		if (range) {
-			return this.tokens.slice(range.start, range.end)
-		} else {
-			return []
+			return tokens;
+		}
+
+		let blockRange = this.getBlockRangeAt(offset);
+
+		if (blockRange) {
+			return expandBlock(this.tokens.slice(blockRange.start, blockRange.end));
 		}
+
+		return [];
 	}
 
 	getReferenceByOffset(offset: number) {
@@ -227,7 +218,7 @@ export default class Statement {
 	getCTEReferences(): CTEReference[] {
 		if (this.type !== StatementType.With) return [];
 
-		const withBlocks = SQLTokeniser.createBlocks(this.tokens.slice(0));
+		const withBlocks = this.blockTokens;
 
 		let cteList: CTEReference[] = [];