From 37d044763b45d43fc8e0599cd3926f82125a0360 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Wed, 24 May 2023 09:54:10 -0700
Subject: [PATCH] Define types for useState/useRef

Defines 4 new types:

* Return type of `useState()`, which has properties "0" and "1" to allow us to
infer the types when destructuring

* Type of useState() set state function

* Return type of `useRef()` so we know what is a ref

* Type of ref.current, so we know what is a ref *value*

Example:

<img width="1670" alt="Screenshot 2023-05-24 at 9 59 37 AM"
src="https://github.com/facebook/react-forget/assets/6425824/3ee7d04a-fda3-4b7b-89b7-d205d9a6fd0d">
---
 compiler/forget/src/HIR/Globals.ts     |  6 +++--
 compiler/forget/src/HIR/ObjectShape.ts | 34 ++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/compiler/forget/src/HIR/Globals.ts b/compiler/forget/src/HIR/Globals.ts
index 16b39c4769..5a38510e7e 100644
--- a/compiler/forget/src/HIR/Globals.ts
+++ b/compiler/forget/src/HIR/Globals.ts
@@ -9,6 +9,8 @@ import { Effect, ValueKind } from "./HIR";
 import {
   BUILTIN_SHAPES,
   BuiltInArrayId,
+  BuiltInUseRefId,
+  BuiltInUseStateId,
   ShapeRegistry,
   addFunction,
   addHook,
@@ -248,7 +250,7 @@ const BUILTIN_HOOKS: Array<[string, FunctionType]> = [
     addHook(DEFAULT_SHAPES, [], {
       positionalParams: [],
       restParam: Effect.Freeze,
-      returnType: { kind: "Poly" },
+      returnType: { kind: "Object", shapeId: BuiltInUseStateId },
       calleeEffect: Effect.Read,
       hookKind: "useState",
       returnValueKind: ValueKind.Frozen,
@@ -259,7 +261,7 @@ const BUILTIN_HOOKS: Array<[string, FunctionType]> = [
     addHook(DEFAULT_SHAPES, [], {
       positionalParams: [],
       restParam: Effect.Capture,
-      returnType: { kind: "Poly" },
+      returnType: { kind: "Object", shapeId: BuiltInUseRefId },
       calleeEffect: Effect.Read,
       hookKind: "useRef",
       returnValueKind: ValueKind.Mutable,
diff --git a/compiler/forget/src/HIR/ObjectShape.ts b/compiler/forget/src/HIR/ObjectShape.ts
index 0ba10a973a..a33e808c69 100644
--- a/compiler/forget/src/HIR/ObjectShape.ts
+++ b/compiler/forget/src/HIR/ObjectShape.ts
@@ -40,9 +40,10 @@ function createAnonId(): string {
 export function addFunction(
   registry: ShapeRegistry,
   properties: Iterable<[string, BuiltInType | PolyType]>,
-  fn: Omit<FunctionSignature, "hookKind">
+  fn: Omit<FunctionSignature, "hookKind">,
+  id: string | null = null
 ): FunctionType {
-  const shapeId = createAnonId();
+  const shapeId = id ?? createAnonId();
   addShape(registry, shapeId, properties, {
     ...fn,
     hookKind: null,
@@ -157,6 +158,10 @@ export type ObjectShape = {
 export type ShapeRegistry = Map<string, ObjectShape>;
 export const BuiltInArrayId = "BuiltInArray";
 export const BuiltInObjectId = "BuiltInObject";
+export const BuiltInUseStateId = "BuiltInUseState";
+export const BuiltInSetStateId = "BuiltInSetState";
+export const BuiltInUseRefId = "BuiltInUseRefId";
+export const BuiltInRefValueId = "BuiltInRefValue";
 
 /**
  * ShapeRegistry with default definitions for built-ins.
@@ -218,6 +223,31 @@ addObject(BUILTIN_SHAPES, BuiltInObjectId, [
   // hasOwnProperty, isPrototypeOf, propertyIsEnumerable, toLocaleString, valueOf
 ]);
 
+addObject(BUILTIN_SHAPES, BuiltInUseStateId, [
+  ["0", { kind: "Poly" }],
+  [
+    "1",
+    addFunction(
+      BUILTIN_SHAPES,
+      [],
+      {
+        positionalParams: [],
+        restParam: Effect.Freeze,
+        returnType: PRIMITIVE_TYPE,
+        calleeEffect: Effect.Read,
+        returnValueKind: ValueKind.Immutable,
+      },
+      BuiltInSetStateId
+    ),
+  ],
+]);
+
+addObject(BUILTIN_SHAPES, BuiltInUseRefId, [
+  ["current", { kind: "Object", shapeId: BuiltInRefValueId }],
+]);
+
+addObject(BUILTIN_SHAPES, BuiltInRefValueId, []);
+
 export const DefaultMutatingHook = addHook(BUILTIN_SHAPES, [], {
   positionalParams: [],
   restParam: Effect.Mutate,