From 29cb62ad6ca2e3d491df038e1cc8cbe882fa31a4 Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Fri, 2 Feb 2024 10:56:44 -0500
Subject: [PATCH] Add type information for jsx/runtime

---
 .../src/HIR/Globals.ts                        | 84 +++++++++++--------
 .../fixtures/compiler/jsx-freeze.expect.md    | 55 ++++++++++++
 .../__tests__/fixtures/compiler/jsx-freeze.js | 17 ++++
 3 files changed, 120 insertions(+), 36 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.js

diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/Globals.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/Globals.ts
index c431998efa..ba4a9140e0 100644
--- a/compiler/packages/babel-plugin-react-forget/src/HIR/Globals.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/HIR/Globals.ts
@@ -346,42 +346,54 @@ const BUILTIN_HOOKS: Array<[string, BuiltInType]> = [
   ],
 ];
 
-TYPED_GLOBALS.push([
-  "React",
-  addObject(DEFAULT_SHAPES, null, [
-    ...BUILTIN_HOOKS,
-    [
-      "createElement",
-      addFunction(DEFAULT_SHAPES, [], {
-        positionalParams: [],
-        restParam: Effect.Freeze,
-        returnType: { kind: "Poly" },
-        calleeEffect: Effect.Read,
-        returnValueKind: ValueKind.Frozen,
-      }),
-    ],
-    [
-      "cloneElement",
-      addFunction(DEFAULT_SHAPES, [], {
-        positionalParams: [],
-        restParam: Effect.Freeze,
-        returnType: { kind: "Poly" },
-        calleeEffect: Effect.Read,
-        returnValueKind: ValueKind.Frozen,
-      }),
-    ],
-    [
-      "createRef",
-      addFunction(DEFAULT_SHAPES, [], {
-        positionalParams: [],
-        restParam: Effect.Capture, // createRef takes no paramters
-        returnType: { kind: "Object", shapeId: BuiltInUseRefId },
-        calleeEffect: Effect.Read,
-        returnValueKind: ValueKind.Mutable,
-      }),
-    ],
-  ]),
-]);
+TYPED_GLOBALS.push(
+  [
+    "React",
+    addObject(DEFAULT_SHAPES, null, [
+      ...BUILTIN_HOOKS,
+      [
+        "createElement",
+        addFunction(DEFAULT_SHAPES, [], {
+          positionalParams: [],
+          restParam: Effect.Freeze,
+          returnType: { kind: "Poly" },
+          calleeEffect: Effect.Read,
+          returnValueKind: ValueKind.Frozen,
+        }),
+      ],
+      [
+        "cloneElement",
+        addFunction(DEFAULT_SHAPES, [], {
+          positionalParams: [],
+          restParam: Effect.Freeze,
+          returnType: { kind: "Poly" },
+          calleeEffect: Effect.Read,
+          returnValueKind: ValueKind.Frozen,
+        }),
+      ],
+      [
+        "createRef",
+        addFunction(DEFAULT_SHAPES, [], {
+          positionalParams: [],
+          restParam: Effect.Capture, // createRef takes no paramters
+          returnType: { kind: "Object", shapeId: BuiltInUseRefId },
+          calleeEffect: Effect.Read,
+          returnValueKind: ValueKind.Mutable,
+        }),
+      ],
+    ]),
+  ],
+  [
+    "_jsx",
+    addFunction(DEFAULT_SHAPES, [], {
+      positionalParams: [],
+      restParam: Effect.Freeze,
+      returnType: { kind: "Poly" },
+      calleeEffect: Effect.Read,
+      returnValueKind: ValueKind.Frozen,
+    }),
+  ]
+);
 
 export type Global = BuiltInType | PolyType;
 export type GlobalRegistry = Map<string, Global>;
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.expect.md
new file mode 100644
index 0000000000..a4a0079cf2
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.expect.md
@@ -0,0 +1,55 @@
+
+## Input
+
+```javascript
+import { jsx as _jsx } from "react/jsx-runtime";
+import { shallowCopy } from "shared-runtime";
+
+function Component(props) {
+  const childprops = { style: { width: props.width } };
+  const element = _jsx("div", {
+    childprops: childprops,
+    children: '"hello world"',
+  });
+  shallowCopy(childprops); // function that in theory could mutate, we assume not bc createElement freezes
+  return element;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+import { jsx as _jsx } from "react/jsx-runtime";
+import { shallowCopy } from "shared-runtime";
+
+function Component(props) {
+  const $ = useMemoCache(2);
+  let element;
+  if ($[0] !== props.width) {
+    const childprops = { style: { width: props.width } };
+    element = _jsx("div", { childprops, children: '"hello world"' });
+    shallowCopy(childprops);
+    $[0] = props.width;
+    $[1] = element;
+  } else {
+    element = $[1];
+  }
+  return element;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div childprops="[object Object]">"hello world"</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.js
new file mode 100644
index 0000000000..55dea91f6a
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/jsx-freeze.js
@@ -0,0 +1,17 @@
+import { jsx as _jsx } from "react/jsx-runtime";
+import { shallowCopy } from "shared-runtime";
+
+function Component(props) {
+  const childprops = { style: { width: props.width } };
+  const element = _jsx("div", {
+    childprops: childprops,
+    children: '"hello world"',
+  });
+  shallowCopy(childprops); // function that in theory could mutate, we assume not bc createElement freezes
+  return element;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};