[hir] Add Load/StoreContext (1/n)

--- This PR adds LoadContext and StoreContext to handle reading and writing to context variables. A context variable is any variable that is declared within a Forget-compiled function and reassigned within a closure. Conceptually, we want to treat these variables as attributes of a `EnvironmentContext` variable (as most javascript VMs do). - context variables currently do not participate in type inference (i.e. we do not produce type equations for loads from context variables). In the future, we can try typing this as `Phi(assignment1Type, assignment2Type, ...)`. - context variables are always treated as `Effect.Mutable`. - context variables do not participate in SSA, or certain optimizing passes (e.g. dead code elimination, constant propagation, etc). There is some still follow ups: - From my understanding, we should introduce a `DeclareContext` instruction. - currently, declaring a context variable (without initializing it) is broken. This is because the declaration lowers to `DeclareLocal`, which assumes it is storing to a SSA-fied identifier. ```js let x; x = 4; () => { x = {}; }; ``` - DeclareContext will also make some initialization logic easier. In this PR, I added some hack-y code to handle initializing effects / mutable ranges / other inference state for the first StoreContext. - Handle or bail on stores to context variables through destructuring assignment - ~~Next PR:~~ - ~~Change closures to track reassigned identifiers (to extend mutable range of primitives)~~
2025-11-01 09:12:30 +00:00 · 2023-05-12 12:17:05 -04:00
parent f5bdf462b2
commit 314a5cfca5
28 changed files with 549 additions and 57 deletions
@@ -18,6 +18,7 @@ import {
  validateUnconditionalHooks,
 } from "./HIR";
 import { Environment, EnvironmentConfig } from "./HIR/Environment";
