Construct LoadGlobal; consume hook info from types

Updates BuildHIR to produce LoadGlobal instructions for references to globals. 
Note that this breaks our previous strategy of finding hook calls: that relied 
on looking at the callee of a CallExpression and checking its name, which relied 
on the callee not being lowered to a temporary. By lowering the name (eg 
`useState`) to a temporary first, we now no longer see the name at the callsite. 

Thankfully @gsathya solved this for us already by teaching type inference about 
hooks, and more generally implementing type inference. I updated this so that we 
infer the type of a LoadGlobal if the name is a hook: the type inference picks 
this up and propagates the type forward correctly. So now, all places that 
needed to check for a hook can just look at the type and everything works. 

This is much more robust than before - you can now reassign a hook to a local 
variable and we'll still detect that when you call it, you're calling a hook.

compiler/forget/src/HIR/BuildHIR.ts

@@ -79,6 +79,14 @@ export function lower(
   func.get("params").forEach((param) => {
     if (param.isIdentifier()) {
       const identifier = builder.resolveIdentifier(param);
+      if (identifier === null) {
+        builder.errors.push({
+          reason: `(BuildHIR::lower) Could not find binding for param '${param.node.name}'`,
+          severity: ErrorSeverity.Invariant,
+          nodePath: param,
+        });
+        return;
+      }
       const place: Place = {
         kind: "Identifier",
         identifier,
@@ -1591,14 +1599,7 @@ function lowerJsxElementName(
   }
   const tag: string = exprPath.node.name;
   if (tag.match(/^[A-Z]/)) {
-    const identifier = builder.resolveIdentifier(exprPath);
-    const place: Place = {
-      kind: "Identifier",
-      identifier: identifier,
-      effect: Effect.Unknown,
-      loc: exprLoc,
-    };
-    return place;
+    return lowerIdentifier(builder, exprPath);
   } else {
     const place: Place = buildTemporaryPlace(builder, exprLoc);
     builder.push({
@@ -1728,11 +1729,34 @@ function lowerExpressionToVoid(
 
 function lowerIdentifier(
   builder: HIRBuilder,
-  exprPath: NodePath<t.Identifier>
+  exprPath: NodePath<t.Identifier | t.JSXIdentifier>
 ): Place {
   const exprNode = exprPath.node;
   const exprLoc = exprNode.loc ?? GeneratedSource;
   const identifier = builder.resolveIdentifier(exprPath);
+  if (identifier === null) {
+    const place = buildTemporaryPlace(
+      builder,
+      exprPath.node.loc ?? GeneratedSource
+    );
+    const global = builder.resolveGlobal(exprPath);
+    let value: InstructionValue;
+    if (global !== null) {
+      value = { kind: "LoadGlobal", name: global.name, loc: place.loc };
+    } else {
+      value = { kind: "UnsupportedNode", node: exprPath.node, loc: place.loc };
+    }
+    builder.push({
+      id: makeInstructionId(0),
+      value,
+      loc: place.loc,
+      lvalue: {
+        place: { ...place },
+        kind: InstructionKind.Const,
+      },
+    });
+    return place;
+  }
   const place: Place = {
     kind: "Identifier",
     identifier: identifier,

compiler/forget/src/HIR/Globals.ts

@@ -13,14 +13,17 @@ const GLOBALS: Map<string, t.Identifier> = new Map([
   ["Math", t.identifier("Math")],
 ]);
 
+export type Global = {
+  name: string;
+};
+
 // TODO: This will work as a stopgap but it isn't really correct. We need proper handling of globals
 // and module-scoped variables, which means understanding module constants and imports.
-export function getOrAddGlobal(identifierName: string): t.Identifier {
+export function getGlobalDeclaration(identifierName: string): Global | null {
   const ident = GLOBALS.get(identifierName);
   if (ident != null) {
     return ident;
   }
-  const newIdent = t.identifier(identifierName);
-  GLOBALS.set(identifierName, newIdent);
-  return newIdent;
+  // TODO: return null if not explicitly configured by the user
+  return { name: identifierName };
 }

compiler/forget/src/HIR/HIR.ts

@@ -8,6 +8,7 @@
 import * as t from "@babel/types";
 import invariant from "invariant";
 import { Environment } from "./Environment";
+import { Hook } from "./Hooks";
 
 // *******************************************************************************************
 // *******************************************************************************************
@@ -641,7 +642,7 @@ export type FunctionType = {
 };
 export type HookType = {
   kind: "Hook";
-  name: string;
+  definition: Hook;
 };
 export type ObjectType = { kind: "Object" };
 export type TypeVar = {
@@ -720,7 +721,9 @@ function funcTypeEquals(tA: Type, tB: Type): boolean {
 }
 
 function hookTypeEquals(tA: Type, tB: Type): boolean {
-  return tA.kind === "Hook" && tB.kind === "Hook" && tA.name === tB.name;
+  return (
+    tA.kind === "Hook" && tB.kind === "Hook" && tA.definition === tB.definition
+  );
 }
 
 function phiTypeEquals(tA: Type, tB: Type): boolean {

compiler/forget/src/HIR/HIRBuilder.ts

@@ -12,7 +12,7 @@ import { CompilerError } from "../CompilerError";
 import { logHIR } from "../Utils/logger";
 import { assertExhaustive } from "../Utils/utils";
 import { Environment } from "./Environment";
-import { getOrAddGlobal } from "./Globals";
+import { getGlobalDeclaration, Global } from "./Globals";
 import {
   BasicBlock,
   BlockId,
@@ -136,6 +136,10 @@ export default class HIRBuilder {
     };
   }
 
+  resolveGlobal(path: NodePath<t.Identifier | t.JSXIdentifier>): Global | null {
+    return getGlobalDeclaration(path.node.name);
+  }
+
   /**
    * Maps an Identifier (or JSX identifier) Babel node to an internal `Identifier`
    * which represents the variable being referenced, according to the JS scoping rules.
@@ -168,15 +172,16 @@ export default class HIRBuilder {
    */
   resolveIdentifier(
     path: NodePath<t.Identifier | t.JSXIdentifier>
-  ): Identifier {
+  ): Identifier | null {
     const originalName = path.node.name;
-    const node =
-      path.scope.getBindingIdentifier(originalName) ??
-      getOrAddGlobal(originalName);
+    const node = path.scope.getBindingIdentifier(originalName);
+    if (node == null) {
+      return null;
+    }
     return this.resolveBinding(node);
   }
 
-  resolveBinding(node: t.Identifier) {
+  resolveBinding(node: t.Identifier): Identifier {
     const originalName = node.name;
     let name = originalName;
     let index = 0;

compiler/forget/src/Inference/DropMemoCalls.ts

@@ -1,11 +1,4 @@
-import invariant from "invariant";
-import {
-  Effect,
-  HIRFunction,
-  HookType,
-  InstructionValue,
-  isHookType,
-} from "../HIR";
+import { Effect, HIRFunction, HookType, isHookType } from "../HIR";
 
 export default function (func: HIRFunction) {
   for (const [_, block] of func.body.blocks) {
@@ -13,7 +6,8 @@ export default function (func: HIRFunction) {
       switch (instr.value.kind) {
         case "CallExpression": {
           if (isHookType(instr.value.callee.identifier)) {
-            const name = (instr.value.callee.identifier.type as HookType).name;
+            const name = (instr.value.callee.identifier.type as HookType)
+              .definition.name;
             if (name === "useMemo") {
               const [fn] = instr.value.args;
 

compiler/forget/src/Inference/InferReferenceEffects.ts

@@ -589,8 +589,8 @@ function inferBlock(
         valueKind = ValueKind.Mutable;
         effectKind = Effect.Mutate;
         const hook =
-          instrValue.callee.identifier.name !== null
-            ? env.getHookDeclaration(instrValue.callee.identifier.name)
+          instrValue.callee.identifier.type.kind === "Hook"
+            ? instrValue.callee.identifier.type.definition
             : null;
         if (hook !== null) {
           effectKind = hook.effectKind;

compiler/forget/src/ReactiveScopes/FlattenScopesWithHooks.ts

@@ -6,8 +6,8 @@
  */
 
 import {
-  Environment,
   InstructionId,
+  isHookType,
   ReactiveFunction,
   ReactiveScopeBlock,
   ReactiveStatement,
@@ -31,19 +31,12 @@ import {
  * to ensure the hook call does not inadvertently become conditional.
  */
 export function flattenScopesWithHooks(fn: ReactiveFunction): void {
-  visitReactiveFunction(fn, new Transform(fn.env), { hasHook: false });
+  visitReactiveFunction(fn, new Transform(), { hasHook: false });
 }
 
 type State = { hasHook: boolean };
 
 class Transform extends ReactiveFunctionTransform<State> {
-  env: Environment;
-
-  constructor(env: Environment) {
-    super();
-    this.env = env;
-  }
-
   override transformScope(
     scope: ReactiveScopeBlock,
     outerState: State
@@ -65,12 +58,9 @@ class Transform extends ReactiveFunctionTransform<State> {
   ): void {
     if (
       value.kind === "CallExpression" &&
-      value.callee.identifier.name !== null
+      isHookType(value.callee.identifier)
     ) {
-      const hook = this.env.getHookDeclaration(value.callee.identifier.name);
-      if (hook !== null) {
-        state.hasHook = true;
-      }
+      state.hasHook = true;
     }
   }
 }

compiler/forget/src/ReactiveScopes/InferReactiveIdentifiers.ts

@@ -6,10 +6,10 @@
  */
 
 import { CompilerError } from "../CompilerError";
-import { Environment } from "../HIR";
 import {
   Effect,
   IdentifierId,
+  isHookType,
   ReactiveFunction,
   ReactiveInstruction,
 } from "../HIR/HIR";
@@ -22,13 +22,6 @@ import {
 
 type IdentifierReactivity = Map<IdentifierId, boolean>;
 class Visitor extends ReactiveFunctionVisitor<IdentifierReactivity> {
-  env: Environment;
-
-  constructor(env: Environment) {
-    super();
-    this.env = env;
-  }
-
   override visitInstruction(
     instr: ReactiveInstruction,
     reactivityMap: IdentifierReactivity
@@ -52,19 +45,14 @@ class Visitor extends ReactiveFunctionVisitor<IdentifierReactivity> {
     if (
       !hasReactiveInput &&
       instr.value.kind === "CallExpression" &&
-      instr.value.callee.identifier.name !== null
+      isHookType(instr.value.callee.identifier)
     ) {
       // Hooks cannot be memoized. Even if they do not accept any reactive inputs,
       // they are not guaranteed to memoize their return value, and their result
       // must be assumed to be reactive.
       // TODO: use types or an opt-in registry of custom hook information to
       // allow treating safe hooks as non-reactive.
-      const hook = this.env.getHookDeclaration(
-        instr.value.callee.identifier.name
-      );
-      if (hook !== null) {
-        hasReactiveInput = true;
-      }
+      hasReactiveInput = true;
     }
     reactivityMap.set(lval.place.identifier.id, hasReactiveInput);
 
@@ -143,7 +131,7 @@ class Visitor extends ReactiveFunctionVisitor<IdentifierReactivity> {
 export function inferReactiveIdentifiers(
   fn: ReactiveFunction
 ): Set<IdentifierId> {
-  const visitor = new Visitor(fn.env);
+  const visitor = new Visitor();
   const reactivityMap: IdentifierReactivity = new Map();
   for (const param of fn.params) {
     reactivityMap.set(param.identifier.id, true);

compiler/forget/src/TypeInference/InferTypes.ts

@@ -121,6 +121,15 @@ function* generateInstructionTypes(
       break;
     }
 
+    case "LoadGlobal": {
+      const hook = env.getHookDeclaration(value.name);
+      if (hook !== null) {
+        const type: Type = { kind: "Hook", definition: hook };
+        yield equation(left, type);
+      }
+      break;
+    }
+
     case "CallExpression": {
       const hook =
         value.callee.identifier.name !== null
@@ -128,7 +137,7 @@ function* generateInstructionTypes(
           : null;
       let type: Type;
       if (hook !== null) {
-        type = { kind: "Hook", name: hook.name };
+        type = { kind: "Hook", definition: hook };
       } else {
         type = { kind: "Function" };
       }

compiler/forget/src/__tests__/fixtures/hir/useRef-rename-mutable.expect.md

@@ -0,0 +1,23 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const x = useRef;
+  const ref = x(null);
+  return ref.current;
+}
+
+```
+
+## Code
+
+```javascript
+function Component(props) {
+  const x = useRef;
+  const ref = x(null);
+  return ref.current;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/useRef-rename-mutable.js

@@ -0,0 +1,5 @@
+function Component(props) {
+  const x = useRef;
+  const ref = x(null);
+  return ref.current;
+}