From d456e4a78ccf1cf607a46e802c67eb1442657300 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 22 May 2023 09:20:16 -0700
Subject: [PATCH] Memoize based on inferred return type, not signature

---
 .../InferReactiveScopeVariables.ts            | 28 +++--------------
 .../compiler/infer-global-object.expect.md    | 31 ++++++-------------
 .../compiler/type-test-polymorphic.expect.md  | 21 +++++--------
 3 files changed, 22 insertions(+), 58 deletions(-)

diff --git a/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts b/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
index 98d5ea57f8..aa4e6d1ab4 100644
--- a/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
+++ b/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
@@ -10,7 +10,6 @@ import {
   HIRFunction,
   Identifier,
   Instruction,
-  InstructionValue,
   makeInstructionId,
   makeScopeId,
   Place,
@@ -21,7 +20,6 @@ import {
   eachInstructionOperand,
   eachPatternOperand,
 } from "../HIR/visitors";
-import { getFunctionCallSignature } from "../Inference/InferReferenceEffects";
 import DisjointSet from "../Utils/DisjointSet";
 import { assertExhaustive } from "../Utils/utils";
 
@@ -105,7 +103,7 @@ export function inferReactiveScopeVariables(fn: HIRFunction): void {
     for (const instr of block.instructions) {
       const operands: Array<Identifier> = [];
       const range = instr.lvalue.identifier.mutableRange;
-      if (range.end > range.start + 1 || mayAllocate(fn.env, instr.value)) {
+      if (range.end > range.start + 1 || mayAllocate(fn.env, instr)) {
         operands.push(instr.lvalue!.identifier);
       }
       if (
@@ -219,7 +217,8 @@ function isMutable({ id }: Instruction, place: Place): boolean {
   return id >= range.start && id < range.end;
 }
 
-function mayAllocate(env: Environment, value: InstructionValue): boolean {
+function mayAllocate(env: Environment, instruction: Instruction): boolean {
+  const { value } = instruction;
   switch (value.kind) {
     case "Destructure": {
       return doesPatternContainSpreadElement(value.lvalue.pattern);
@@ -247,27 +246,10 @@ function mayAllocate(env: Environment, value: InstructionValue): boolean {
     case "Debugger": {
       return false;
     }
-    case "CallExpression": {
-      if (env.enableFunctionCallSignatureOptimizations) {
-        const signature = getFunctionCallSignature(
-          env,
-          value.callee.identifier.type
-        );
-        if (signature !== null) {
-          return signature.returnType.kind !== "Primitive";
-        }
-      }
-      return true;
-    }
+    case "CallExpression":
     case "MethodCall": {
       if (env.enableFunctionCallSignatureOptimizations) {
-        const signature = getFunctionCallSignature(
-          env,
-          value.receiver.identifier.type
-        );
-        if (signature !== null) {
-          return signature.returnType.kind !== "Primitive";
-        }
+        return instruction.lvalue.identifier.type.kind !== "Primitive";
       }
       return true;
     }
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/infer-global-object.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/infer-global-object.expect.md
index f69e8eac44..c82f9d4289 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/infer-global-object.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/infer-global-object.expect.md
@@ -24,7 +24,7 @@ function Component(props) {
 import { unstable_useMemoCache as useMemoCache } from "react"; // Check that we correctly resolve type and effect lookups on the javascript
 // global object.
 function Component(props) {
-  const $ = useMemoCache(7);
+  const $ = useMemoCache(4);
   const c_0 = $[0] !== props.b;
   let t0;
   if (c_0) {
@@ -35,32 +35,21 @@ function Component(props) {
     t0 = $[1];
   }
   const neverAliasedOrMutated = t0;
-  const c_2 = $[2] !== props.a;
-  const c_3 = $[3] !== neverAliasedOrMutated;
-  let t1;
-  if (c_2 || c_3) {
-    t1 = Math.max(props.a, neverAliasedOrMutated);
-    $[2] = props.a;
-    $[3] = neverAliasedOrMutated;
-    $[4] = t1;
-  } else {
-    t1 = $[4];
-  }
-  const primitiveVal1 = t1;
+  const primitiveVal1 = Math.max(props.a, neverAliasedOrMutated);
 
   const primitiveVal3 = globaThis.globalThis.NaN;
 
   foo(primitiveVal1, Infinity, primitiveVal3);
-  const c_5 = $[5] !== primitiveVal1;
-  let t2;
-  if (c_5) {
-    t2 = { primitiveVal1, primitiveVal2: Infinity, primitiveVal3 };
-    $[5] = primitiveVal1;
-    $[6] = t2;
+  const c_2 = $[2] !== primitiveVal1;
+  let t1;
+  if (c_2) {
+    t1 = { primitiveVal1, primitiveVal2: Infinity, primitiveVal3 };
+    $[2] = primitiveVal1;
+    $[3] = t1;
   } else {
-    t2 = $[6];
+    t1 = $[3];
   }
-  return t2;
+  return t1;
 }
 
 ```
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/type-test-polymorphic.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/type-test-polymorphic.expect.md
index 0006510875..321cef23e2 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/type-test-polymorphic.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/type-test-polymorphic.expect.md
@@ -24,33 +24,26 @@ function component() {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function component() {
-  const $ = useMemoCache(3);
+  const $ = useMemoCache(2);
+  const p = makePrimitive();
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = makePrimitive();
+    t0 = {};
     $[0] = t0;
   } else {
     t0 = $[0];
   }
-  const p = t0;
-  let t1;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = {};
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
-  const o = t1;
+  const o = t0;
   let x;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
     x = {};
 
     x.t = p;
 
     x.t = o;
-    $[2] = x;
+    $[1] = x;
   } else {
-    x = $[2];
+    x = $[1];
   }
   const y = x.t;
   return y;