From ec44c4789b6db4f8d06f23dffc808a38d46a8d85 Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Thu, 10 Aug 2023 17:24:21 -0400
Subject: [PATCH] Add ComponentDeclaration util

---
 .../src/Entrypoint/Program.ts                  |  5 ++---
 .../src/Utils/ComponentDeclaration.ts          | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)
 create mode 100644 compiler/forget/packages/babel-plugin-react-forget/src/Utils/ComponentDeclaration.ts

diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
index 1b707f80f5..f760bd5659 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
@@ -14,6 +14,7 @@ import {
   ErrorSeverity,
 } from "../CompilerError";
 import { GeneratedSource } from "../HIR";
+import { isComponentDeclaration } from "../Utils/ComponentDeclaration";
 import { getOrInsertDefault } from "../Utils/utils";
 import { addInstrumentForget } from "./Instrumentation";
 import { ExternalFunction, PluginOptions, parsePluginOptions } from "./Options";
@@ -330,9 +331,7 @@ function shouldVisitNode(
   pass: CompilerPass
 ): boolean {
   if (pass.opts.enableOnlyOnReactScript) {
-    // Typescript gets angry about lookup the magic prop
-    let fnAny = fn as any;
-    if (!fnAny.node.__componentDeclaration) {
+    if (!isComponentDeclaration(fn.node)) {
       return false;
     }
   }
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Utils/ComponentDeclaration.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Utils/ComponentDeclaration.ts
new file mode 100644
index 0000000000..4b50f4efee
--- /dev/null
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/Utils/ComponentDeclaration.ts
@@ -0,0 +1,18 @@
+import * as t from "@babel/types";
+
+type FunctionDeclOrExpr = t.FunctionDeclaration | t.ArrowFunctionExpression;
+type ComponentDeclaration = FunctionDeclOrExpr & {
+  __componentDeclaration: boolean;
+};
+
+export function isComponentDeclaration(
+  node: FunctionDeclOrExpr
+): node is ComponentDeclaration {
+  return Object.prototype.hasOwnProperty.call(node, "__componentDeclaration");
+}
+
+export function parseComponentDeclaration(
+  node: FunctionDeclOrExpr
+): ComponentDeclaration | null {
+  return isComponentDeclaration(node) ? node : null;
+}