Fix candidate collection for query with many Statements containing WITHs

packages/language-support/src/autocompletion/completionCoreCompletions.ts

@@ -10,6 +10,7 @@ import CypherParser, {
   Expression2Context,
 } from '../generated-parser/CypherCmdParser';
 import {
+  findLastClause,
   findParent,
   findPreviousNonSpace,
   resolveCypherVersion,
@@ -24,7 +25,7 @@ import {
 
 import { getMethodName, ParsedStatement } from '../parserWrapper';
 
-import type { CandidateRule } from '../../../../vendor/antlr4-c3/dist/esm/index.js';
+import type { CandidateRule, ParserRuleContext } from '../../../../vendor/antlr4-c3/dist/esm/index.js';
 import {
   CandidatesCollection,
   CodeCompletionCore,
@@ -543,7 +544,9 @@ export function completionCoreCompletion(
       .map(([token]) => Number(token)),
   );
 
-  const candidates = codeCompletion.collectCandidates(caretIndex);
+  const lastClause = findLastClause(parsingResult.ctx);
+
+  const candidates = codeCompletion.collectCandidates(caretIndex, lastClause as ParserRuleContext);
 
   const ruleCompletions = Array.from(candidates.rules.entries()).flatMap(
     (candidate): CompletionItem[] => {

packages/language-support/src/helpers.ts

@@ -9,8 +9,11 @@ import antlrDefaultExport, {
 import { DbSchema } from './dbSchema';
 import CypherLexer from './generated-parser/CypherCmdLexer';
 import CypherParser, {
+  ClauseContext,
   NodePatternContext,
+  QueryWithLocalDefinitionsContext,
   RelationshipPatternContext,
+  SingleQueryContext,
   StatementsOrCommandsContext,
 } from './generated-parser/CypherCmdParser';
 import { ParsedStatement, ParsingResult } from './parserWrapper';
@@ -209,6 +212,69 @@ export function resolveCypherVersion(
   return cypherVersion;
 }
 
+/**
+ * Takes a rule context checks if the current (rightmost) rule is within a clause.
+ * If so, checks how many other clauses are inside the corresponding SingleQueryContext and returns the current
+ * ClauseContext if the count is >= 15. This is so we can use the clause as context for candidate collection, 
+ * which can get slow when there are multiple clauses in a single query.
+ */
+export function findLastClause(ctx: ParserRuleContext) {
+  let current: ParserRuleContext = ctx;
+  //ctx can vary from rule inside latest clause statement containing it
+  //Here we go up to the outer singlequery (which is not empty)
+  while(!(current instanceof SingleQueryContext && current.getText()) && current.parentCtx) {
+    current = current.parentCtx;
+  };
+
+  const newCtx = checkNumClauses(current);
+
+  return newCtx;
+}
+
+/**
+ * Takes a parser rule, and returns the number of clauses under the rightmost SingleQueryContext
+ * If the rightmost child of said SingleQueryContext is a QueryWithLocalDefinitionsContext:s 
+ * it can contain more SingleQueryContext:s and we thus recursively search this.
+ * @param ctx 
+ * @returns 
+ */
+export function checkNumClauses(ctx: ParserRuleContext): ParserRuleContext {
+  let current = ctx;
+  let lastCurrent: ParserRuleContext = undefined;
+  while(lastCurrent !== current && current.children) {
+    lastCurrent = current;
+    if (current instanceof SingleQueryContext) {
+      const foundClauseCount: boolean = false;
+      let lastCandidate = 1;
+      while (!foundClauseCount && lastCandidate <= current.children.length){
+        const lastChild = current.children.at(-1*lastCandidate);
+        if (lastChild instanceof ClauseContext) {
+          const clauseChildren = current.children.filter(x => x instanceof ClauseContext);
+          const numClauses = clauseChildren.length;
+          if (numClauses >= 15) {
+            return lastChild;
+          } else {
+            return undefined;
+          }
+        } else if (lastChild instanceof QueryWithLocalDefinitionsContext) {
+          return checkNumClauses(lastChild);
+        } else {
+          lastCandidate ++;
+        }
+      }
+      return undefined;
+    }
+    for (let i = current.children.length-1; i >= 0; i--) {
+      const candidate = current.children[i];
+      if (candidate instanceof ParserRuleContext) {
+        current = candidate;
+        break;
+      }
+    }
+  }
+  return undefined;
+}
+
 export const rulesDefiningVariables = [
   CypherParser.RULE_returnItem,
   CypherParser.RULE_unwindClause,

packages/language-support/src/syntaxValidation/completionCoreErrors.ts

@@ -1,6 +1,6 @@
 import { Token } from 'antlr4';
 import { distance } from 'fastest-levenshtein';
-import { CodeCompletionCore } from '../../../../vendor/antlr4-c3/dist/esm/index.js';
+import { CodeCompletionCore, ParserRuleContext } from '../../../../vendor/antlr4-c3/dist/esm/index.js';
 import CypherLexer from '../generated-parser/CypherCmdLexer';
 import CypherParser from '../generated-parser/CypherCmdParser';
 import {
@@ -9,6 +9,7 @@ import {
   lexerSymbols,
   tokenNames,
 } from '../lexerSymbols';
+import { findLastClause } from '../helpers.js';
 
 /*
 We ask for 0.7 similarity (number between 0 and 1) for 
@@ -63,7 +64,9 @@ export function completionCoreErrormessage(
 
   const errorText = currentToken.text;
 
-  const candidates = codeCompletion.collectCandidates(caretIndex);
+  const lastClause = findLastClause(parser._ctx);
+
+  const candidates = codeCompletion.collectCandidates(caretIndex, lastClause as ParserRuleContext);
 
   const ruleCandidates = Array.from(candidates.rules.keys());