Lookup hook declarations on environment

Changes from calling the global parseHookCall() function to looking up the hook 
declaration on the environment.

compiler/forget/src/HIR/Environment.ts

@@ -1,4 +1,4 @@
-import { IdentifierId, makeIdentifierId } from "./HIR";
+import { Effect, IdentifierId, makeIdentifierId, ValueKind } from "./HIR";
 import { BUILTIN_HOOKS, Hook } from "./Hooks";
 
 const HOOK_PATTERN = /^_?use/;
@@ -35,6 +35,11 @@ export class Environment {
     if (hook !== undefined) {
       return hook;
     }
-    return null;
+    return {
+      kind: "Custom",
+      name,
+      effectKind: Effect.Mutate,
+      valueKind: ValueKind.Mutable,
+    };
   }
 }

compiler/forget/src/HIR/HIR.ts

@@ -49,6 +49,7 @@ export type ReactiveFunction = {
   generator: boolean;
   async: boolean;
   body: ReactiveBlock;
+  env: Environment;
 };
 
 export type ReactiveScopeBlock = {

compiler/forget/src/HIR/Hooks.ts

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import { Effect, Place, ValueKind } from "./HIR";
+import { Effect, ValueKind } from "./HIR";
 
 export const BUILTIN_HOOKS: Map<string, Hook> = new Map([
   [
@@ -62,20 +62,3 @@ export type Hook = {
   effectKind: Effect;
   valueKind: ValueKind;
 };
-
-export function parseHookCall(place: Place): Hook | null {
-  const name = place.identifier.name;
-  if (name === null || !name.match(/^_?use/)) {
-    return null;
-  }
-  const hook = BUILTIN_HOOKS.get(name);
-  if (hook != null) {
-    return hook;
-  }
-  return {
-    kind: "Custom",
-    name: place.identifier.name ?? "",
-    effectKind: Effect.Mutate,
-    valueKind: ValueKind.Mutable,
-  };
-}

compiler/forget/src/Inference/InferReferenceEffects.ts

@@ -7,6 +7,7 @@
 
 import invariant from "invariant";
 import { CompilerError } from "../CompilerError";
+import { Environment } from "../HIR";
 import {
   BasicBlock,
   BlockId,
@@ -19,7 +20,6 @@ import {
   Place,
   ValueKind,
 } from "../HIR/HIR";
-import { parseHookCall } from "../HIR/Hooks";
 import {
   printMixedHIR,
   printPlace,
@@ -151,7 +151,7 @@ export default function inferReferenceEffects(fn: HIRFunction) {
 
       statesByBlock.set(blockId, incomingState);
       const state = incomingState.clone();
-      inferBlock(state, block);
+      inferBlock(fn.env, state, block);
 
       for (const nextBlockId of eachTerminalSuccessor(block.terminal)) {
         queue(nextBlockId, state);
@@ -551,7 +551,12 @@ function mergeValues(a: ValueKind, b: ValueKind): ValueKind {
  * Iterates over the given @param block, defining variables and
  * recording references on the @param state according to JS semantics.
  */
-function inferBlock(state: InferenceState, block: BasicBlock) {
+function inferBlock(
+  env: Environment,
+
+  state: InferenceState,
+  block: BasicBlock
+) {
   for (const phi of block.phis) {
     state.inferPhi(phi);
   }
@@ -583,7 +588,10 @@ function inferBlock(state: InferenceState, block: BasicBlock) {
       case "CallExpression": {
         valueKind = ValueKind.Mutable;
         effectKind = Effect.Mutate;
-        const hook = parseHookCall(instrValue.callee);
+        const hook =
+          instrValue.callee.identifier.name !== null
+            ? env.getHookDeclaration(instrValue.callee.identifier.name)
+            : null;
         if (hook !== null) {
           effectKind = hook.effectKind;
           valueKind = hook.valueKind;

compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts

@@ -48,6 +48,7 @@ export function buildReactiveFunction(fn: HIRFunction): ReactiveFunction {
     generator: fn.generator,
     async: fn.async,
     body,
+    env: fn.env,
   };
 }
 

compiler/forget/src/ReactiveScopes/FlattenScopesWithHooks.ts

@@ -6,13 +6,13 @@
  */
 
 import {
+  Environment,
   InstructionId,
   ReactiveFunction,
   ReactiveScopeBlock,
   ReactiveStatement,
   ReactiveValue,
 } from "../HIR";
-import { parseHookCall } from "../HIR/Hooks";
 import {
   ReactiveFunctionTransform,
   Transformed,
@@ -31,12 +31,19 @@ import {
  * to ensure the hook call does not inadvertently become conditional.
  */
 export function flattenScopesWithHooks(fn: ReactiveFunction): void {
-  visitReactiveFunction(fn, new Transform(), { hasHook: false });
+  visitReactiveFunction(fn, new Transform(fn.env), { 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
@@ -56,8 +63,11 @@ class Transform extends ReactiveFunctionTransform<State> {
     value: ReactiveValue,
     state: State
   ): void {
-    if (value.kind === "CallExpression") {
-      const hook = parseHookCall(value.callee);
+    if (
+      value.kind === "CallExpression" &&
+      value.callee.identifier.name !== null
+    ) {
+      const hook = this.env.getHookDeclaration(value.callee.identifier.name);
       if (hook !== null) {
         state.hasHook = true;
       }

compiler/forget/src/ReactiveScopes/InferReactiveIdentifiers.ts

@@ -6,13 +6,13 @@
  */
 
 import { CompilerError } from "../CompilerError";
+import { Environment } from "../HIR";
 import {
   Effect,
   IdentifierId,
   ReactiveFunction,
   ReactiveInstruction,
 } from "../HIR/HIR";
-import { parseHookCall } from "../HIR/Hooks";
 import { assertExhaustive } from "../Utils/utils";
 import {
   eachReactiveValueOperand,
@@ -21,7 +21,14 @@ import {
 } from "./visitors";
 
 type IdentifierReactivity = Map<IdentifierId, boolean>;
-class Environment extends ReactiveFunctionVisitor<IdentifierReactivity> {
+class Visitor extends ReactiveFunctionVisitor<IdentifierReactivity> {
+  env: Environment;
+
+  constructor(env: Environment) {
+    super();
+    this.env = env;
+  }
+
   override visitInstruction(
     instr: ReactiveInstruction,
     reactivityMap: IdentifierReactivity
@@ -42,13 +49,19 @@ class Environment extends ReactiveFunctionVisitor<IdentifierReactivity> {
         break;
       }
     }
-    if (!hasReactiveInput && instr.value.kind === "CallExpression") {
+    if (
+      !hasReactiveInput &&
+      instr.value.kind === "CallExpression" &&
+      instr.value.callee.identifier.name !== null
+    ) {
       // 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 = parseHookCall(instr.value.callee);
+      const hook = this.env.getHookDeclaration(
+        instr.value.callee.identifier.name
+      );
       if (hook !== null) {
         hasReactiveInput = true;
       }
@@ -130,7 +143,7 @@ class Environment extends ReactiveFunctionVisitor<IdentifierReactivity> {
 export function inferReactiveIdentifiers(
   fn: ReactiveFunction
 ): Set<IdentifierId> {
-  const visitor = new Environment();
+  const visitor = new Visitor(fn.env);
   const reactivityMap: IdentifierReactivity = new Map();
   for (const param of fn.params) {
     reactivityMap.set(param.identifier.id, true);

compiler/forget/src/TypeInference/InferTypes.ts

@@ -1,5 +1,6 @@
 import * as t from "@babel/types";
 import invariant from "invariant";
+import { Environment } from "../HIR";
 import {
   HIRFunction,
   Instruction,
@@ -8,7 +9,6 @@ import {
   TypeId,
   TypeVar,
 } from "../HIR/HIR";
-import { parseHookCall } from "../HIR/Hooks";
 import { eachInstructionOperand } from "../HIR/visitors";
 
 function isPrimitiveBinaryOp(op: t.BinaryExpression["operator"]) {
@@ -83,12 +83,13 @@ function* generate(
     }
 
     for (const instr of block.instructions) {
-      yield* generateInstructionTypes(instr);
+      yield* generateInstructionTypes(func.env, instr);
     }
   }
 }
 
 function* generateInstructionTypes(
+  env: Environment,
   instr: Instruction
 ): Generator<TypeEquation, void, undefined> {
   const { lvalue, value } = instr;
@@ -121,7 +122,10 @@ function* generateInstructionTypes(
     }
 
     case "CallExpression": {
-      const hook = parseHookCall(value.callee);
+      const hook =
+        value.callee.identifier.name !== null
+          ? env.getHookDeclaration(value.callee.identifier.name)
+          : null;
       let type: Type;
       if (hook !== null) {
         type = { kind: "Hook", name: hook.name };