From 989b0cccc215e5c4692552b0cc01d23938dcf99b Mon Sep 17 00:00:00 2001
From: lauren <poteto@users.noreply.github.com>
Date: Fri, 7 Feb 2025 16:04:46 -0500
Subject: [PATCH] [compiler] Add simple walltime measurement (#32331)

Adds a new Timing logger event to the compiler which currently only
records the walltime of running the compiler from the time the babel
plugin's Program visitor enters to the time it exits.

To enable, run the compiler with `ENABLE_REACT_COMPILER_TIMINGS=1 ...`
or `export ENABLE_REACT_COMPILER_TIMINGS=1` to set it by default.
---
 .../src/Babel/BabelPlugin.ts                  | 93 +++++++++++++------
 .../src/Entrypoint/Options.ts                 |  4 +
 2 files changed, 67 insertions(+), 30 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Babel/BabelPlugin.ts b/compiler/packages/babel-plugin-react-compiler/src/Babel/BabelPlugin.ts
index c648c66043..aa49bda22b 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Babel/BabelPlugin.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Babel/BabelPlugin.ts
@@ -6,12 +6,15 @@
  */
 
 import type * as BabelCore from '@babel/core';
-import {compileProgram, parsePluginOptions} from '../Entrypoint';
+import {compileProgram, Logger, parsePluginOptions} from '../Entrypoint';
 import {
   injectReanimatedFlag,
   pipelineUsesReanimatedPlugin,
 } from '../Entrypoint/Reanimated';
 
+const ENABLE_REACT_COMPILER_TIMINGS =
+  process.env['ENABLE_REACT_COMPILER_TIMINGS'] === '1';
+
 /*
  * The React Forget Babel Plugin
  * @param {*} _babel
@@ -28,35 +31,65 @@ export default function BabelPluginReactCompiler(
        * prior to B, if A does not have a Program visitor and B does, B will run first. We always
        * want Forget to run true to source as possible.
        */
-      Program(prog, pass): void {
-        let opts = parsePluginOptions(pass.opts);
-        const isDev =
-          (typeof __DEV__ !== 'undefined' && __DEV__ === true) ||
-          process.env['NODE_ENV'] === 'development';
-        if (
-          opts.enableReanimatedCheck === true &&
-          pipelineUsesReanimatedPlugin(pass.file.opts.plugins)
-        ) {
-          opts = injectReanimatedFlag(opts);
-        }
-        if (
-          opts.environment.enableResetCacheOnSourceFileChanges !== false &&
-          isDev
-        ) {
-          opts = {
-            ...opts,
-            environment: {
-              ...opts.environment,
-              enableResetCacheOnSourceFileChanges: true,
-            },
-          };
-        }
-        compileProgram(prog, {
-          opts,
-          filename: pass.filename ?? null,
-          comments: pass.file.ast.comments ?? [],
-          code: pass.file.code,
-        });
+      Program: {
+        enter(prog, pass): void {
+          const filename = pass.filename ?? 'unknown';
+          if (ENABLE_REACT_COMPILER_TIMINGS === true) {
+            performance.mark(`${filename}:start`, {
+              detail: 'BabelPlugin:Program:start',
+            });
+          }
+          let opts = parsePluginOptions(pass.opts);
+          const isDev =
+            (typeof __DEV__ !== 'undefined' && __DEV__ === true) ||
+            process.env['NODE_ENV'] === 'development';
+          if (
+            opts.enableReanimatedCheck === true &&
+            pipelineUsesReanimatedPlugin(pass.file.opts.plugins)
+          ) {
+            opts = injectReanimatedFlag(opts);
+          }
+          if (
+            opts.environment.enableResetCacheOnSourceFileChanges !== false &&
+            isDev
+          ) {
+            opts = {
+              ...opts,
+              environment: {
+                ...opts.environment,
+                enableResetCacheOnSourceFileChanges: true,
+              },
+            };
+          }
+          compileProgram(prog, {
+            opts,
+            filename: pass.filename ?? null,
+            comments: pass.file.ast.comments ?? [],
+            code: pass.file.code,
+          });
+          if (ENABLE_REACT_COMPILER_TIMINGS === true) {
+            performance.mark(`${filename}:end`, {
+              detail: 'BabelPlugin:Program:end',
+            });
+          }
+        },
+        exit(_, pass): void {
+          if (ENABLE_REACT_COMPILER_TIMINGS === true) {
+            const filename = pass.filename ?? 'unknown';
+            const measurement = performance.measure(filename, {
+              start: `${filename}:start`,
+              end: `${filename}:end`,
+              detail: 'BabelPlugin:Program',
+            });
+            if ('logger' in pass.opts && pass.opts.logger != null) {
+              const logger: Logger = pass.opts.logger as Logger;
+              logger.logEvent(filename, {
+                kind: 'Timing',
+                measurement,
+              });
+            }
+          }
+        },
       },
     },
   };
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Options.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Options.ts
index fb951d25c5..35c2c4134e 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Options.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Options.ts
@@ -206,6 +206,10 @@ export type LoggerEvent =
       kind: 'PipelineError';
       fnLoc: t.SourceLocation | null;
       data: string;
+    }
+  | {
+      kind: 'Timing';
+      measurement: PerformanceMeasure;
     };
 
 export type Logger = {