[forgive] Polish decorations

Polishes up decorations.

Co-authored-by: Jordan Brown <jmbrown@meta.com>

compiler/packages/babel-plugin-react-compiler/src/Inference/InferEffectDependencies.ts

@@ -261,6 +261,7 @@ export function inferEffectDependencies(fn: HIRFunction): void {
         } else if (
           value.args.length >= 2 &&
           value.args.length - 1 === autodepFnLoads.get(callee.identifier.id) &&
+          value.args[0] != null &&
           value.args[0].kind === 'Identifier'
         ) {
           const penultimateArg = value.args[value.args.length - 2];

compiler/packages/react-forgive/client/src/autodeps.ts

@@ -0,0 +1,99 @@
+import * as vscode from 'vscode';
+import {
+  LanguageClient,
+  RequestType,
+  type Position,
+} from 'vscode-languageclient/node';
+import {positionLiteralToVSCodePosition, positionsToRange} from './mapping';
+
+export type AutoDepsDecorationsLSPEvent = {
+  useEffectCallExpr: [Position, Position];
+  decorations: Array<[Position, Position]>;
+};
+
+export interface AutoDepsDecorationsParams {
+  position: Position;
+}
+
+export namespace AutoDepsDecorationsRequest {
+  export const type = new RequestType<
+    AutoDepsDecorationsParams,
+    AutoDepsDecorationsLSPEvent | null,
+    void
+  >('react/autodeps_decorations');
+}
+
+const inferredEffectDepDecoration =
+  vscode.window.createTextEditorDecorationType({
+    // TODO: make configurable?
+    borderColor: new vscode.ThemeColor('diffEditor.move.border'),
+    borderStyle: 'solid',
+    borderWidth: '0 0 4px 0',
+  });
+
+let currentlyDecoratedAutoDepFnLoc: vscode.Range | null = null;
+export function getCurrentlyDecoratedAutoDepFnLoc(): vscode.Range | null {
+  return currentlyDecoratedAutoDepFnLoc;
+}
+export function setCurrentlyDecoratedAutoDepFnLoc(range: vscode.Range): void {
+  currentlyDecoratedAutoDepFnLoc = range;
+}
+export function clearCurrentlyDecoratedAutoDepFnLoc(): void {
+  currentlyDecoratedAutoDepFnLoc = null;
+}
+
+let decorationRequestId = 0;
+export type AutoDepsDecorationsOptions = {
+  shouldUpdateCurrent: boolean;
+};
+export function requestAutoDepsDecorations(
+  client: LanguageClient,
+  position: vscode.Position,
+  options: AutoDepsDecorationsOptions,
+) {
+  const id = ++decorationRequestId;
+  client
+    .sendRequest(AutoDepsDecorationsRequest.type, {position})
+    .then(response => {
+      if (response !== null) {
+        const {
+          decorations,
+          useEffectCallExpr: [start, end],
+        } = response;
+        // Maintain ordering
+        if (decorationRequestId === id) {
+          if (options.shouldUpdateCurrent) {
+            setCurrentlyDecoratedAutoDepFnLoc(positionsToRange(start, end));
+          }
+          drawInferredEffectDepDecorations(decorations);
+        }
+      } else {
+        clearCurrentlyDecoratedAutoDepFnLoc();
+        clearDecorations(inferredEffectDepDecoration);
+      }
+    });
+}
+
+export function drawInferredEffectDepDecorations(
+  decorations: Array<[Position, Position]>,
+): void {
+  const decorationOptions = decorations.map(([start, end]) => {
+    return {
+      range: new vscode.Range(
+        positionLiteralToVSCodePosition(start),
+        positionLiteralToVSCodePosition(end),
+      ),
+      hoverMessage: 'Inferred as an effect dependency',
+    };
+  });
+  vscode.window.activeTextEditor?.setDecorations(
+    inferredEffectDepDecoration,
+    decorationOptions,
+  );
+}
+
+export function clearDecorations(
+  decorationType: vscode.TextEditorDecorationType,
+) {
+  vscode.window.activeTextEditor?.setDecorations(decorationType, []);
+}

compiler/packages/react-forgive/client/src/extension.ts

@@ -5,29 +5,16 @@ import {
   LanguageClient,
   LanguageClientOptions,
   type Position,
-  RequestType,
   ServerOptions,
   TransportKind,
 } from 'vscode-languageclient/node';
-import {WHITE} from './colors';
+import {positionLiteralToVSCodePosition} from './mapping';
+import {
+  getCurrentlyDecoratedAutoDepFnLoc,
+  requestAutoDepsDecorations,
+} from './autodeps';
 
 let client: LanguageClient;
-const inferredEffectDepDecoration =
-  vscode.window.createTextEditorDecorationType({
-    backgroundColor: WHITE.toAlphaString(0.3),
-  });
-
-type Range = [Position, Position];
-interface AutoDepsDecorationsParams {
-  position: Position;
-}
-namespace AutoDepsDecorationsRequest {
-  export const type = new RequestType<
-    AutoDepsDecorationsParams,
-    Array<Range> | null,
-    void
-  >('react/autodeps_decorations');
-}
 
 export function activate(context: vscode.ExtensionContext) {
   const serverModule = context.asAbsolutePath(path.join('dist', 'server.js'));
@@ -71,31 +58,31 @@ export function activate(context: vscode.ExtensionContext) {
 
   vscode.languages.registerHoverProvider(documentSelector, {
     provideHover(_document, position, _token) {
-      client
-        .sendRequest(AutoDepsDecorationsRequest.type, {position})
-        .then(decorations => {
-          if (Array.isArray(decorations)) {
-            const decorationOptions = decorations.map(([start, end]) => {
-              return {
-                range: new vscode.Range(
-                  new vscode.Position(start.line, start.character),
-                  new vscode.Position(end.line, end.character),
-                ),
-                hoverMessage: 'Inferred as an effect dependency',
-              };
-            });
-            vscode.window.activeTextEditor?.setDecorations(
-              inferredEffectDepDecoration,
-              decorationOptions,
-            );
-          } else {
-            clearDecorations(inferredEffectDepDecoration);
-          }
-        });
+      requestAutoDepsDecorations(client, position, {shouldUpdateCurrent: true});
       return null;
     },
   });
 
+  vscode.workspace.onDidChangeTextDocument(async _e => {
+    const currentlyDecoratedAutoDepFnLoc = getCurrentlyDecoratedAutoDepFnLoc();
+    if (currentlyDecoratedAutoDepFnLoc !== null) {
+      requestAutoDepsDecorations(client, currentlyDecoratedAutoDepFnLoc.start, {
+        shouldUpdateCurrent: false,
+      });
+    }
+  });
+
+  vscode.commands.registerCommand(
+    'react.requestAutoDepsDecorations',
+    (position: Position) => {
+      requestAutoDepsDecorations(
+        client,
+        positionLiteralToVSCodePosition(position),
+        {shouldUpdateCurrent: true},
+      );
+    },
+  );
+
   client.registerProposedFeatures();
   client.start();
 }
@@ -106,9 +93,3 @@ export function deactivate(): Thenable<void> | undefined {
   }
   return;
 }
-
-export function clearDecorations(
-  decorationType: vscode.TextEditorDecorationType,
-) {
-  vscode.window.activeTextEditor?.setDecorations(decorationType, []);
-}

compiler/packages/react-forgive/client/src/mapping.ts

@@ -0,0 +1,15 @@
+import * as vscode from 'vscode';
+import {Position} from 'vscode-languageclient/node';
+
+export function positionLiteralToVSCodePosition(
+  position: Position,
+): vscode.Position {
+  return new vscode.Position(position.line, position.character);
+}
+
+export function positionsToRange(start: Position, end: Position): vscode.Range {
+  return new vscode.Range(
+    positionLiteralToVSCodePosition(start),
+    positionLiteralToVSCodePosition(end),
+  );
+}

compiler/packages/react-forgive/server/src/index.ts

@@ -10,6 +10,7 @@ import {
   CodeAction,
   CodeActionKind,
   CodeLens,
+  Command,
   createConnection,
   type InitializeParams,
   type InitializeResult,
@@ -96,6 +97,7 @@ connection.onInitialize((_params: InitializeParams) => {
     logger: {
       logEvent(_filename: string | null, event: LoggerEvent) {
         connection.console.info(`Received event: ${event.kind}`);
+        connection.console.debug(JSON.stringify(event, null, 2));
         if (event.kind === 'CompileSuccess') {
           compiledFns.add(event);
         }
@@ -191,7 +193,6 @@ connection.onCodeLensResolve(lens => {
 
 connection.onCodeAction(params => {
   connection.console.log('onCodeAction');
-  connection.console.log(JSON.stringify(params, null, 2));
   const codeActions: Array<CodeAction> = [];
   for (const codeActionEvent of codeActionEvents) {
     if (
@@ -200,39 +201,41 @@ connection.onCodeAction(params => {
         codeActionEvent.anchorRange,
       )
     ) {
-      codeActions.push(
-        CodeAction.create(
-          codeActionEvent.title,
-          {
-            changes: {
-              [params.textDocument.uri]: [
-                {
-                  newText: codeActionEvent.newText,
-                  range: codeActionEvent.editRange,
-                },
-              ],
-            },
+      const codeAction = CodeAction.create(
+        codeActionEvent.title,
+        {
+          changes: {
+            [params.textDocument.uri]: [
+              {
+                newText: codeActionEvent.newText,
+                range: codeActionEvent.editRange,
+              },
+            ],
           },
-          codeActionEvent.kind,
-        ),
+        },
+        codeActionEvent.kind,
       );
+      // After executing a codeaction, we want to draw autodep decorations again
+      codeAction.command = Command.create(
+        'Request autodeps decorations',
+        'react.requestAutoDepsDecorations',
+        codeActionEvent.anchorRange[0],
+      );
+      codeActions.push(codeAction);
     }
   }
   return codeActions;
 });
 
-connection.onCodeActionResolve(codeAction => {
-  connection.console.log('onCodeActionResolve');
-  connection.console.log(JSON.stringify(codeAction, null, 2));
-  return codeAction;
-});
-
+/**
+ * The client can request the server to compute autodeps decorations based on a currently selected
+ * position if the selected position is within an autodep eligible function call.
+ */
 connection.onRequest(AutoDepsDecorationsRequest.type, async params => {
   const position = params.position;
-  connection.console.debug('Client hovering on: ' + JSON.stringify(position));
-  for (const dec of autoDepsDecorations) {
-    if (isPositionWithinRange(position, dec.useEffectCallExpr)) {
-      return dec.decorations;
+  for (const decoration of autoDepsDecorations) {
+    if (isPositionWithinRange(position, decoration.useEffectCallExpr)) {
+      return decoration;
     }
   }
   return null;

compiler/packages/react-forgive/server/src/requests/autodepsdecorations.ts

@@ -13,7 +13,7 @@ export interface AutoDepsDecorationsParams {
 export namespace AutoDepsDecorationsRequest {
   export const type = new RequestType<
     AutoDepsDecorationsParams,
-    Array<Range> | null,
+    AutoDepsDecorationsLSPEvent,
     void
   >('react/autodeps_decorations');
 }