Cycle completions (#32)

* Add ability to cycle through completions

* Don’t store originalDocument

* Add cycle widget

* Switch back to ctrl

Author: tmcw

demo/demo.css

@@ -54,3 +54,32 @@ main {
 .cm-ghostText:hover {
   background: #eee;
 }
+
+.cm-codeium-cycle {
+  font-size: 9px;
+  background-color: #eee;
+  padding: 2px;
+  border-radius: 2px;
+  display: inline-block;
+}
+
+.cm-codeium-cycle-key {
+  font-size: 9px;
+  font-family: monospace;
+  display: inline-block;
+  padding: 2px;
+  border-radius: 2px;
+  border: none;
+  background-color: lightblue;
+  margin-left: 5px;
+}
+
+.cm-codeium-cycle-key:hover {
+  background-color: deepskyblue;
+}
+
+.cm-codeium-cycle-explanation {
+  font-family: monospace;
+  display: inline-block;
+  padding: 2px;
+}

demo/index.ts

@@ -8,13 +8,7 @@ import {
 import { python } from "@codemirror/lang-python";
 
 new EditorView({
-  doc: `let hasAnError: string = 10;
-
-function increment(num: number) {
-  return num + 1;
-}
-
-increment('not a number');`,
+  doc: "// Factorial function",
   extensions: [
     basicSetup,
     javascript({

src/commands.ts

@@ -1,8 +1,12 @@
-import { Transaction, EditorSelection } from "@codemirror/state";
-import type { EditorView } from "@codemirror/view";
+import { Transaction, EditorSelection, ChangeSet } from "@codemirror/state";
+import type { Command, EditorView } from "@codemirror/view";
 import { copilotEvent, copilotIgnore } from "./annotations.js";
 import { completionDecoration } from "./completionDecoration.js";
-import { acceptSuggestion, clearSuggestion } from "./effects.js";
+import {
+  acceptSuggestion,
+  addSuggestions,
+  clearSuggestion,
+} from "./effects.js";
 
 /**
  * Accepting a suggestion: we remove the ghost text, which
@@ -53,6 +57,79 @@ export function acceptSuggestionCommand(view: EditorView) {
   return true;
 }
 
+/**
+ * Cycle through suggested AI completed code.
+ */
+export const nextSuggestionCommand: Command = (view: EditorView) => {
+  const { state } = view;
+  // We delete the suggestion, then carry through with the original keypress
+  const stateField = state.field(completionDecoration);
+
+  if (!stateField) {
+    return false;
+  }
+
+  const { changeSpecs } = stateField;
+
+  if (changeSpecs.length < 2) {
+    return false;
+  }
+
+  // Loop through next suggestion.
+  const index = (stateField.index + 1) % changeSpecs.length;
+  const nextSpec = changeSpecs.at(index);
+  if (!nextSpec) {
+    return false;
+  }
+
+  /**
+   * First, get the original document, by applying the stored
+   * reverse changeset against the currently-displayed document.
+   */
+  const originalDocument = stateField.reverseChangeSet.apply(state.doc);
+
+  /**
+   * Get the changeset that we will apply that will
+   *
+   * 1. Reverse the currently-displayed suggestion, to get us back to
+   *    the original document
+   * 2. Apply the next suggestion.
+   *
+   * It does both in the same changeset, so there is no flickering.
+   */
+  const reverseCurrentSuggestionAndApplyNext = ChangeSet.of(
+    stateField.reverseChangeSet,
+    state.doc.length,
+  ).compose(ChangeSet.of(nextSpec, originalDocument.length));
+
+  /**
+   * Generate the next changeset
+   */
+  const nextSet = ChangeSet.of(nextSpec, originalDocument.length);
+  const reverseChangeSet = nextSet.invert(originalDocument);
+
+  view.dispatch({
+    changes: reverseCurrentSuggestionAndApplyNext,
+    effects: addSuggestions.of({
+      index,
+      reverseChangeSet,
+      changeSpecs,
+    }),
+    annotations: [
+      // Tell upstream integrations to ignore this
+      // change.
+      copilotIgnore.of(null),
+      // Tell ourselves not to request a completion
+      // because of this change.
+      copilotEvent.of(null),
+      // Don't add this to history.
+      Transaction.addToHistory.of(false),
+    ],
+  });
+
+  return true;
+};
+
 /**
  * Rejecting a suggestion: this looks at the currently-shown suggestion
  * and reverses it, clears the suggestion, and makes sure

src/completionDecoration.ts

@@ -10,6 +10,7 @@ import {
   clearSuggestion,
 } from "./effects.js";
 import type { CompletionState } from "./types.js";
+import { codeiumConfig } from "./config.js";
 
 const ghostMark = Decoration.mark({ class: "cm-ghostText" });
 
@@ -24,22 +25,36 @@ export const completionDecoration = StateField.define<CompletionState>({
     return null;
   },
   update(state: CompletionState, transaction: Transaction) {
+    const config = transaction.state.facet(codeiumConfig);
     for (const effect of transaction.effects) {
       if (effect.is(addSuggestions)) {
         const { changeSpecs, index } = effect.value;
 
         // NOTE: here we're adjusting the decoration range
         // to refer to locations in the document _after_ we've
         // inserted the text.
-        const decorations = Decoration.set(
-          changeSpecs[index]!.map((suggestionRange) => {
-            const range = ghostMark.range(
-              suggestionRange.absoluteStartPos,
-              suggestionRange.absoluteEndPos,
-            );
-            return range;
-          }),
-        );
+        const ranges = changeSpecs[index]!.map((suggestionRange) => {
+          const range = ghostMark.range(
+            suggestionRange.absoluteStartPos,
+            suggestionRange.absoluteEndPos,
+          );
+          return range;
+        });
+        const widgetPos = ranges.at(-1)?.to;
+
+        const decorations = Decoration.set([
+          ...ranges,
+          ...(widgetPos !== undefined &&
+          changeSpecs.length > 1 &&
+          config.widgetClass
+            ? [
+                Decoration.widget({
+                  widget: new config.widgetClass(index, changeSpecs.length),
+                  side: 1,
+                }).range(widgetPos),
+              ]
+            : []),
+        ]);
 
         return {
           index,

src/config.ts

@@ -3,6 +3,7 @@ import { Language } from "./api/proto/exa/codeium_common_pb/codeium_common_pb.js
 import type { Document } from "./api/proto/exa/language_server_pb/language_server_pb.js";
 import type { PartialMessage } from "@bufbuild/protobuf";
 import type { CompletionContext } from "@codemirror/autocomplete";
+import { DefaultCycleWidget } from "./defaultCycleWidget.js";
 
 export interface CodeiumConfig {
   /**
@@ -30,6 +31,12 @@ export interface CodeiumConfig {
    * autocomplete sources.
    */
   shouldComplete?: (context: CompletionContext) => boolean;
+
+  /**
+   * The class for the widget that is shown at the end a suggestion
+   * when there are multiple suggestions to cycle through.
+   */
+  widgetClass?: typeof DefaultCycleWidget | null;
 }
 
 export const codeiumConfig = Facet.define<
@@ -42,6 +49,7 @@ export const codeiumConfig = Facet.define<
       {
         language: Language.TYPESCRIPT,
         timeout: 150,
+        widgetClass: DefaultCycleWidget,
       },
       {},
     );

src/defaultCycleWidget.ts

@@ -0,0 +1,34 @@
+import { WidgetType } from "@codemirror/view";
+
+/**
+ * Shown at the end of a suggestion if there are multiple
+ * suggestions to cycle through.
+ */
+export class DefaultCycleWidget extends WidgetType {
+  index: number;
+  total: number;
+
+  constructor(index: number, total: number) {
+    super();
+    this.index = index;
+    this.total = total;
+  }
+
+  toDOM() {
+    const wrap = document.createElement("span");
+    wrap.setAttribute("aria-hidden", "true");
+    wrap.className = "cm-codeium-cycle";
+    const words = wrap.appendChild(document.createElement("span"));
+    words.className = "cm-codeium-cycle-explanation";
+    words.innerText = `${this.index + 1}/${this.total}`;
+    const key = wrap.appendChild(document.createElement("button"));
+    key.className = "cm-codeium-cycle-key";
+    key.innerText = "→ (Ctrl ])";
+    key.dataset.action = "codeium-cycle";
+    return wrap;
+  }
+
+  ignoreEvent() {
+    return false;
+  }
+}

src/plugin.ts

@@ -1,11 +1,12 @@
-import { EditorView, type ViewUpdate } from "@codemirror/view";
+import { EditorView, keymap, type ViewUpdate } from "@codemirror/view";
 import { type Extension, Prec } from "@codemirror/state";
 import { completionDecoration } from "./completionDecoration.js";
 import { completionRequester } from "./completionRequester.js";
 import {
   sameKeyCommand,
   rejectSuggestionCommand,
   acceptSuggestionCommand,
+  nextSuggestionCommand,
 } from "./commands.js";
 import {
   type CodeiumConfig,
@@ -45,6 +46,13 @@ function isDecorationClicked(view: EditorView): boolean {
 function completionPlugin() {
   return EditorView.domEventHandlers({
     keydown(event, view) {
+      // Ideally, we handle infighting between
+      // the nextSuggestionCommand and this handler
+      // by using precedence, but I can't get that to work
+      // yet.
+      if (event.key === "]" && event.ctrlKey) {
+        return false;
+      }
       if (
         event.key !== "Shift" &&
         event.key !== "Control" &&
@@ -55,7 +63,17 @@ function completionPlugin() {
       }
       return false;
     },
-    mouseup(_event, view) {
+    mouseup(event, view) {
+      const target = event.target as HTMLElement;
+      if (
+        target.nodeName === "BUTTON" &&
+        target.dataset.action === "codeium-cycle"
+      ) {
+        nextSuggestionCommand(view);
+        event.stopPropagation();
+        event.preventDefault();
+        return true;
+      }
       if (isDecorationClicked(view)) {
         return acceptSuggestionCommand(view);
       }
@@ -64,6 +82,18 @@ function completionPlugin() {
   });
 }
 
+/**
+ * Next completion map
+ */
+function nextCompletionPlugin() {
+  return keymap.of([
+    {
+      key: "Ctrl-]",
+      run: nextSuggestionCommand,
+    },
+  ]);
+}
+
 /**
  * Changing the editor's focus - blurring it by clicking outside -
  * rejects the suggestion
@@ -81,6 +111,7 @@ export {
   copilotIgnore,
   codeiumConfig,
   codeiumOtherDocumentsConfig,
+  nextSuggestionCommand,
   type CodeiumOtherDocumentsConfig,
   type CodeiumConfig,
 };
@@ -93,8 +124,9 @@ export function copilotPlugin(config: CodeiumConfig): Extension {
   return [
     codeiumConfig.of(config),
     completionDecoration,
-    Prec.highest(completionPlugin()),
+    Prec.highest(nextCompletionPlugin()),
     Prec.highest(viewCompletionPlugin()),
+    Prec.high(completionPlugin()),
     completionRequester(),
   ];
 }

src/types.ts

@@ -1,5 +1,5 @@
-import type { ChangeSet } from "@codemirror/state";
-import type { DecorationSet } from "@codemirror/view";
+import type { Range, ChangeSet } from "@codemirror/state";
+import type { Decoration, DecorationSet } from "@codemirror/view";
 
 /**
  * We dispatch an effect that updates the CompletionState.