[forgive] Add basic codelens provider (facebook#32476)

Adds a first codelens provider for successfully compiled functions. A
later PR will add an actual command that will fire when the codelens is
clicked

![Screenshot 2025-02-25 at 6 40
20 PM](https://github.com/user-attachments/assets/924586e0-f70a-45d1-b0e6-a89af9371c8d)

Author: poteto

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

@@ -176,41 +176,48 @@ export type CompilationMode = z.infer<typeof CompilationModeSchema>;
  *   babel or other unhandled exceptions).
  */
 export type LoggerEvent =
-  | {
-      kind: 'CompileError';
-      fnLoc: t.SourceLocation | null;
-      detail: CompilerErrorDetailOptions;
-    }
-  | {
-      kind: 'CompileDiagnostic';
-      fnLoc: t.SourceLocation | null;
-      detail: Omit<Omit<CompilerErrorDetailOptions, 'severity'>, 'suggestions'>;
-    }
-  | {
-      kind: 'CompileSkip';
-      fnLoc: t.SourceLocation | null;
-      reason: string;
-      loc: t.SourceLocation | null;
-    }
-  | {
-      kind: 'CompileSuccess';
-      fnLoc: t.SourceLocation | null;
-      fnName: string | null;
-      memoSlots: number;
-      memoBlocks: number;
-      memoValues: number;
-      prunedMemoBlocks: number;
-      prunedMemoValues: number;
-    }
-  | {
-      kind: 'PipelineError';
-      fnLoc: t.SourceLocation | null;
-      data: string;
-    }
-  | {
-      kind: 'Timing';
-      measurement: PerformanceMeasure;
-    };
+  | CompileSuccessEvent
+  | CompileErrorEvent
+  | CompileDiagnosticEvent
+  | CompileSkipEvent
+  | PipelineErrorEvent
+  | TimingEvent;
+
+export type CompileErrorEvent = {
+  kind: 'CompileError';
+  fnLoc: t.SourceLocation | null;
+  detail: CompilerErrorDetailOptions;
+};
+export type CompileDiagnosticEvent = {
+  kind: 'CompileDiagnostic';
+  fnLoc: t.SourceLocation | null;
+  detail: Omit<Omit<CompilerErrorDetailOptions, 'severity'>, 'suggestions'>;
+};
+export type CompileSuccessEvent = {
+  kind: 'CompileSuccess';
+  fnLoc: t.SourceLocation | null;
+  fnName: string | null;
+  memoSlots: number;
+  memoBlocks: number;
+  memoValues: number;
+  prunedMemoBlocks: number;
+  prunedMemoValues: number;
+};
+export type CompileSkipEvent = {
+  kind: 'CompileSkip';
+  fnLoc: t.SourceLocation | null;
+  reason: string;
+  loc: t.SourceLocation | null;
+};
+export type PipelineErrorEvent = {
+  kind: 'PipelineError';
+  fnLoc: t.SourceLocation | null;
+  data: string;
+};
+export type TimingEvent = {
+  kind: 'Timing';
+  measurement: PerformanceMeasure;
+};
 
 export type Logger = {
   logEvent: (filename: string | null, event: LoggerEvent) => void;

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

@@ -0,0 +1,22 @@
+import {SourceLocation} from 'babel-plugin-react-compiler/src';
+import {type Range} from 'vscode-languageserver';
+
+export function babelLocationToRange(loc: SourceLocation): Range | null {
+  if (typeof loc === 'symbol') {
+    return null;
+  }
+  return {
+    start: {line: loc.start.line - 1, character: loc.start.column},
+    end: {line: loc.end.line - 1, character: loc.end.column},
+  };
+}
+
+/**
+ * Refine range to only the first character.
+ */
+export function getRangeFirstCharacter(range: Range): Range {
+  return {
+    start: range.start,
+    end: range.start,
+  };
+}

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

@@ -11,10 +11,12 @@ import BabelPluginReactCompiler, {
   type PluginOptions,
 } from 'babel-plugin-react-compiler/src';
 import * as babelParser from 'prettier/plugins/babel.js';
-import * as estreeParser from 'prettier/plugins/estree';
+import estreeParser from 'prettier/plugins/estree';
 import * as typescriptParser from 'prettier/plugins/typescript';
 import * as prettier from 'prettier/standalone';
 
+export let lastResult: BabelCore.BabelFileResult | null = null;
+
 type CompileOptions = {
   text: string;
   file: string;
@@ -24,14 +26,17 @@ export async function compile({
   text,
   file,
   options,
-}: CompileOptions): Promise<BabelCore.BabelFileResult> {
+}: CompileOptions): Promise<BabelCore.BabelFileResult | null> {
   const ast = await parseAsync(text, {
     sourceFileName: file,
     parserOpts: {
       plugins: ['typescript', 'jsx'],
     },
     sourceType: 'module',
   });
+  if (ast == null) {
+    return null;
+  }
   const plugins =
     options != null
       ? [[BabelPluginReactCompiler, options]]
@@ -54,5 +59,8 @@ export async function compile({
     parser: 'babel-ts',
     plugins: [babelParser, estreeParser, typescriptParser],
   });
+  if (result.code != null) {
+    lastResult = result;
+  }
   return result;
 }

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

@@ -7,17 +7,23 @@
 
 import {TextDocument} from 'vscode-languageserver-textdocument';
 import {
+  CodeLens,
   createConnection,
   type InitializeParams,
   type InitializeResult,
   ProposedFeatures,
   TextDocuments,
   TextDocumentSyncKind,
 } from 'vscode-languageserver/node';
-import {compile} from './compiler';
+import {compile, lastResult} from './compiler';
 import {type PluginOptions} from 'babel-plugin-react-compiler/src';
 import {resolveReactConfig} from './compiler/options';
-import {type BabelFileResult} from '@babel/core';
+import {
+  CompileSuccessEvent,
+  defaultOptions,
+  LoggerEvent,
+} from 'babel-plugin-react-compiler/src/Entrypoint/Options';
+import {babelLocationToRange, getRangeFirstCharacter} from './compiler/compat';
 
 const SUPPORTED_LANGUAGE_IDS = new Set([
   'javascript',
@@ -30,11 +36,21 @@ const connection = createConnection(ProposedFeatures.all);
 const documents = new TextDocuments(TextDocument);
 
 let compilerOptions: PluginOptions | null = null;
-let lastResult: BabelFileResult | null = null;
+let compiledFns: Set<CompileSuccessEvent> = new Set();
 
 connection.onInitialize((_params: InitializeParams) => {
   // TODO(@poteto) get config fr
-  compilerOptions = resolveReactConfig('.');
+  compilerOptions = resolveReactConfig('.') ?? defaultOptions;
+  compilerOptions = {
+    ...compilerOptions,
+    logger: {
+      logEvent(_filename: string | null, event: LoggerEvent) {
+        if (event.kind === 'CompileSuccess') {
+          compiledFns.add(event);
+        }
+      },
+    },
+  };
   const result: InitializeResult = {
     capabilities: {
       textDocumentSync: TextDocumentSyncKind.Full,
@@ -48,44 +64,58 @@ connection.onInitialized(() => {
   connection.console.log('initialized');
 });
 
-documents.onDidOpen(async event => {
-  if (SUPPORTED_LANGUAGE_IDS.has(event.document.languageId)) {
-    const text = event.document.getText();
-    const result = await compile({
-      text,
-      file: event.document.uri,
-      options: compilerOptions,
-    });
-    if (result.code != null) {
-      lastResult = result;
-    }
-  }
-});
-
 documents.onDidChangeContent(async event => {
+  connection.console.info(`Changed: ${event.document.uri}`);
+  compiledFns.clear();
   if (SUPPORTED_LANGUAGE_IDS.has(event.document.languageId)) {
     const text = event.document.getText();
-    const result = await compile({
+    await compile({
       text,
       file: event.document.uri,
       options: compilerOptions,
     });
-    if (result.code != null) {
-      lastResult = result;
-    }
   }
 });
 
 connection.onDidChangeWatchedFiles(change => {
+  compiledFns.clear();
   connection.console.log(
     change.changes.map(c => `File changed: ${c.uri}`).join('\n'),
   );
 });
 
 connection.onCodeLens(params => {
-  connection.console.log('lastResult: ' + JSON.stringify(lastResult, null, 2));
-  connection.console.log('params: ' + JSON.stringify(params, null, 2));
-  return [];
+  connection.console.info(`Handling codelens for: ${params.textDocument.uri}`);
+  if (compiledFns.size === 0) {
+    return;
+  }
+  const lenses: Array<CodeLens> = [];
+  for (const compiled of compiledFns) {
+    if (compiled.fnLoc != null) {
+      const fnLoc = babelLocationToRange(compiled.fnLoc);
+      if (fnLoc === null) continue;
+      const lens = CodeLens.create(
+        getRangeFirstCharacter(fnLoc),
+        compiled.fnLoc,
+      );
+      if (lastResult?.code != null) {
+        lens.command = {
+          title: 'Optimized by React Compiler',
+          command: 'todo',
+        };
+      }
+      lenses.push(lens);
+    }
+  }
+  return lenses;
+});
+
+connection.onCodeLensResolve(lens => {
+  connection.console.info(`Resolving codelens for: ${JSON.stringify(lens)}`);
+  if (lastResult?.code != null) {
+    connection.console.log(lastResult.code);
+  }
+  return lens;
 });
 
 documents.listen(connection);