VSCode extensions: add completion cache. (#128)

clients/vscode/package.json

@@ -107,6 +107,7 @@
     "axios": "^1.3.4",
     "events": "^3.3.0",
     "form-data": "^4.0.0",
+    "linked-list-typescript": "^1.0.15",
     "process": "^0.11.10"
   }
 }

clients/vscode/src/CompletionCache.ts

@@ -0,0 +1,75 @@
+import { LinkedList } from "linked-list-typescript";
+import { CompletionResponse, Choice } from "./generated";
+
+type Range = {
+  start: number;
+  end: number;
+};
+
+export type CompletionCacheEntry = {
+  documentId: any;
+  promptRange: Range;
+  prompt: string;
+  completion: CompletionResponse;
+};
+
+export class CompletionCache {
+  public static cacheSize = 10;
+  private cache = new LinkedList<CompletionCacheEntry>();
+
+  constructor() {}
+
+  private evict() {
+    while (this.cache.length > CompletionCache.cacheSize) {
+      this.cache.removeTail();
+    }
+  }
+
+  private pop(entry: CompletionCacheEntry) {
+    this.cache.remove(entry);
+    this.cache.prepend(entry);
+  }
+
+  public add(entry: CompletionCacheEntry) {
+    this.evict();
+    this.cache.prepend(entry);
+  }
+
+  public findCompatible(documentId: any, text: string, cursor: number): CompletionResponse | null {
+    let hit: { entry: CompletionCacheEntry; compatibleChoices: Choice[] } | null = null;
+    for (const entry of this.cache) {
+      if (entry.documentId !== documentId) {
+        continue;
+      }
+      // Check if text in prompt range has not changed
+      if (text.slice(entry.promptRange.start, entry.promptRange.end) !== entry.prompt) {
+        continue;
+      }
+      // Filter choices that start with inputed text after prompt
+      const compatibleChoices = entry.completion.choices
+        .filter((choice) => choice.text.startsWith(text.slice(entry.promptRange.end, cursor)))
+        .map((choice) => {
+          return {
+            index: choice.index,
+            text: choice.text.substring(cursor - entry.promptRange.end),
+          };
+        });
+      if (compatibleChoices.length > 0) {
+        hit = {
+          entry,
+          compatibleChoices,
+        };
+        break;
+      }
+    }
+    if (hit) {
+      this.pop(hit.entry);
+      return {
+        id: hit.entry.completion.id,
+        created: hit.entry.completion.created,
+        choices: hit.compatibleChoices,
+      };
+    }
+    return null;
+  }
+}

clients/vscode/src/TabbyCompletionProvider.ts

@@ -12,6 +12,7 @@ import {
 } from "vscode";
 import { CompletionResponse, EventType, ChoiceEvent, ApiError, CancelablePromise, CancelError } from "./generated";
 import { TabbyClient } from "./TabbyClient";
+import { CompletionCache } from "./CompletionCache";
 import { sleep } from "./utils";
 
 export class TabbyCompletionProvider implements InlineCompletionItemProvider {
@@ -20,6 +21,7 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
   private pendingCompletion: CancelablePromise<CompletionResponse> | null = null;
 
   private tabbyClient = TabbyClient.getInstance();
+  private completionCache = new CompletionCache();
   // User Settings
   private enabled: boolean = true;
   private suggestionDelay: number = 150;
@@ -42,7 +44,8 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
       return emptyResponse;
     }
 
-    const prompt = this.getPrompt(document, position);
+    const promptRange = this.calculatePromptRange(position);
+    const prompt = document.getText(promptRange);
     if (this.isNil(prompt)) {
       console.debug("Prompt is empty, skipping");
       return emptyResponse;
@@ -56,6 +59,15 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
       return emptyResponse;
     }
 
+    const replaceRange = this.calculateReplaceRange(document, position);
+
+    const compatibleCache = this.completionCache.findCompatible(document.uri, document.getText(), document.offsetAt(position));
+    if (compatibleCache) {
+      const completions = this.toInlineCompletions(compatibleCache, replaceRange);
+      console.debug("Use cached completions: ", compatibleCache);
+      return Promise.resolve(completions);
+    }
+
     console.debug(
       "Requesting: ",
       {
@@ -86,15 +98,14 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
     });
     this.pendingCompletion = null;
 
-    const hasSuffixParen = this.hasSuffixParen(document, position);
-    const replaceRange = hasSuffixParen
-      ? new Range(
-          position.line,
-          position.character,
-          position.line,
-          position.character + 1
-        )
-      : new Range(position, position);
+    if (completion) {
+      this.completionCache.add({
+        documentId: document.uri,
+        promptRange: { start: document.offsetAt(promptRange.start), end: document.offsetAt(promptRange.end) },
+        prompt,
+        completion,
+      });
+    }
     const completions = this.toInlineCompletions(completion, replaceRange);
     console.debug("Result completions: ", completions);
     return Promise.resolve(completions);
@@ -106,13 +117,6 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
     this.suggestionDelay = configuration.get("suggestionDelay", 150);
   }
 
-  private getPrompt(document: TextDocument, position: Position): string | undefined {
-    const maxLines = 20;
-    const firstLine = Math.max(position.line - maxLines, 0);
-
-    return document.getText(new Range(firstLine, 0, position.line, position.character));
-  }
-
   private isNil(value: string | undefined | null): boolean {
     return value === undefined || value === null || value.length === 0;
   }
@@ -140,4 +144,19 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
     );
     return ")]}".indexOf(suffix) > -1;
   }
+
+  private calculatePromptRange(position: Position): Range {
+    const maxLines = 20;
+    const firstLine = Math.max(position.line - maxLines, 0);
+    return new Range(firstLine, 0, position.line, position.character);
+  }
+
+  private calculateReplaceRange(document: TextDocument, position: Position): Range {
+    const hasSuffixParen = this.hasSuffixParen(document, position);
+    if (hasSuffixParen) {
+      return new Range(position.line, position.character, position.line, position.character + 1);
+    } else {
+      return new Range(position, position);
+    }
+  }
 }

clients/vscode/yarn.lock

@@ -1740,6 +1740,11 @@ lie@~3.3.0:
   dependencies:
     immediate "~3.0.5"
 
+linked-list-typescript@^1.0.15:
+  version "1.0.15"
+  resolved "https://registry.yarnpkg.com/linked-list-typescript/-/linked-list-typescript-1.0.15.tgz#faeed93cf9203f102e2158c29edcddda320abe82"
+  integrity sha512-RIyUu9lnJIyIaMe63O7/aFv/T2v3KsMFuXMBbUQCHX+cgtGro86ETDj5ed0a8gQL2+DFjzYYsgVG4I36/cUwgw==
+
 linkify-it@^3.0.1:
   version "3.0.3"
   resolved "https://registry.npmmirror.com/linkify-it/-/linkify-it-3.0.3.tgz#a98baf44ce45a550efb4d49c769d07524cc2fa2e"