{"commit":{"authorEmail":"mofei.zhang@utexas.edu","authorName":"Mofei Zhang","revision":"b76ee1ab32201143b73d8792ad1fd6d1e1207bc3","authoredAt":"2023-05-23T13:55:10-04:00","subject":"[types] Consolidate Hook and Function types: Hooks are functions","shortRevision":"b76ee1ab32"},"parents":["47e9a2b8a5d2161324edfbd2f805227e88fb4313"],"body":"","files":[{"status":"M","path":"compiler\/forget\/packages\/playground\/components\/Editor\/index.tsx"},{"status":"M","path":"compiler\/forget\/packages\/snap\/src\/compiler-worker.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/Environment.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/Globals.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/HIR.ts"},{"status":"D","path":"compiler\/forget\/src\/HIR\/Hooks.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/ObjectShape.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/Types.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/ValidateHooksUsage.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/ValidateUnconditionalHooks.ts"},{"status":"M","path":"compiler\/forget\/src\/HIR\/index.ts"},{"status":"M","path":"compiler\/forget\/src\/Inference\/DropMemoCalls.ts"},{"status":"M","path":"compiler\/forget\/src\/Inference\/InferReferenceEffects.ts"},{"status":"M","path":"compiler\/forget\/src\/ReactiveScopes\/FlattenScopesWithHooks.ts"},{"status":"M","path":"compiler\/forget\/src\/ReactiveScopes\/InferReactiveIdentifiers.ts"},{"status":"M","path":"compiler\/forget\/src\/ReactiveScopes\/PruneNonEscapingScopes.ts"}],"patch":"diff --git a\/compiler\/forget\/packages\/playground\/components\/Editor\/index.tsx b\/compiler\/forget\/packages\/playground\/components\/Editor\/index.tsx\nindex 1d19de18f3..5714bd8720 100644\n--- a\/compiler\/forget\/packages\/playground\/components\/Editor\/index.tsx\n+++ b\/compiler\/forget\/packages\/playground\/components\/Editor\/index.tsx\n@@ -74,8 +74,6 @@ const COMMON_HOOKS: Array<[string, Hook]> = [\n [\n \"useFragment\",\n {\n- name: \"useFragment\",\n- kind: \"Custom\",\n valueKind: ValueKind.Frozen,\n effectKind: Effect.Freeze,\n },\n@@ -83,8 +81,6 @@ const COMMON_HOOKS: Array<[string, Hook]> = [\n [\n \"usePaginationFragment\",\n {\n- name: \"usePaginationFragment\",\n- kind: \"Custom\",\n valueKind: ValueKind.Frozen,\n effectKind: Effect.Freeze,\n },\n@@ -92,8 +88,6 @@ const COMMON_HOOKS: Array<[string, Hook]> = [\n [\n \"useRefetchableFragment\",\n {\n- name: \"useRefetchableFragment\",\n- kind: \"Custom\",\n valueKind: ValueKind.Frozen,\n effectKind: Effect.Freeze,\n },\n@@ -101,8 +95,6 @@ const COMMON_HOOKS: Array<[string, Hook]> = [\n [\n \"useLazyLoadQuery\",\n {\n- name: \"useLazyLoadQuery\",\n- kind: \"Custom\",\n valueKind: ValueKind.Frozen,\n effectKind: Effect.Freeze,\n },\n@@ -110,8 +102,6 @@ const COMMON_HOOKS: Array<[string, Hook]> = [\n [\n \"usePreloadedQuery\",\n {\n- name: \"usePreloadedQuery\",\n- kind: \"Custom\",\n valueKind: ValueKind.Frozen,\n effectKind: Effect.Freeze,\n },\ndiff --git a\/compiler\/forget\/packages\/snap\/src\/compiler-worker.ts b\/compiler\/forget\/packages\/snap\/src\/compiler-worker.ts\nindex 7f3bb45bf8..b70f784693 100644\n--- a\/compiler\/forget\/packages\/snap\/src\/compiler-worker.ts\n+++ b\/compiler\/forget\/packages\/snap\/src\/compiler-worker.ts\n@@ -134,8 +134,6 @@ export async function compile(\n [\n \"useFreeze\",\n {\n- name: \"useFreeze\",\n- kind: \"Custom\",\n valueKind: \"frozen\",\n effectKind: \"freeze\",\n },\ndiff --git a\/compiler\/forget\/src\/HIR\/Environment.ts b\/compiler\/forget\/src\/HIR\/Environment.ts\nindex e60e9e4b66..1a78bde8d0 100644\n--- a\/compiler\/forget\/src\/HIR\/Environment.ts\n+++ b\/compiler\/forget\/src\/HIR\/Environment.ts\n@@ -17,15 +17,27 @@ import {\n import {\n BlockId,\n BuiltInType,\n+ Effect,\n FunctionType,\n IdentifierId,\n ObjectType,\n PolyType,\n+ ValueKind,\n makeBlockId,\n makeIdentifierId,\n } from \".\/HIR\";\n-import { Hook } from \".\/Hooks\";\n-import { FunctionSignature, ShapeRegistry } from \".\/ObjectShape\";\n+import {\n+ DefaultMutatingHook,\n+ DefaultNonmutatingHook,\n+ FunctionSignature,\n+ ShapeRegistry,\n+ addHook,\n+} from \".\/ObjectShape\";\n+\n+export type Hook = {\n+ effectKind: Effect;\n+ valueKind: ValueKind;\n+};\n \n \/\/ TODO(mofeiZ): User defined global types (with corresponding shapes).\n \/\/ User defined global types should have inline ObjectShapes instead of directly\n@@ -107,7 +119,7 @@ export class Environment {\n config: EnvironmentConfig | null,\n contextIdentifiers: Set\n ) {\n- this.#shapes = DEFAULT_SHAPES;\n+ this.#shapes = new Map(DEFAULT_SHAPES);\n \n if (config?.customHooks) {\n this.#globals = new Map(DEFAULT_GLOBALS);\n@@ -116,10 +128,17 @@ export class Environment {\n !this.#globals.has(hookName),\n `[Globals] Found existing definition in global registry for custom hook ${hookName}`\n );\n- this.#globals.set(hookName, {\n- kind: \"Hook\",\n- definition: hook,\n- });\n+ this.#globals.set(\n+ hookName,\n+ addHook(this.#shapes, [], {\n+ positionalParams: [],\n+ restParam: hook.effectKind,\n+ returnType: { kind: \"Poly\" },\n+ returnValueKind: hook.valueKind,\n+ calleeEffect: Effect.Read,\n+ hookKind: \"Custom\",\n+ })\n+ );\n }\n } else {\n this.#globals = DEFAULT_GLOBALS;\n@@ -150,10 +169,11 @@ export class Environment {\n if (resolvedGlobal === null) {\n \/\/ Hack, since we don't track module level declarations and imports\n if (isHookName(name)) {\n- return {\n- kind: \"Hook\",\n- definition: null,\n- };\n+ if (this.enableAssumeHooksFollowRulesOfReact) {\n+ return DefaultNonmutatingHook;\n+ } else {\n+ return DefaultMutatingHook;\n+ }\n } else {\n log(() => `Undefined global '${name}'`);\n }\ndiff --git a\/compiler\/forget\/src\/HIR\/Globals.ts b\/compiler\/forget\/src\/HIR\/Globals.ts\nindex 12e17b0f56..16b39c4769 100644\n--- a\/compiler\/forget\/src\/HIR\/Globals.ts\n+++ b\/compiler\/forget\/src\/HIR\/Globals.ts\n@@ -6,15 +6,15 @@\n *\/\n \n import { Effect, ValueKind } from \".\/HIR\";\n-import { Hook } from \".\/Hooks\";\n import {\n BUILTIN_SHAPES,\n BuiltInArrayId,\n ShapeRegistry,\n addFunction,\n+ addHook,\n addObject,\n } from \".\/ObjectShape\";\n-import { BuiltInType, HookType, PolyType } from \".\/Types\";\n+import { BuiltInType, FunctionType, PolyType } from \".\/Types\";\n \n \/**\n * This file exports types and defaults for JavaScript global objects.\n@@ -229,85 +229,91 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [\n \/\/ TODO: rest of Global objects\n ];\n \n-const BUILTIN_HOOKS: Array<[string, Hook]> = [\n+\/\/ TODO(mofeiZ): We currently only store rest param effects for hooks\n+\/\/ until FeatureFlag `enableTreatHooksAsFunctions` is removed\n+const BUILTIN_HOOKS: Array<[string, FunctionType]> = [\n [\n \"useContext\",\n- {\n- kind: \"State\",\n- name: \"useContext\",\n- effectKind: Effect.Read,\n- valueKind: ValueKind.Mutable,\n- },\n+ addHook(DEFAULT_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Read,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"useContext\",\n+ returnValueKind: ValueKind.Mutable,\n+ }),\n ],\n [\n \"useState\",\n- {\n- kind: \"State\",\n- name: \"useState\",\n- effectKind: Effect.Freeze,\n- valueKind: ValueKind.Frozen,\n- },\n+ addHook(DEFAULT_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Freeze,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"useState\",\n+ returnValueKind: ValueKind.Frozen,\n+ }),\n ],\n [\n \"useRef\",\n- {\n- kind: \"Ref\",\n- name: \"useRef\",\n- effectKind: Effect.Capture,\n- valueKind: ValueKind.Mutable,\n- },\n+ addHook(DEFAULT_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Capture,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"useRef\",\n+ returnValueKind: ValueKind.Mutable,\n+ }),\n ],\n [\n \"useMemo\",\n- {\n- kind: \"Memo\",\n- name: \"useMemo\",\n- effectKind: Effect.Freeze,\n- valueKind: ValueKind.Frozen,\n- },\n+ addHook(DEFAULT_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Freeze,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"useMemo\",\n+ returnValueKind: ValueKind.Frozen,\n+ }),\n ],\n [\n \"useCallback\",\n- {\n- kind: \"Memo\",\n- name: \"useCallback\",\n- effectKind: Effect.Freeze,\n- valueKind: ValueKind.Frozen,\n- },\n+ addHook(DEFAULT_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Freeze,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"useCallback\",\n+ returnValueKind: ValueKind.Frozen,\n+ }),\n ],\n [\n \"useEffect\",\n- {\n- kind: \"Memo\",\n- name: \"useEffect\",\n- effectKind: Effect.Freeze,\n- valueKind: ValueKind.Frozen,\n- },\n+ addHook(DEFAULT_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Freeze,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"useEffect\",\n+ returnValueKind: ValueKind.Frozen,\n+ }),\n ],\n [\n \"useLayoutEffect\",\n- {\n- kind: \"Memo\",\n- name: \"useLayoutEffect\",\n- effectKind: Effect.Freeze,\n- valueKind: ValueKind.Frozen,\n- },\n+ addHook(DEFAULT_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Freeze,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"useLayoutEffect\",\n+ returnValueKind: ValueKind.Frozen,\n+ }),\n ],\n ];\n \n-export type Global = BuiltInType | HookType | PolyType;\n+export type Global = BuiltInType | PolyType;\n export type GlobalRegistry = Map;\n-export const DEFAULT_GLOBALS: GlobalRegistry = new Map(\n- BUILTIN_HOOKS.map(([hookName, hook]) => {\n- return [\n- hookName,\n- {\n- kind: \"Hook\",\n- definition: hook,\n- },\n- ];\n- })\n-);\n+export const DEFAULT_GLOBALS: GlobalRegistry = new Map(BUILTIN_HOOKS);\n \n \/\/ Hack until we add ObjectShapes for all globals\n for (const name of UNTYPED_GLOBALS) {\ndiff --git a\/compiler\/forget\/src\/HIR\/HIR.ts b\/compiler\/forget\/src\/HIR\/HIR.ts\nindex 59a3988932..ec23776565 100644\n--- a\/compiler\/forget\/src\/HIR\/HIR.ts\n+++ b\/compiler\/forget\/src\/HIR\/HIR.ts\n@@ -8,6 +8,7 @@\n import * as t from \"@babel\/types\";\n import invariant from \"invariant\";\n import { Environment } from \".\/Environment\";\n+import { HookKind } from \".\/ObjectShape\";\n import { Type } from \".\/Types\";\n \n \/\/ *******************************************************************************************\n@@ -951,8 +952,13 @@ export function isPrimitiveType(id: Identifier): boolean {\n return id.type.kind === \"Primitive\";\n }\n \n-export function isHookType(id: Identifier): boolean {\n- return id.type.kind === \"Hook\";\n+export function getHookKind(env: Environment, id: Identifier): HookKind | null {\n+ const idType = id.type;\n+ if (idType.kind === \"Function\") {\n+ const signature = env.getFunctionSignature(idType);\n+ return signature?.hookKind ?? null;\n+ }\n+ return null;\n }\n \n export * from \".\/Types\";\ndiff --git a\/compiler\/forget\/src\/HIR\/Hooks.ts b\/compiler\/forget\/src\/HIR\/Hooks.ts\ndeleted file mode 100644\nindex b8bb0936c2..0000000000\n--- a\/compiler\/forget\/src\/HIR\/Hooks.ts\n+++ \/dev\/null\n@@ -1,16 +0,0 @@\n-\/**\n- * Copyright (c) Meta Platforms, Inc. and affiliates.\n- *\n- * This source code is licensed under the MIT license found in the\n- * LICENSE file in the root directory of this source tree.\n- *\/\n-\n-import { Effect, ValueKind } from \".\/HIR\";\n-\n-export type HookKind = \"State\" | \"Ref\" | \"Custom\" | \"Memo\";\n-export type Hook = {\n- kind: HookKind;\n- name: string;\n- effectKind: Effect;\n- valueKind: ValueKind;\n-};\ndiff --git a\/compiler\/forget\/src\/HIR\/ObjectShape.ts b\/compiler\/forget\/src\/HIR\/ObjectShape.ts\nindex 873140f75c..0ba10a973a 100644\n--- a\/compiler\/forget\/src\/HIR\/ObjectShape.ts\n+++ b\/compiler\/forget\/src\/HIR\/ObjectShape.ts\n@@ -33,14 +33,36 @@ function createAnonId(): string {\n }\n \n \/**\n- * Add a function to an existing ShapeRegistry.\n+ * Add a non-hook function to an existing ShapeRegistry.\n *\n * @returns a {@link FunctionType} representing the added function.\n *\/\n export function addFunction(\n registry: ShapeRegistry,\n properties: Iterable<[string, BuiltInType | PolyType]>,\n- fn: FunctionSignature\n+ fn: Omit\n+): FunctionType {\n+ const shapeId = createAnonId();\n+ addShape(registry, shapeId, properties, {\n+ ...fn,\n+ hookKind: null,\n+ });\n+ return {\n+ kind: \"Function\",\n+ return: fn.returnType,\n+ shapeId,\n+ };\n+}\n+\n+\/**\n+ * Add a hook to an existing ShapeRegistry.\n+ *\n+ * @returns a {@link FunctionType} representing the added hook function.\n+ *\/\n+export function addHook(\n+ registry: ShapeRegistry,\n+ properties: Iterable<[string, BuiltInType | PolyType]>,\n+ fn: FunctionSignature & { hookKind: HookKind }\n ): FunctionType {\n const shapeId = createAnonId();\n addShape(registry, shapeId, properties, fn);\n@@ -88,6 +110,16 @@ function addShape(\n return shape;\n }\n \n+export type HookKind =\n+ | \"useContext\"\n+ | \"useState\"\n+ | \"useRef\"\n+ | \"useEffect\"\n+ | \"useLayoutEffect\"\n+ | \"useMemo\"\n+ | \"useCallback\"\n+ | \"Custom\";\n+\n \/**\n * Call signature of a function, used for type and effect inference.\n *\n@@ -102,6 +134,7 @@ export type FunctionSignature = {\n returnType: BuiltInType | PolyType;\n returnValueKind: ValueKind;\n calleeEffect: Effect;\n+ hookKind: HookKind | null;\n };\n \n \/**\n@@ -184,3 +217,21 @@ addObject(BUILTIN_SHAPES, BuiltInObjectId, [\n \/\/ TODO:\n \/\/ hasOwnProperty, isPrototypeOf, propertyIsEnumerable, toLocaleString, valueOf\n ]);\n+\n+export const DefaultMutatingHook = addHook(BUILTIN_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Mutate,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"Custom\",\n+ returnValueKind: ValueKind.Mutable,\n+});\n+\n+export const DefaultNonmutatingHook = addHook(BUILTIN_SHAPES, [], {\n+ positionalParams: [],\n+ restParam: Effect.Freeze,\n+ returnType: { kind: \"Poly\" },\n+ calleeEffect: Effect.Read,\n+ hookKind: \"Custom\",\n+ returnValueKind: ValueKind.Frozen,\n+});\ndiff --git a\/compiler\/forget\/src\/HIR\/Types.ts b\/compiler\/forget\/src\/HIR\/Types.ts\nindex 14083d107f..8d26868bdc 100644\n--- a\/compiler\/forget\/src\/HIR\/Types.ts\n+++ b\/compiler\/forget\/src\/HIR\/Types.ts\n@@ -6,22 +6,11 @@\n *\/\n \n import invariant from \"invariant\";\n-import { Hook } from \".\/Hooks\";\n \n export type BuiltInType = PrimitiveType | FunctionType | ObjectType;\n \n-export type Type =\n- | BuiltInType\n- | HookType\n- | PhiType\n- | TypeVar\n- | PolyType\n- | PropType;\n+export type Type = BuiltInType | PhiType | TypeVar | PolyType | PropType;\n export type PrimitiveType = { kind: \"Primitive\" };\n-export type HookType = {\n- kind: \"Hook\";\n- definition: Hook | null;\n-};\n \n \/**\n * An {@link FunctionType} or {@link ObjectType} (also a JS object) may be associated with an\n@@ -93,7 +82,6 @@ export function typeEquals(tA: Type, tB: Type): boolean {\n return (\n typeVarEquals(tA, tB) ||\n funcTypeEquals(tA, tB) ||\n- hookTypeEquals(tA, tB) ||\n objectTypeEquals(tA, tB) ||\n primitiveTypeEquals(tA, tB) ||\n polyTypeEquals(tA, tB) ||\n@@ -131,12 +119,6 @@ function funcTypeEquals(tA: Type, tB: Type): boolean {\n return typeEquals(tA.return, tB.return);\n }\n \n-function hookTypeEquals(tA: Type, tB: Type): boolean {\n- return (\n- tA.kind === \"Hook\" && tB.kind === \"Hook\" && tA.definition === tB.definition\n- );\n-}\n-\n function phiTypeEquals(tA: Type, tB: Type): boolean {\n if (tA.kind === \"Phi\" && tB.kind === \"Phi\") {\n if (tA.operands.length !== tB.operands.length) {\ndiff --git a\/compiler\/forget\/src\/HIR\/ValidateHooksUsage.ts b\/compiler\/forget\/src\/HIR\/ValidateHooksUsage.ts\nindex f1134e9edc..c1a588d126 100644\n--- a\/compiler\/forget\/src\/HIR\/ValidateHooksUsage.ts\n+++ b\/compiler\/forget\/src\/HIR\/ValidateHooksUsage.ts\n@@ -11,7 +11,7 @@ import {\n ErrorSeverity,\n } from \"..\/CompilerError\";\n import { hasBackEdge } from \"..\/Optimization\/DeadCodeElimination\";\n-import { HIRFunction, IdentifierId, Place, isHookType } from \".\/HIR\";\n+import { HIRFunction, IdentifierId, Place, getHookKind } from \".\/HIR\";\n import { eachInstructionValueOperand, eachTerminalOperand } from \".\/visitors\";\n \n \/**\n@@ -56,7 +56,7 @@ export function validateHooksUsage(fn: HIRFunction): void {\n for (const instr of block.instructions) {\n if (\n instr.value.kind === \"LoadGlobal\" &&\n- isHookType(instr.lvalue.identifier)\n+ getHookKind(fn.env, instr.lvalue.identifier) != null\n ) {\n hooks.add(instr.lvalue.identifier.id);\n } else if (instr.value.kind === \"CallExpression\") {\ndiff --git a\/compiler\/forget\/src\/HIR\/ValidateUnconditionalHooks.ts b\/compiler\/forget\/src\/HIR\/ValidateUnconditionalHooks.ts\nindex a80eebed11..77d5656b87 100644\n--- a\/compiler\/forget\/src\/HIR\/ValidateUnconditionalHooks.ts\n+++ b\/compiler\/forget\/src\/HIR\/ValidateUnconditionalHooks.ts\n@@ -13,7 +13,7 @@ import {\n import { findBlocksWithBackEdges } from \"..\/Optimization\/DeadCodeElimination\";\n import { Err, Ok, Result } from \"..\/Utils\/Result\";\n import { PostDominator, computePostDominatorTree } from \".\/Dominator\";\n-import { BlockId, HIRFunction, isHookType } from \".\/HIR\";\n+import { BlockId, HIRFunction, getHookKind } from \".\/HIR\";\n \n \/**\n * Validates that the function honors the [Rules of Hooks](https:\/\/react.dev\/warnings\/invalid-hook-call-warning)\n@@ -85,7 +85,7 @@ export function validateUnconditionalHooks(\n for (const instr of block.instructions) {\n if (\n instr.value.kind === \"CallExpression\" &&\n- isHookType(instr.value.callee.identifier)\n+ getHookKind(fn.env, instr.value.callee.identifier) != null\n ) {\n const loc = instr.loc;\n \/\/ TODO: the current ESLint rule has different error messages for code that is called conditionally, in a loop, etc.\ndiff --git a\/compiler\/forget\/src\/HIR\/index.ts b\/compiler\/forget\/src\/HIR\/index.ts\nindex 0a06c43cfd..47e296ed83 100644\n--- a\/compiler\/forget\/src\/HIR\/index.ts\n+++ b\/compiler\/forget\/src\/HIR\/index.ts\n@@ -7,7 +7,7 @@\n \n export { lower } from \".\/BuildHIR\";\n export { computeDominatorTree, computePostDominatorTree } from \".\/Dominator\";\n-export { Environment } from \".\/Environment\";\n+export { Environment, Hook } from \".\/Environment\";\n export * from \".\/HIR\";\n export {\n markInstructionIds,\n@@ -15,7 +15,6 @@ export {\n removeUnreachableFallthroughs,\n reversePostorderBlocks,\n } from \".\/HIRBuilder\";\n-export { Hook } from \".\/Hooks\";\n export { mergeConsecutiveBlocks } from \".\/MergeConsecutiveBlocks\";\n export { printFunction, printHIR } from \".\/PrintHIR\";\n export { validateConsistentIdentifiers } from \".\/ValidateConsistentIdentifiers\";\ndiff --git a\/compiler\/forget\/src\/Inference\/DropMemoCalls.ts b\/compiler\/forget\/src\/Inference\/DropMemoCalls.ts\nindex 4b1669f7d3..6267b7c385 100644\n--- a\/compiler\/forget\/src\/Inference\/DropMemoCalls.ts\n+++ b\/compiler\/forget\/src\/Inference\/DropMemoCalls.ts\n@@ -5,17 +5,16 @@\n * LICENSE file in the root directory of this source tree.\n *\/\n \n-import { Effect, HIRFunction, HookType, isHookType } from \"..\/HIR\";\n+import { Effect, HIRFunction, getHookKind } from \"..\/HIR\";\n \n export default function (func: HIRFunction): void {\n for (const [_, block] of func.body.blocks) {\n for (const instr of block.instructions) {\n switch (instr.value.kind) {\n case \"CallExpression\": {\n- if (isHookType(instr.value.callee.identifier)) {\n- const name = (instr.value.callee.identifier.type as HookType)\n- .definition?.name;\n- if (name === \"useMemo\") {\n+ const hookKind = getHookKind(func.env, instr.value.callee.identifier);\n+ if (hookKind != null) {\n+ if (hookKind === \"useMemo\") {\n const [fn] = instr.value.args;\n \n \/\/ TODO(gsn): Consider inlining the function passed to useMemo,\n@@ -38,7 +37,7 @@ export default function (func: HIRFunction): void {\n loc: instr.value.loc,\n };\n }\n- } else if (name === \"useCallback\") {\n+ } else if (hookKind === \"useCallback\") {\n const [fn] = instr.value.args;\n \n \/\/ Instead of a Call, just alias the callback directly.\ndiff --git a\/compiler\/forget\/src\/Inference\/InferReferenceEffects.ts b\/compiler\/forget\/src\/Inference\/InferReferenceEffects.ts\nindex edfc37dbc9..a76cd1bdce 100644\n--- a\/compiler\/forget\/src\/Inference\/InferReferenceEffects.ts\n+++ b\/compiler\/forget\/src\/Inference\/InferReferenceEffects.ts\n@@ -24,7 +24,11 @@ import {\n Type,\n ValueKind,\n } from \"..\/HIR\/HIR\";\n-import { FunctionSignature } from \"..\/HIR\/ObjectShape\";\n+import {\n+ DefaultMutatingHook,\n+ DefaultNonmutatingHook,\n+ FunctionSignature,\n+} from \"..\/HIR\/ObjectShape\";\n import {\n printIdentifier,\n printMixedHIR,\n@@ -708,25 +712,33 @@ function inferBlock(\n continue;\n }\n case \"CallExpression\": {\n- if (instrValue.callee.identifier.type.kind === \"Hook\") {\n- const definition = instrValue.callee.identifier.type.definition;\n- if (definition !== null) {\n- effectKind = definition.effectKind;\n- valueKind = definition.valueKind;\n- } else if (env.enableAssumeHooksFollowRulesOfReact) {\n- effectKind = Effect.Freeze;\n- valueKind = ValueKind.Frozen;\n- } else {\n- effectKind = Effect.Mutate;\n- valueKind = ValueKind.Mutable;\n- }\n+ let signature = getFunctionCallSignature(\n+ env,\n+ instrValue.callee.identifier.type\n+ );\n+ signature =\n+ env.enableFunctionCallSignatureOptimizations ||\n+ signature?.hookKind != null\n+ ? signature\n+ : null;\n+\n+ if (\n+ signature &&\n+ signature.hookKind != null &&\n+ !env.enableTreatHooksAsFunctions\n+ ) {\n+ effectKind = signature.restParam;\n+ valueKind = signature.returnValueKind;\n break;\n }\n \n- const signature = env.enableFunctionCallSignatureOptimizations\n- ? getFunctionCallSignature(env, instrValue.callee.identifier.type)\n- : null;\n-\n+ \/\/ We currently always check reference effects of typed functions\n+ \/\/ (i.e. call `referenceAndCheckError`). However, default custom hooks\n+ \/\/ should not assert reference effects, since their signatures are only\n+ \/\/ assumptions \/ defaults.\n+ const isDefaultCustomHook =\n+ instrValue.callee.identifier.type === DefaultMutatingHook ||\n+ instrValue.callee.identifier.type === DefaultNonmutatingHook;\n const effects =\n signature !== null ? getFunctionEffects(instrValue, signature) : null;\n const returnValueKind =\n@@ -735,9 +747,13 @@ function inferBlock(\n const arg = instrValue.args[i];\n const place = arg.kind === \"Identifier\" ? arg : arg.place;\n if (effects !== null) {\n- \/\/ If effects are inferred for an argument, we should fail invalid\n- \/\/ mutating effects\n- state.referenceAndCheckError(place, effects[i]);\n+ if (isDefaultCustomHook) {\n+ state.reference(place, effects[i]);\n+ } else {\n+ \/\/ If effects are inferred for an argument, we should fail invalid\n+ \/\/ mutating effects\n+ state.referenceAndCheckError(place, effects[i]);\n+ }\n } else {\n state.reference(place, Effect.Mutate);\n }\ndiff --git a\/compiler\/forget\/src\/ReactiveScopes\/FlattenScopesWithHooks.ts b\/compiler\/forget\/src\/ReactiveScopes\/FlattenScopesWithHooks.ts\nindex 6b6432a799..14fa5fe11b 100644\n--- a\/compiler\/forget\/src\/ReactiveScopes\/FlattenScopesWithHooks.ts\n+++ b\/compiler\/forget\/src\/ReactiveScopes\/FlattenScopesWithHooks.ts\n@@ -6,12 +6,13 @@\n *\/\n \n import {\n+ Environment,\n InstructionId,\n- isHookType,\n ReactiveFunction,\n ReactiveScopeBlock,\n ReactiveStatement,\n ReactiveValue,\n+ getHookKind,\n } from \"..\/HIR\";\n import {\n ReactiveFunctionTransform,\n@@ -31,17 +32,26 @@ import {\n * to ensure the hook call does not inadvertently become conditional.\n *\/\n export function flattenScopesWithHooks(fn: ReactiveFunction): void {\n- visitReactiveFunction(fn, new Transform(), { hasHook: false });\n+ visitReactiveFunction(fn, new Transform(), {\n+ env: fn.env,\n+ hasHook: false,\n+ });\n }\n \n-type State = { hasHook: boolean };\n+type State = {\n+ env: Environment;\n+ hasHook: boolean;\n+};\n \n class Transform extends ReactiveFunctionTransform {\n override transformScope(\n scope: ReactiveScopeBlock,\n outerState: State\n ): Transformed {\n- const innerState: State = { hasHook: false };\n+ const innerState: State = {\n+ env: outerState.env,\n+ hasHook: false,\n+ };\n this.visitScope(scope, innerState);\n outerState.hasHook ||= innerState.hasHook;\n if (innerState.hasHook) {\n@@ -59,7 +69,7 @@ class Transform extends ReactiveFunctionTransform {\n this.traverseValue(id, value, state);\n if (\n value.kind === \"CallExpression\" &&\n- isHookType(value.callee.identifier)\n+ getHookKind(state.env, value.callee.identifier) != null\n ) {\n state.hasHook = true;\n }\ndiff --git a\/compiler\/forget\/src\/ReactiveScopes\/InferReactiveIdentifiers.ts b\/compiler\/forget\/src\/ReactiveScopes\/InferReactiveIdentifiers.ts\nindex da749fd527..6c7c9ec54b 100644\n--- a\/compiler\/forget\/src\/ReactiveScopes\/InferReactiveIdentifiers.ts\n+++ b\/compiler\/forget\/src\/ReactiveScopes\/InferReactiveIdentifiers.ts\n@@ -6,26 +6,32 @@\n *\/\n \n import { CompilerError } from \"..\/CompilerError\";\n+import { Environment } from \"..\/HIR\";\n import {\n Effect,\n IdentifierId,\n- isHookType,\n ReactiveFunction,\n ReactiveInstruction,\n+ getHookKind,\n } from \"..\/HIR\/HIR\";\n import { eachInstructionLValue } from \"..\/HIR\/visitors\";\n import { assertExhaustive } from \"..\/Utils\/utils\";\n import {\n- eachReactiveValueOperand,\n ReactiveFunctionVisitor,\n+ eachReactiveValueOperand,\n visitReactiveFunction,\n } from \".\/visitors\";\n \n type IdentifierReactivity = Map;\n \n class State {\n+ env: Environment;\n reactivityMap: IdentifierReactivity = new Map();\n temporaries: Map = new Map();\n+\n+ constructor(env: Environment) {\n+ this.env = env;\n+ }\n }\n \n class Visitor extends ReactiveFunctionVisitor {\n@@ -65,7 +71,7 @@ class Visitor extends ReactiveFunctionVisitor {\n if (\n !hasReactiveInput &&\n instr.value.kind === \"CallExpression\" &&\n- isHookType(instr.value.callee.identifier)\n+ getHookKind(state.env, instr.value.callee.identifier) != null\n ) {\n \/\/ Hooks cannot be memoized. Even if they do not accept any reactive inputs,\n \/\/ they are not guaranteed to memoize their return value, and their result\n@@ -169,7 +175,7 @@ export function inferReactiveIdentifiers(\n fn: ReactiveFunction\n ): Set {\n const visitor = new Visitor();\n- const state = new State();\n+ const state = new State(fn.env);\n for (const param of fn.params) {\n state.reactivityMap.set(param.identifier.id, true);\n }\ndiff --git a\/compiler\/forget\/src\/ReactiveScopes\/PruneNonEscapingScopes.ts b\/compiler\/forget\/src\/ReactiveScopes\/PruneNonEscapingScopes.ts\nindex eb0e82ca09..0db85aa470 100644\n--- a\/compiler\/forget\/src\/ReactiveScopes\/PruneNonEscapingScopes.ts\n+++ b\/compiler\/forget\/src\/ReactiveScopes\/PruneNonEscapingScopes.ts\n@@ -9,10 +9,9 @@ import invariant from \"invariant\";\n import prettyFormat from \"pretty-format\";\n import { CompilerError } from \"..\/CompilerError\";\n import {\n+ Environment,\n IdentifierId,\n InstructionId,\n- isHookType,\n- isMutableEffect,\n Pattern,\n Place,\n ReactiveFunction,\n@@ -23,6 +22,8 @@ import {\n ReactiveTerminalStatement,\n ReactiveValue,\n ScopeId,\n+ getHookKind,\n+ isMutableEffect,\n } from \"..\/HIR\";\n import { eachInstructionValueOperand } from \"..\/HIR\/visitors\";\n import { log } from \"..\/Utils\/logger\";\n@@ -30,10 +31,10 @@ import { assertExhaustive } from \"..\/Utils\/utils\";\n import { getPlaceScope } from \".\/BuildReactiveBlocks\";\n import { printReactiveFunction } from \".\/PrintReactiveFunction\";\n import {\n- eachReactiveValueOperand,\n ReactiveFunctionTransform,\n ReactiveFunctionVisitor,\n Transformed,\n+ eachReactiveValueOperand,\n visitReactiveFunction,\n } from \".\/visitors\";\n \n@@ -116,7 +117,7 @@ export function pruneNonEscapingScopes(\n ): void {\n \/\/ First build up a map of which instructions are involved in creating which values,\n \/\/ and which values are returned.\n- const state = new State();\n+ const state = new State(fn.env);\n if (fn.id !== null) {\n state.declare(fn.id.id);\n }\n@@ -200,6 +201,7 @@ type ScopeNode = {\n \n \/\/ Stores the identifier and scope graphs, set of returned identifiers, etc\n class State {\n+ env: Environment;\n \/\/ Maps lvalues for LoadLocal to the identifier being loaded, to resolve indirections\n \/\/ in subsequent lvalues\/rvalues\n definitions: Map = new Map();\n@@ -208,6 +210,10 @@ class State {\n scopes: Map = new Map();\n escapingValues: Set = new Set();\n \n+ constructor(env: Environment) {\n+ this.env = env;\n+ }\n+\n \/**\n * Declare a new identifier, used for function id and params\n *\/\n@@ -715,7 +721,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor {\n );\n } else if (instruction.value.kind === \"CallExpression\") {\n const callee = instruction.value.callee;\n- if (isHookType(callee.identifier)) {\n+ if (getHookKind(state.env, callee.identifier)) {\n for (const operand of eachInstructionValueOperand(instruction.value)) {\n state.escapingValues.add(operand.identifier.id);\n }\n"}

dillon1000/react

Files changed