diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
index 753ae91beb..6a23e532ba 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
@@ -431,7 +431,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       const deps = instrValue.dependencies
         .map((dep) => printPlace(dep))
         .join(",");
-      value = `Function @deps[${deps}]:\n${fn}`;
+      const context = instrValue.loweredFunc.context
+        .map((dep) => printPlace(dep))
+        .join(",");
+      value = `Function @deps[${deps}] @context[${context}]:\n${fn}`;
       break;
     }
     case "TaggedTemplateExpression": {
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts
index 37020eb4dd..2ecfe2e8c3 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts
@@ -46,8 +46,8 @@ class SSABuilder {
   #unknown: Set<Identifier> = new Set();
   #context: Set<Identifier> = new Set();
 
-  constructor(env: Environment, blocks: Map<BlockId, BasicBlock>) {
-    this.#blocks = blocks;
+  constructor(env: Environment, blocks: ReadonlyMap<BlockId, BasicBlock>) {
+    this.#blocks = new Map(blocks);
     this.#env = env;
   }
 
@@ -55,6 +55,18 @@ class SSABuilder {
     return this.#env.nextIdentifierId;
   }
 
+  defineFunction(func: HIRFunction): void {
+    for (const [id, block] of func.body.blocks) {
+      this.#blocks.set(id, block);
+    }
+  }
+
+  enter(fn: () => void): void {
+    const current = this.#current;
+    fn();
+    this.#current = current;
+  }
+
   state(): State {
     invariant(
       this.#current !== null,
@@ -125,7 +137,9 @@ class SSABuilder {
     if (block.preds.size == 0) {
       // We're at the entry block and haven't found our defintion yet.
       // console.log(
-      //   `Unable to find "${printIdentifier(oldId)}", assuming it's a global`
+      //   `Unable to find "${printIdentifier(
+      //     oldId
+      //   )}" in bb${blockId}, assuming it's a global`
       // );
       this.#unknown.add(oldId);
       return oldId;
@@ -213,8 +227,16 @@ class SSABuilder {
 }
 
 export default function enterSSA(func: HIRFunction): void {
-  const visitedBlocks: Set<BasicBlock> = new Set();
   const builder = new SSABuilder(func.env, func.body.blocks);
+  enterSSAImpl(func, builder, func.body.entry);
+}
+
+function enterSSAImpl(
+  func: HIRFunction,
+  builder: SSABuilder,
+  rootEntry: BlockId
+): void {
+  const visitedBlocks: Set<BasicBlock> = new Set();
   for (const [blockId, block] of func.body.blocks) {
     invariant(
       !visitedBlocks.has(block),
@@ -224,8 +246,14 @@ export default function enterSSA(func: HIRFunction): void {
 
     builder.startBlock(block);
 
-    if (func.body.entry === blockId) {
-      func.context = func.context.map((p) => builder.defineContext(p));
+    if (blockId === rootEntry) {
+      // NOTE: func.context should be empty for the root function
+      if (func.context.length !== 0) {
+        CompilerError.invariant(
+          `Expected function context to be empty for outer function declarations`,
+          func.loc
+        );
+      }
       func.params = func.params.map((p) => builder.definePlace(p));
     }
 
@@ -234,7 +262,24 @@ export default function enterSSA(func: HIRFunction): void {
       mapInstructionOperands(instr, (place) => builder.getPlace(place));
 
       if (instr.value.kind === "FunctionExpression") {
-        enterSSA(instr.value.loweredFunc);
+        const loweredFunc = instr.value.loweredFunc;
+        const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!;
+        invariant(
+          entry.preds.size === 0,
+          "Expected function expression entry block to have zero predecessors"
+        );
+        entry.preds.add(blockId);
+        builder.defineFunction(loweredFunc);
+        builder.enter(() => {
+          loweredFunc.context = loweredFunc.context.map((p) =>
+            builder.getPlace(p)
+          );
+          loweredFunc.params = loweredFunc.params.map((p) =>
+            builder.definePlace(p)
+          );
+          enterSSAImpl(loweredFunc, builder, rootEntry);
+        });
+        entry.preds.clear();
       }
     }
 
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
index 0b8acae0be..0f42adab0b 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
@@ -65,8 +65,11 @@ function apply(func: HIRFunction, unifier: Unifier): void {
       for (const place of eachInstructionOperand(instr)) {
         place.identifier.type = unifier.get(place.identifier.type);
       }
-      const { lvalue } = instr;
+      const { lvalue, value } = instr;
       lvalue.identifier.type = unifier.get(lvalue.identifier.type);
+      if (value.kind === "FunctionExpression") {
+        apply(value.loweredFunc, unifier);
+      }
     }
   }
 }
@@ -125,14 +128,13 @@ function* generateInstructionTypes(
       break;
     }
 
-    // For now, we won't infer types for context variables
-    case "StoreContext": {
-      break;
-    }
+    // We intentionally do not infer types for context variables
+    case "DeclareContext":
+    case "StoreContext":
     case "LoadContext": {
-      yield equation(left, value.place.identifier.type);
       break;
     }
+
     case "StoreLocal": {
       yield equation(left, value.value.identifier.type);
       yield equation(
@@ -245,12 +247,11 @@ function* generateInstructionTypes(
     }
 
     case "FunctionExpression": {
-      inferTypes(value.loweredFunc);
+      yield* generate(value.loweredFunc);
       break;
     }
 
     case "DeclareLocal":
-    case "DeclareContext":
     case "NewExpression":
     case "JsxExpression":
     case "JsxFragment":
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md
index 330045c423..eb63670c41 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md
@@ -20,41 +20,42 @@ function Component(props) {
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
   const $ = useMemoCache(7);
-  const c_0 = $[0] !== props.x;
-  let t0;
+  const t0 = props.x;
+  const c_0 = $[0] !== t0;
+  let t1;
   if (c_0) {
-    t0 = foo(props.x);
-    $[0] = props.x;
-    $[1] = t0;
+    t1 = foo(t0);
+    $[0] = t0;
+    $[1] = t1;
   } else {
-    t0 = $[1];
+    t1 = $[1];
   }
-  const x = t0;
+  const x = t1;
   const c_2 = $[2] !== props;
   const c_3 = $[3] !== x;
-  let t1;
+  let t2;
   if (c_2 || c_3) {
-    t1 = function () {
+    t2 = function () {
       const arr = [...bar(props)];
       return arr.at(x);
     };
     $[2] = props;
     $[3] = x;
-    $[4] = t1;
+    $[4] = t2;
   } else {
-    t1 = $[4];
+    t2 = $[4];
   }
-  const fn = t1;
+  const fn = t2;
   const c_5 = $[5] !== fn;
-  let t2;
+  let t3;
   if (c_5) {
-    t2 = fn();
+    t3 = fn();
     $[5] = fn;
-    $[6] = t2;
+    $[6] = t3;
   } else {
-    t2 = $[6];
+    t3 = $[6];
   }
-  const fnResult = t2;
+  const fnResult = t3;
   return fnResult;
 }
 
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md
similarity index 100%
rename from compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md
rename to compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js
similarity index 100%
rename from compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js
rename to compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
index 431e23cdd7..536652b111 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
@@ -38,8 +38,8 @@ function Component(props) {
   if (c_0) {
     const allUrls = [];
 
-    const { media: t85, comments, urls } = post;
-    media = t85;
+    const { media: t84, comments, urls } = post;
+    media = t84;
     const c_3 = $[3] !== comments.length;
     let t0;
     if (c_3) {