compiler: Known hooks/nonescaping scopes dont count as pruned

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

@@ -331,29 +331,8 @@ class CountMemoBlockVisitor extends ReactiveFunctionVisitor<void> {
     scopeBlock: PrunedReactiveScopeBlock,
     state: void
   ): void {
-    let isHookOnlyMemoBlock = false;
-    if (
-      scopeBlock.instructions.length === 1 &&
-      scopeBlock.instructions[0].kind === "instruction"
-    ) {
-      const instr = scopeBlock.instructions[0]!.instruction;
-      if (
-        instr.value.kind === "MethodCall" ||
-        instr.value.kind === "CallExpression"
-      ) {
-        const callee =
-          instr.value.kind === "MethodCall"
-            ? instr.value.property
-            : instr.value.callee;
-        if (getHookKind(this.env, callee.identifier) != null) {
-          isHookOnlyMemoBlock = true;
-        }
-      }
-    }
-    if (!isHookOnlyMemoBlock) {
-      this.prunedMemoBlocks += 1;
-      this.prunedMemoValues += scopeBlock.scope.declarations.size;
-    }
+    this.prunedMemoBlocks += 1;
+    this.prunedMemoValues += scopeBlock.scope.declarations.size;
     this.traversePrunedScope(scopeBlock, state);
   }
 }

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

