From ab87e8eb65184e827a9e8d966d4d794b6c47a3ef Mon Sep 17 00:00:00 2001
From: Mike Vitousek <mmvitousek@gmail.com>
Date: Tue, 11 Jun 2024 13:09:01 -0700
Subject: [PATCH] [compiler] Expect component props annotations to be potential
 objects

Summary: We now expect that candidate components that have Flow or TS type annotations on their first parameters have annotations that are potentially objects--this lets us reject compiling functions that explicitly take e.g. `number` as a parameter.

ghstack-source-id: 03622a68d62bb7a4029dd8cf228e62b290b6dd7b
Pull Request resolved: https://github.com/facebook/react/pull/29866
---
 .../src/Entrypoint/Program.ts                 | 55 ++++++++++++++++++-
 .../infer-no-component-annot.expect.md        | 39 +++++++++++++
 .../compiler/infer-no-component-annot.ts      | 12 ++++
 3 files changed, 104 insertions(+), 2 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.ts

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
index 79d337e65216e..6ca9d2aafdb46 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
@@ -651,18 +651,69 @@ function isMemoCallback(path: NodePath<t.Expression>): boolean {
   );
 }
 
+function isValidPropsAnnotation(
+  annot: t.TypeAnnotation | t.TSTypeAnnotation | t.Noop | null | undefined
+): boolean {
+  if (annot == null) {
+    return true;
+  } else if (annot.type === "TSTypeAnnotation") {
+    switch (annot.typeAnnotation.type) {
+      case "TSArrayType":
+      case "TSBigIntKeyword":
+      case "TSBooleanKeyword":
+      case "TSConstructorType":
+      case "TSFunctionType":
+      case "TSLiteralType":
+      case "TSNeverKeyword":
+      case "TSNumberKeyword":
+      case "TSStringKeyword":
+      case "TSSymbolKeyword":
+      case "TSTupleType":
+        return false;
+    }
+    return true;
+  } else if (annot.type === "TypeAnnotation") {
+    switch (annot.typeAnnotation.type) {
+      case "ArrayTypeAnnotation":
+      case "BooleanLiteralTypeAnnotation":
+      case "BooleanTypeAnnotation":
+      case "EmptyTypeAnnotation":
+      case "FunctionTypeAnnotation":
+      case "NumberLiteralTypeAnnotation":
+      case "NumberTypeAnnotation":
+      case "StringLiteralTypeAnnotation":
+      case "StringTypeAnnotation":
+      case "SymbolTypeAnnotation":
+      case "ThisTypeAnnotation":
+      case "TupleTypeAnnotation":
+        return false;
+    }
+    return true;
+  } else if (annot.type === "Noop") {
+    return true;
+  } else {
+    assertExhaustive(annot, `Unexpected annotation node \`${annot}\``);
+  }
+}
+
 function isValidComponentParams(
   params: Array<NodePath<t.Identifier | t.Pattern | t.RestElement>>
 ): boolean {
   if (params.length === 0) {
     return true;
   } else if (params.length === 1) {
-    return !params[0].isRestElement();
+    return (
+      isValidPropsAnnotation(params[0].node.typeAnnotation) &&
+      !params[0].isRestElement()
+    );
   } else if (params.length === 2) {
     // check if second param might be a ref
     if (params[1].isIdentifier()) {
       const { name } = params[1].node;
-      return name.includes("ref") || name.includes("Ref");
+      return (
+        isValidPropsAnnotation(params[0].node.typeAnnotation) &&
+        (name.includes("ref") || name.includes("Ref"))
+      );
     }
     /**
      * Otherwise, avoid helper functions that take more than one argument.
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.expect.md
new file mode 100644
index 0000000000000..d47c57d0d63db
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.expect.md
@@ -0,0 +1,39 @@
+
+## Input
+
+```javascript
+// @compilationMode(infer)
+import { useIdentity, identity } from "shared-runtime";
+
+function Component(fakeProps: number) {
+  const x = useIdentity(fakeProps);
+  return identity(x);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [42],
+};
+
+```
+
+## Code
+
+```javascript
+// @compilationMode(infer)
+import { useIdentity, identity } from "shared-runtime";
+
+function Component(fakeProps: number) {
+  const x = useIdentity(fakeProps);
+  return identity(x);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [42],
+};
+
+```
+      
+### Eval output
+(kind: ok) 42
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.ts
new file mode 100644
index 0000000000000..425c2f7fadb09
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-no-component-annot.ts
@@ -0,0 +1,12 @@
+// @compilationMode(infer)
+import { useIdentity, identity } from "shared-runtime";
+
+function Component(fakeProps: number) {
+  const x = useIdentity(fakeProps);
+  return identity(x);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [42],
+};