Promote enableOptimizeFunctionExpressions to default/only option

We’ve had this feature turned on internally for a while with only a couple minor 
bugs and no fundamental issues. It’s a necessary change for simplifying function 
expression dependencies and context variables, so let’s remove the feature flag 
and fix forward for any issues.

compiler/forget/apps/playground/components/Editor/index.tsx

@@ -115,7 +115,6 @@ function parsePragma(pragma: string) {
   let disableAllMemoization = false;
   let validateRefAccessDuringRender = true;
   let enableEmitFreeze = null;
-  let enableOptimizeFunctionExpressions = true;
   let inlineUseMemo = true;
   let validateHooksUsage = true;
   let enableFunctionCallSignatureOptimizations = true;
@@ -137,9 +136,6 @@ function parsePragma(pragma: string) {
   if (pragma.includes("@validateRefAccessDuringRender false")) {
     validateRefAccessDuringRender = false;
   }
-  if (pragma.includes("@enableOptimizeFunctionExpressions false")) {
-    enableOptimizeFunctionExpressions = false;
-  }
   if (pragma.includes("@enableEmitFreeze")) {
     enableEmitFreeze = {
       source: "react-forget-runtime",
@@ -173,7 +169,6 @@ function parsePragma(pragma: string) {
     validateRefAccessDuringRender,
     validateFrozenLambdas,
     enableEmitFreeze,
-    enableOptimizeFunctionExpressions,
     assertValidMutableRanges,
   };
 }

compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts

@@ -159,15 +159,6 @@ export type EnvironmentConfig = Partial<{
    */
   enableEmitFreeze: ExternalFunction | null;
 
-  /**
-   * When enabled, function expression codegen uses a subset of the compiler pipeline
-   * to transform and optimize their contents. When disabled, function expression
-   * codegen uses the original, un-transformed function body.
-   *
-   * Defaults to false (use the un-transformed function body).
-   */
-  enableOptimizeFunctionExpressions: boolean;
-
   /**
    * Enable validation of mutable ranges
    *
@@ -190,7 +181,6 @@ export class Environment {
   enableTreatHooksAsFunctions: boolean;
   disableAllMemoization: boolean;
   enableEmitFreeze: ExternalFunction | null;
-  enableOptimizeFunctionExpressions: boolean;
   assertValidMutableRanges: boolean;
 
   #contextIdentifiers: Set<t.Identifier>;
@@ -237,8 +227,6 @@ export class Environment {
       config?.enableTreatHooksAsFunctions ?? true;
     this.disableAllMemoization = config?.disableAllMemoization ?? false;
     this.enableEmitFreeze = config?.enableEmitFreeze ?? null;
-    this.enableOptimizeFunctionExpressions =
-      config?.enableOptimizeFunctionExpressions ?? true;
     this.assertValidMutableRanges = config?.assertValidMutableRanges ?? false;
     this.validateNoSetStateInRender =
       config?.validateNoSetStateInRender ?? false;

compiler/forget/packages/babel-plugin-react-forget/src/HIR/MergeConsecutiveBlocks.ts

@@ -31,11 +31,9 @@ import { mapOptionalFallthroughs } from "./visitors";
 export function mergeConsecutiveBlocks(fn: HIRFunction): void {
   const merged = new MergedBlocks();
   for (const [, block] of fn.body.blocks) {
-    if (fn.env.enableOptimizeFunctionExpressions) {
-      for (const instr of block.instructions) {
-        if (instr.value.kind === "FunctionExpression") {
-          mergeConsecutiveBlocks(instr.value.loweredFunc);
-        }
+    for (const instr of block.instructions) {
+      if (instr.value.kind === "FunctionExpression") {
+        mergeConsecutiveBlocks(instr.value.loweredFunc);
       }
     }
 

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

@@ -14,14 +14,12 @@ import {
   isRefValueType,
   isSetStateType,
   isUseRefType,
-  mergeConsecutiveBlocks,
   Place,
   ReactiveScopeDependency,
 } from "../HIR";
-import { constantPropagation, deadCodeElimination } from "../Optimization";
+import { deadCodeElimination } from "../Optimization";
 import { inferReactiveScopeVariables } from "../ReactiveScopes";
-import { eliminateRedundantPhi, enterSSA, leaveSSA } from "../SSA";
-import { inferTypes } from "../TypeInference";
+import { leaveSSA } from "../SSA";
 import { logHIRFunction } from "../Utils/logger";
 import { inferMutableContextVariables } from "./InferMutableContextVariables";
 import { inferMutableRanges } from "./InferMutableRanges";
@@ -104,14 +102,6 @@ export default function analyseFunctions(func: HIRFunction): void {
 }
 
 function lower(func: HIRFunction): void {
-  if (!func.env.enableOptimizeFunctionExpressions) {
-    mergeConsecutiveBlocks(func);
-    enterSSA(func);
-    eliminateRedundantPhi(func);
-    constantPropagation(func);
-    inferTypes(func);
-  }
-
   analyseFunctions(func);
   inferReferenceEffects(func, { isFunctionExpression: true });
   deadCodeElimination(func);

compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts

@@ -131,12 +131,6 @@ function applyConstantPropagation(
         continue;
       }
       const instr = block.instructions[i]!;
-      if (!fn.env.enableOptimizeFunctionExpressions) {
-        // Don't propagate constants used as function expression dependencies
-        if (functionDependencies.has(instr.lvalue.identifier.id)) {
-          continue;
-        }
-      }
       const value = evaluateInstruction(fn.env, constants, instr);
       if (value !== null) {
         constants.set(instr.lvalue.identifier.id, value);
@@ -438,9 +432,7 @@ function evaluateInstruction(
       return placeValue;
     }
     case "FunctionExpression": {
-      if (env.enableOptimizeFunctionExpressions) {
-        constantPropagationImpl(value.loweredFunc, constants);
-      }
+      constantPropagationImpl(value.loweredFunc, constants);
       return null;
     }
     default: {

compiler/forget/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -993,34 +993,30 @@ function codegenInstructionValue(
       break;
     }
     case "FunctionExpression": {
-      if (cx.env.enableOptimizeFunctionExpressions) {
-        const loweredFunc = instrValue.loweredFunc;
-        const reactiveFunction = buildReactiveFunction(loweredFunc);
-        pruneUnusedLabels(reactiveFunction);
-        pruneUnusedLValues(reactiveFunction);
-        renameVariables(reactiveFunction);
-        const fn = codegenReactiveFunction(reactiveFunction).unwrap();
-        if (instrValue.expr.type === "ArrowFunctionExpression") {
-          let body: t.BlockStatement | t.Expression = fn.body;
-          if (body.body.length === 1) {
-            const stmt = body.body[0]!;
-            if (stmt.type === "ReturnStatement" && stmt.argument != null) {
-              body = stmt.argument;
-            }
+      const loweredFunc = instrValue.loweredFunc;
+      const reactiveFunction = buildReactiveFunction(loweredFunc);
+      pruneUnusedLabels(reactiveFunction);
+      pruneUnusedLValues(reactiveFunction);
+      renameVariables(reactiveFunction);
+      const fn = codegenReactiveFunction(reactiveFunction).unwrap();
+      if (instrValue.expr.type === "ArrowFunctionExpression") {
+        let body: t.BlockStatement | t.Expression = fn.body;
+        if (body.body.length === 1) {
+          const stmt = body.body[0]!;
+          if (stmt.type === "ReturnStatement" && stmt.argument != null) {
+            body = stmt.argument;
           }
-          value = t.arrowFunctionExpression(fn.params, body, fn.async);
-        } else {
-          value = t.functionExpression(
-            fn.id ??
-              (instrValue.name != null ? t.identifier(instrValue.name) : null),
-            fn.params,
-            fn.body,
-            fn.generator,
-            fn.async
-          );
         }
+        value = t.arrowFunctionExpression(fn.params, body, fn.async);
       } else {
-        value = t.cloneNode(instrValue.expr, true, false);
+        value = t.functionExpression(
+          fn.id ??
+            (instrValue.name != null ? t.identifier(instrValue.name) : null),
+          fn.params,
+          fn.body,
+          fn.generator,
+          fn.async
+        );
       }
       break;
     }

compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts

@@ -106,10 +106,7 @@ export function eliminateRedundantPhi(
           rewritePlace(place, rewrites);
         }
 
-        if (
-          instr.value.kind === "FunctionExpression" &&
-          fn.env.enableOptimizeFunctionExpressions
-        ) {
+        if (instr.value.kind === "FunctionExpression") {
           const { context } = instr.value.loweredFunc;
           for (const place of context) {
             rewritePlace(place, rewrites);

compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts

@@ -250,16 +250,12 @@ function enterSSAImpl(
 
     if (blockId === rootEntry) {
       // NOTE: func.context should be empty for the root function
-      if (func.env.enableOptimizeFunctionExpressions) {
-        CompilerError.invariant(func.context.length === 0, {
-          reason: `Expected function context to be empty for outer function declarations`,
-          description: null,
-          loc: func.loc,
-          suggestions: null,
-        });
-      } else {
-        func.context = func.context.map((p) => builder.defineContext(p));
-      }
+      CompilerError.invariant(func.context.length === 0, {
+        reason: `Expected function context to be empty for outer function declarations`,
+        description: null,
+        loc: func.loc,
+        suggestions: null,
+      });
       func.params = func.params.map((p) => builder.definePlace(p));
     }
 
@@ -267,10 +263,7 @@ function enterSSAImpl(
       mapInstructionOperands(instr, (place) => builder.getPlace(place));
       mapInstructionLValues(instr, (lvalue) => builder.definePlace(lvalue));
 
-      if (
-        instr.value.kind === "FunctionExpression" &&
-        func.env.enableOptimizeFunctionExpressions
-      ) {
+      if (instr.value.kind === "FunctionExpression") {
         const loweredFunc = instr.value.loweredFunc;
         const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!;
         CompilerError.invariant(entry.preds.size === 0, {

compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts

@@ -68,10 +68,7 @@ function apply(func: HIRFunction, unifier: Unifier): void {
       const { lvalue, value } = instr;
       lvalue.identifier.type = unifier.get(lvalue.identifier.type);
 
-      if (
-        value.kind === "FunctionExpression" &&
-        func.env.enableOptimizeFunctionExpressions
-      ) {
+      if (value.kind === "FunctionExpression") {
         apply(value.loweredFunc, unifier);
       }
     }
@@ -259,9 +256,7 @@ function* generateInstructionTypes(
     }
 
     case "FunctionExpression": {
-      if (env.enableOptimizeFunctionExpressions) {
-        yield* generate(value.loweredFunc);
-      }
+      yield* generate(value.loweredFunc);
       break;
     }
 

compiler/forget/packages/fixture-test-utils/src/compiler-utils.ts

@@ -32,7 +32,6 @@ export function transformFixtureInput(
   let validateRefAccessDuringRender = true;
   let validateNoSetStateInRender = true;
   let enableEmitFreeze = null;
-  let enableOptimizeFunctionExpressions = true;
   let enableOnlyOnReactScript = false;
 
   if (firstLine.indexOf("@forgetDirective") !== -1) {
@@ -71,9 +70,6 @@ export function transformFixtureInput(
   if (firstLine.includes("@validateNoSetStateInRender false")) {
     validateNoSetStateInRender = false;
   }
-  if (firstLine.includes("@enableOptimizeFunctionExpressions false")) {
-    enableOptimizeFunctionExpressions = false;
-  }
   if (firstLine.includes("@enableEmitFreeze")) {
     enableEmitFreeze = {
       source: "react-forget-runtime",
@@ -111,7 +107,6 @@ export function transformFixtureInput(
         validateFrozenLambdas: true,
         validateNoSetStateInRender,
         enableEmitFreeze,
-        enableOptimizeFunctionExpressions,
         assertValidMutableRanges: true,
       },
       enableOnlyOnUseForgetDirective,

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

@@ -5,8 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import type { runReactForgetBabelPlugin as RunReactForgetBabelPlugin } from 'babel-plugin-react-forget/src/Babel/RunReactForgetBabelPlugin';
-import { TestFixture, transformFixtureInput, writeOutputToString } from 'fixture-test-utils';
+import type { runReactForgetBabelPlugin as RunReactForgetBabelPlugin } from "babel-plugin-react-forget/src/Babel/RunReactForgetBabelPlugin";
+import {
+  TestFixture,
+  transformFixtureInput,
+  writeOutputToString,
+} from "fixture-test-utils";
 import fs from "fs/promises";
 
 const originalConsoleError = console.error;
@@ -47,13 +51,10 @@ export async function compile(
     clearRequireCache();
   }
   version = compilerVersion;
-  const { inputPath, inputExists, outputPath, outputExists, basename } = fixture;
-  const input = inputExists
-    ? await fs.readFile(inputPath, "utf8")
-    : null;
-  const expected = outputExists
-    ? await fs.readFile(outputPath, "utf8")
-    : null;
+  const { inputPath, inputExists, outputPath, outputExists, basename } =
+    fixture;
+  const input = inputExists ? await fs.readFile(inputPath, "utf8") : null;
+  const expected = outputExists ? await fs.readFile(outputPath, "utf8") : null;
 
   // Input will be null if the input file did not exist, in which case the output file
   // is stale
@@ -72,14 +73,20 @@ export async function compile(
   try {
     // NOTE: we intentionally require lazily here so that we can clear the require cache
     // and load fresh versions of the compiler when `compilerVersion` changes.
-    const { runReactForgetBabelPlugin } = require(compilerPath) as { runReactForgetBabelPlugin: typeof RunReactForgetBabelPlugin };
+    const { runReactForgetBabelPlugin } = require(compilerPath) as {
+      runReactForgetBabelPlugin: typeof RunReactForgetBabelPlugin;
+    };
     const { toggleLogging } = require(loggerPath);
 
     // only try logging if we filtered out all but one fixture,
     // since console log order is non-deterministic
     const shouldLogPragma = input.split("\n")[0].includes("@debug");
     toggleLogging(isOnlyFixture && shouldLogPragma);
-    code = transformFixtureInput(input, basename, runReactForgetBabelPlugin).code;
+    code = transformFixtureInput(
+      input,
+      basename,
+      runReactForgetBabelPlugin
+    ).code;
   } catch (e) {
     error = e;
   }