[compiler] Infer return types of function expressions

Uses the returnIdentifier added in the previous PR to provide a stable identifier for which we can infer a return type for functions, then wires up the equations in InferTypes to infer the type.

ghstack-source-id: c46c802
Pull Request resolved: #30785

Author: josephsavona

compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts

@@ -72,6 +72,7 @@ export function printFunction(fn: HIRFunction): string {
   if (definition.length !== 0) {
     output.push(definition);
   }
+  output.push(printType(fn.returnIdentifier.type));
   output.push(printHIR(fn.body));
   output.push(...fn.directives);
   return output.join('\n');
@@ -555,7 +556,8 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
             }
           })
           .join(', ') ?? '';
-      value = `${kind} ${name} @deps[${deps}] @context[${context}] @effects[${effects}]:\n${fn}`;
+      const type = printType(instrValue.loweredFunc.func.returnIdentifier.type).trim();
+      value = `${kind} ${name} @deps[${deps}] @context[${context}] @effects[${effects}]${type!== '' ? ` return${type}` : ''}:\n${fn}`;
       break;
     }
     case 'TaggedTemplateExpression': {

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

@@ -88,6 +88,7 @@ function apply(func: HIRFunction, unifier: Unifier): void {
       }
     }
   }
+  func.returnIdentifier.type = unifier.get(func.returnIdentifier.type);
 }
 
 type TypeEquation = {
@@ -122,6 +123,7 @@ function* generate(
   }
 
   const names = new Map();
+  const returnTypes: Array<Type> = [];
   for (const [_, block] of func.body.blocks) {
     for (const phi of block.phis) {
       yield equation(phi.type, {
@@ -133,6 +135,18 @@ function* generate(
     for (const instr of block.instructions) {
       yield* generateInstructionTypes(func.env, names, instr);
     }
+    const terminal = block.terminal;
+    if (terminal.kind === 'return') {
+      returnTypes.push(terminal.value.identifier.type);
+    }
+  }
+  if (returnTypes.length > 1) {
+    yield equation(func.returnIdentifier.type, {
+      kind: 'Phi',
+      operands: returnTypes,
+    });
+  } else if (returnTypes.length === 1) {
+    yield equation(func.returnIdentifier.type, returnTypes[0]!);
   }
 }
 
@@ -346,7 +360,7 @@ function* generateInstructionTypes(
 
     case 'FunctionExpression': {
       yield* generate(value.loweredFunc.func);
-      yield equation(left, {kind: 'Object', shapeId: BuiltInFunctionId});
+      yield equation(left, {kind: 'Function', shapeId: BuiltInFunctionId, return: value.loweredFunc.func.returnIdentifier.type});
       break;
     }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md

@@ -41,7 +41,7 @@ import { c as _c } from "react/compiler-runtime";
 import { useFragment } from "shared-runtime";
 
 function Component(props) {
-  const $ = _c(4);
+  const $ = _c(8);
   const post = useFragment(
     graphql`
       fragment F on T {
@@ -50,33 +50,45 @@ function Component(props) {
     `,
     props.post,
   );
-  let t0;
+  let media;
+  let onClick;
   if ($[0] !== post) {
     const allUrls = [];
 
-    const { media, comments, urls } = post;
+    const { media: t0, comments, urls } = post;
+    media = t0;
     let t1;
-    if ($[2] !== comments.length) {
+    if ($[3] !== comments.length) {
       t1 = (e) => {
         if (!comments.length) {
           return;
         }
 
         console.log(comments.length);
       };
-      $[2] = comments.length;
-      $[3] = t1;
+      $[3] = comments.length;
+      $[4] = t1;
     } else {
-      t1 = $[3];
+      t1 = $[4];
     }
-    const onClick = t1;
+    onClick = t1;
 
     allUrls.push(...urls);
-    t0 = <Media media={media} onClick={onClick} />;
     $[0] = post;
-    $[1] = t0;
+    $[1] = media;
+    $[2] = onClick;
+  } else {
+    media = $[1];
+    onClick = $[2];
+  }
+  let t0;
+  if ($[5] !== media || $[6] !== onClick) {
+    t0 = <Media media={media} onClick={onClick} />;
+    $[5] = media;
+    $[6] = onClick;
+    $[7] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[7];
   }
   return t0;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/flag-enable-emit-hook-guards.expect.md

@@ -49,7 +49,7 @@ import {
 
 const MyContext = createContext("my context value");
 function Component(t0) {
-  const $ = _c(4);
+  const $ = _c(5);
   try {
     $dispatcherGuard(0);
     const { value } = t0;
@@ -71,21 +71,24 @@ function Component(t0) {
     })();
     print(value, state);
     let t2;
-    let t3;
     if ($[1] !== state) {
       t2 = () => {
         if (state === 4) {
           setState(5);
         }
       };
-
-      t3 = [state];
       $[1] = state;
       $[2] = t2;
-      $[3] = t3;
     } else {
       t2 = $[2];
-      t3 = $[3];
+    }
+    let t3;
+    if ($[3] !== state) {
+      t3 = [state];
+      $[3] = state;
+      $[4] = t3;
+    } else {
+      t3 = $[4];
     }
     (function () {
       try {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-simple-const-declaration.expect.md

@@ -25,18 +25,17 @@ export const FIXTURE_ENTRYPOINT = {
 import { c as _c } from "react/compiler-runtime";
 function hoisting() {
   const $ = _c(1);
-  let t0;
+  let foo;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    const foo = () => bar + baz;
+    foo = () => bar + baz;
 
     const bar = 3;
     const baz = 2;
-    t0 = foo();
-    $[0] = t0;
+    $[0] = foo;
   } else {
-    t0 = $[0];
+    foo = $[0];
   }
-  return t0;
+  return foo();
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-simple-let-declaration.expect.md

@@ -25,18 +25,17 @@ export const FIXTURE_ENTRYPOINT = {
 import { c as _c } from "react/compiler-runtime";
 function hoisting() {
   const $ = _c(1);
-  let t0;
+  let foo;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    const foo = () => bar + baz;
+    foo = () => bar + baz;
 
     let bar = 3;
     let baz = 2;
-    t0 = foo();
-    $[0] = t0;
+    $[0] = foo;
   } else {
-    t0 = $[0];
+    foo = $[0];
   }
-  return t0;
+  return foo();
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-ternary-local-variable.expect.md

@@ -24,17 +24,25 @@ import { c as _c } from "react/compiler-runtime";
 import { RenderPropAsChild, StaticText1, StaticText2 } from "shared-runtime";
 
 function Component(props) {
-  const $ = _c(2);
+  const $ = _c(4);
   const Foo = props.showText1 ? StaticText1 : StaticText2;
   let t0;
   if ($[0] !== Foo) {
-    t0 = <RenderPropAsChild items={[() => <Foo key="0" />]} />;
+    t0 = () => <Foo key="0" />;
     $[0] = Foo;
     $[1] = t0;
   } else {
     t0 = $[1];
   }
-  return t0;
+  let t1;
+  if ($[2] !== t0) {
+    t1 = <RenderPropAsChild items={[t0]} />;
+    $[2] = t0;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-deps-subset-of-decls.expect.md

@@ -33,23 +33,55 @@ import { c as _c } from "react/compiler-runtime";
 import { useState } from "react";
 
 function Component() {
-  const $ = _c(2);
+  const $ = _c(11);
   const [count, setCount] = useState(0);
   let t0;
   if ($[0] !== count) {
-    t0 = (
-      <div>
-        <button onClick={() => setCount(count - 1)}>Decrement</button>
-
-        <button onClick={() => setCount(count + 1)}>Increment</button>
-      </div>
-    );
+    t0 = () => setCount(count - 1);
     $[0] = count;
     $[1] = t0;
   } else {
     t0 = $[1];
   }
-  return t0;
+  let t1;
+  if ($[2] !== t0) {
+    t1 = <button onClick={t0}>Decrement</button>;
+    $[2] = t0;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  let t2;
+  if ($[4] !== count) {
+    t2 = () => setCount(count + 1);
+    $[4] = count;
+    $[5] = t2;
+  } else {
+    t2 = $[5];
+  }
+  let t3;
+  if ($[6] !== t2) {
+    t3 = <button onClick={t2}>Increment</button>;
+    $[6] = t2;
+    $[7] = t3;
+  } else {
+    t3 = $[7];
+  }
+  let t4;
+  if ($[8] !== t1 || $[9] !== t3) {
+    t4 = (
+      <div>
+        {t1}
+        {t3}
+      </div>
+    );
+    $[8] = t1;
+    $[9] = t3;
+    $[10] = t4;
+  } else {
+    t4 = $[10];
+  }
+  return t4;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.expect.md

@@ -34,47 +34,59 @@ import { useState } from "react";
 import { Stringify } from "shared-runtime";
 
 function Component() {
-  const $ = _c(7);
+  const $ = _c(10);
   const [state, setState] = useState(0);
   let t0;
-  let t1;
   if ($[0] !== state) {
-    t0 = (
-      <button data-testid="button" onClick={() => setState(state + 1)}>
-        increment
-      </button>
-    );
-    t1 = <span>{state}</span>;
+    t0 = () => setState(state + 1);
     $[0] = state;
     $[1] = t0;
-    $[2] = t1;
   } else {
     t0 = $[1];
-    t1 = $[2];
+  }
+  let t1;
+  if ($[2] !== t0) {
+    t1 = (
+      <button data-testid="button" onClick={t0}>
+        increment
+      </button>
+    );
+    $[2] = t0;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
   }
   let t2;
-  if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
-    t2 = <Stringify text="Counter" />;
-    $[3] = t2;
+  if ($[4] !== state) {
+    t2 = <span>{state}</span>;
+    $[4] = state;
+    $[5] = t2;
   } else {
-    t2 = $[3];
+    t2 = $[5];
   }
   let t3;
-  if ($[4] !== t1 || $[5] !== t0) {
-    t3 = (
+  if ($[6] === Symbol.for("react.memo_cache_sentinel")) {
+    t3 = <Stringify text="Counter" />;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  let t4;
+  if ($[7] !== t2 || $[8] !== t1) {
+    t4 = (
       <div>
+        {t3}
         {t2}
         {t1}
-        {t0}
       </div>
     );
-    $[4] = t1;
-    $[5] = t0;
-    $[6] = t3;
+    $[7] = t2;
+    $[8] = t1;
+    $[9] = t4;
   } else {
-    t3 = $[6];
+    t4 = $[9];
   }
-  return t3;
+  return t4;
 }
 
 export const FIXTURE_ENTRYPOINT = {