Option to disable memoization

Adds a `removeAllMemoization` flag that runs the entire compiler pipeline but 
strips out all memoization. The intent is to be able to compare (in limited 
use-cases) the performance of an existing app with all memoization removed, vs 
the performance with manual memoization, vs the performance with Forget enabled. 

In terms of how this works: we already strip out useMemo/useCallback since 
Forget is more accurate. The new option adds an extra pass that strips out all 
reactive scopes. Collectively this leaves ~zero memoization within components 
(this does leave React.memo, but close enough).

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

@@ -95,6 +95,7 @@ export async function compile(
     let memoizeJsxElements = true;
     let enableAssumeHooksFollowRulesOfReact = false;
     let enableTreatHooksAsFunctions = true;
+    let disableAllMemoization = false;
     if (firstLine.indexOf("@forgetDirective") !== -1) {
       enableOnlyOnUseForgetDirective = true;
     }
@@ -128,6 +129,9 @@ export async function compile(
     if (firstLine.indexOf("@enableTreatHooksAsFunctions false") !== -1) {
       enableTreatHooksAsFunctions = false;
     }
+    if (firstLine.indexOf("@disableAllMemoization true") !== -1) {
+      disableAllMemoization = true;
+    }
 
     const language = parseLanguage(firstLine);
 
@@ -143,12 +147,13 @@ export async function compile(
             },
           ],
         ]),
-        validateHooksUsage: true,
-        enableFunctionCallSignatureOptimizations: true,
         enableAssumeHooksFollowRulesOfReact,
+        enableFunctionCallSignatureOptimizations: true,
+        disableAllMemoization,
+        enableTreatHooksAsFunctions,
         inlineUseMemo: true,
         memoizeJsxElements,
-        enableTreatHooksAsFunctions,
+        validateHooksUsage: true,
       },
       logger: null,
       gating,

compiler/forget/src/CompilerPipeline.ts

@@ -40,6 +40,7 @@ import {
   mergeOverlappingReactiveScopes,
   promoteUsedTemporaries,
   propagateScopeDependencies,
+  pruneAllReactiveScopes,
   pruneNonEscapingScopes,
   pruneNonReactiveDependencies,
   pruneUnusedLabels,
@@ -159,6 +160,15 @@ export function* run(
     value: reactiveFunction,
   });
 
+  if (env.disableAllMemoization) {
+    pruneAllReactiveScopes(reactiveFunction);
+    yield log({
+      kind: "reactive",
+      name: "PruneAllReactiveScopes",
+      value: reactiveFunction,
+    });
+  }
+
   flattenReactiveLoops(reactiveFunction);
   yield log({
     kind: "reactive",

compiler/forget/src/HIR/Environment.ts

@@ -102,6 +102,15 @@ export type EnvironmentConfig = Partial<{
    * Defaults to true
    */
   enableTreatHooksAsFunctions: boolean;
+
+  /**
+   * When enabled, removes *all* memoization from the function: this includes
+   * removing manually added useMemo/useCallback as well as not adding Forget's
+   * usual useMemoCache-based memoization.
+   *
+   * Defaults to false (ie, by default memoization is enabled)
+   */
+  disableAllMemoization: boolean;
 }>;
 
 export class Environment {
@@ -113,6 +122,7 @@ export class Environment {
   enableFunctionCallSignatureOptimizations: boolean;
   enableAssumeHooksFollowRulesOfReact: boolean;
   enableTreatHooksAsFunctions: boolean;
+  disableAllMemoization: boolean;
   #contextIdentifiers: Set<t.Identifier>;
 
   constructor(
@@ -150,6 +160,7 @@ export class Environment {
       config?.enableAssumeHooksFollowRulesOfReact ?? false;
     this.enableTreatHooksAsFunctions =
       config?.enableTreatHooksAsFunctions ?? true;
+    this.disableAllMemoization = config?.disableAllMemoization ?? false;
     this.#contextIdentifiers = contextIdentifiers;
   }
 

compiler/forget/src/ReactiveScopes/PruneAllReactiveScopes.ts

@@ -0,0 +1,36 @@
+/**
+ * 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 {
+  ReactiveFunction,
+  ReactiveScopeBlock,
+  ReactiveStatement,
+} from "../HIR/HIR";
+import {
+  ReactiveFunctionTransform,
+  Transformed,
+  visitReactiveFunction,
+} from "./visitors";
+
+/**
+ * Removes *all* reactive scopes. Intended for experimentation only, to allow
+ * accurately removing memoization using the compiler pipeline to get a baseline
+ * for performance of a product without memoization applied.
+ */
+export function pruneAllReactiveScopes(fn: ReactiveFunction): void {
+  visitReactiveFunction(fn, new Transform(), undefined);
+}
+
+class Transform extends ReactiveFunctionTransform<void> {
+  override transformScope(
+    scopeBlock: ReactiveScopeBlock,
+    state: void
+  ): Transformed<ReactiveStatement> {
+    this.visitScope(scopeBlock, state);
+    return { kind: "replace-many", value: scopeBlock.instructions };
+  }
+}

compiler/forget/src/ReactiveScopes/index.ts

@@ -18,6 +18,7 @@ export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes
 export { printReactiveFunction } from "./PrintReactiveFunction";
 export { promoteUsedTemporaries } from "./PromoteUsedTemporaries";
 export { propagateScopeDependencies } from "./PropagateScopeDependencies";
+export { pruneAllReactiveScopes } from "./PruneAllReactiveScopes";
 export { pruneNonEscapingScopes } from "./PruneNonEscapingScopes";
 export { pruneNonReactiveDependencies } from "./PruneNonReactiveDependencies";
 export { pruneTemporaryLValues as pruneUnusedLValues } from "./PruneTemporaryLValues";

compiler/forget/src/__tests__/fixtures/compiler/remove-memoization-kitchen-sink.expect.md

@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+// @disableAllMemoization true
+function Component(props) {
+  const [x, setX] = useState(() => initializeState(props));
+  const onChange = useCallback((e) => {
+    setX(e.target.value);
+  });
+  const object = { x, onChange };
+  return useMemo(() => {
+    const { x, onChange } = object;
+    return <input value={x} onChange={onChange} />;
+  }, [x]);
+}
+
+```
+
+## Code
+
+```javascript
+// @disableAllMemoization true
+function Component(props) {
+  const [x, setX] = useState(() => initializeState(props));
+  const onChange = (e) => {
+    setX(e.target.value);
+  };
+  const object = { x, onChange };
+
+  const { x: x_0, onChange: onChange_0 } = object;
+  const t44 = <input value={x_0} onChange={onChange_0} />;
+  return t44;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/remove-memoization-kitchen-sink.js

@@ -0,0 +1,12 @@
+// @disableAllMemoization true
+function Component(props) {
+  const [x, setX] = useState(() => initializeState(props));
+  const onChange = useCallback((e) => {
+    setX(e.target.value);
+  });
+  const object = { x, onChange };
+  return useMemo(() => {
+    const { x, onChange } = object;
+    return <input value={x} onChange={onChange} />;
+  }, [x]);
+}