[rfc][env] Hook, object, and function types for Environment globals

Adds `GlobalRegistry`, which holds the names and types of known global objects, 
i.e. 

```js 

type GlobalRegistry = Map<string, PrimitiveType | ObjectType | FunctionType | 
HookType | PolyType>; 

// ... 

globalRegistry.get("NaN"); // {kind: "Primitive"} 

globalRegistry.get("parseInt"); // {kind: "Function", shapeId: "..."} 

globalRegistry.get("Math"); // {kind: "Object", shapeId: "..."} 

``` 

Since we currently do not track module imports and module-level declarations, 
builtin and custom hooks currently also live in GlobalRegistry. 

```js 

globalRegistry.get("useState"); // {kind: "Hook", definition: {...}} 

globalRegistry.get("useFreeze"); // {kind: Hook, definition: {...}} 

``` 

This PR does not allow Forget users to define their own globals. When we add 
this as a configuration, we should not expose `ShapeRegistry` to the user, as a 
user-provided ShapeRegistry may accidentally be not well formed. (i.e. missing 
(1) required shapes (BuiltInArray for [] and BuiltInObject for {}) or (2) some 
recursive shapeIds) 

```js 

export type UserType = UserObject | UserFunction | "Primitive" | "BuiltinObject" 
| ...; 

export type UserObject = { 

kind: "Object", 

properties: Map<string, UserType> 

} 

export type UserFunction = { 

kind: "Function", 

properties: Map<string, UserType>, 

signature: ... 

} 

export type UserGlobals = Map<string, UserType>; 

class Environment { 

constructor(globals: Map<string, UserType>, ...) { 

// ... 

addUserDefinedGlobals(this.#globals, this.#shapes); 

```

compiler/forget/src/HIR/Environment.ts

@@ -1,6 +1,6 @@
 import invariant from "invariant";
 import { log } from "../Utils/logger";
