From 23e0ce02a61e57360b4eed57e4fdc38e43a2efbd Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 13 Mar 2023 13:25:09 -0700
Subject: [PATCH] Make JSX memoization optional (on by default)

Makes JSX memoized by default again, but adds an option to disable memoization
of JSX. Also adds a new test and fixtures directory to test the opt-in
no-jsx-memoization behavior.
---
 compiler/forget/src/CompilerPipeline.ts       |   4 +-
 compiler/forget/src/HIR/Environment.ts        |   2 +
 .../ReactiveScopes/PruneNonEscapingScopes.ts  |  50 +++++---
 .../forget/src/__tests__/compiler-test.ts     |   2 +-
 .../disableMemoizeJsxElements-test.ts         | 110 ++++++++++++++++++
 .../fixtures/compiler/call.expect.md          |  11 +-
 .../capture-func-passed-to-jsx.expect.md      |  13 ++-
 ...turing-function-member-expr-call.expect.md |  15 ++-
 .../fixtures/compiler/component.expect.md     |  34 ++++--
 .../compiler/concise-arrow-expr.expect.md     |  13 ++-
 .../compiler/conditional-on-mutable.expect.md |  30 ++++-
 .../fixtures/compiler/constructor.expect.md   |  11 +-
 .../compiler/controlled-input.expect.md       |  15 ++-
 ...alysis-destructured-rest-element.expect.md |  15 ++-
 .../escape-analysis-jsx-child.expect.md       |  14 ++-
 ...ing-test-export-default-function.expect.md |  24 +++-
 ...test-export-function-and-default.expect.md |  24 +++-
 .../gating-test-export-function.expect.md     |  24 +++-
 .../fixtures/compiler/gating-test.expect.md   |  24 +++-
 .../fixtures/compiler/hook-call.expect.md     |  23 ++--
 .../compiler/independent-across-if.expect.md  |  15 ++-
 .../fixtures/compiler/independent.expect.md   |  15 ++-
 .../interdependent-across-if.expect.md        |  15 ++-
 .../compiler/interdependent.expect.md         |  15 ++-
 .../fixtures/compiler/jsx-fragment.expect.md  |  36 +++++-
 .../fixtures/compiler/jsx-spread.expect.md    |  15 ++-
 ...ng-scopes-shadowing-within-block.expect.md |  14 ++-
 .../compiler/property-assignment.expect.md    |  15 ++-
 .../compiler/reactive-scopes-if.expect.md     |  13 ++-
 .../compiler/reactive-scopes.expect.md        |  13 ++-
 .../reassignment-conditional.expect.md        |  15 ++-
 .../reassignment-separate-scopes.expect.md    |  67 +++++++----
 .../fixtures/compiler/reassignment.expect.md  |  15 ++-
 ...s-dep-and-redeclare-maybe-frozen.expect.md |  85 +++++++++-----
 ...me-variable-as-dep-and-redeclare.expect.md |  81 ++++++++-----
 .../compiler/ssa-call-jsx-2.expect.md         |  11 +-
 .../fixtures/compiler/ssa-call-jsx.expect.md  |  11 +-
 .../compiler/ssa-leave-case.expect.md         |  25 ++--
 .../switch-non-final-default.expect.md        |  27 ++++-
 .../fixtures/compiler/switch.expect.md        |  27 ++++-
 ...emporary-at-start-of-value-block.expect.md |  12 +-
 .../fixtures/compiler/timers.expect.md        |  19 ++-
 .../compiler/use-callback-simple.expect.md    |  13 ++-
 .../compiler/useMemo-simple.expect.md         |  13 ++-
 .../Component.expect.md                       |  41 +++++++
 .../disableMemoizeJsxElements/Component.js    |  11 ++
 46 files changed, 899 insertions(+), 198 deletions(-)
 create mode 100644 compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts
 create mode 100644 compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.expect.md
 create mode 100644 compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.js

diff --git a/compiler/forget/src/CompilerPipeline.ts b/compiler/forget/src/CompilerPipeline.ts
index 6aff1039f7..e628207fc2 100644
--- a/compiler/forget/src/CompilerPipeline.ts
+++ b/compiler/forget/src/CompilerPipeline.ts
@@ -146,7 +146,9 @@ export function* run(
     value: reactiveFunction,
   });
 
