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

/**
 * 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 {TextDocument} from 'vscode-languageserver-textdocument';
import {
  CodeAction,
  CodeActionKind,
  CodeLens,
  createConnection,
  type InitializeParams,
  type InitializeResult,
  Position,
  ProposedFeatures,
  TextDocuments,
  TextDocumentSyncKind,
} from 'vscode-languageserver/node';
import {compile, lastResult} from './compiler';
import {type PluginOptions} from 'babel-plugin-react-compiler/src';
import {resolveReactConfig} from './compiler/options';
import {
  type CompileSuccessEvent,
  type LoggerEvent,
  defaultOptions,
} from 'babel-plugin-react-compiler/src/Entrypoint/Options';
import {babelLocationToRange, getRangeFirstCharacter} from './compiler/compat';
import {
  type AutoDepsDecorationsLSPEvent,
  AutoDepsDecorationsRequest,
  mapCompilerEventToLSPEvent,
} from './requests/autodepsdecorations';
import {
  isPositionWithinRange,
  isRangeWithinRange,
  Range,
  sourceLocationToRange,
} from './utils/range';

const SUPPORTED_LANGUAGE_IDS = new Set([
  'javascript',
  'javascriptreact',
  'typescript',
  'typescriptreact',
]);

const connection = createConnection(ProposedFeatures.all);
const documents = new TextDocuments(TextDocument);

let compilerOptions: PluginOptions | null = null;
let compiledFns: Set<CompileSuccessEvent> = new Set();
let autoDepsDecorations: Array<AutoDepsDecorationsLSPEvent> = [];
let codeActionEvents: Array<CodeActionLSPEvent> = [];

type CodeActionLSPEvent = {
  title: string;
  kind: CodeActionKind;
  newText: string;
  anchorRange: Range;
  editRange: {start: Position; end: Position};
};

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));
        }
        if (event.kind === 'AutoDepsEligible') {
          const depArrayLoc = sourceLocationToRange(event.depArrayLoc);
          codeActionEvents.push({
            title: 'Use React Compiler inferred dependency array',
            kind: CodeActionKind.QuickFix,
            newText: '',
            anchorRange: sourceLocationToRange(event.fnLoc),
            editRange: {start: depArrayLoc[0], end: depArrayLoc[1]},
          });
        }
      },
    },
  };
  const result: InitializeResult = {
    capabilities: {
      textDocumentSync: TextDocumentSyncKind.Full,
      codeLensProvider: {resolveProvider: true},
      codeActionProvider: {resolveProvider: true},
    },
  };
  return result;
});

connection.onInitialized(() => {
  connection.console.log('initialized');
});

documents.onDidChangeContent(async event => {
  connection.console.info(`Changed: ${event.document.uri}`);
  resetState();
  if (SUPPORTED_LANGUAGE_IDS.has(event.document.languageId)) {
    const text = event.document.getText();
    try {
      await compile({
        text,
        file: event.document.uri,
        options: compilerOptions,
      });
    } catch (err) {
      if (err instanceof Error) {
        connection.console.error(err.stack ?? '');
      }
    }
  }
});

connection.onDidChangeWatchedFiles(change => {
  resetState();
  connection.console.log(
    change.changes.map(c => `File changed: ${c.uri}`).join('\n'),
  );
});

connection.onCodeLens(params => {
  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;
});

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 (
      isRangeWithinRange(
        [params.range.start, params.range.end],
        codeActionEvent.anchorRange,
      )
    ) {
      codeActions.push(
        CodeAction.create(
          codeActionEvent.title,
          {
            changes: {
              [params.textDocument.uri]: [
                {
                  newText: codeActionEvent.newText,
                  range: codeActionEvent.editRange,
                },
              ],
            },
          },
          codeActionEvent.kind,
        ),
      );
    }
  }
  return codeActions;
});

connection.onCodeActionResolve(codeAction => {
  connection.console.log('onCodeActionResolve');
  connection.console.log(JSON.stringify(codeAction, null, 2));
  return codeAction;
});

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;
    }
  }
  return null;
});

function resetState() {
  compiledFns.clear();
  autoDepsDecorations = [];
  codeActionEvents = [];
}

documents.listen(connection);
connection.listen();
connection.console.info(`React Analyzer running in node ${process.version}`);