Merge pull request #4161 from ferenci84/diff_cache

Implement caching mechanism for diff snippets based on file save timestamp

Author: sestinj

core/autocomplete/snippets/getAllSnippets.ts

@@ -11,6 +11,8 @@ import {
   AutocompleteSnippetType,
 } from "./types";
 
+const IDE_SNIPPETS_ENABLED = false; // ideSnippets is not used, so it's temporarily disabled
+
 export interface SnippetPayload {
   rootPathSnippets: AutocompleteCodeSnippet[];
   importDefinitionSnippets: AutocompleteCodeSnippet[];
@@ -29,6 +31,27 @@ function racePromise<T>(promise: Promise<T[]>): Promise<T[]> {
   return Promise.race([promise, timeoutPromise]);
 }
 
+class DiffSnippetsCache {
+  private cache: Map<number, any> = new Map();
+  private lastTimestamp: number = 0;
+
+  public set<T>(timestamp: number, value: T): T {
+    // Clear old cache entry if exists
+    if (this.lastTimestamp !== timestamp) {
+      this.cache.clear();
+    }
+    this.lastTimestamp = timestamp;
+    this.cache.set(timestamp, value);
+    return value;
+  }
+
+  public get(timestamp: number): any | undefined {
+    return this.cache.get(timestamp);
+  }
+}
+
+const diffSnippetsCache = new DiffSnippetsCache();
+
 // Some IDEs might have special ways of finding snippets (e.g. JetBrains and VS Code have different "LSP-equivalent" systems,
 // or they might separately track recently edited ranges)
 async function getIdeSnippets(
@@ -90,19 +113,31 @@ const getClipboardSnippets = async (
 const getDiffSnippets = async (
   ide: IDE,
 ): Promise<AutocompleteDiffSnippet[]> => {
+  const currentTimestamp = ide.getLastFileSaveTimestamp ?
+    ide.getLastFileSaveTimestamp() :
+    Math.floor(Date.now() / 10000) * 10000; // Defaults to update once in every 10 seconds
+
+  // Check cache first
+  const cached = diffSnippetsCache.get(currentTimestamp) as AutocompleteDiffSnippet[];
+
+  if (cached) {
+    return cached;
+  }
+
   let diff: string[] = [];
   try {
     diff = await ide.getDiff(true);
   } catch (e) {
     console.error("Error getting diff for autocomplete", e);
   }
 
-  return diff.map((item) => {
+  return diffSnippetsCache.set(currentTimestamp, diff.map((item) => {
     return {
       content: item,
       type: AutocompleteSnippetType.Diff,
     };
-  });
+  }));
+
 };
 
 export const getAllSnippets = async ({
@@ -127,10 +162,8 @@ export const getAllSnippets = async ({
     clipboardSnippets,
   ] = await Promise.all([
     racePromise(contextRetrievalService.getRootPathSnippets(helper)),
-    racePromise(
-      contextRetrievalService.getSnippetsFromImportDefinitions(helper),
-    ),
-    racePromise(getIdeSnippets(helper, ide, getDefinitionsFromLsp)),
+    racePromise(contextRetrievalService.getSnippetsFromImportDefinitions(helper)),
+    IDE_SNIPPETS_ENABLED ? racePromise(getIdeSnippets(helper, ide, getDefinitionsFromLsp)) : [],
     racePromise(getDiffSnippets(ide)),
     racePromise(getClipboardSnippets(ide)),
   ]);

core/index.d.ts

@@ -43,13 +43,13 @@ export interface IndexingProgressUpdate {
   desc: string;
   shouldClearIndexes?: boolean;
   status:
-    | "loading"
-    | "indexing"
-    | "done"
-    | "failed"
-    | "paused"
-    | "disabled"
-    | "cancelled";
+  | "loading"
+  | "indexing"
+  | "done"
+  | "failed"
+  | "paused"
+  | "disabled"
+  | "cancelled";
   debugInfo?: string;
 }
 
@@ -667,12 +667,16 @@ export interface IDE {
   getCurrentFile(): Promise<
     | undefined
     | {
-        isUntitled: boolean;
-        path: string;
-        contents: string;
-      }
+      isUntitled: boolean;
+      path: string;
+      contents: string;
+    }
   >;
 
+  getLastFileSaveTimestamp?(): number;
+
+  updateLastFileSaveTimestamp?(): void;
+
   getPinnedFiles(): Promise<string[]>;
 
   getSearchResults(query: string): Promise<string>;
@@ -849,11 +853,11 @@ export interface CustomCommand {
 export interface Prediction {
   type: "content";
   content:
-    | string
-    | {
-        type: "text";
-        text: string;
-      }[];
+  | string
+  | {
+    type: "text";
+    text: string;
+  }[];
 }
 
 export interface ToolExtras {
@@ -1182,9 +1186,9 @@ export interface Config {
   embeddingsProvider?: EmbeddingsProviderDescription | ILLM;
   /** The model that Continue will use for tab autocompletions. */
   tabAutocompleteModel?:
-    | CustomLLM
-    | ModelDescription
-    | (CustomLLM | ModelDescription)[];
+  | CustomLLM
+  | ModelDescription
+  | (CustomLLM | ModelDescription)[];
   /** Options for tab autocomplete */
   tabAutocompleteOptions?: Partial<TabAutocompleteOptions>;
   /** UI styles customization */
@@ -1275,9 +1279,9 @@ export type PackageDetailsSuccess = PackageDetails & {
 export type PackageDocsResult = {
   packageInfo: ParsedPackageInfo;
 } & (
-  | { error: string; details?: never }
-  | { details: PackageDetailsSuccess; error?: never }
-);
+    | { error: string; details?: never }
+    | { details: PackageDetailsSuccess; error?: never }
+  );
 
 export interface TerminalOptions {
   reuseTerminal?: boolean;

extensions/vscode/src/VsCodeIde.ts

@@ -33,6 +33,7 @@ import { SecretStorage } from "./stubs/SecretStorage";
 class VsCodeIde implements IDE {
   ideUtils: VsCodeIdeUtils;
   secretStorage: SecretStorage;
+  private lastFileSaveTimestamp: number = Date.now();
 
   constructor(
     private readonly vscodeWebviewProtocolPromise: Promise<VsCodeWebviewProtocol>,
@@ -42,6 +43,14 @@ class VsCodeIde implements IDE {
     this.secretStorage = new SecretStorage(context);
   }
 
+  public updateLastFileSaveTimestamp(): void {
+    this.lastFileSaveTimestamp = Date.now();
+  }
+
+  public getLastFileSaveTimestamp(): number {
+    return this.lastFileSaveTimestamp;
+  }
+
   async readSecrets(keys: string[]): Promise<Record<string, string>> {
     const secretValuePromises = keys.map((key) => this.secretStorage.get(key));
     const secretValues = await Promise.all(secretValuePromises);

extensions/vscode/src/autocomplete/lsp.ts

@@ -33,7 +33,7 @@ function gotoInputKey(input: GotoInput) {
   return `${input.name}${input.uri.toString()}${input.line}${input.character}`;
 }
 
-const MAX_CACHE_SIZE = 50;
+const MAX_CACHE_SIZE = 500;
 const gotoCache = new Map<string, RangeInFile[]>();
 
 export async function executeGotoProvider(
@@ -151,28 +151,32 @@ async function crawlTypes(
   identifierNodes.forEach((node) => searchedLabels.add(node.text));
 
   // Use LSP to get the definitions of those types
-  const definitions = await Promise.all(
-    identifierNodes.map(async (node) => {
-      const [typeDef] = await executeGotoProvider({
-        uri: vscode.Uri.parse(rif.filepath),
-        // TODO: tree-sitter is zero-indexed, but there seems to be an off-by-one
-        // error at least with the .ts parser sometimes
-        line:
-          rif.range.start.line +
-          Math.min(node.startPosition.row, astLineCount - 1),
-        character: rif.range.start.character + node.startPosition.column,
-        name: "vscode.executeDefinitionProvider",
-      });
+  const definitions = [];
+
+  for (const node of identifierNodes) {
+    const [typeDef] = await executeGotoProvider({
+      uri: vscode.Uri.parse(rif.filepath),
+      // TODO: tree-sitter is zero-indexed, but there seems to be an off-by-one
+      // error at least with the .ts parser sometimes
+      line:
+        rif.range.start.line +
+        Math.min(node.startPosition.row, astLineCount - 1),
+      character: rif.range.start.character + node.startPosition.column,
+      name: "vscode.executeDefinitionProvider",
+    });
+
+    if (!typeDef) {
+      definitions.push(undefined);
+      continue;
+    }
 
-      if (!typeDef) {
-        return undefined;
-      }
-      return {
-        ...typeDef,
-        contents: await ide.readRangeInFile(typeDef.filepath, typeDef.range),
-      };
-    }),
-  );
+    const contents = await ide.readRangeInFile(typeDef.filepath, typeDef.range);
+
+    definitions.push({
+      ...typeDef,
+      contents,
+    });
+  }
 
   // TODO: Filter out if not in our code?
 

extensions/vscode/src/extension/VsCodeExtension.ts

@@ -267,6 +267,7 @@ export class VsCodeExtension {
     });
 
     vscode.workspace.onDidSaveTextDocument(async (event) => {
+      this.ide.updateLastFileSaveTimestamp();
       this.core.invoke("files/changed", {
         uris: [event.uri.toString()],
       });
@@ -346,8 +347,7 @@ export class VsCodeExtension {
 
     // Register a content provider for the readonly virtual documents
     const documentContentProvider = new (class
-      implements vscode.TextDocumentContentProvider
-    {
+      implements vscode.TextDocumentContentProvider {
       // emitter and its event
       onDidChangeEmitter = new vscode.EventEmitter<vscode.Uri>();
       onDidChange = this.onDidChangeEmitter.event;

extensions/vscode/src/util/ideUtils.ts

@@ -394,6 +394,27 @@ export class VsCodeIdeUtils {
     }
   }
 
+  private _getRepositories(): Repository[] | undefined {
+    const extension =
+      vscode.extensions.getExtension<GitExtension>("vscode.git");
+    if (
+      typeof extension === "undefined" ||
+      !extension.isActive ||
+      typeof vscode.workspace.workspaceFolders === "undefined"
+    ) {
+      return undefined;
+    }
+
+    try {
+      const git = extension.exports.getAPI(1);
+      return git.repositories;
+    } catch (e) {
+      this._repoWasNone = true;
+      console.warn("Git not found: ", e);
+      return undefined;
+    }
+  }
+
   private _repoWasNone: boolean = false;
   private repoCache: Map<string, Repository> = new Map();
   private static secondsToWaitForGitToLoad =
@@ -475,20 +496,28 @@ export class VsCodeIdeUtils {
   async getDiff(includeUnstaged: boolean): Promise<string[]> {
     const diffs: string[] = [];
 
-    for (const dir of this.getWorkspaceDirectories()) {
-      const repo = await this.getRepo(dir);
-      if (!repo) {
-        continue;
-      }
+    const repos = this._getRepositories();
+
+    try {
+      if (repos) {
+        for (const repo of repos) {
+
+          const staged = await repo.diff(true);
 
-      const staged = await repo.diff(true);
-      diffs.push(staged);
-      if (includeUnstaged) {
-        const unstaged = await repo.diff(false);
-        diffs.push(unstaged);
+          diffs.push(staged);
+          if (includeUnstaged) {
+            const unstaged = await repo.diff(false);
+            diffs.push(unstaged);
+          }
+        }
       }
+
+      return diffs.flatMap((diff) => this.splitDiff(diff));
+
+    } catch (e) {
+      console.error(e);
+      return [];
     }
 
-    return diffs.flatMap((diff) => this.splitDiff(diff));
   }
 }