-  pruneNonEscapingScopes(reactiveFunction);
+  pruneNonEscapingScopes(reactiveFunction, {
+    memoizeJsxElements: options?.memoizeJsxElements ?? true,
+  });
   yield log({
     kind: "reactive",
     name: "PruneNonEscapingDependencies",
diff --git a/compiler/forget/src/HIR/Environment.ts b/compiler/forget/src/HIR/Environment.ts
index 05c062a50a..b8806ee80d 100644
--- a/compiler/forget/src/HIR/Environment.ts
+++ b/compiler/forget/src/HIR/Environment.ts
@@ -8,11 +8,13 @@ const HOOK_PATTERN = /^_?use/;
 export type EnvironmentOptions = {
   customHooks: Map<string, Hook>;
   globals: Set<string>;
+  memoizeJsxElements: boolean;
 };
 
 const DEFAULT_OPTIONS: EnvironmentOptions = {
   customHooks: new Map(),
   globals: DEFAULT_GLOBALS,
+  memoizeJsxElements: true,
 };
 
 export function mergeOptions(
diff --git a/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts
index 7f8a2de5ac..f2303a409c 100644
--- a/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts
+++ b/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts
@@ -102,7 +102,10 @@ import {
  *    conditional node with an aliased dep promotes to aliased).
  * 4. Finally we prune scopes whose outputs weren't marked.
  */
-export function pruneNonEscapingScopes(fn: ReactiveFunction): void {
+export function pruneNonEscapingScopes(
+  fn: ReactiveFunction,
+  options: MemoizationOptions
+): void {
   // First build up a map of which instructions are involved in creating which values,
   // and which values are returned.
   const state = new State();
@@ -112,7 +115,7 @@ export function pruneNonEscapingScopes(fn: ReactiveFunction): void {
   for (const param of fn.params) {
     state.declare(param.identifier.id);
   }
-  visitReactiveFunction(fn, new CollectDependenciesVisitor(), state);
+  visitReactiveFunction(fn, new CollectDependenciesVisitor(options), state);
 
   log(() => prettyFormat(state));
 
@@ -128,6 +131,10 @@ export function pruneNonEscapingScopes(fn: ReactiveFunction): void {
   visitReactiveFunction(fn, new PruneScopesTransform(), memoized);
 }
 
+export type MemoizationOptions = {
+  memoizeJsxElements: boolean;
+};
+
 // Describes how to determine whether a value should be memoized, relative to dependees and dependencies
 enum MemoizationLevel {
   // The value should be memoized if it escapes
@@ -316,7 +323,8 @@ type LValueMemoization = {
  */
 function computeMemoizationInputs(
   value: ReactiveValue,
-  lvalue: Place | null
+  lvalue: Place | null,
+  options: MemoizationOptions
 ): {
   // can optionally return a custom set of lvalues per instruction
   lvalues: Array<LValueMemoization>;
@@ -332,8 +340,8 @@ function computeMemoizationInputs(
             : [],
         rvalues: [
           // Conditionals do not alias their test value.
-          ...computeMemoizationInputs(value.consequent, null).rvalues,
-          ...computeMemoizationInputs(value.alternate, null).rvalues,
+          ...computeMemoizationInputs(value.consequent, null, options).rvalues,
+          ...computeMemoizationInputs(value.alternate, null, options).rvalues,
         ],
       };
     }
@@ -345,8 +353,8 @@ function computeMemoizationInputs(
             ? [{ place: lvalue, level: MemoizationLevel.Conditional }]
             : [],
         rvalues: [
-          ...computeMemoizationInputs(value.left, null).rvalues,
-          ...computeMemoizationInputs(value.right, null).rvalues,
+          ...computeMemoizationInputs(value.left, null, options).rvalues,
+          ...computeMemoizationInputs(value.right, null, options).rvalues,
         ],
       };
     }
@@ -360,7 +368,7 @@ function computeMemoizationInputs(
         // Only the final value of the sequence is a true rvalue:
         // values from the sequence's instructions are evaluated
         // as separate nodes
-        rvalues: computeMemoizationInputs(value.value, null).rvalues,
+        rvalues: computeMemoizationInputs(value.value, null, options).rvalues,
       };
     }
     case "JsxExpression": {
@@ -378,24 +386,24 @@ function computeMemoizationInputs(
           operands.push(child);
         }
       }
+      const level = options.memoizeJsxElements
+        ? MemoizationLevel.Memoized
+        : MemoizationLevel.Unmemoized;
       return {
         // JSX elements themselves are not memoized unless forced to
         // avoid breaking downstream memoization
-        lvalues:
-          lvalue !== null
-            ? [{ place: lvalue, level: MemoizationLevel.Unmemoized }]
-            : [],
+        lvalues: lvalue !== null ? [{ place: lvalue, level }] : [],
         rvalues: operands,
       };
     }
     case "JsxFragment": {
+      const level = options.memoizeJsxElements
+        ? MemoizationLevel.Memoized
+        : MemoizationLevel.Unmemoized;
       return {
         // JSX elements themselves are not memoized unless forced to
         // avoid breaking downstream memoization
-        lvalues:
-          lvalue !== null
-            ? [{ place: lvalue, level: MemoizationLevel.Unmemoized }]
-            : [],
+        lvalues: lvalue !== null ? [{ place: lvalue, level }] : [],
         rvalues: value.children,
       };
     }
@@ -564,6 +572,13 @@ function computePatternLValues(pattern: Pattern): Array<LValueMemoization> {
  * identifier's and scope's dependencies.
  */
 class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
+  options: MemoizationOptions;
+
+  constructor(options: MemoizationOptions) {
+    super();
+    this.options = options;
+  }
+
   override visitInstruction(
     instruction: ReactiveInstruction,
     state: State
@@ -573,7 +588,8 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
     // Determe the level of memoization for this value and the lvalues/rvalues
     const aliasing = computeMemoizationInputs(
       instruction.value,
-      instruction.lvalue
+      instruction.lvalue,
+      this.options
     );
 
     // Associate all the rvalues with the instruction's scope if it has one
diff --git a/compiler/forget/src/__tests__/compiler-test.ts b/compiler/forget/src/__tests__/compiler-test.ts
index 6a3b6d4eb4..f259489f52 100644
--- a/compiler/forget/src/__tests__/compiler-test.ts
+++ b/compiler/forget/src/__tests__/compiler-test.ts
@@ -28,7 +28,7 @@ wasmFolder(
   path.join(__dirname, "..", "..", "node_modules", "@hpcc-js", "wasm", "dist")
 );
 
-describe("React Forget (HIR version)", () => {
+describe("React Forget", () => {
   generateTestsFromFixtures(
     path.join(__dirname, "fixtures", "compiler"),
     (input, file, options) => {
diff --git a/compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts b/compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts
new file mode 100644
index 0000000000..a0c36d19ae
--- /dev/null
+++ b/compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts
@@ -0,0 +1,110 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+"use strict";
+
+import path from "path";
+import runReactForgetBabelPlugin from "../Babel/RunReactForgetBabelPlugin";
+import { Effect, ValueKind } from "../index";
+import { toggleLogging } from "../Utils/logger";
+import generateTestsFromFixtures from "./test-utils/generateTestsFromFixtures";
+
+type TestOutput = {
+  js: string;
+};
+
+function wrapWithTripleBackticks(s: string, ext?: string) {
+  return `\`\`\`${ext ?? ""}
+${s}
+\`\`\``;
+}
+
+describe("React Forget (Disable memoization of JSX elements)", () => {
+  generateTestsFromFixtures(
+    path.join(__dirname, "fixtures", "disableMemoizeJsxElements"),
+    (input, file, options) => {
+      let items: Array<TestOutput> = [];
+      let error: Error | null = null;
+      if (options.debug) {
+        toggleLogging(options.debug);
+      }
+      try {
+        items.push({
+          js: runReactForgetBabelPlugin(input, file, options.language, {
+            enableOnlyOnUseForgetDirective:
+              options.enableOnlyOnUseForgetDirective,
+            environment: {
+              memoizeJsxElements: false,
+              customHooks: new Map([
+                [
+                  "useFreeze",
+                  {
+                    name: "useFreeze",
+                    kind: "Custom",
+                    valueKind: ValueKind.Frozen,
+                    effectKind: Effect.Freeze,
+                  },
+                ],
+              ]),
+            },
+            logger: null,
+            gatingModule: options.gatingModule,
+          }).code,
+        });
+      } catch (e) {
+        error = e;
+      }
+      let outputs: Array<string>;
+
+      const expectError = file.startsWith("error.");
+      if (expectError) {
+        if (error === null) {
+          throw new Error(
+            `Expected an error to be thrown for fixture: '${file}', remove the 'error.' prefix if an error is not expected.`
+          );
+        } else {
+          outputs = [formatErrorOutput(error)];
+        }
+      } else {
+        if (error !== null) {
+          error.message = `Expected fixture '${file}' to succeed but it failed with error:\n\n${error.message}`;
+          throw error;
+        }
+        if (items === null || items.length === 0) {
+          throw new Error(`Expected at least one output for file '${file}'.`);
+        }
+        outputs = formatOutput(items);
+      }
+      return `
+## Input
+
+${wrapWithTripleBackticks(input, "javascript")}
+
+${outputs.join("\n")}
+      `;
+    }
+  );
+});
+
+function formatErrorOutput(error: Error): string {
+  error.message = error.message.replace(/^\/.*?:\s/, "");
+  return `
+## Error
+
+${wrapWithTripleBackticks(error.message)}
+          `;
+}
+
+function formatOutput(items: Array<TestOutput>): Array<string> {
+  return items.map(({ js }) => {
+    return `
+## Code
+
+${wrapWithTripleBackticks(js, "javascript")}
+        `.trim();
+  });
+}
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/call.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/call.expect.md
index 60377a0666..7af3132878 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/call.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/call.expect.md
@@ -21,7 +21,7 @@ function Component(props) {
 function foo() {}
 
 function Component(props) {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(3);
   let a;
   let b;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -36,7 +36,14 @@ function Component(props) {
     a = $[0];
     b = $[1];
   }
-  return <div a={a} b={b}></div>;
+  let t0;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div a={a} b={b}></div>;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/capture-func-passed-to-jsx.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/capture-func-passed-to-jsx.expect.md
index be0199a9b0..acae2e4956 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/capture-func-passed-to-jsx.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/capture-func-passed-to-jsx.expect.md
@@ -20,7 +20,7 @@ function component(a, b) {
 
 ```javascript
 function component(a, b) {
-  const $ = React.unstable_useMemoCache(7);
+  const $ = React.unstable_useMemoCache(9);
   const c_0 = $[0] !== b;
   let t0;
   if (c_0) {
@@ -56,7 +56,16 @@ function component(a, b) {
     t2 = $[6];
   }
   const x = t2;
-  const t = <Foo x={x}></Foo>;
+  const c_7 = $[7] !== x;
+  let t3;
+  if (c_7) {
+    t3 = <Foo x={x}></Foo>;
+    $[7] = x;
+    $[8] = t3;
+  } else {
+    t3 = $[8];
+  }
+  const t = t3;
   mutate(x);
   return t;
 }
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/capturing-function-member-expr-call.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/capturing-function-member-expr-call.expect.md
index f727b5e205..2dc57b8057 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/capturing-function-member-expr-call.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/capturing-function-member-expr-call.expect.md
@@ -20,7 +20,7 @@ function component({ mutator }) {
 
 ```javascript
 function component(t27) {
-  const $ = React.unstable_useMemoCache(4);
+  const $ = React.unstable_useMemoCache(7);
   const { mutator } = t27;
   const c_0 = $[0] !== mutator;
   let t0;
@@ -46,7 +46,18 @@ function component(t27) {
     t1 = $[3];
   }
   const hide = t1;
-  return <Foo poke={poke} hide={hide}></Foo>;
+  const c_4 = $[4] !== poke;
+  const c_5 = $[5] !== hide;
+  let t2;
+  if (c_4 || c_5) {
+    t2 = <Foo poke={poke} hide={hide}></Foo>;
+    $[4] = poke;
+    $[5] = hide;
+    $[6] = t2;
+  } else {
+    t2 = $[6];
+  }
+  return t2;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/component.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/component.expect.md
index 2518ff4e3d..4ca1f899fe 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/component.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/component.expect.md
@@ -35,7 +35,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(8);
   const items = props.items;
   const maxItems = props.maxItems;
   const c_0 = $[0] !== maxItems;
@@ -65,12 +65,32 @@ function Component(props) {
   }
 
   const count = renderedItems.length;
-  return (
-    <div>
-      {<h1>{count} Items</h1>}
-      {renderedItems}
-    </div>
-  );
+  const c_3 = $[3] !== count;
+  let t0;
+  if (c_3) {
+    t0 = <h1>{count} Items</h1>;
+    $[3] = count;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
+  }
+  const c_5 = $[5] !== t0;
+  const c_6 = $[6] !== renderedItems;
+  let t1;
+  if (c_5 || c_6) {
+    t1 = (
+      <div>
+        {t0}
+        {renderedItems}
+      </div>
+    );
+    $[5] = t0;
+    $[6] = renderedItems;
+    $[7] = t1;
+  } else {
+    t1 = $[7];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/concise-arrow-expr.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/concise-arrow-expr.expect.md
index 006ec5f533..517503efd5 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/concise-arrow-expr.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/concise-arrow-expr.expect.md
@@ -14,7 +14,7 @@ function component() {
 
 ```javascript
 function component() {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(4);
   const [x, setX] = useState(0);
   const c_0 = $[0] !== setX;
   let t0;
@@ -26,7 +26,16 @@ function component() {
     t0 = $[1];
   }
   const handler = t0;
-  return <Foo handler={handler}></Foo>;
+  const c_2 = $[2] !== handler;
+  let t1;
+  if (c_2) {
+    t1 = <Foo handler={handler}></Foo>;
+    $[2] = handler;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/conditional-on-mutable.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/conditional-on-mutable.expect.md
index 990b07d7f5..f80aacf62e 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/conditional-on-mutable.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/conditional-on-mutable.expect.md
@@ -35,7 +35,7 @@ function mayMutate() {}
 
 ```javascript
 function ComponentA(props) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(6);
   const c_0 = $[0] !== props;
   let a;
   let b;
@@ -55,11 +55,22 @@ function ComponentA(props) {
     a = $[1];
     b = $[2];
   }
-  return <Foo a={a} b={b}></Foo>;
+  const c_3 = $[3] !== a;
+  const c_4 = $[4] !== b;
+  let t0;
+  if (c_3 || c_4) {
+    t0 = <Foo a={a} b={b}></Foo>;
+    $[3] = a;
+    $[4] = b;
+    $[5] = t0;
+  } else {
+    t0 = $[5];
+  }
+  return t0;
 }
 
 function ComponentB(props) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(6);
   const c_0 = $[0] !== props;
   let a;
   let b;
@@ -79,7 +90,18 @@ function ComponentB(props) {
     a = $[1];
     b = $[2];
   }
-  return <Foo a={a} b={b}></Foo>;
+  const c_3 = $[3] !== a;
+  const c_4 = $[4] !== b;
+  let t0;
+  if (c_3 || c_4) {
+    t0 = <Foo a={a} b={b}></Foo>;
+    $[3] = a;
+    $[4] = b;
+    $[5] = t0;
+  } else {
+    t0 = $[5];
+  }
+  return t0;
 }
 
 function Foo() {}
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/constructor.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/constructor.expect.md
index 3c0d0d9f28..08b57f01d8 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/constructor.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/constructor.expect.md
@@ -21,7 +21,7 @@ function Component(props) {
 function Foo() {}
 
 function Component(props) {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(3);
   let a;
   let b;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -35,7 +35,14 @@ function Component(props) {
     a = $[0];
     b = $[1];
   }
-  return <div a={a} b={b}></div>;
+  let t0;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div a={a} b={b}></div>;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/controlled-input.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/controlled-input.expect.md
index dfdf8c1b09..417dcef789 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/controlled-input.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/controlled-input.expect.md
@@ -14,7 +14,7 @@ function component() {
 
 ```javascript
 function component() {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(5);
   const [x, setX] = useState(0);
   const c_0 = $[0] !== setX;
   let t0;
@@ -26,7 +26,18 @@ function component() {
     t0 = $[1];
   }
   const handler = t0;
-  return <input onChange={handler} value={x}></input>;
+  const c_2 = $[2] !== handler;
+  const c_3 = $[3] !== x;
+  let t1;
+  if (c_2 || c_3) {
+    t1 = <input onChange={handler} value={x}></input>;
+    $[2] = handler;
+    $[3] = x;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-destructured-rest-element.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-destructured-rest-element.expect.md
index d9fde0f28d..14682a35a5 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-destructured-rest-element.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-destructured-rest-element.expect.md
@@ -16,7 +16,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(4);
+  const $ = React.unstable_useMemoCache(7);
   const c_0 = $[0] !== props.a;
   let b;
   if (c_0) {
@@ -35,7 +35,18 @@ function Component(props) {
   } else {
     d = $[3];
   }
-  return <div b={b} d={d}></div>;
+  const c_4 = $[4] !== b;
+  const c_5 = $[5] !== d;
+  let t0;
+  if (c_4 || c_5) {
+    t0 = <div b={b} d={d}></div>;
+    $[4] = b;
+    $[5] = d;
+    $[6] = t0;
+  } else {
+    t0 = $[6];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-jsx-child.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-jsx-child.expect.md
index 7f41bd83cd..8ed5b42f64 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-jsx-child.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-jsx-child.expect.md
@@ -20,7 +20,7 @@ function foo(a, b, c) {
 
 ```javascript
 function foo(a, b, c) {
-  const $ = React.unstable_useMemoCache(7);
+  const $ = React.unstable_useMemoCache(9);
   const c_0 = $[0] !== a;
   const c_1 = $[1] !== b;
   const c_2 = $[2] !== c;
@@ -42,8 +42,16 @@ function foo(a, b, c) {
       } else {
         y = $[6];
       }
-
-      x.push(<div>{y}</div>);
+      const c_7 = $[7] !== y;
+      let t0;
+      if (c_7) {
+        t0 = <div>{y}</div>;
+        $[7] = y;
+        $[8] = t0;
+      } else {
+        t0 = $[8];
+      }
+      x.push(t0);
     }
     $[0] = a;
     $[1] = b;
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md
index 6518388ae8..8ce6c365f5 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md
@@ -28,7 +28,17 @@ function Bar_uncompiled(props) {
   return <div>{props.bar}</div>;
 }
 function Bar_forget(props) {
-  return <div>{props.bar}</div>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <div>{props.bar}</div>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 const Bar = isForgetEnabled ? Bar_forget : Bar_uncompiled;
 export default Bar;
@@ -42,7 +52,17 @@ function Foo_uncompiled(props) {
   return <Foo>{props.bar}</Foo>;
 }
 function Foo_forget(props) {
-  return <Foo>{props.bar}</Foo>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <Foo>{props.bar}</Foo>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 const Foo = isForgetEnabled ? Foo_forget : Foo_uncompiled;
 
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md
index 6c12dcde2c..e4bfcfe42e 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md
@@ -28,7 +28,17 @@ function Bar_uncompiled(props) {
   return <div>{props.bar}</div>;
 }
 function Bar_forget(props) {
-  return <div>{props.bar}</div>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <div>{props.bar}</div>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 const Bar = isForgetEnabled ? Bar_forget : Bar_uncompiled;
 export default Bar;
@@ -42,7 +52,17 @@ function Foo_uncompiled(props) {
   return <Foo>{props.bar}</Foo>;
 }
 function Foo_forget(props) {
-  return <Foo>{props.bar}</Foo>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <Foo>{props.bar}</Foo>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 export const Foo = isForgetEnabled ? Foo_forget : Foo_uncompiled;
 
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md
index 36363fb738..81b8dd9e8e 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md
@@ -28,7 +28,17 @@ function Bar_uncompiled(props) {
   return <div>{props.bar}</div>;
 }
 function Bar_forget(props) {
-  return <div>{props.bar}</div>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <div>{props.bar}</div>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 export const Bar = isForgetEnabled ? Bar_forget : Bar_uncompiled;
 
@@ -41,7 +51,17 @@ function Foo_uncompiled(props) {
   return <Foo>{props.bar}</Foo>;
 }
 function Foo_forget(props) {
-  return <Foo>{props.bar}</Foo>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <Foo>{props.bar}</Foo>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 export const Foo = isForgetEnabled ? Foo_forget : Foo_uncompiled;
 
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/gating-test.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/gating-test.expect.md
index 624514d94e..19260edb97 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/gating-test.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/gating-test.expect.md
@@ -28,7 +28,17 @@ function Bar_uncompiled(props) {
   return <div>{props.bar}</div>;
 }
 function Bar_forget(props) {
-  return <div>{props.bar}</div>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <div>{props.bar}</div>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 const Bar = isForgetEnabled ? Bar_forget : Bar_uncompiled;
 
@@ -41,7 +51,17 @@ function Foo_uncompiled(props) {
   return <Foo>{props.bar}</Foo>;
 }
 function Foo_forget(props) {
-  return <Foo>{props.bar}</Foo>;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props.bar;
+  let t0;
+  if (c_0) {
+    t0 = <Foo>{props.bar}</Foo>;
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 const Foo = isForgetEnabled ? Foo_forget : Foo_uncompiled;
 
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/hook-call.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/hook-call.expect.md
index b7a764d207..4ecc4cc7fa 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/hook-call.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/hook-call.expect.md
@@ -26,7 +26,7 @@ function useFreeze() {}
 function foo() {}
 
 function Component(props) {
-  const $ = React.unstable_useMemoCache(1);
+  const $ = React.unstable_useMemoCache(3);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = [];
@@ -37,12 +37,21 @@ function Component(props) {
   const x = t0;
   const y = useFreeze(x);
   foo(y, x);
-  return (
-    <Component>
-      {x}
-      {y}
-    </Component>
-  );
+  const c_1 = $[1] !== y;
+  let t1;
+  if (c_1) {
+    t1 = (
+      <Component>
+        {x}
+        {y}
+      </Component>
+    );
+    $[1] = y;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/independent-across-if.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/independent-across-if.expect.md
index 68f656e264..34633c17d8 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/independent-across-if.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/independent-across-if.expect.md
@@ -56,7 +56,7 @@ function Foo() {}
  *   return = <Foo a={a} b={b} />
  */
 function Component(props) {
-  const $ = React.unstable_useMemoCache(5);
+  const $ = React.unstable_useMemoCache(8);
   const c_0 = $[0] !== props.a;
   const c_1 = $[1] !== props.b;
   const c_2 = $[2] !== props.c;
@@ -78,7 +78,18 @@ function Component(props) {
     a = $[3];
     b = $[4];
   }
-  return <Foo a={a} b={b}></Foo>;
+  const c_5 = $[5] !== a;
+  const c_6 = $[6] !== b;
+  let t0;
+  if (c_5 || c_6) {
+    t0 = <Foo a={a} b={b}></Foo>;
+    $[5] = a;
+    $[6] = b;
+    $[7] = t0;
+  } else {
+    t0 = $[7];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/independent.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/independent.expect.md
index c9cc9692f4..9a344bd0d9 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/independent.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/independent.expect.md
@@ -38,7 +38,7 @@ function Foo() {}
  *   return = <Foo a={a} b={b} />
  */
 function Component(props) {
-  const $ = React.unstable_useMemoCache(4);
+  const $ = React.unstable_useMemoCache(7);
   const c_0 = $[0] !== props.a;
   let t0;
   if (c_0) {
@@ -59,7 +59,18 @@ function Component(props) {
     t1 = $[3];
   }
   const b = t1;
-  return <Foo a={a} b={b}></Foo>;
+  const c_4 = $[4] !== a;
+  const c_5 = $[5] !== b;
+  let t2;
+  if (c_4 || c_5) {
+    t2 = <Foo a={a} b={b}></Foo>;
+    $[4] = a;
+    $[5] = b;
+    $[6] = t2;
+  } else {
+    t2 = $[6];
+  }
+  return t2;
 }
 
 function compute() {}
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/interdependent-across-if.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/interdependent-across-if.expect.md
index 9339072294..0d0f6cab93 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/interdependent-across-if.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/interdependent-across-if.expect.md
@@ -45,7 +45,7 @@ function Foo() {}
  *   return = <Foo a={a} b={b} />
  */
 function Component(props) {
-  const $ = React.unstable_useMemoCache(5);
+  const $ = React.unstable_useMemoCache(8);
   const c_0 = $[0] !== props.a;
   const c_1 = $[1] !== props.b;
   const c_2 = $[2] !== props.c;
@@ -66,7 +66,18 @@ function Component(props) {
     a = $[3];
     b = $[4];
   }
-  return <Foo a={a} b={b}></Foo>;
+  const c_5 = $[5] !== a;
+  const c_6 = $[6] !== b;
+  let t0;
+  if (c_5 || c_6) {
+    t0 = <Foo a={a} b={b}></Foo>;
+    $[5] = a;
+    $[6] = b;
+    $[7] = t0;
+  } else {
+    t0 = $[7];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/interdependent.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/interdependent.expect.md
index fa63c5ee4f..9c7bc4c7c3 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/interdependent.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/interdependent.expect.md
@@ -37,7 +37,7 @@ function Foo() {}
  *   return = <Foo a={a} b={b} />
  */
 function Component(props) {
-  const $ = React.unstable_useMemoCache(4);
+  const $ = React.unstable_useMemoCache(7);
   const c_0 = $[0] !== props.a;
   const c_1 = $[1] !== props.b;
   let a;
@@ -54,7 +54,18 @@ function Component(props) {
     a = $[2];
     b = $[3];
   }
-  return <Foo a={a} b={b}></Foo>;
+  const c_4 = $[4] !== a;
+  const c_5 = $[5] !== b;
+  let t0;
+  if (c_4 || c_5) {
+    t0 = <Foo a={a} b={b}></Foo>;
+    $[4] = a;
+    $[5] = b;
+    $[6] = t0;
+  } else {
+    t0 = $[6];
+  }
+  return t0;
 }
 
 function compute() {}
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/jsx-fragment.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/jsx-fragment.expect.md
index 896216a662..09b15eefae 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/jsx-fragment.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/jsx-fragment.expect.md
@@ -19,12 +19,36 @@ function Foo(props) {
 
 ```javascript
 function Foo(props) {
-  return (
-    <>
-      Hello {props.greeting}
-      {<div>{<>Text</>}</div>}
-    </>
-  );
+  const $ = React.unstable_useMemoCache(4);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <>Text</>;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = <div>{t0}</div>;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const c_2 = $[2] !== props.greeting;
+  let t2;
+  if (c_2) {
+    t2 = (
+      <>
+        Hello {props.greeting}
+        {t1}
+      </>
+    );
+    $[2] = props.greeting;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  return t2;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/jsx-spread.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/jsx-spread.expect.md
index 2dfe91d80a..0dae3baf71 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/jsx-spread.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/jsx-spread.expect.md
@@ -14,7 +14,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(5);
 
   const t0 = props.cond ? props.foo : props.bar;
   const c_0 = $[0] !== t0;
@@ -26,7 +26,18 @@ function Component(props) {
   } else {
     t1 = $[1];
   }
-  return <Component {...props} {...t1}></Component>;
+  const c_2 = $[2] !== props;
+  const c_3 = $[3] !== t1;
+  let t2;
+  if (c_2 || c_3) {
+    t2 = <Component {...props} {...t1}></Component>;
+    $[2] = props;
+    $[3] = t1;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  return t2;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/overlapping-scopes-shadowing-within-block.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/overlapping-scopes-shadowing-within-block.expect.md
index 8a4627755b..d24c7ef2f9 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/overlapping-scopes-shadowing-within-block.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/overlapping-scopes-shadowing-within-block.expect.md
@@ -21,7 +21,7 @@ function foo(a, b, c) {
 
 ```javascript
 function foo(a, b, c) {
-  const $ = React.unstable_useMemoCache(7);
+  const $ = React.unstable_useMemoCache(9);
   const c_0 = $[0] !== a;
   const c_1 = $[1] !== b;
   const c_2 = $[2] !== c;
@@ -43,8 +43,16 @@ function foo(a, b, c) {
       } else {
         y = $[6];
       }
-
-      x.push(<div>{y}</div>);
+      const c_7 = $[7] !== y;
+      let t0;
+      if (c_7) {
+        t0 = <div>{y}</div>;
+        $[7] = y;
+        $[8] = t0;
+      } else {
+        t0 = $[8];
+      }
+      x.push(t0);
     }
     $[0] = a;
     $[1] = b;
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/property-assignment.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/property-assignment.expect.md
index 643cbcac2e..0589203765 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/property-assignment.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/property-assignment.expect.md
@@ -17,7 +17,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(6);
   const c_0 = $[0] !== props.p0;
   let x;
   let child;
@@ -35,7 +35,18 @@ function Component(props) {
     x = $[1];
     child = $[2];
   }
-  return <Component data={x}>{child}</Component>;
+  const c_3 = $[3] !== x;
+  const c_4 = $[4] !== child;
+  let t0;
+  if (c_3 || c_4) {
+    t0 = <Component data={x}>{child}</Component>;
+    $[3] = x;
+    $[4] = child;
+    $[5] = t0;
+  } else {
+    t0 = $[5];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes-if.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes-if.expect.md
index 1ec13b7c14..3faf4cc51f 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes-if.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes-if.expect.md
@@ -20,7 +20,7 @@ function foo(a, b, c) {
 
 ```javascript
 function foo(a, b, c) {
-  const $ = React.unstable_useMemoCache(6);
+  const $ = React.unstable_useMemoCache(8);
   const c_0 = $[0] !== a;
   const c_1 = $[1] !== b;
   const c_2 = $[2] !== c;
@@ -38,7 +38,16 @@ function foo(a, b, c) {
       } else {
         y = $[5];
       }
-      x.push(<div>{y}</div>);
+      const c_6 = $[6] !== y;
+      let t0;
+      if (c_6) {
+        t0 = <div>{y}</div>;
+        $[6] = y;
+        $[7] = t0;
+      } else {
+        t0 = $[7];
+      }
+      x.push(t0);
     } else {
       x.push(c);
     }
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes.expect.md
index d70fbe3f80..006beeecae 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/reactive-scopes.expect.md
@@ -19,7 +19,7 @@ function f(a, b) {
 
 ```javascript
 function f(a, b) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(5);
   const c_0 = $[0] !== a.length;
   const c_1 = $[1] !== b;
   let x;
@@ -36,7 +36,16 @@ function f(a, b) {
   } else {
     x = $[2];
   }
-  return <div>{x}</div>;
+  const c_3 = $[3] !== x;
+  let t0;
+  if (c_3) {
+    t0 = <div>{x}</div>;
+    $[3] = x;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/reassignment-conditional.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/reassignment-conditional.expect.md
index 3eaac56f77..c7552085ba 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/reassignment-conditional.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/reassignment-conditional.expect.md
@@ -23,7 +23,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(6);
   const c_0 = $[0] !== props.p0;
   let x;
   if (c_0) {
@@ -47,7 +47,18 @@ function Component(props) {
   }
 
   y.push(props.p2);
-  return <Component x={x} y={y}></Component>;
+  const c_3 = $[3] !== x;
+  const c_4 = $[4] !== y;
+  let t1;
+  if (c_3 || c_4) {
+    t1 = <Component x={x} y={y}></Component>;
+    $[3] = x;
+    $[4] = y;
+    $[5] = t1;
+  } else {
+    t1 = $[5];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/reassignment-separate-scopes.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/reassignment-separate-scopes.expect.md
index 315c9f16d1..79d36dd99e 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/reassignment-separate-scopes.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/reassignment-separate-scopes.expect.md
@@ -34,7 +34,7 @@ function foo(a, b, c) {
 
 ```javascript
 function foo(a, b, c) {
-  const $ = React.unstable_useMemoCache(6);
+  const $ = React.unstable_useMemoCache(11);
   const c_0 = $[0] !== a;
   let x;
   if (c_0) {
@@ -47,39 +47,58 @@ function foo(a, b, c) {
   } else {
     x = $[1];
   }
-
-  const y = <div>{x}</div>;
+  const c_2 = $[2] !== x;
+  let t0;
+  if (c_2) {
+    t0 = <div>{x}</div>;
+    $[2] = x;
+    $[3] = t0;
+  } else {
+    t0 = $[3];
+  }
+  const y = t0;
   bb3: switch (b) {
     case 0: {
-      const c_2 = $[2] !== b;
-      if (c_2) {
-        x = [];
-        x.push(b);
-        $[2] = b;
-        $[3] = x;
-      } else {
-        x = $[3];
-      }
-      break bb3;
-    }
-    default: {
-      const c_4 = $[4] !== c;
+      const c_4 = $[4] !== b;
       if (c_4) {
         x = [];
-        x.push(c);
-        $[4] = c;
+        x.push(b);
+        $[4] = b;
         $[5] = x;
       } else {
         x = $[5];
       }
+      break bb3;
+    }
+    default: {
+      const c_6 = $[6] !== c;
+      if (c_6) {
+        x = [];
+        x.push(c);
+        $[6] = c;
+        $[7] = x;
+      } else {
+        x = $[7];
+      }
     }
   }
-  return (
-    <div>
-      {y}
-      {x}
-    </div>
-  );
+  const c_8 = $[8] !== y;
+  const c_9 = $[9] !== x;
+  let t1;
+  if (c_8 || c_9) {
+    t1 = (
+      <div>
+        {y}
+        {x}
+      </div>
+    );
+    $[8] = y;
+    $[9] = x;
+    $[10] = t1;
+  } else {
+    t1 = $[10];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/reassignment.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/reassignment.expect.md
index 82cd2284a5..0b3212b74d 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/reassignment.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/reassignment.expect.md
@@ -21,7 +21,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(5);
+  const $ = React.unstable_useMemoCache(8);
   const c_0 = $[0] !== props.p0;
   const c_1 = $[1] !== props.p1;
   let x;
@@ -48,7 +48,18 @@ function Component(props) {
     x = $[2];
     y = $[3];
   }
-  return <Component x={x} y={y}></Component>;
+  const c_5 = $[5] !== x;
+  const c_6 = $[6] !== y;
+  let t1;
+  if (c_5 || c_6) {
+    t1 = <Component x={x} y={y}></Component>;
+    $[5] = x;
+    $[6] = y;
+    $[7] = t1;
+  } else {
+    t1 = $[7];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare-maybe-frozen.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare-maybe-frozen.expect.md
index 055923ee83..c4f6052007 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare-maybe-frozen.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare-maybe-frozen.expect.md
@@ -46,7 +46,7 @@ function foo(props) {
 // note: comments are for the ideal scopes, not what is currently
 // emitted
 function foo(props) {
-  const $ = React.unstable_useMemoCache(7);
+  const $ = React.unstable_useMemoCache(16);
   const c_0 = $[0] !== props.a;
   let x;
   if (c_0) {
@@ -57,39 +57,70 @@ function foo(props) {
   } else {
     x = $[1];
   }
-
-  const header = props.showHeader ? <div>{x}</div> : null;
-  const c_2 = $[2] !== x;
-  const c_3 = $[3] !== props.b;
-  const c_4 = $[4] !== props.c;
+  const c_2 = $[2] !== props.showHeader;
+  const c_3 = $[3] !== x;
+  let t0;
+  if (c_2 || c_3) {
+    t0 = props.showHeader ? <div>{x}</div> : null;
+    $[2] = props.showHeader;
+    $[3] = x;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
+  }
+  const header = t0;
+  const c_5 = $[5] !== x;
+  const c_6 = $[6] !== props.b;
+  const c_7 = $[7] !== props.c;
   let y;
-  if (c_2 || c_3 || c_4) {
+  if (c_5 || c_6 || c_7) {
     y = [x];
     x = [];
     y.push(props.b);
     x.push(props.c);
-    $[2] = x;
-    $[3] = props.b;
-    $[4] = props.c;
-    $[5] = y;
-    $[6] = x;
+    $[5] = x;
+    $[6] = props.b;
+    $[7] = props.c;
+    $[8] = y;
+    $[9] = x;
   } else {
-    y = $[5];
-    x = $[6];
+    y = $[8];
+    x = $[9];
   }
-
-  const content = (
-    <div>
-      {x}
-      {y}
-    </div>
-  );
-  return (
-    <>
-      {header}
-      {content}
-    </>
-  );
+  const c_10 = $[10] !== x;
+  const c_11 = $[11] !== y;
+  let t1;
+  if (c_10 || c_11) {
+    t1 = (
+      <div>
+        {x}
+        {y}
+      </div>
+    );
+    $[10] = x;
+    $[11] = y;
+    $[12] = t1;
+  } else {
+    t1 = $[12];
+  }
+  const content = t1;
+  const c_13 = $[13] !== header;
+  const c_14 = $[14] !== content;
+  let t2;
+  if (c_13 || c_14) {
+    t2 = (
+      <>
+        {header}
+        {content}
+      </>
+    );
+    $[13] = header;
+    $[14] = content;
+    $[15] = t2;
+  } else {
+    t2 = $[15];
+  }
+  return t2;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare.expect.md
index fa6a6227b1..dc4c090b37 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/same-variable-as-dep-and-redeclare.expect.md
@@ -46,7 +46,7 @@ function foo(props) {
 // note: comments are for the ideal scopes, not what is currently
 // emitted
 function foo(props) {
-  const $ = React.unstable_useMemoCache(7);
+  const $ = React.unstable_useMemoCache(15);
   const c_0 = $[0] !== props.a;
   let x;
   if (c_0) {
@@ -57,39 +57,68 @@ function foo(props) {
   } else {
     x = $[1];
   }
-
-  const header = <div>{x}</div>;
   const c_2 = $[2] !== x;
-  const c_3 = $[3] !== props.b;
-  const c_4 = $[4] !== props.c;
+  let t0;
+  if (c_2) {
+    t0 = <div>{x}</div>;
+    $[2] = x;
+    $[3] = t0;
+  } else {
+    t0 = $[3];
+  }
+  const header = t0;
+  const c_4 = $[4] !== x;
+  const c_5 = $[5] !== props.b;
+  const c_6 = $[6] !== props.c;
   let y;
-  if (c_2 || c_3 || c_4) {
+  if (c_4 || c_5 || c_6) {
     y = [x];
     x = [];
     y.push(props.b);
     x.push(props.c);
-    $[2] = x;
-    $[3] = props.b;
-    $[4] = props.c;
-    $[5] = y;
-    $[6] = x;
+    $[4] = x;
+    $[5] = props.b;
+    $[6] = props.c;
+    $[7] = y;
+    $[8] = x;
   } else {
-    y = $[5];
-    x = $[6];
+    y = $[7];
+    x = $[8];
   }
-
-  const content = (
-    <div>
-      {x}
-      {y}
-    </div>
-  );
-  return (
-    <>
-      {header}
-      {content}
-    </>
-  );
+  const c_9 = $[9] !== x;
+  const c_10 = $[10] !== y;
+  let t1;
+  if (c_9 || c_10) {
+    t1 = (
+      <div>
+        {x}
+        {y}
+      </div>
+    );
+    $[9] = x;
+    $[10] = y;
+    $[11] = t1;
+  } else {
+    t1 = $[11];
+  }
+  const content = t1;
+  const c_12 = $[12] !== header;
+  const c_13 = $[13] !== content;
+  let t2;
+  if (c_12 || c_13) {
+    t2 = (
+      <>
+        {header}
+        {content}
+      </>
+    );
+    $[12] = header;
+    $[13] = content;
+    $[14] = t2;
+  } else {
+    t2 = $[14];
+  }
+  return t2;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx-2.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx-2.expect.md
index 8d33b08ebb..e658014990 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx-2.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx-2.expect.md
@@ -25,7 +25,7 @@ function Component(props) {
 function foo() {}
 
 function Component(props) {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(3);
   let a;
   let b;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -42,7 +42,14 @@ function Component(props) {
     a = $[0];
     b = $[1];
   }
-  return <div a={a} b={b}></div>;
+  let t0;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div a={a} b={b}></div>;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx.expect.md
index 8a59efbd00..e786729615 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/ssa-call-jsx.expect.md
@@ -21,7 +21,7 @@ function Component(props) {
 function foo() {}
 
 function Component(props) {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(3);
   let a;
   let b;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -36,7 +36,14 @@ function Component(props) {
     a = $[0];
     b = $[1];
   }
-  return <div a={a} b={b}></div>;
+  let t0;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div a={a} b={b}></div>;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/ssa-leave-case.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/ssa-leave-case.expect.md
index a9fe1e2ee3..72c1e52d44 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/ssa-leave-case.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/ssa-leave-case.expect.md
@@ -23,7 +23,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(6);
   const c_0 = $[0] !== props;
   let x;
   let y;
@@ -41,12 +41,23 @@ function Component(props) {
     x = $[1];
     y = $[2];
   }
-  return (
-    <Component>
-      {x}
-      {y}
-    </Component>
-  );
+  const c_3 = $[3] !== x;
+  const c_4 = $[4] !== y;
+  let t0;
+  if (c_3 || c_4) {
+    t0 = (
+      <Component>
+        {x}
+        {y}
+      </Component>
+    );
+    $[3] = x;
+    $[4] = y;
+    $[5] = t0;
+  } else {
+    t0 = $[5];
+  }
+  return t0;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/switch-non-final-default.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/switch-non-final-default.expect.md
index edf573ba87..617fbc88fb 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/switch-non-final-default.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/switch-non-final-default.expect.md
@@ -32,7 +32,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(4);
+  const $ = React.unstable_useMemoCache(9);
   const c_0 = $[0] !== props;
   let x;
   let y;
@@ -69,10 +69,29 @@ function Component(props) {
     x = $[1];
     y = $[2];
   }
-
-  const child = <Component data={x}></Component>;
+  const c_4 = $[4] !== x;
+  let t1;
+  if (c_4) {
+    t1 = <Component data={x}></Component>;
+    $[4] = x;
+    $[5] = t1;
+  } else {
+    t1 = $[5];
+  }
+  const child = t1;
   y.push(props.p4);
-  return <Component data={y}>{child}</Component>;
+  const c_6 = $[6] !== y;
+  const c_7 = $[7] !== child;
+  let t2;
+  if (c_6 || c_7) {
+    t2 = <Component data={y}>{child}</Component>;
+    $[6] = y;
+    $[7] = child;
+    $[8] = t2;
+  } else {
+    t2 = $[8];
+  }
+  return t2;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/switch.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/switch.expect.md
index beccdd160d..1ab3f13f01 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/switch.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/switch.expect.md
@@ -27,7 +27,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(8);
   const c_0 = $[0] !== props;
   let x;
   let y;
@@ -50,10 +50,29 @@ function Component(props) {
     x = $[1];
     y = $[2];
   }
-
-  const child = <Component data={x}></Component>;
+  const c_3 = $[3] !== x;
+  let t0;
+  if (c_3) {
+    t0 = <Component data={x}></Component>;
+    $[3] = x;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
+  }
+  const child = t0;
   y.push(props.p4);
-  return <Component data={y}>{child}</Component>;
+  const c_5 = $[5] !== y;
+  const c_6 = $[6] !== child;
+  let t1;
+  if (c_5 || c_6) {
+    t1 = <Component data={y}>{child}</Component>;
+    $[5] = y;
+    $[6] = child;
+    $[7] = t1;
+  } else {
+    t1 = $[7];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/temporary-at-start-of-value-block.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/temporary-at-start-of-value-block.expect.md
index 7aadfcb57b..29f79e4370 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/temporary-at-start-of-value-block.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/temporary-at-start-of-value-block.expect.md
@@ -14,7 +14,17 @@ function component(props) {
 
 ```javascript
 function component(props) {
-  const x = isMenuShown ? <Bar> {props.a ? props.b : props.c}</Bar> : null;
+  const $ = React.unstable_useMemoCache(2);
+  const c_0 = $[0] !== props;
+  let t0;
+  if (c_0) {
+    t0 = isMenuShown ? <Bar> {props.a ? props.b : props.c}</Bar> : null;
+    $[0] = props;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const x = t0;
   return x;
 }
 
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/timers.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/timers.expect.md
index ff5ed5784f..398278c470 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/timers.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/timers.expect.md
@@ -19,7 +19,7 @@ function Component(props) {
 
 ```javascript
 function Component(props) {
-  const $ = React.unstable_useMemoCache(1);
+  const $ = React.unstable_useMemoCache(2);
   const start = performance.now();
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -30,11 +30,18 @@ function Component(props) {
   }
   const now = t0;
   const time = performance.now() - start;
-  return (
-    <div>
-      rendering took {time} at {now}
-    </div>
-  );
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = (
+      <div>
+        rendering took {time} at {now}
+      </div>
+    );
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/use-callback-simple.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/use-callback-simple.expect.md
index 4228e8a584..27e04507db 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/use-callback-simple.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/use-callback-simple.expect.md
@@ -15,7 +15,7 @@ function component() {
 
 ```javascript
 function component() {
-  const $ = React.unstable_useMemoCache(3);
+  const $ = React.unstable_useMemoCache(5);
   const [count, setCount] = useState(0);
   const c_0 = $[0] !== setCount;
   const c_1 = $[1] !== count;
@@ -29,7 +29,16 @@ function component() {
     t0 = $[2];
   }
   const increment = t0;
-  return <Foo onClick={increment}></Foo>;
+  const c_3 = $[3] !== increment;
+  let t1;
+  if (c_3) {
+    t1 = <Foo onClick={increment}></Foo>;
+    $[3] = increment;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md
index 182a54a307..ef1baed8c5 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md
@@ -13,7 +13,7 @@ function component(a) {
 
 ```javascript
 function component(a) {
-  const $ = React.unstable_useMemoCache(2);
+  const $ = React.unstable_useMemoCache(4);
   const c_0 = $[0] !== a;
   let t0;
   if (c_0) {
@@ -24,7 +24,16 @@ function component(a) {
     t0 = $[1];
   }
   const x = t0;
-  return <Foo x={x}></Foo>;
+  const c_2 = $[2] !== x;
+  let t1;
+  if (c_2) {
+    t1 = <Foo x={x}></Foo>;
+    $[2] = x;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.expect.md b/compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.expect.md
new file mode 100644
index 0000000000..9a6b460671
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.expect.md
@@ -0,0 +1,41 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const [name, setName] = useState(null);
+  const onChange = function (e) {
+    setName(e.target.value);
+  };
+  return (
+    <form>
+      <input onChange={onChange} value={name} />
+    </form>
+  );
+}
+
+```
+
+## Code
+
+```javascript
+function Component(props) {
+  const $ = React.unstable_useMemoCache(2);
+  const [name, setName] = useState(null);
+  const c_0 = $[0] !== setName;
+  let t0;
+  if (c_0) {
+    t0 = function (e) {
+      setName(e.target.value);
+    };
+    $[0] = setName;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const onChange = t0;
+  return <form>{<input onChange={onChange} value={name}></input>}</form>;
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.js b/compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.js
new file mode 100644
index 0000000000..98a0b104c5
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/disableMemoizeJsxElements/Component.js
@@ -0,0 +1,11 @@
+function Component(props) {
+  const [name, setName] = useState(null);
+  const onChange = function (e) {
+    setName(e.target.value);
+  };
+  return (
+    <form>
+      <input onChange={onChange} value={name} />
+    </form>
+  );
+}