-import { DEFAULT_GLOBALS, Global } from "./Globals";
+import { DEFAULT_GLOBALS, Global, GlobalRegistry } from "./Globals";
 import {
   BuiltInType,
   Effect,
@@ -10,7 +10,7 @@ import {
   ObjectType,
   ValueKind,
 } from "./HIR";
-import { BUILTIN_HOOKS, Hook } from "./Hooks";
+import { Hook } from "./Hooks";
 import {
   BUILTIN_SHAPES,
   FunctionSignature,
@@ -19,21 +19,41 @@ import {
 
 const HOOK_PATTERN = /^_?use/;
 
+// TODO(mofeiZ): User defined global types (with corresponding shapes).
+// User defined global types should have inline ObjectShapes instead of directly
+// using ObjectShapes.ShapeRegistry, as a user-provided ShapeRegistry may be
+// accidentally be not well formed.
+// i.e.
+//   missing required shapes (BuiltInArray for [] and BuiltInObject for {})
+//   missing some recursive Object / Function shapeIds
 export type EnvironmentConfig = Partial<{
   customHooks: Map<string, Hook>;
   memoizeJsxElements: boolean;
 }>;
 
 export class Environment {
-  #customHooks: Map<string, Hook>;
-  #globals: Set<string>;
+  #globals: GlobalRegistry;
   #shapes: ShapeRegistry;
   #nextIdentifer: number = 0;
 
   constructor(config: EnvironmentConfig | null) {
-    this.#customHooks = config?.customHooks ?? new Map();
     this.#shapes = BUILTIN_SHAPES;
-    this.#globals = DEFAULT_GLOBALS;
+
+    if (config?.customHooks) {
+      this.#globals = new Map(DEFAULT_GLOBALS);
+      for (const [hookName, hook] of config.customHooks) {
+        invariant(
+          !this.#globals.has(hookName),
+          `[Globals] Found existing definition in global registry for custom hook ${hookName}`
+        );
+        this.#globals.set(hookName, {
+          kind: "Hook",
+          definition: hook,
+        });
+      }
+    } else {
+      this.#globals = DEFAULT_GLOBALS;
+    }
   }
 
   get nextIdentifierId(): IdentifierId {
@@ -41,26 +61,24 @@ export class Environment {
   }
 
   getGlobalDeclaration(name: string): Global | null {
-    if (!this.#globals.has(name)) {
-      log(() => `Undefined global '${name}'`);
+    let resolvedGlobal: Global | null = this.#globals.get(name) ?? null;
+    if (resolvedGlobal === null) {
+      // Hack, since we don't track module level declarations and imports
+      if (name.match(HOOK_PATTERN)) {
+        return {
+          kind: "Hook",
+          definition: {
+            kind: "Custom",
+            name,
+            effectKind: Effect.Mutate,
+            valueKind: ValueKind.Mutable,
+          },
+        };
+      } else {
+        log(() => `Undefined global '${name}'`);
+      }
     }
-    return { name };
-  }
-
-  getHookDeclaration(name: string): Hook | null {
-    if (!name.match(HOOK_PATTERN)) {
-      return null;
-    }
-    const hook = BUILTIN_HOOKS.get(name) ?? this.#customHooks.get(name);
-    if (hook !== undefined) {
-      return hook;
-    }
-    return {
-      kind: "Custom",
-      name,
-      effectKind: Effect.Mutate,
-      valueKind: ValueKind.Mutable,
-    };
+    return resolvedGlobal;
   }
 
   getPropertyType(

compiler/forget/src/HIR/Globals.ts

@@ -5,7 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-export const DEFAULT_GLOBALS: Set<string> = new Set([
+import { Effect, ValueKind } from "./HIR";
+import { Hook } from "./Hooks";
+import { BuiltInType, HookType, PolyType } from "./Types";
+
+// Hack until we add ObjectShapes for all globals
+const UNTYPED_GLOBALS: Set<string> = new Set([
   "String",
   "Object",
   "Function",
@@ -50,6 +55,62 @@ export const DEFAULT_GLOBALS: Set<string> = new Set([
   "decodeURIComponent",
 ]);
 
-export type Global = {
-  name: string;
-};
+const BUILTIN_HOOKS: Array<[string, Hook]> = [
+  [
+    "useState",
+    {
+      kind: "State",
+      name: "useState",
+      effectKind: Effect.Freeze,
+      valueKind: ValueKind.Frozen,
+    },
+  ],
+  [
+    "useRef",
+    {
+      kind: "Ref",
+      name: "useRef",
+      effectKind: Effect.Capture,
+      valueKind: ValueKind.Mutable,
+    },
+  ],
+  [
+    "useMemo",
+    {
+      kind: "Memo",
+      name: "useMemo",
+      effectKind: Effect.Freeze,
+      valueKind: ValueKind.Frozen,
+    },
+  ],
+  [
+    "useCallback",
+    {
+      kind: "Memo",
+      name: "useCallback",
+      effectKind: Effect.Freeze,
+      valueKind: ValueKind.Frozen,
+    },
+  ],
+];
+
+export type Global = BuiltInType | HookType | PolyType;
+export type GlobalRegistry = Map<string, Global>;
+export const DEFAULT_GLOBALS: GlobalRegistry = new Map(
+  BUILTIN_HOOKS.map(([hookName, hook]) => {
+    return [
+      hookName,
+      {
+        kind: "Hook",
+        definition: hook,
+      },
+    ];
+  })
+);
+
+// Hack until we add ObjectShapes for all globals
+for (const name of UNTYPED_GLOBALS) {
+  DEFAULT_GLOBALS.set(name, {
+    kind: "Poly",
+  });
+}

compiler/forget/src/HIR/HIRBuilder.ts

@@ -143,8 +143,23 @@ export default class HIRBuilder {
     };
   }
 
-  resolveGlobal(path: NodePath<t.Identifier | t.JSXIdentifier>): Global | null {
-    return this.#env.getGlobalDeclaration(path.node.name);
+  resolveGlobal(
+    path: NodePath<t.Identifier | t.JSXIdentifier>
+  ): (Global & { name: string }) | null {
+    const name = path.node.name;
+    const resolvedGlobal = this.#env.getGlobalDeclaration(name);
+    if (resolvedGlobal) {
+      return {
+        ...resolvedGlobal,
+        name,
+      };
+    } else {
+      // if env records no global with the given name, load it as an unknown type
+      return {
+        kind: "Poly",
+        name,
+      };
+    }
   }
 
   /**

compiler/forget/src/HIR/Hooks.ts

@@ -7,45 +7,6 @@
 
 import { Effect, ValueKind } from "./HIR";
 
-export const BUILTIN_HOOKS: Map<string, Hook> = new Map([
-  [
-    "useState",
-    {
-      kind: "State",
-      name: "useState",
-      effectKind: Effect.Freeze,
-      valueKind: ValueKind.Frozen,
-    },
-  ],
-  [
-    "useRef",
-    {
-      kind: "Ref",
-      name: "useRef",
-      effectKind: Effect.Capture,
-      valueKind: ValueKind.Mutable,
-    },
-  ],
-  [
-    "useMemo",
-    {
-      kind: "Memo",
-      name: "useMemo",
-      effectKind: Effect.Freeze,
-      valueKind: ValueKind.Frozen,
-    },
-  ],
-  [
-    "useCallback",
-    {
-      kind: "Memo",
-      name: "useCallback",
-      effectKind: Effect.Freeze,
-      valueKind: ValueKind.Frozen,
-    },
-  ],
-]);
-
 export type HookKind = "State" | "Ref" | "Custom" | "Memo";
 export type Hook = {
   kind: HookKind;

compiler/forget/src/TypeInference/InferTypes.ts

@@ -135,10 +135,12 @@ function* generateInstructionTypes(
     }
 
     case "LoadGlobal": {
-      const hook = env.getHookDeclaration(value.name);
-      if (hook !== null) {
-        const type: Type = { kind: "Hook", definition: hook };
-        yield equation(left, type);
+      const globalType = env.getGlobalDeclaration(value.name);
+      if (globalType) {
+        if (globalType.kind === "Hook") {
+          yield equation(left, globalType);
+        }
+        // TODO(mofeiZ): add type inference for other globals
       }
       break;
     }