+import { findContextIdentifiers } from "./HIR/FindContextIdentifiers";
 import {
  analyseFunctions,
  dropMemoCalls,
@@ -60,7 +61,8 @@ export function* run(
  func: NodePath<t.FunctionDeclaration>,
  config?: EnvironmentConfig | null
 ): Generator<CompilerPipelineValue, t.FunctionDeclaration> {
-  const env = new Environment(config ?? null);
+  const contextIdentifiers = findContextIdentifiers(func);
+  const env = new Environment(config ?? null, contextIdentifiers);
  const hir = lower(func, env).unwrap();
  yield log({ kind: "hir", name: "HIR", value: hir });

@@ -947,7 +947,7 @@ function lowerExpression(
      const expr = exprPath as NodePath<t.Identifier>;
      const place = lowerIdentifier(builder, expr);
      return {
-        kind: "LoadLocal",
+        kind: getLoadKind(builder, expr),
        place,
        loc: exprLoc,
      };
@@ -1367,7 +1367,7 @@ function lowerExpression(
            loc: exprLoc,
          });
          lowerValueToTemporary(builder, {
-            kind: "StoreLocal",
+            kind: getStoreKind(builder, leftExpr),
            lvalue: {
              place: { ...identifier },
              kind: InstructionKind.Reassign,
@@ -1701,7 +1701,7 @@ function lowerExpression(
        loc: exprLoc,
      });
      lowerValueToTemporary(builder, {
-        kind: "StoreLocal",
+        kind: getStoreKind(builder, argument),
        lvalue: { place: { ...identifier }, kind: InstructionKind.Reassign },
        value: { ...temp },
        loc: exprLoc,
@@ -2404,6 +2404,22 @@ function buildTemporaryPlace(builder: HIRBuilder, loc: SourceLocation): Place {
  return place;
 }

+function getStoreKind(
+  builder: HIRBuilder,
+  identifier: NodePath<t.Identifier>
+): "StoreLocal" | "StoreContext" {
+  const isContext = builder.isContextIdentifier(identifier);
+  return isContext ? "StoreContext" : "StoreLocal";
+}
+
+function getLoadKind(
+  builder: HIRBuilder,
+  identifier: NodePath<t.Identifier>
+): "LoadLocal" | "LoadContext" {
+  const isContext = builder.isContextIdentifier(identifier);
+  return isContext ? "LoadContext" : "LoadLocal";
+}
+
 function lowerAssignment(
  builder: HIRBuilder,
  loc: SourceLocation,
@@ -2446,7 +2462,7 @@ function lowerAssignment(
        loc: lvalue.node.loc ?? GeneratedSource,
      };
      const temporary = lowerValueToTemporary(builder, {
-        kind: "StoreLocal",
+        kind: getStoreKind(builder, lvalue),
        lvalue: { place: { ...place }, kind },
        value,
        loc,
@@ -2501,6 +2517,7 @@ function lowerAssignment(
      }
    }
    case "ArrayPattern": {
+      // TODO
      const lvalue = lvaluePath as NodePath<t.ArrayPattern>;
      const elements = lvalue.get("elements");
      const items: ArrayPattern["items"] = [];
@@ -5,6 +5,7 @@
 * LICENSE file in the root directory of this source tree.
 */

+import * as t from "@babel/types";
 import invariant from "invariant";
 import { log } from "../Utils/logger";
 import {
@@ -47,8 +48,12 @@ export class Environment {
  #nextIdentifer: number = 0;
  #nextBlock: number = 0;
  validateHooksUsage: boolean;
+  #contextIdentifiers: Set<t.Identifier>;

-  constructor(config: EnvironmentConfig | null) {
+  constructor(
+    config: EnvironmentConfig | null,
+    contextIdentifiers: Set<t.Identifier>
+  ) {
    this.#shapes = DEFAULT_SHAPES;

    if (config?.customHooks) {
@@ -67,6 +72,7 @@ export class Environment {
      this.#globals = DEFAULT_GLOBALS;
    }
    this.validateHooksUsage = config?.validateHooksUsage ?? false;
+    this.#contextIdentifiers = contextIdentifiers;
  }

  get nextIdentifierId(): IdentifierId {
@@ -76,6 +82,9 @@ export class Environment {
  get nextBlockId(): BlockId {
    return makeBlockId(this.#nextBlock++);
  }
+  isContextIdentifier(node: t.Identifier): boolean {
+    return this.#contextIdentifiers.has(node);
+  }

  getGlobalDeclaration(name: string): Global | null {
    let resolvedGlobal: Global | null = this.#globals.get(name) ?? null;
@@ -0,0 +1,154 @@
+import { NodePath } from "@babel/traverse";
+import * as t from "@babel/types";
+import { CompilerError } from "../CompilerError";
+import { GeneratedSource } from "./HIR";
+
+type FindContextIdentifierState = {
+  inLambda: number;
+  currentLambda: Array<
+    NodePath<t.FunctionExpression> | NodePath<t.ArrowFunctionExpression>
+  >;
+  contextIdentifiers: Set<t.Identifier>;
+};
+
+export function findContextIdentifiers(
+  func: NodePath<t.Function>
+): Set<t.Identifier> {
+  const state: FindContextIdentifierState = {
+    inLambda: 0,
+    currentLambda: [],
+    contextIdentifiers: new Set(),
+  };
+
+  func.traverse<FindContextIdentifierState>(
+    {
+      FunctionExpression: {
+        enter(
+          fn: NodePath<t.FunctionExpression>,
+          state: FindContextIdentifierState
+        ): void {
+          state.currentLambda.push(fn);
+        },
+        exit(
+          fn: NodePath<t.FunctionExpression>,
+          state: FindContextIdentifierState
+        ): void {
+          state.currentLambda.pop();
+        },
+      },
+
+      ArrowFunctionExpression: {
+        enter(
+          fn: NodePath<t.ArrowFunctionExpression>,
+          state: FindContextIdentifierState
+        ): void {
+          state.currentLambda.push(fn);
+        },
+        exit(
+          fn: NodePath<t.ArrowFunctionExpression>,
+          state: FindContextIdentifierState
+        ): void {
+          state.currentLambda.pop();
+        },
+      },
+      AssignmentExpression(
+        path: NodePath<t.AssignmentExpression>,
+        state: FindContextIdentifierState
+      ): void {
+        const currentLambda = state.currentLambda.at(-1);
+        if (currentLambda) {
+          const left = path.get("left");
+          handleAssignment(currentLambda, state.contextIdentifiers, left);
+        }
+      },
+    },
+    state
+  );
+  return state.contextIdentifiers;
+}
+
+function handleAssignment(
+  currentLambda:
+    | NodePath<t.FunctionExpression>
+    | NodePath<t.ArrowFunctionExpression>,
+  contextIdentifiers: Set<t.Identifier>,
+  lvalPath: NodePath<t.LVal>
+): void {
+  // Find all reassignments to identifiers declared outside of currentLambda
+  // This closely follows destructuring assignment assumptions and logic in BuildHIR
+  const lvalNode = lvalPath.node;
+  switch (lvalNode.type) {
+    case "Identifier": {
+      const path = lvalPath as NodePath<t.Identifier>;
+      const name = path.node.name;
+      const ownBinding = path.scope.getBinding(name);
+      const bindingAboveLambdaScope =
+        currentLambda.scope.parent.getBinding(name);
+
+      if (ownBinding != null && ownBinding === bindingAboveLambdaScope) {
+        contextIdentifiers.add(ownBinding.identifier);
+      }
+      break;
+    }
+    case "ArrayPattern": {
+      const path = lvalPath as NodePath<t.ArrayPattern>;
+      for (const element of path.get("elements")) {
+        if (nonNull(element)) {
+          handleAssignment(currentLambda, contextIdentifiers, element);
+        }
+      }
+      break;
+    }
+    case "ObjectPattern": {
+      const path = lvalPath as NodePath<t.ObjectPattern>;
+      for (const property of path.get("properties")) {
+        if (property.isObjectProperty()) {
+          const valuePath = property.get("value");
+          if (!valuePath.isLVal()) {
+            CompilerError.invariant(
+              `[FindContextIdentifiers] Expected object property value to be an LVal, got: ${valuePath.type}`,
+              valuePath.node.loc ?? GeneratedSource
+            );
+          }
+          handleAssignment(currentLambda, contextIdentifiers, valuePath);
+        } else {
+          if (!property.isRestElement()) {
+            CompilerError.invariant(
+              `[FindContextIdentifiers] Invalid assumptions for babel types.`,
+              property.node.loc ?? GeneratedSource
+            );
+          }
+          handleAssignment(currentLambda, contextIdentifiers, property);
+        }
+      }
+      break;
+    }
+    case "AssignmentPattern": {
+      const path = lvalPath as NodePath<t.AssignmentPattern>;
+      const left = path.get("left");
+      handleAssignment(currentLambda, contextIdentifiers, left);
+      break;
+    }
+    case "RestElement": {
+      const path = lvalPath as NodePath<t.RestElement>;
+      handleAssignment(currentLambda, contextIdentifiers, path.get("argument"));
+      break;
+    }
+    case "MemberExpression": {
+      // Interior mutability (not a reassign)
+      break;
+    }
+    default: {
+      CompilerError.todo(
+        `[FindContextIdentifiers] Cannot handle Object destructuring assignment target ${lvalNode.type}`,
+        lvalNode.loc ?? GeneratedSource
+      );
+    }
+  }
+}
+
+function nonNull<T extends NonNullable<t.Node>>(
+  t: NodePath<T | null>
+): t is NodePath<T> {
+  return t.node != null;
+}
@@ -552,6 +552,11 @@ export type InstructionValue =
      place: Place;
      loc: SourceLocation;
    }
+  | {
+      kind: "LoadContext";
+      place: Place;
+      loc: SourceLocation;
+    }
  | {
      kind: "DeclareLocal";
      lvalue: LValue;
@@ -563,6 +568,12 @@ export type InstructionValue =
      value: Place;
      loc: SourceLocation;
    }
+  | {
+      kind: "StoreContext";
+      lvalue: LValue;
+      value: Place;
+      loc: SourceLocation;
+    }
  | {
      kind: "Destructure";
      lvalue: LValuePattern;
@@ -5,7 +5,7 @@
 * LICENSE file in the root directory of this source tree.
 */

-import { NodePath } from "@babel/traverse";
+import { Binding, NodePath } from "@babel/traverse";
 import * as t from "@babel/types";
 import invariant from "invariant";
 import { CompilerError } from "../CompilerError";
@@ -22,10 +22,10 @@ import {
  Identifier,
  IdentifierId,
  Instruction,
+  Terminal,
  makeBlockId,
  makeInstructionId,
  makeType,
-  Terminal,
 } from "./HIR";
 import { printInstruction } from "./PrintHIR";
 import {
@@ -164,6 +164,35 @@ export default class HIRBuilder {
    }
  }

+  #resolveBabelBinding(
+    path: NodePath<t.Identifier | t.JSXIdentifier>
+  ): Binding | null {
+    const originalName = path.node.name;
+    const binding = path.scope.getBinding(originalName);
+    if (binding == null) {
+      return null;
+    }
+    // If the binding is from the parent function's outer scope, then
+    // we treat it equivalently to a global.
+    //
+    // TODO: remove the exception that resolves references to the
+    // parent function itself. We don't need to support self-recursion,
+    // so we can treat such references as globals.
+    const outerBinding =
+      this.parentFunction.scope.parent.getBinding(originalName);
+    if (binding === outerBinding) {
+      const func = this.parentFunction;
+      const isParentFunctionReference =
+        func.isFunctionDeclaration() &&
+        func.get("id").node != null &&
+        func.get("id").node!.name === originalName;
+      if (!isParentFunctionReference) {
+        return null;
+      }
+    }
+    return binding;
+  }
+
  /**
   * Maps an Identifier (or JSX identifier) Babel node to an internal `Identifier`
   * which represents the variable being referenced, according to the JS scoping rules.
@@ -198,36 +227,26 @@ export default class HIRBuilder {
    path: NodePath<t.Identifier | t.JSXIdentifier>
  ): Identifier | null {
    const originalName = path.node.name;
-    const binding = path.scope.getBinding(originalName);
-    if (binding == null) {
+    const babelBinding = this.#resolveBabelBinding(path);
+    if (babelBinding == null) {
      return null;
    }
-    // If the binding is from the parent function's outer scope, then
-    // we treat it equivalently to a global.
-    //
-    // TODO: remove the exception that resolves references to the
-    // parent function itself. We don't need to support self-recursion,
-    // so we can treat such references as globals.
-    const outerBinding =
-      this.parentFunction.scope.parent.getBinding(originalName);
-    if (binding === outerBinding) {
-      const func = this.parentFunction;
-      const isParentFunctionReference =
-        func.isFunctionDeclaration() &&
-        func.get("id").node != null &&
-        func.get("id").node!.name === originalName;
-      if (!isParentFunctionReference) {
-        return null;
-      }
-    }
-
-    const resolvedBinding = this.resolveBinding(binding.identifier);
+    const resolvedBinding = this.resolveBinding(babelBinding.identifier);
    if (resolvedBinding.name && resolvedBinding.name !== originalName) {
-      binding.scope.rename(originalName, resolvedBinding.name);
+      babelBinding.scope.rename(originalName, resolvedBinding.name);
    }
    return resolvedBinding;
  }

+  isContextIdentifier(path: NodePath<t.Identifier>): boolean {
+    const binding = this.#resolveBabelBinding(path);
+    if (binding) {
+      return this.#env.isContextIdentifier(binding.identifier);
+    } else {
+      return false;
+    }
+  }
+
  resolveBinding(node: t.Identifier): Identifier {
    const originalName = node.name;
    let name = originalName;
@@ -117,7 +117,7 @@ export function printInstruction(instr: ReactiveInstruction): string {
  }
 }

-function printPhi(phi: Phi): string {
+export function printPhi(phi: Phi): string {
  const items = [];
  items.push(printIdentifier(phi.id));
  items.push(printMutableRange(phi.id));
@@ -360,6 +360,16 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
      )} = ${printPlace(instrValue.value)}`;
      break;
    }
+    case "LoadContext": {
+      value = `LoadContext ${printPlace(instrValue.place)}`;
+      break;
+    }
+    case "StoreContext": {
+      value = `StoreContext ${instrValue.lvalue.kind} ${printPlace(
+        instrValue.lvalue.place
+      )} = ${printPlace(instrValue.value)}`;
+      break;
+    }
    case "Destructure": {
      value = `Destructure ${instrValue.lvalue.kind} ${printPattern(
        instrValue.lvalue.pattern
@@ -64,7 +64,8 @@ export function* eachInstructionValueOperand(
    case "DeclareLocal": {
      break;
    }
-    case "LoadLocal": {
+    case "LoadLocal":
+    case "LoadContext": {
      yield instrValue.place;
      break;
    }
@@ -72,6 +73,11 @@ export function* eachInstructionValueOperand(
      yield instrValue.value;
      break;
    }
+    case "StoreContext": {
+      yield instrValue.lvalue.place;
+      yield instrValue.value;
+      break;
+    }
    case "Destructure": {
      yield instrValue.value;
      break;
@@ -348,7 +354,8 @@ export function mapInstructionOperands(
    case "DeclareLocal": {
      break;
    }
-    case "LoadLocal": {
+    case "LoadLocal":
+    case "LoadContext": {
      instrValue.place = fn(instrValue.place);
      break;
    }
@@ -356,6 +363,11 @@ export function mapInstructionOperands(
      instrValue.value = fn(instrValue.value);
      break;
    }
+    case "StoreContext": {
+      instrValue.lvalue.place = fn(instrValue.lvalue.place);
+      instrValue.value = fn(instrValue.value);
+      break;
+    }
    case "Destructure": {
      instrValue.value = fn(instrValue.value);
      break;
@@ -75,7 +75,8 @@ export default function analyseFunctions(func: HIRFunction): void {
          state.declareProperty(instr.lvalue, instr.value.object, "");
          break;
        }
-        case "LoadLocal": {
+        case "LoadLocal":
+        case "LoadContext": {
          if (instr.lvalue.identifier.name === null) {
            state.declareTemporary(instr.lvalue, instr.value.place);
          }
@@ -34,14 +34,16 @@ function inferInstr(
  const { lvalue, value: instrValue } = instr;
  let alias: Place | null = null;
  switch (instrValue.kind) {
-    case "LoadLocal": {
+    case "LoadLocal":
+    case "LoadContext": {
      if (isPrimitiveType(instrValue.place.identifier)) {
        return;
      }
      alias = instrValue.place;
      break;
    }
-    case "StoreLocal": {
+    case "StoreLocal":
+    case "StoreContext": {
      alias = instrValue.value;
      break;
    }
@@ -118,6 +118,16 @@ export function inferMutableLifetimes(
    }

    for (const instr of block.instructions) {
+      if (instr.value.kind === "StoreContext") {
+        const id = instr.value.lvalue.place.identifier;
+        // Context variables do not participate in SSA and are not generally considered
+        // lvalues (). This hack tries to initialize a mutable range the first time we
+        // visit an context variable assignment.
+        if (id.mutableRange.start === 0 && id.mutableRange.end === 0) {
+          id.mutableRange.start = instr.id;
+          id.mutableRange.end = makeInstructionId(instr.id + 1);
+        }
+      }
      for (const operand of eachInstructionLValue(instr)) {
        const lvalueId = operand.identifier;

@@ -282,6 +282,12 @@ class InferenceState {
  #referenceImpl(place: Place, effectKind: Effect, shouldError: boolean): void {
    const values = this.#variables.get(place.identifier.id);
    if (values === undefined) {
+      if (effectKind === Effect.Store) {
+        CompilerError.invariant(
+          "[InferReferenceEffects] Unhandled store reference effect",
+          place.loc
+        );
+      }
      place.effect = effectKind === Effect.Mutate ? Effect.Mutate : Effect.Read;
      return;
    }
@@ -849,6 +855,19 @@ function inferBlock(
        state.alias(lvalue, instrValue.place);
        continue;
      }
+      case "LoadContext": {
+        state.reference(instrValue.place, Effect.Capture);
+        const lvalue = instr.lvalue;
+        lvalue.effect = Effect.Mutate;
+        const valueKind = state.kind(instrValue.place);
+        invariant(
+          valueKind === ValueKind.Mutable || valueKind === ValueKind.Context,
+          "[InferReferenceEffects] Context variables are always mutable."
+        );
+        state.initialize(instrValue, valueKind);
+        state.define(lvalue, instrValue);
+        continue;
+      }
      case "DeclareLocal": {
        const value: InstructionValue = {
          kind: "Primitive",
@@ -876,6 +895,30 @@ function inferBlock(
        state.reference(instrValue.lvalue.place, Effect.Store);
        continue;
      }
+      case "StoreContext": {
+        state.reference(instrValue.value, Effect.Mutate);
+        state.reference(instrValue.lvalue.place, Effect.Mutate);
+
+        const lvalue = instr.lvalue;
+        state.alias(lvalue, instrValue.value);
+        // this logic is really awkward
+        // Essentially, we want to say that
+        // 1. instr.lvalue (the value produced by the instruction itself) has a
+        //    ValueKind of the rhs.
+        //     - this is for chained assignment
+        // 2. instr.value.lvalue (the store location) has a ValueKind of Mutable
+
+        // As an alternative, we could insert a CreateContextVariable instruction
+        // before the initial StoreContext
+        const storeLValue = instrValue.lvalue.place;
+        if (!state.isDefined(storeLValue)) {
+          const instrCopy = { ...instrValue };
+          state.initialize(instrCopy, ValueKind.Mutable);
+          state.define(storeLValue, instrCopy);
+        }
+        lvalue.effect = Effect.Store;
+        continue;
+      }
      case "Destructure": {
        let effect: Effect = Effect.Capture;
        for (const place of eachPatternOperand(instrValue.lvalue.pattern)) {
@@ -222,6 +222,10 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
      // another StoreLocal or Destructure instruction, but conceptually we can't prune
      return false;
    }
+    case "LoadContext":
+    case "StoreContext": {
+      return false;
+    }
    case "RegExpLiteral":
    case "LoadGlobal":
    case "ArrayExpression":
@@ -441,13 +441,17 @@ function codegenInstructionNullable(
 ): t.Statement | null {
  if (
    instr.value.kind === "StoreLocal" ||
+    instr.value.kind === "StoreContext" ||
    instr.value.kind === "Destructure" ||
    instr.value.kind === "DeclareLocal"
  ) {
    let kind: InstructionKind = instr.value.lvalue.kind;
    let lvalue;
    let value: t.Expression | null;
-    if (instr.value.kind === "StoreLocal") {
+    if (
+      instr.value.kind === "StoreLocal" ||
+      instr.value.kind === "StoreContext"
+    ) {
      kind = cx.hasDeclared(instr.value.lvalue.place.identifier)
        ? InstructionKind.Reassign
        : kind;
@@ -899,7 +903,8 @@ function codegenInstructionValue(
      );
      break;
    }
-    case "LoadLocal": {
+    case "LoadLocal":
+    case "LoadContext": {
      value = codegenPlace(cx, instrValue.place);
      break;
    }
@@ -1013,7 +1018,8 @@ function codegenInstructionValue(
    case "Debugger":
    case "DeclareLocal":
    case "Destructure":
-    case "StoreLocal": {
+    case "StoreLocal":
+    case "StoreContext": {
      CompilerError.invariant(
        `Unexpected ${instrValue.kind} in codegenInstructionValue`,
        instrValue.loc
@@ -106,7 +106,10 @@ export function inferReactiveScopeVariables(fn: HIRFunction): void {
      if (range.end > range.start + 1 || mayAllocate(instr.value)) {
        operands.push(instr.lvalue!.identifier);
      }
-      if (instr.value.kind === "StoreLocal") {
+      if (
+        instr.value.kind === "StoreLocal" ||
+        instr.value.kind === "StoreContext"
+      ) {
        if (
          instr.value.lvalue.place.identifier.mutableRange.end >
          instr.value.lvalue.place.identifier.mutableRange.start + 1
@@ -226,6 +229,8 @@ function mayAllocate(value: InstructionValue): boolean {
    case "TypeCastExpression":
    case "BinaryExpression":
    case "LoadLocal":
+    case "LoadContext":
+    case "StoreContext":
    case "PropertyLoad":
    case "PropertyDelete":
    case "ComputedLoad":
@@ -91,6 +91,7 @@ class FindPromotedTemporaries extends ReactiveFunctionVisitor<TemporariesUsedOut
    }
    switch (instruction.value.kind) {
      case "LoadLocal":
+      case "LoadContext":
      case "PropertyLoad": {
        state.declarations.set(instruction.lvalue.identifier.id, scope);
        break;
@@ -512,7 +513,10 @@ class PropagationVisitor extends ReactiveFunctionVisitor<Context> {
    value: ReactiveValue,
    lvalue: Place | null
  ): void {
-    if (value.kind === "LoadLocal" && lvalue !== null) {
+    if (
+      (value.kind === "LoadLocal" || value.kind === "LoadContext") &&
+      lvalue !== null
+    ) {
      if (
        value.place.identifier.name !== null &&
        lvalue.identifier.name === null &&
@@ -528,7 +532,7 @@ class PropagationVisitor extends ReactiveFunctionVisitor<Context> {
      } else {
        context.visitProperty(value.object, value.property);
      }
-    } else if (value.kind === "StoreLocal") {
+    } else if (value.kind === "StoreLocal" || value.kind === "StoreContext") {
      context.visitOperand(value.value);
      if (value.lvalue.kind === InstructionKind.Reassign) {
        context.visitReassignment(value.lvalue.place);
@@ -461,6 +461,16 @@ function computeMemoizationInputs(
        rvalues: [value.place],
      };
    }
+    case "LoadContext": {
+      return {
+        // Should never be pruned
+        lvalues:
+          lvalue !== null
+            ? [{ place: lvalue, level: MemoizationLevel.Conditional }]
+            : [],
+        rvalues: [value.place],
+      };
+    }
    case "DeclareLocal": {
      const lvalues = [
        { place: value.lvalue.place, level: MemoizationLevel.Unmemoized },
@@ -486,6 +496,20 @@ function computeMemoizationInputs(
        rvalues: [value.value],
      };
    }
+    case "StoreContext": {
+      // Should never be pruned
+      const lvalues = [
+        { place: value.lvalue.place, level: MemoizationLevel.Memoized },
+      ];
+      if (lvalue !== null) {
+        lvalues.push({ place: lvalue, level: MemoizationLevel.Conditional });
+      }
+
+      return {
+        lvalues,
+        rvalues: [value.value],
+      };
+    }
    case "Destructure": {
      // Indirection for the inner value, memoized if the value is
      const lvalues = [];
@@ -124,6 +124,14 @@ function* generateInstructionTypes(
      break;
    }

+    // For now, we won't infer types for context variables
+    case "StoreContext": {
+      break;
+    }
+    case "LoadContext": {
+      yield equation(left, value.place.identifier.type);
+      break;
+    }
    case "StoreLocal": {
      yield equation(left, value.value.identifier.type);
      yield equation(
@@ -17,14 +17,24 @@ function component(a) {
 ## Code

 ```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
 function component(a) {
-  const x = { a };
-
-  (function () {
-    y = x;
-  })();
-  mutate(1);
-  return 1;
+  const $ = useMemoCache(2);
+  const c_0 = $[0] !== a;
+  let y;
+  if (c_0) {
+    const x = { a };
+    y = 1;
+    (function () {
+      y = x;
+    })();
+    mutate(y);
+    $[0] = a;
+    $[1] = y;
+  } else {
+    y = $[1];
+  }
+  return y;
 }

 ```
@@ -0,0 +1,35 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let a;
+  [a, b] = props.value;
+
+  return [a, b];
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(2);
+
+  const [a] = props.value;
+  const c_0 = $[0] !== a;
+  let t0;
+  if (c_0) {
+    t0 = [a, b];
+    $[0] = a;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+```
+      
@@ -0,0 +1,6 @@
+function Component(props) {
+  let a;
+  [a, b] = props.value;
+
+  return [a, b];
+}
@@ -22,17 +22,26 @@ function Component() {
 ## Code

 ```javascript
-// writing to primitives is not a 'mutate' or 'store' to context references,
+import { unstable_useMemoCache as useMemoCache } from "react"; // writing to primitives is not a 'mutate' or 'store' to context references,
 // under current analysis in AnalyzeFunctions.
 // <unknown> $23:TFunction = Function @deps[<unknown>
 //   $21:TPrimitive,<unknown> $22:TPrimitive]:

 function Component() {
-  const fn = function () {
-    x = x + 1;
-  };
-  fn();
-  return 40;
+  const $ = useMemoCache(1);
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x = 40;
+
+    const fn = function () {
+      x = x + 1;
+    };
+    fn();
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+  return x;
 }

 ```
@@ -31,6 +31,7 @@ function Component() {
  }
  const x = t0;

+  let x_0 = 56;
  const fn = function () {
    x_0 = 42;
  };
@@ -28,7 +28,7 @@ function bar(a, b) {
  if (c_0 || c_1) {
    const x = [a, b];
    y = {};
-    const t = {};
+    let t = {};
    (function () {
      y = x[0][1];
      t = x[1][0];
@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = [];
+  let foo = () => {
+    x = {};
+  };
+  foo();
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(1);
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x = [];
+    const foo = () => {
+      x = {};
+    };
+    foo();
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+  return x;
+}
+
+```
+      
@@ -0,0 +1,8 @@
+function Component(props) {
+  let x = [];
+  let foo = () => {
+    x = {};
+  };
+  foo();
+  return x;
+}
@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = 5;
+  let foo = () => {
+    x = {};
+  };
+  foo();
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(1);
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x = 5;
+    const foo = () => {
+      x = {};
+    };
+    foo();
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+  return x;
+}
+
+```
+      
@@ -0,0 +1,8 @@
+function Component(props) {
+  let x = 5;
+  let foo = () => {
+    x = {};
+  };
+  foo();
+  return x;
+}