Update on "[compiler][wip] Inline single return JSX for known function components"

See comments, still WIP bc i'm getting a weird error with Babel and line numbers, one of the source locations must be off.

[ghstack-poisoned]

compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineSingleReturnJSX.ts

@@ -17,6 +17,7 @@ import {
   SpreadPattern,
 } from '../HIR';
 import {createTemporaryPlace, markInstructionIds} from '../HIR/HIRBuilder';
+import {BuiltInPropsId} from '../HIR/ObjectShape';
 
 export function inlineSingleReturnJSX(fn: HIRFunction): void {
   if (fn.fnType !== 'Component') {
@@ -65,6 +66,7 @@ export function inlineSingleReturnJSX(fn: HIRFunction): void {
     nextInstructions ??= returnBlock.instructions.slice(0, i);
     const propsTemp = createTemporaryPlace(fn.env, value.loc);
     propsTemp.effect = Effect.Freeze;
+    propsTemp.identifier.type = {kind: 'Object', shapeId: BuiltInPropsId};
 
     const properties: Array<ObjectProperty | SpreadPattern> = value.props.map(
       attr => {

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

@@ -1475,7 +1475,7 @@ function withLoc<T extends (...args: Array<any>) => t.Node>(
     ...args: Parameters<T>
   ): ReturnType<T> => {
     const node = fn(...args);
-    if (loc != null && loc != GeneratedSource) {
+    if (loc != null && typeof loc !== 'symbol') {
       node.loc = loc;
     }
     return node as ReturnType<T>;
@@ -1550,7 +1550,7 @@ function createCallExpression(
   isHook: boolean,
 ): t.CallExpression {
   const callExpr = t.callExpression(callee, args);
-  if (loc != null && loc != GeneratedSource) {
+  if (loc != null && typeof loc !== 'symbol') {
     callExpr.loc = loc;
   }
 
@@ -2555,7 +2555,7 @@ function codegenPlace(cx: Context, place: Place): t.Expression | t.JSXText {
     suggestions: null,
   });
   const identifier = convertIdentifier(place.identifier);
-  identifier.loc = place.loc as any;
+  identifier.loc = typeof place.loc !== 'symbol' ? (place.loc as any) : null;
   return identifier;
 }
 

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

@@ -17,6 +17,7 @@ import {
   ReactiveScope,
   makeInstructionId,
 } from '../HIR/HIR';
+import {BuiltInPropsId} from '../HIR/ObjectShape';
 import {
   doesPatternContainSpreadElement,
   eachInstructionOperand,
@@ -227,6 +228,15 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
     case 'StoreGlobal': {
       return false;
     }
+    case 'ObjectExpression': {
+      const type = instruction.lvalue.identifier.type;
+      if (type.kind === 'Object' && type.shapeId === BuiltInPropsId) {
+        // If this is an object literal for inlined JSX, we don't need to memoize
+        // it. The props object is always assumed to have changed.
+        return false;
+      }
+      return true;
+    }
     case 'TaggedTemplateExpression':
     case 'CallExpression':
     case 'MethodCall': {
@@ -239,7 +249,6 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
     case 'JsxExpression':
     case 'JsxFragment':
     case 'NewExpression':
-    case 'ObjectExpression':
     case 'UnsupportedNode':
     case 'ObjectMethod':
     case 'FunctionExpression': {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-single-return-jsx/inline-single-return-jsx.expect.md

@@ -0,0 +1,105 @@
+
+## Input
+
+```javascript
+// @enableInlineSingleReturnJSX
+function Component({a, b}) {
+  return (
+    <Child value={a}>
+      <div>{b}</div>
+    </Child>
+  );
+}
+
+function Child({value, children}) {
+  return (
+    <div>
+      <span>{value}</span>
+      {children}
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 0, b: 1},
+    {a: 0, b: 0},
+    {a: 1, b: 1},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableInlineSingleReturnJSX
+function Component(t0) {
+  const $ = _c(2);
+  const { a, b } = t0;
+  let t1;
+  if ($[0] !== b) {
+    t1 = <div>{b}</div>;
+    $[0] = b;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return Child({ value: a, children: t1 });
+}
+
+function Child(t0) {
+  const $ = _c(5);
+  const { value, children } = t0;
+  let t1;
+  if ($[0] !== value) {
+    t1 = <span>{value}</span>;
+    $[0] = value;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  let t2;
+  if ($[2] !== t1 || $[3] !== children) {
+    t2 = (
+      <div>
+        {t1}
+        {children}
+      </div>
+    );
+    $[2] = t1;
+    $[3] = children;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 1, b: 0 },
+    { a: 1, b: 1 },
+    { a: 0, b: 1 },
+    { a: 0, b: 0 },
+    { a: 1, b: 1 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div><span>0</span><div>0</div></div>
+<div><span>1</span><div>0</div></div>
+<div><span>1</span><div>1</div></div>
+<div><span>0</span><div>1</div></div>
+<div><span>0</span><div>0</div></div>
+<div><span>1</span><div>1</div></div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-single-return-jsx/inline-single-return-jsx.js

@@ -1,9 +1,30 @@
 // @enableInlineSingleReturnJSX
 function Component({a, b}) {
-  const c = [a, b];
   return (
-    <Child value={c}>
-      <div />
+    <Child value={a}>
+      <div>{b}</div>
     </Child>
   );
 }
+
+function Child({value, children}) {
+  return (
+    <div>
+      <span>{value}</span>
+      {children}
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 0, b: 1},
+    {a: 0, b: 0},
+    {a: 1, b: 1},
+  ],
+};