[forgive] Hacky first pass at adding decorations for inferred deps (#32998)

Draws basic decorations for inferred deps on hover.

Co-authored-by: Jordan Brown <jmbrown@meta.com>
---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/32998).
* #33002
* #33001
* #33000
* #32999
* __->__ #32998

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

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Options.ts

@@ -222,8 +222,8 @@ export type TimingEvent = {
 };
 export type AutoDepsDecorations = {
   kind: 'AutoDepsDecorations';
-  useEffectCallExpr: t.SourceLocation | null;
-  decorations: Array<t.SourceLocation | null>;
+  useEffectCallExpr: t.SourceLocation;
+  decorations: Array<t.SourceLocation>;
 };
 
 export type Logger = {

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

@@ -1,3 +1,11 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import * as t from '@babel/types';
 import {CompilerError, SourceLocation} from '..';
 import {
   ArrayExpression,
@@ -212,14 +220,20 @@ export function inferEffectDependencies(fn: HIRFunction): void {
             }
 
             // For LSP autodeps feature.
-            fn.env.logger?.logEvent(fn.env.filename, {
-              kind: 'AutoDepsDecorations',
-              useEffectCallExpr:
-                typeof value.loc !== 'symbol' ? value.loc : null,
-              decorations: collectDepUsages(usedDeps, fnExpr.value).map(loc =>
-                typeof loc !== 'symbol' ? loc : null,
-              ),
-            });
+            const decorations: Array<t.SourceLocation> = [];
+            for (const loc of collectDepUsages(usedDeps, fnExpr.value)) {
+              if (typeof loc === 'symbol') {
+                continue;
+              }
+              decorations.push(loc);
+            }
+            if (typeof value.loc !== 'symbol') {
+              fn.env.logger?.logEvent(fn.env.filename, {
+                kind: 'AutoDepsDecorations',
+                useEffectCallExpr: value.loc,
+                decorations,
+              });
+            }
 
             newInstructions.push({
               id: makeInstructionId(0),

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

@@ -1,17 +1,22 @@
 import * as path from 'path';
-import {ExtensionContext, window as Window} from 'vscode';
+import * as vscode from 'vscode';
 
 import {
   LanguageClient,
   LanguageClientOptions,
+  Position,
   ServerOptions,
   TransportKind,
 } from 'vscode-languageclient/node';
 
 let client: LanguageClient;
 
-export function activate(context: ExtensionContext) {
+export function activate(context: vscode.ExtensionContext) {
   const serverModule = context.asAbsolutePath(path.join('dist', 'server.js'));
+  const documentSelector = [
+    {scheme: 'file', language: 'javascriptreact'},
+    {scheme: 'file', language: 'typescriptreact'},
+  ];
 
   // If the extension is launched in debug mode then the debug server options are used
   // Otherwise the run options are used
@@ -27,10 +32,7 @@ export function activate(context: ExtensionContext) {
   };
 
   const clientOptions: LanguageClientOptions = {
-    documentSelector: [
-      {scheme: 'file', language: 'javascriptreact'},
-      {scheme: 'file', language: 'typescriptreact'},
-    ],
+    documentSelector,
     progressOnInitialization: true,
   };
 
@@ -43,12 +45,38 @@ export function activate(context: ExtensionContext) {
       clientOptions,
     );
   } catch {
-    Window.showErrorMessage(
+    vscode.window.showErrorMessage(
       `React Analyzer couldn't be started. See the output channel for details.`,
     );
     return;
   }
 
+  vscode.languages.registerHoverProvider(documentSelector, {
+    provideHover(_document, position, _token) {
+      client
+        .sendRequest('react/autodepsdecorations', position)
+        .then((decorations: Array<[Position, Position]>) => {
+          for (const [start, end] of decorations) {
+            const range = new vscode.Range(
+              new vscode.Position(start.line, start.character),
+              new vscode.Position(end.line, end.character),
+            );
+            const vscodeDecoration =
+              vscode.window.createTextEditorDecorationType({
+                backgroundColor: 'red',
+              });
+            vscode.window.activeTextEditor?.setDecorations(vscodeDecoration, [
+              {
+                range,
+                hoverMessage: 'hehe',
+              },
+            ]);
+          }
+        });
+      return null;
+    },
+  });
+
   client.registerProposedFeatures();
   client.start();
 }

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

@@ -0,0 +1,18 @@
+import {AutoDepsDecorations} from 'babel-plugin-react-compiler/src/Entrypoint';
+import {Position} from 'vscode-languageserver-textdocument';
+import {sourceLocationToRange} from '../utils/lsp-adapter';
+
+export type Range = [Position, Position];
+export type AutoDepsDecorationsLSPEvent = {
+  useEffectCallExpr: Range;
+  decorations: Array<Range>;
+};
+
+export function mapCompilerEventToLSPEvent(
+  event: AutoDepsDecorations,
+): AutoDepsDecorationsLSPEvent {
+  return {
+    useEffectCallExpr: sourceLocationToRange(event.useEffectCallExpr),
+    decorations: event.decorations.map(sourceLocationToRange),
+  };
+}

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

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {TextDocument} from 'vscode-languageserver-textdocument';
+import {Position, TextDocument} from 'vscode-languageserver-textdocument';
 import {
   CodeLens,
   createConnection,
@@ -24,6 +24,10 @@ import {
   LoggerEvent,
 } from 'babel-plugin-react-compiler/src/Entrypoint/Options';
 import {babelLocationToRange, getRangeFirstCharacter} from './compiler/compat';
+import {
+  AutoDepsDecorationsLSPEvent,
+  mapCompilerEventToLSPEvent,
+} from './custom-requests/autodepsdecorations';
 
 const SUPPORTED_LANGUAGE_IDS = new Set([
   'javascript',
@@ -37,17 +41,48 @@ const documents = new TextDocuments(TextDocument);
 
 let compilerOptions: PluginOptions | null = null;
 let compiledFns: Set<CompileSuccessEvent> = new Set();
+let autoDepsDecorations: Array<AutoDepsDecorationsLSPEvent> = [];
 
 connection.onInitialize((_params: InitializeParams) => {
   // TODO(@poteto) get config fr
   compilerOptions = resolveReactConfig('.') ?? defaultOptions;
   compilerOptions = {
     ...compilerOptions,
+    environment: {
+      ...compilerOptions.environment,
+      inferEffectDependencies: [
+        {
+          function: {
+            importSpecifierName: 'useEffect',
+            source: 'react',
+          },
+          numRequiredArgs: 1,
+        },
+        {
+          function: {
+            importSpecifierName: 'useSpecialEffect',
+            source: 'shared-runtime',
+          },
+          numRequiredArgs: 2,
+        },
+        {
+          function: {
+            importSpecifierName: 'default',
+            source: 'useEffectWrapper',
+          },
+          numRequiredArgs: 1,
+        },
+      ],
+    },
     logger: {
       logEvent(_filename: string | null, event: LoggerEvent) {
+        connection.console.info(`Received event: ${event.kind}`);
         if (event.kind === 'CompileSuccess') {
           compiledFns.add(event);
         }
+        if (event.kind === 'AutoDepsDecorations') {
+          autoDepsDecorations.push(mapCompilerEventToLSPEvent(event));
+        }
       },
     },
   };
@@ -67,6 +102,7 @@ connection.onInitialized(() => {
 documents.onDidChangeContent(async event => {
   connection.console.info(`Changed: ${event.document.uri}`);
   compiledFns.clear();
+  autoDepsDecorations = [];
   if (SUPPORTED_LANGUAGE_IDS.has(event.document.languageId)) {
     const text = event.document.getText();
     await compile({
@@ -79,6 +115,7 @@ documents.onDidChangeContent(async event => {
 
 connection.onDidChangeWatchedFiles(change => {
   compiledFns.clear();
+  autoDepsDecorations = [];
   connection.console.log(
     change.changes.map(c => `File changed: ${c.uri}`).join('\n'),
   );
@@ -118,6 +155,25 @@ connection.onCodeLensResolve(lens => {
   return lens;
 });
 
+connection.onRequest('react/autodepsdecorations', (position: Position) => {
+  connection.console.log('Client hovering on: ' + JSON.stringify(position));
+  connection.console.log(JSON.stringify(autoDepsDecorations, null, 2));
+
+  for (const dec of autoDepsDecorations) {
+    // TODO: extract to helper
+    if (
+      position.line >= dec.useEffectCallExpr[0].line &&
+      position.line <= dec.useEffectCallExpr[1].line
+    ) {
+      connection.console.log(
+        'found decoration: ' + JSON.stringify(dec.decorations),
+      );
+      return dec.decorations;
+    }
+  }
+  return null;
+});
+
 documents.listen(connection);
 connection.listen();
 connection.console.info(`React Analyzer running in node ${process.version}`);

compiler/packages/react-forgive/server/src/utils/lsp-adapter.ts

@@ -0,0 +1,11 @@
+import * as t from '@babel/types';
+import {Position} from 'vscode-languageserver/node';
+
+export function sourceLocationToRange(
+  loc: t.SourceLocation,
+): [Position, Position] {
+  return [
+    {line: loc.start.line - 1, character: loc.start.column},
+    {line: loc.end.line - 1, character: loc.end.column},
+  ];
+}