[compiler][be] Logger based debug printing in test runner

Avoid mutable logging enabled state and writing to `process.stdout` within our babel transform.

compiler/packages/babel-plugin-react-compiler/package.json

@@ -42,9 +42,7 @@
     "babel-jest": "^29.0.3",
     "babel-plugin-fbt": "^1.0.0",
     "babel-plugin-fbt-runtime": "^1.0.0",
-    "chalk": "4",
     "eslint": "^8.57.1",
-    "glob": "^7.1.6",
     "invariant": "^2.2.4",
     "jest": "^29.0.3",
     "jest-environment-jsdom": "^29.0.3",

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

@@ -79,13 +79,6 @@ import {
   rewriteInstructionKindsBasedOnReassignment,
 } from '../SSA';
 import {inferTypes} from '../TypeInference';
-import {
-  logCodegenFunction,
-  logDebug,
-  logHIRFunction,
-  logReactiveFunction,
-} from '../Utils/logger';
-import {assertExhaustive} from '../Utils/utils';
 import {
   validateContextVariableLValues,
   validateHooksUsage,
@@ -139,13 +132,7 @@ function run(
     name: 'EnvironmentConfig',
     value: prettyFormat(env.config),
   });
-  printLog({
-    kind: 'debug',
-    name: 'EnvironmentConfig',
-    value: prettyFormat(env.config),
-  });
-  const ast = runWithEnvironment(func, env);
-  return ast;
+  return runWithEnvironment(func, env);
 }
 
 /*
@@ -158,10 +145,8 @@ function runWithEnvironment(
   >,
   env: Environment,
 ): CodegenFunction {
-  const log = (value: CompilerPipelineValue): CompilerPipelineValue => {
-    printLog(value);
+  const log = (value: CompilerPipelineValue): void => {
     env.logger?.debugLogIRs?.(value);
-    return value;
   };
   const hir = lower(func, env).unwrap();
   log({kind: 'hir', name: 'HIR', value: hir});
@@ -545,28 +530,3 @@ export function compileFn(
     code,
   );
 }
-
-function printLog(value: CompilerPipelineValue): CompilerPipelineValue {
-  switch (value.kind) {
-    case 'ast': {
-      logCodegenFunction(value.name, value.value);
-      break;
-    }
-    case 'hir': {
-      logHIRFunction(value.name, value.value);
-      break;
-    }
-    case 'reactive': {
-      logReactiveFunction(value.name, value.value);
-      break;
-    }
-    case 'debug': {
-      logDebug(value.name, value.value);
-      break;
-    }
-    default: {
-      assertExhaustive(value, 'Unexpected compilation kind');
-    }
-  }
-  return value;
-}

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

@@ -19,7 +19,6 @@ import {
 import {deadCodeElimination} from '../Optimization';
 import {inferReactiveScopeVariables} from '../ReactiveScopes';
 import {rewriteInstructionKindsBasedOnReassignment} from '../SSA';
-import {logHIRFunction} from '../Utils/logger';
 import {inferMutableContextVariables} from './InferMutableContextVariables';
 import {inferMutableRanges} from './InferMutableRanges';
 import inferReferenceEffects from './InferReferenceEffects';
@@ -112,7 +111,11 @@ function lower(func: HIRFunction): void {
   rewriteInstructionKindsBasedOnReassignment(func);
   inferReactiveScopeVariables(func);
   inferMutableContextVariables(func);
-  logHIRFunction('AnalyseFunction (inner)', func);
+  func.env.logger?.debugLogIRs?.({
+    kind: 'hir',
+    name: 'AnalyseFunction (inner)',
+    value: func,
+  });
 }
 
 function infer(

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -25,7 +25,6 @@ import {
   eachPatternOperand,
 } from '../HIR/visitors';
 import DisjointSet from '../Utils/DisjointSet';
-import {logHIRFunction} from '../Utils/logger';
 import {assertExhaustive} from '../Utils/utils';
 
 /*
@@ -156,7 +155,11 @@ export function inferReactiveScopeVariables(fn: HIRFunction): void {
       scope.range.end > maxInstruction + 1
     ) {
       // Make it easier to debug why the error occurred
-      logHIRFunction('InferReactiveScopeVariables (invalid scope)', fn);
+      fn.env.logger?.debugLogIRs?.({
+        kind: 'hir',
+        name: 'InferReactiveScopeVariables (invalid scope)',
+        value: fn,
+      });
       CompilerError.invariant(false, {
         reason: `Invalid mutable range for scope`,
         loc: GeneratedSource,

compiler/packages/babel-plugin-react-compiler/src/Utils/logger.ts

@@ -1,110 +0,0 @@
-/**
- * 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 generate from '@babel/generator';
-import * as t from '@babel/types';
-import chalk from 'chalk';
-import {HIR, HIRFunction, ReactiveFunction} from '../HIR/HIR';
-import {printFunctionWithOutlined, printHIR} from '../HIR/PrintHIR';
-import {CodegenFunction} from '../ReactiveScopes';
-import {printReactiveFunctionWithOutlined} from '../ReactiveScopes/PrintReactiveFunction';
-
-let ENABLED: boolean = false;
-
-let lastLogged: string;
-
-export function toggleLogging(enabled: boolean): void {
-  ENABLED = enabled;
-}
-
-export function logDebug(step: string, value: string): void {
-  if (ENABLED) {
-    process.stdout.write(`${chalk.green(step)}:\n${value}\n\n`);
-  }
-}
-
-export function logHIR(step: string, ir: HIR): void {
-  if (ENABLED) {
-    const printed = printHIR(ir);
-    if (printed !== lastLogged) {
-      lastLogged = printed;
-      process.stdout.write(`${chalk.green(step)}:\n${printed}\n\n`);
-    } else {
-      process.stdout.write(`${chalk.blue(step)}: (no change)\n\n`);
-    }
-  }
-}
-
-export function logCodegenFunction(step: string, fn: CodegenFunction): void {
-  if (ENABLED) {
-    let printed: string | null = null;
-    try {
-      const node = t.functionDeclaration(
-        fn.id,
-        fn.params,
-        fn.body,
-        fn.generator,
-        fn.async,
-      );
-      const ast = generate(node);
-      printed = ast.code;
-    } catch (e) {
-      let errMsg: string;
-      if (
-        typeof e === 'object' &&
-        e != null &&
-        'message' in e &&
-        typeof e.message === 'string'
-      ) {
-        errMsg = e.message.toString();
-      } else {
-        errMsg = '[empty]';
-      }
-      console.log('Error formatting AST: ' + errMsg);
-    }
-    if (printed === null) {
-      return;
-    }
-    if (printed !== lastLogged) {
-      lastLogged = printed;
-      process.stdout.write(`${chalk.green(step)}:\n${printed}\n\n`);
-    } else {
-      process.stdout.write(`${chalk.blue(step)}: (no change)\n\n`);
-    }
-  }
-}
-
-export function logHIRFunction(step: string, fn: HIRFunction): void {
-  if (ENABLED) {
-    const printed = printFunctionWithOutlined(fn);
-    if (printed !== lastLogged) {
-      lastLogged = printed;
-      process.stdout.write(`${chalk.green(step)}:\n${printed}\n\n`);
-    } else {
-      process.stdout.write(`${chalk.blue(step)}: (no change)\n\n`);
-    }
-  }
-}
-
-export function logReactiveFunction(step: string, fn: ReactiveFunction): void {
-  if (ENABLED) {
-    const printed = printReactiveFunctionWithOutlined(fn);
-    if (printed !== lastLogged) {
-      lastLogged = printed;
-      process.stdout.write(`${chalk.green(step)}:\n${printed}\n\n`);
-    } else {
-      process.stdout.write(`${chalk.blue(step)}: (no change)\n\n`);
-    }
-  }
-}
-
-export function log(fn: () => string): void {
-  if (ENABLED) {
-    const message = fn();
-    process.stdout.write(message.trim() + '\n\n');
-  }
-}

compiler/packages/snap/src/compiler.ts

@@ -19,6 +19,7 @@ import type {
   PanicThresholdOptions,
   PluginOptions,
   CompilerReactTarget,
+  CompilerPipelineValue,
 } from 'babel-plugin-react-compiler/src/Entrypoint';
 import type {Effect, ValueKind} from 'babel-plugin-react-compiler/src/HIR';
 import type {
@@ -45,6 +46,7 @@ export function parseLanguage(source: string): 'flow' | 'typescript' {
 function makePluginOptions(
   firstLine: string,
   parseConfigPragmaFn: typeof ParseConfigPragma,
+  debugIRLogger: (value: CompilerPipelineValue) => void,
   EffectEnum: typeof Effect,
   ValueKindEnum: typeof ValueKind,
 ): [PluginOptions, Array<{filename: string | null; event: LoggerEvent}>] {
@@ -182,15 +184,15 @@ function makePluginOptions(
       .filter(s => s.length > 0);
   }
 
-  let logs: Array<{filename: string | null; event: LoggerEvent}> = [];
-  let logger: Logger | null = null;
-  if (firstLine.includes('@logger')) {
-    logger = {
-      logEvent(filename: string | null, event: LoggerEvent): void {
-        logs.push({filename, event});
-      },
-    };
-  }
+  const logs: Array<{filename: string | null; event: LoggerEvent}> = [];
+  const logger: Logger = {
+    logEvent: firstLine.includes('@logger')
+      ? (filename, event) => {
+          logs.push({filename, event});
+        }
+      : () => {},
+    debugLogIRs: debugIRLogger,
+  };
 
   const config = parseConfigPragmaFn(firstLine);
   const options = {
@@ -338,6 +340,7 @@ export async function transformFixtureInput(
   parseConfigPragmaFn: typeof ParseConfigPragma,
   plugin: BabelCore.PluginObj,
   includeEvaluator: boolean,
+  debugIRLogger: (value: CompilerPipelineValue) => void,
   EffectEnum: typeof Effect,
   ValueKindEnum: typeof ValueKind,
 ): Promise<{kind: 'ok'; value: TransformResult} | {kind: 'err'; msg: string}> {
@@ -365,6 +368,7 @@ export async function transformFixtureInput(
   const [options, logs] = makePluginOptions(
     firstLine,
     parseConfigPragmaFn,
+    debugIRLogger,
     EffectEnum,
     ValueKindEnum,
   );

compiler/packages/snap/src/constants.ts

@@ -18,11 +18,17 @@ export const COMPILER_PATH = path.join(
   'BabelPlugin.js',
 );
 export const COMPILER_INDEX_PATH = path.join(process.cwd(), 'dist', 'index');
-export const LOGGER_PATH = path.join(
+export const PRINT_HIR_PATH = path.join(
   process.cwd(),
   'dist',
-  'Utils',
-  'logger.js',
+  'HIR',
+  'PrintHIR.js',
+);
+export const PRINT_REACTIVE_IR_PATH = path.join(
+  process.cwd(),
+  'dist',
+  'ReactiveScopes',
+  'PrintReactiveFunction.js',
 );
 export const PARSE_CONFIG_PRAGMA_PATH = path.join(
   process.cwd(),

compiler/packages/snap/src/runner-worker.ts

@@ -8,16 +8,21 @@
 import {codeFrameColumns} from '@babel/code-frame';
 import type {PluginObj} from '@babel/core';
 import type {parseConfigPragmaForTests as ParseConfigPragma} from 'babel-plugin-react-compiler/src/HIR/Environment';
+import type {printFunctionWithOutlined as PrintFunctionWithOutlined} from 'babel-plugin-react-compiler/src/HIR/PrintHIR';
+import type {printReactiveFunctionWithOutlined as PrintReactiveFunctionWithOutlined} from 'babel-plugin-react-compiler/src/ReactiveScopes/PrintReactiveFunction';
 import {TransformResult, transformFixtureInput} from './compiler';
 import {
   COMPILER_PATH,
   COMPILER_INDEX_PATH,
-  LOGGER_PATH,
   PARSE_CONFIG_PRAGMA_PATH,
+  PRINT_HIR_PATH,
+  PRINT_REACTIVE_IR_PATH,
 } from './constants';
 import {TestFixture, getBasename, isExpectError} from './fixture-utils';
 import {TestResult, writeOutputToString} from './reporter';
 import {runSprout} from './sprout';
+import {CompilerPipelineValue} from 'babel-plugin-react-compiler/src';
+import chalk from 'chalk';
 
 const originalConsoleError = console.error;
 
@@ -64,20 +69,56 @@ async function compile(
     const {Effect: EffectEnum, ValueKind: ValueKindEnum} = require(
       COMPILER_INDEX_PATH,
     );
-    const {toggleLogging} = require(LOGGER_PATH);
+    const {printFunctionWithOutlined} = require(PRINT_HIR_PATH) as {
+      printFunctionWithOutlined: typeof PrintFunctionWithOutlined;
+    };
+    const {printReactiveFunctionWithOutlined} = require(
+      PRINT_REACTIVE_IR_PATH,
+    ) as {
+      printReactiveFunctionWithOutlined: typeof PrintReactiveFunctionWithOutlined;
+    };
+
+    let lastLogged: string | null = null;
+    const debugIRLogger = shouldLog
+      ? (value: CompilerPipelineValue) => {
+          let printed: string;
+          switch (value.kind) {
+            case 'hir':
+              printed = printFunctionWithOutlined(value.value);
+              break;
+            case 'reactive':
+              printed = printReactiveFunctionWithOutlined(value.value);
+              break;
+            case 'debug':
+              printed = value.value;
+              break;
+            case 'ast':
+              // skip printing ast as we already write fixture output JS
+              printed = '(ast)';
+              break;
+          }
+
+          if (printed !== lastLogged) {
+            lastLogged = printed;
+            console.log(`${chalk.green(value.name)}:\n ${printed}\n`);
+          } else {
+            console.log(`${chalk.blue(value.name)}: (no change)\n`);
+          }
+        }
+      : () => {};
     const {parseConfigPragmaForTests} = require(PARSE_CONFIG_PRAGMA_PATH) as {
       parseConfigPragmaForTests: typeof ParseConfigPragma;
     };
 
     // only try logging if we filtered out all but one fixture,
     // since console log order is non-deterministic
-    toggleLogging(shouldLog);
     const result = await transformFixtureInput(
       input,
       fixturePath,
       parseConfigPragmaForTests,
       BabelPluginReactCompiler,
       includeEvaluator,
+      debugIRLogger,
       EffectEnum,
       ValueKindEnum,
     );