From 19963f1c9df9a85215570880c42b7660191f0c26 Mon Sep 17 00:00:00 2001
From: Sathya Gunsasekaran <gsathya.ceg@gmail.com>
Date: Wed, 31 Jul 2024 14:35:30 +0100
Subject: [PATCH] [compiler] Add typing for useContext hook

In the future, we can use this to identify useContext calls.

[ghstack-poisoned]
---
 .../src/HIR/Globals.ts                        | 23 +++++++++++--------
 .../src/HIR/HIR.ts                            |  6 +++++
 .../src/HIR/ObjectShape.ts                    |  1 +
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts
index 9b3cfebb9af..5e90401c688 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts
@@ -10,6 +10,7 @@ import {
   BUILTIN_SHAPES,
   BuiltInArrayId,
   BuiltInUseActionStateId,
+  BuiltInUseContextHookId,
   BuiltInUseEffectHookId,
   BuiltInUseInsertionEffectHookId,
   BuiltInUseLayoutEffectHookId,
@@ -245,15 +246,19 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
 const REACT_APIS: Array<[string, BuiltInType]> = [
   [
     'useContext',
-    addHook(DEFAULT_SHAPES, {
-      positionalParams: [],
-      restParam: Effect.Read,
-      returnType: {kind: 'Poly'},
-      calleeEffect: Effect.Read,
-      hookKind: 'useContext',
-      returnValueKind: ValueKind.Frozen,
-      returnValueReason: ValueReason.Context,
-    }),
+    addHook(
+      DEFAULT_SHAPES,
+      {
+        positionalParams: [],
+        restParam: Effect.Read,
+        returnType: {kind: 'Poly'},
+        calleeEffect: Effect.Read,
+        hookKind: 'useContext',
+        returnValueKind: ValueKind.Frozen,
+        returnValueReason: ValueReason.Context,
+      },
+      BuiltInUseContextHookId,
+    ),
   ],
   [
     'useState',
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
index 4578a9b874a..f0ada2d3708 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
@@ -1599,6 +1599,12 @@ export function isUseInsertionEffectHookType(id: Identifier): boolean {
   );
 }
 
+export function isUseContextHookType(id: Identifier): boolean {
+  return (
+    id.type.kind === 'Function' && id.type.shapeId === 'BuiltInUseContextHook'
+  );
+}
+
 export function getHookKind(env: Environment, id: Identifier): HookKind | null {
   return getHookKindForType(env, id.type);
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts
index 63c4d7f6f9e..3d377dba59d 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts
@@ -208,6 +208,7 @@ export const BuiltInUseInsertionEffectHookId = 'BuiltInUseInsertionEffectHook';
 export const BuiltInUseOperatorId = 'BuiltInUseOperator';
 export const BuiltInUseReducerId = 'BuiltInUseReducer';
 export const BuiltInDispatchId = 'BuiltInDispatch';
+export const BuiltInUseContextHookId = 'BuiltInUseContextHook';
 
 // ShapeRegistry with default definitions for built-ins.
 export const BUILTIN_SHAPES: ShapeRegistry = new Map();