From 57550bbb7add096fb6289d789c763c72e795bd31 Mon Sep 17 00:00:00 2001
From: Mike Vitousek <mmvitousek@gmail.com>
Date: Tue, 11 Jun 2024 13:42:31 -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: e2c23348265b7ef651232b962ed7be7f6fed1930
Pull Request resolved: https://github.com/facebook/react/pull/29866
---
 .../src/Entrypoint/Program.ts                 | 71 ++++++++++++++++---
 .../infer-no-component-annot.expect.md        | 39 ++++++++++
 .../compiler/infer-no-component-annot.ts      | 12 ++++
 3 files changed, 112 insertions(+), 10 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 8e5b00e6cd160..adee97238b2b7 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
@@ -651,24 +651,75 @@ 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();
-  } else if (params.length === 2) {
-    // check if second param might be a ref
-    if (params[1].isIdentifier()) {
+  } else if (params.length > 0 && params.length <= 2) {
+    if (!isValidPropsAnnotation(params[0].node.typeAnnotation)) {
+      return false;
+    }
+
+    if (params.length === 1) {
+      return !params[0].isRestElement();
+    } else if (params[1].isIdentifier()) {
+      // check if second param might be a ref
       const { name } = params[1].node;
       return name.includes("ref") || name.includes("Ref");
+    } else {
+      /**
+       * Otherwise, avoid helper functions that take more than one argument.
+       * Helpers are _usually_ named with lowercase, but some code may
+       * violate this rule
+       */
+      return false;
     }
-    /**
-     * Otherwise, avoid helper functions that take more than one argument.
-     * Helpers are _usually_ named with lowercase, but some code may
-     * violate this rule
-     */
   }
   return false;
 }
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],
+};