[forgive] Refactor inferred deps

Refactor.

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

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

@@ -183,7 +183,7 @@ export type LoggerEvent =
   | CompileSkipEvent
   | PipelineErrorEvent
   | TimingEvent
-  | AutoDepsDecorations;
+  | AutoDepsDecorationsEvent;
 
 export type CompileErrorEvent = {
   kind: 'CompileError';
@@ -220,7 +220,7 @@ export type TimingEvent = {
   kind: 'Timing';
   measurement: PerformanceMeasure;
 };
-export type AutoDepsDecorations = {
+export type AutoDepsDecorationsEvent = {
   kind: 'AutoDepsDecorations';
   useEffectCallExpr: t.SourceLocation;
   decorations: Array<t.SourceLocation>;

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

@@ -380,7 +380,10 @@ function collectDepUsages(
 
   for (const [, block] of fnExpr.loweredFunc.func.body.blocks) {
     for (const instr of block.instructions) {
-      if (instr.value.kind === 'LoadLocal') {
+      if (
+        instr.value.kind === 'LoadLocal' &&
+        identifiers.has(instr.value.place.identifier.id)
+      ) {
         loadedDeps.add(instr.lvalue.identifier.id);
       }
       for (const place of eachInstructionOperand(instr)) {

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

@@ -0,0 +1,73 @@
+type RGB = [number, number, number];
+
+const int = Math.floor;
+
+export class Color {
+  constructor(
+    private r: number,
+    private g: number,
+    private b: number,
+  ) {}
+
+  toAlphaString(a: number) {
+    return this.toCssString(a);
+  }
+  toString() {
+    return this.toCssString(1);
+  }
+
+  /**
+   * Adjust the color by a multiplier to lighten (`> 1.0`) or darken (`< 1.0`) the color. Returns a new
+   * instance.
+   */
+  adjusted(mult: number) {
+    const adjusted = Color.redistribute([
+      this.r * mult,
+      this.g * mult,
+      this.b * mult,
+    ]);
+    return new Color(...adjusted);
+  }
+
+  private toCssString(a: number) {
+    return `rgba(${this.r},${this.g},${this.b},${a})`;
+  }
+  /**
+   * Redistributes rgb, maintaing hue until its clamped.
+   * https://stackoverflow.com/a/141943
+   */
+  private static redistribute([r, g, b]: RGB): RGB {
+    const threshold = 255.999;
+    const max = Math.max(r, g, b);
+    if (max <= threshold) {
+      return [int(r), int(g), int(b)];
+    }
+    const total = r + g + b;
+    if (total >= 3 * threshold) {
+      return [int(threshold), int(threshold), int(threshold)];
+    }
+    const x = (3 * threshold - total) / (3 * max - total);
+    const gray = threshold - x * max;
+    return [int(gray + x * r), int(gray + x * g), int(gray + x * b)];
+  }
+}
+
+export const BLACK = new Color(0, 0, 0);
+export const WHITE = new Color(255, 255, 255);
+
+const COLOR_POOL = [
+  new Color(249, 65, 68),
+  new Color(243, 114, 44),
+  new Color(248, 150, 30),
+  new Color(249, 132, 74),
+  new Color(249, 199, 79),
+  new Color(144, 190, 109),
+  new Color(67, 170, 139),
+  new Color(77, 144, 142),
+  new Color(87, 117, 144),
+  new Color(39, 125, 161),
+];
+
+export function getColorFor(index: number): Color {
+  return COLOR_POOL[Math.abs(index) % COLOR_POOL.length]!;
+}

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

@@ -4,12 +4,30 @@ import * as vscode from 'vscode';
 import {
   LanguageClient,
   LanguageClientOptions,
-  Position,
+  type Position,
+  RequestType,
   ServerOptions,
   TransportKind,
 } from 'vscode-languageclient/node';
+import {WHITE} from './colors';
 
 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'));
@@ -54,23 +72,24 @@ export function activate(context: vscode.ExtensionContext) {
   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),
+        .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,
             );
-            const vscodeDecoration =
-              vscode.window.createTextEditorDecorationType({
-                backgroundColor: 'red',
-              });
-            vscode.window.activeTextEditor?.setDecorations(vscodeDecoration, [
-              {
-                range,
-                hoverMessage: 'hehe',
-              },
-            ]);
+          } else {
+            clearDecorations(inferredEffectDepDecoration);
           }
         });
       return null;
@@ -85,4 +104,11 @@ export function deactivate(): Thenable<void> | undefined {
   if (client !== undefined) {
     return client.stop();
   }
+  return;
+}
+
+export function clearDecorations(
+  decorationType: vscode.TextEditorDecorationType,
+) {
+  vscode.window.activeTextEditor?.setDecorations(decorationType, []);
 }

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

@@ -1,18 +0,0 @@
-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 {Position, TextDocument} from 'vscode-languageserver-textdocument';
+import {TextDocument} from 'vscode-languageserver-textdocument';
 import {
   CodeLens,
   createConnection,
@@ -19,15 +19,17 @@ import {compile, lastResult} from './compiler';
 import {type PluginOptions} from 'babel-plugin-react-compiler/src';
 import {resolveReactConfig} from './compiler/options';
 import {
-  CompileSuccessEvent,
+  type CompileSuccessEvent,
+  type LoggerEvent,
   defaultOptions,
-  LoggerEvent,
 } from 'babel-plugin-react-compiler/src/Entrypoint/Options';
 import {babelLocationToRange, getRangeFirstCharacter} from './compiler/compat';
 import {
-  AutoDepsDecorationsLSPEvent,
+  type AutoDepsDecorationsLSPEvent,
+  AutoDepsDecorationsRequest,
   mapCompilerEventToLSPEvent,
-} from './custom-requests/autodepsdecorations';
+} from './requests/autodepsdecorations';
+import {isPositionWithinRange} from './utils/range';
 
 const SUPPORTED_LANGUAGE_IDS = new Set([
   'javascript',
@@ -155,19 +157,11 @@ 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));
-
+connection.onRequest(AutoDepsDecorationsRequest.type, async params => {
+  const position = params.position;
+  connection.console.debug('Client hovering on: ' + JSON.stringify(position));
   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),
-      );
+    if (isPositionWithinRange(position, dec.useEffectCallExpr)) {
       return dec.decorations;
     }
   }

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

@@ -0,0 +1,28 @@
+import {type AutoDepsDecorationsEvent} from 'babel-plugin-react-compiler/src/Entrypoint';
+import {type Position} from 'vscode-languageserver-textdocument';
+import {RequestType} from 'vscode-languageserver/node';
+import {type Range, sourceLocationToRange} from '../utils/range';
+
+export type AutoDepsDecorationsLSPEvent = {
+  useEffectCallExpr: Range;
+  decorations: Array<Range>;
+};
+export interface AutoDepsDecorationsParams {
+  position: Position;
+}
+export namespace AutoDepsDecorationsRequest {
+  export const type = new RequestType<
+    AutoDepsDecorationsParams,
+    Array<Range> | null,
+    void
+  >('react/autodeps_decorations');
+}
+
+export function mapCompilerEventToLSPEvent(
+  event: AutoDepsDecorationsEvent,
+): AutoDepsDecorationsLSPEvent {
+  return {
+    useEffectCallExpr: sourceLocationToRange(event.useEffectCallExpr),
+    decorations: event.decorations.map(sourceLocationToRange),
+  };
+}

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

@@ -1,11 +0,0 @@
-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},
-  ];
-}

compiler/packages/react-forgive/server/src/utils/range.ts

@@ -0,0 +1,19 @@
+import * as t from '@babel/types';
+import {type Position} from 'vscode-languageserver/node';
+
+export type Range = [Position, Position];
+export function isPositionWithinRange(
+  position: Position,
+  [start, end]: Range,
+): boolean {
+  return position.line >= start.line && position.line <= end.line;
+}
+
+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},
+  ];
+}

compiler/packages/react-forgive/tsconfig.json

@@ -0,0 +1,16 @@
+{
+  "extends": "@tsconfig/strictest/tsconfig.json",
+  "compilerOptions": {
+    "module": "Node16",
+    "moduleResolution": "Node16",
+    "rootDir": "../",
+    "noEmit": true,
+    "jsx": "react-jsxdev",
+    "lib": ["ES2022"],
+
+    "target": "ES2022",
+    "importsNotUsedAsValues": "remove",
+  },
+  "exclude": ["node_modules"],
+  "include": ["server/src/**/*.ts", "client/src/**/*.ts"],
+}