From 746c1dbed581ad128f0ea22b8fe20267391e1c2f Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Tue, 18 Apr 2023 15:29:17 -0700
Subject: [PATCH] Add common (Relay) hooks to playground environment

This aligns the playground configuration with our internal compiler
configuration to make it easier to repro compilation issues on playground. There
is a bug that doesn't repro right now and i suspect it's because of different
hooks being configured.

Test plan:

Before: playground output has no obvious bugs, but is different than internal
compilation output where the bug occurs

After: playground output matches internal compilation output w the bug
---
 .../playground/components/Editor/index.tsx    | 55 ++++++++++++++++++-
 .../forget/packages/playground/next-env.d.ts  |  7 ---
 compiler/forget/src/HIR/index.ts              |  1 +
 compiler/forget/src/index.ts                  |  4 +-
 4 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/compiler/forget/packages/playground/components/Editor/index.tsx b/compiler/forget/packages/playground/components/Editor/index.tsx
index 0a14b29683..97b762c4b0 100644
--- a/compiler/forget/packages/playground/components/Editor/index.tsx
+++ b/compiler/forget/packages/playground/components/Editor/index.tsx
@@ -9,9 +9,12 @@ import { parse, ParserPlugin } from "@babel/parser";
 import traverse, { NodePath } from "@babel/traverse";
 import * as t from "@babel/types";
 import {
+  Effect,
+  Hook,
   printHIR,
   printReactiveFunction,
   run,
+  ValueKind,
 } from "babel-plugin-react-forget";
 import clsx from "clsx";
 import invariant from "invariant";
@@ -67,6 +70,54 @@ function parseFunctions(
   return items;
 }
 
+const COMMON_HOOKS: Array<[string, Hook]> = [
+  [
+    "useFragment",
+    {
+      name: "useFragment",
+      kind: "Custom",
+      valueKind: ValueKind.Frozen,
+      effectKind: Effect.Freeze,
+    },
+  ],
+  [
+    "usePaginationFragment",
+    {
+      name: "usePaginationFragment",
+      kind: "Custom",
+      valueKind: ValueKind.Frozen,
+      effectKind: Effect.Freeze,
+    },
+  ],
+  [
+    "useRefetchableFragment",
+    {
+      name: "useRefetchableFragment",
+      kind: "Custom",
+      valueKind: ValueKind.Frozen,
+      effectKind: Effect.Freeze,
+    },
+  ],
+  [
+    "useLazyLoadQuery",
+    {
+      name: "useLazyLoadQuery",
+      kind: "Custom",
+      valueKind: ValueKind.Frozen,
+      effectKind: Effect.Freeze,
+    },
+  ],
+  [
+    "usePreloadedQuery",
+    {
+      name: "usePreloadedQuery",
+      kind: "Custom",
+      valueKind: ValueKind.Frozen,
+      effectKind: Effect.Freeze,
+    },
+  ],
+];
+
 function compile(source: string): CompilerOutput {
   const results = new Map<string, PrintedCompilerPipelineValue[]>();
   const upsert = (result: PrintedCompilerPipelineValue) => {
@@ -79,7 +130,9 @@ function compile(source: string): CompilerOutput {
   };
   try {
     for (const fn of parseFunctions(source)) {
-      for (const result of run(fn)) {
+      for (const result of run(fn, {
+        customHooks: new Map([...COMMON_HOOKS]),
+      })) {
         const fnName = fn.node.id?.name ?? null;
         switch (result.kind) {
           case "ast": {
diff --git a/compiler/forget/packages/playground/next-env.d.ts b/compiler/forget/packages/playground/next-env.d.ts
index 811cbc51ad..4f11a03dc6 100644
--- a/compiler/forget/packages/playground/next-env.d.ts
+++ b/compiler/forget/packages/playground/next-env.d.ts
@@ -1,10 +1,3 @@
-/**
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
 /// <reference types="next" />
 /// <reference types="next/image-types/global" />
 
diff --git a/compiler/forget/src/HIR/index.ts b/compiler/forget/src/HIR/index.ts
index 8834ae1bd2..92d345d71e 100644
--- a/compiler/forget/src/HIR/index.ts
+++ b/compiler/forget/src/HIR/index.ts
@@ -15,6 +15,7 @@ export {
   reversePostorderBlocks,
   shrink,
 } from "./HIRBuilder";
+export { Hook } from "./Hooks";
 export { mergeConsecutiveBlocks } from "./MergeConsecutiveBlocks";
 export { printFunction, printHIR } from "./PrintHIR";
 export { validateConsistentIdentifiers } from "./ValidateConsistentIdentifiers";
diff --git a/compiler/forget/src/index.ts b/compiler/forget/src/index.ts
index e8a89302fc..deb0786942 100644
--- a/compiler/forget/src/index.ts
+++ b/compiler/forget/src/index.ts
@@ -11,8 +11,8 @@ export {
   CompilerErrorDetail,
   ErrorSeverity,
 } from "./CompilerError";
-export { compile, CompilerPipelineValue, run } from "./CompilerPipeline";
-export { Effect, printHIR, ValueKind } from "./HIR";
+export { CompilerPipelineValue, compile, run } from "./CompilerPipeline";
+export { Effect, Hook, ValueKind, printHIR } from "./HIR";
 export { printReactiveFunction } from "./ReactiveScopes";
 
 declare global {