@@ -66,6 +66,20 @@ class Transform extends ReactiveFunctionTransform<State> {
     this.visitScope(scope, innerState);
     outerState.hasHook ||= innerState.hasHook;
     if (innerState.hasHook) {
+      if (scope.instructions.length === 1) {
+        /*
+         * This was a scope just for a hook call, which doesn't need memoization.
+         * flatten it away
+         */
+        return {
+          kind: "replace-many",
+          value: scope.instructions,
+        };
+      }
+      /*
+       * else this scope had multiple instructions and produced some other value:
+       * mark it as pruned
+       */
       return {
         kind: "replace",
         value: {

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

@@ -951,12 +951,8 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
     } else {
       this.prunedScopes.add(scopeBlock.scope.id);
       return {
-        kind: "replace",
-        value: {
-          kind: "pruned-scope",
-          scope: scopeBlock.scope,
-          instructions: scopeBlock.instructions,
-        },
+        kind: "replace-many",
+        value: scopeBlock.instructions,
       };
     }
   }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/log-pruned-memoization.expect.md

@@ -3,10 +3,17 @@
 
 ```javascript
 // @logger
-import { useState } from "react";
-import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+import { createContext, use, useState } from "react";
+import {
+  Stringify,
+  identity,
+  makeObject_Primitives,
+  useHook,
+} from "shared-runtime";
 
 function Component() {
+  const w = use(Context);
+
   // The scopes for x and x2 are interleaved, so this is one scope with two values
   const x = makeObject_Primitives();
   const x2 = makeObject_Primitives();
@@ -26,11 +33,21 @@ function Component() {
   }
 
   // Overall we expect two pruned scopes (for x+x2, and obj), with 3 pruned scope values.
-  return [x, x2, y, z];
+  return <Stringify items={[w, x, x2, y, z]} />;
+}
+
+const Context = createContext();
+
+function Wrapper() {
+  return (
+    <Context value={42}>
+      <Component />
+    </Context>
+  );
 }
 
 export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
+  fn: Wrapper,
   params: [{}],
 };
 
@@ -40,11 +57,17 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @logger
-import { useState } from "react";
-import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+import { createContext, use, useState } from "react";
+import {
+  Stringify,
+  identity,
+  makeObject_Primitives,
+  useHook,
+} from "shared-runtime";
 
 function Component() {
-  const $ = _c(5);
+  const $ = _c(6);
+  const w = use(Context);
 
   const x = makeObject_Primitives();
   const x2 = makeObject_Primitives();
@@ -65,20 +88,39 @@ function Component() {
     z = $[0];
   }
   let t0;
-  if ($[1] !== x || $[2] !== x2 || $[3] !== y) {
-    t0 = [x, x2, y, z];
-    $[1] = x;
-    $[2] = x2;
-    $[3] = y;
-    $[4] = t0;
+  if ($[1] !== w || $[2] !== x || $[3] !== x2 || $[4] !== y) {
+    t0 = <Stringify items={[w, x, x2, y, z]} />;
+    $[1] = w;
+    $[2] = x;
+    $[3] = x2;
+    $[4] = y;
+    $[5] = t0;
+  } else {
+    t0 = $[5];
+  }
+  return t0;
+}
+
+const Context = createContext();
+
+function Wrapper() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = (
+      <Context value={42}>
+        <Component />
+      </Context>
+    );
+    $[0] = t0;
   } else {
-    t0 = $[4];
+    t0 = $[0];
   }
   return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
+  fn: Wrapper,
   params: [{}],
 };
 
@@ -87,8 +129,9 @@ export const FIXTURE_ENTRYPOINT = {
 ## Logs
 
 ```
-{"kind":"CompileSuccess","fnLoc":{"start":{"line":5,"column":0,"index":121},"end":{"line":26,"column":1,"index":813},"filename":"log-pruned-memoization.ts"},"fnName":"Component","memoSlots":5,"memoBlocks":2,"memoValues":2,"prunedMemoBlocks":2,"prunedMemoValues":3}
+{"kind":"CompileSuccess","fnLoc":{"start":{"line":10,"column":0,"index":161},"end":{"line":33,"column":1,"index":905},"filename":"log-pruned-memoization.ts"},"fnName":"Component","memoSlots":6,"memoBlocks":2,"memoValues":2,"prunedMemoBlocks":2,"prunedMemoValues":3}
+{"kind":"CompileSuccess","fnLoc":{"start":{"line":37,"column":0,"index":941},"end":{"line":43,"column":1,"index":1039},"filename":"log-pruned-memoization.ts"},"fnName":"Wrapper","memoSlots":1,"memoBlocks":1,"memoValues":1,"prunedMemoBlocks":0,"prunedMemoValues":0}
 ```
 
 ### Eval output
-(kind: ok) [{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},[{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true}]]
+(kind: ok) <div>{"items":[42,{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},[{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true}]]}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/log-pruned-memoization.js

@@ -1,8 +1,15 @@
 // @logger
-import { useState } from "react";
-import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+import { createContext, use, useState } from "react";
+import {
+  Stringify,
+  identity,
+  makeObject_Primitives,
+  useHook,
+} from "shared-runtime";
 
 function Component() {
+  const w = use(Context);
+
   // The scopes for x and x2 are interleaved, so this is one scope with two values
   const x = makeObject_Primitives();
   const x2 = makeObject_Primitives();
@@ -22,10 +29,20 @@ function Component() {
   }
 
   // Overall we expect two pruned scopes (for x+x2, and obj), with 3 pruned scope values.
-  return [x, x2, y, z];
+  return <Stringify items={[w, x, x2, y, z]} />;
+}
+
+const Context = createContext();
+
+function Wrapper() {
+  return (
+    <Context value={42}>
+      <Component />
+    </Context>
+  );
 }
 
 export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
+  fn: Wrapper,
   params: [{}],
 };

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nonreactive-noescaping-dependency-can-inline-into-consuming-scope.expect.md

@@ -0,0 +1,46 @@
+
+## Input
+
+```javascript
+// @flow
+function Component() {
+  return (
+    <div
+      className={stylex(
+        // this value is a) in its own scope, b) non-reactive, and c) non-escaping
+        // its scope gets pruned bc it's non-escaping, but this doesn't mean we need to
+        // create a temporary for it
+        flags.feature("feature-name") ? styles.featureNameStyle : null
+      )}
+    ></div>
+  );
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = (
+      <div
+        className={stylex(
+          flags.feature("feature-name") ? styles.featureNameStyle : null,
+        )}
+      />
+    );
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+```
+
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nonreactive-noescaping-dependency-can-inline-into-consuming-scope.js

@@ -0,0 +1,13 @@
+// @flow
+function Component() {
+  return (
+    <div
+      className={stylex(
+        // this value is a) in its own scope, b) non-reactive, and c) non-escaping
+        // its scope gets pruned bc it's non-escaping, but this doesn't mean we need to
+        // create a temporary for it
+        flags.feature("feature-name") ? styles.featureNameStyle : null
+      )}
+    ></div>
+  );
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-control-dependency-from-interleaved-reactivity-for-of.expect.md

@@ -39,7 +39,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(2);
+  const $ = _c(1);
 
   const a = [];
   const b = [];
@@ -53,12 +53,11 @@ function Component(props) {
     x = 1;
   }
   let t0;
-  if ($[0] !== x) {
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = [x];
-    $[0] = x;
-    $[1] = t0;
+    $[0] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[0];
   }
   return t0;
 }