[hir] add DeclareContext (2/n)

--- 

In `lower`, we now ensure that all context variables are declared by a 
`DeclareContext` instruction. `DeclareContext` always produces a `let` 
declaration, and `StoreContext` is always a reassign. There are a few reasons we 
need `DeclareContext`: 

- DeclareLocal assumes it is storing to a SSA-fied identifier (which always 
stores an immutable primitive). This does not fit context variables. 

- Without DeclareContext, we need custom logic in some passes to initialize 
identifier / context state (e.g. `MutableRange`, ValueKind, etc) for the 
`StoreContext` that declares the context. 

This PR stack models context variables as concrete identifiers (with references 
to context variables modeled by `Place` referencing the context variable 
identifier). @josephsavona pointed out that this is abusing the notion of 
Identifier/Place, as context variables are essentially interior properties of a 
ContextEnvironment. Since we are not modeling `ContextEnvironment` implicitly or 
explicitly, all inference for context variables is essentially pointer analysis.

compiler/forget/src/HIR/BuildHIR.ts

@@ -647,19 +647,38 @@ function lowerStatement(
               nodePath: id,
             });
           } else {
-            lowerValueToTemporary(builder, {
-              kind: "DeclareLocal",
-              lvalue: {
-                kind,
-                place: {
-                  effect: Effect.Unknown,
-                  identifier,
-                  kind: "Identifier",
-                  loc: id.node.loc ?? GeneratedSource,
-                },
-              },
+            const place: Place = {
+              effect: Effect.Unknown,
+              identifier,
+              kind: "Identifier",
               loc: id.node.loc ?? GeneratedSource,
-            });
+            };
+            if (builder.isContextIdentifier(id)) {
+              if (kind === InstructionKind.Const) {
+                builder.errors.push({
+                  reason: `(BuildHIR::lowerAssignment) Invalid declaration kind (const) for variable later reassigned.`,
+                  severity: ErrorSeverity.InvalidInput,
+                  nodePath: id,
+                });
+              }
+              lowerValueToTemporary(builder, {
+                kind: "DeclareContext",
+                lvalue: {
+                  kind: InstructionKind.Let,
+                  place,
+                },
+                loc: id.node.loc ?? GeneratedSource,
+              });
+            } else {
+              lowerValueToTemporary(builder, {
+                kind: "DeclareLocal",
+                lvalue: {
+                  kind,
+                  place,
+                },
+                loc: id.node.loc ?? GeneratedSource,
+              });
+            }
           }
         } else {
           builder.errors.push({
@@ -2461,12 +2480,41 @@ function lowerAssignment(
         effect: Effect.Unknown,
         loc: lvalue.node.loc ?? GeneratedSource,
       };
-      const temporary = lowerValueToTemporary(builder, {
-        kind: getStoreKind(builder, lvalue),
-        lvalue: { place: { ...place }, kind },
-        value,
-        loc,
-      });
+
+      let temporary;
+      if (builder.isContextIdentifier(lvalue)) {
+        if (kind !== InstructionKind.Reassign) {
+          if (kind === InstructionKind.Const) {
+            builder.errors.push({
+              reason: `(BuildHIR::lowerAssignment) Invalid declaration kind (const) for variable later reassigned.`,
+              severity: ErrorSeverity.InvalidInput,
+              nodePath: lvalue,
+            });
+          }
+          lowerValueToTemporary(builder, {
+            kind: "DeclareContext",
+            lvalue: {
+              kind: InstructionKind.Let,
+              place: { ...place },
+            },
+            loc: place.loc,
+          });
+        }
+
+        temporary = lowerValueToTemporary(builder, {
+          kind: "StoreContext",
+          lvalue: { place: { ...place }, kind: InstructionKind.Reassign },
+          value,
+          loc,
+        });
+      } else {
+        temporary = lowerValueToTemporary(builder, {
+          kind: "StoreLocal",
+          lvalue: { place: { ...place }, kind },
+          value,
+          loc,
+        });
+      }
       return { kind: "LoadLocal", place: temporary, loc: temporary.loc };
     }
     case "MemberExpression": {

compiler/forget/src/HIR/HIR.ts

@@ -562,6 +562,14 @@ export type InstructionValue =
       lvalue: LValue;
       loc: SourceLocation;
     }
+  | {
+      kind: "DeclareContext";
+      lvalue: {
+        kind: InstructionKind.Let;
+        place: Place;
+      };
+      loc: SourceLocation;
+    }
   | {
       kind: "StoreLocal";
       lvalue: LValue;
@@ -570,7 +578,10 @@ export type InstructionValue =
     }
   | {
       kind: "StoreContext";
-      lvalue: LValue;
+      lvalue: {
+        kind: InstructionKind.Reassign;
+        place: Place;
+      };
       value: Place;
       loc: SourceLocation;
     }

compiler/forget/src/HIR/PrintHIR.ts

@@ -354,6 +354,12 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       )}`;
       break;
     }
+    case "DeclareContext": {
+      value = `DeclareContext ${instrValue.lvalue.kind} ${printPlace(
+        instrValue.lvalue.place
+      )}`;
+      break;
+    }
     case "StoreLocal": {
       value = `StoreLocal ${instrValue.lvalue.kind} ${printPlace(
         instrValue.lvalue.place

compiler/forget/src/HIR/visitors.ts

@@ -26,6 +26,7 @@ export function* eachInstructionLValue(
   }
   switch (instr.value.kind) {
     case "DeclareLocal":
+    case "DeclareContext":
     case "StoreLocal": {
       yield instr.value.lvalue.place;
       break;
@@ -61,6 +62,7 @@ export function* eachInstructionValueOperand(
       yield* eachCallArgument(instrValue.args);
       break;
     }
+    case "DeclareContext":
     case "DeclareLocal": {
       break;
     }
@@ -351,6 +353,7 @@ export function mapInstructionOperands(
       instrValue.value = fn(instrValue.value);
       break;
     }
+    case "DeclareContext":
     case "DeclareLocal": {
       break;
     }

compiler/forget/src/Inference/InferMutableLifetimes.ts

@@ -118,16 +118,6 @@ 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;
 

compiler/forget/src/Inference/InferReferenceEffects.ts

@@ -876,8 +876,11 @@ function inferBlock(
         };
         state.initialize(value, ValueKind.Immutable);
         state.define(instrValue.lvalue.place, value);
-        state.alias(instr.lvalue, instrValue.lvalue.place);
-        instr.lvalue.effect = Effect.Mutate;
+        continue;
+      }
+      case "DeclareContext": {
+        state.initialize(instrValue, ValueKind.Mutable);
+        state.define(instrValue.lvalue.place, instrValue);
         continue;
       }
       case "StoreLocal": {
@@ -901,21 +904,6 @@ function inferBlock(
 
         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;
       }

compiler/forget/src/Optimization/DeadCodeElimination.ts

@@ -223,6 +223,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
       return false;
     }
     case "LoadContext":
+    case "DeclareContext":
     case "StoreContext": {
       return false;
     }

compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -443,21 +443,25 @@ function codegenInstructionNullable(
     instr.value.kind === "StoreLocal" ||
     instr.value.kind === "StoreContext" ||
     instr.value.kind === "Destructure" ||
-    instr.value.kind === "DeclareLocal"
+    instr.value.kind === "DeclareLocal" ||
+    instr.value.kind === "DeclareContext"
   ) {
     let kind: InstructionKind = instr.value.lvalue.kind;
     let lvalue;
     let value: t.Expression | null;
-    if (
-      instr.value.kind === "StoreLocal" ||
-      instr.value.kind === "StoreContext"
-    ) {
+    if (instr.value.kind === "StoreLocal") {
       kind = cx.hasDeclared(instr.value.lvalue.place.identifier)
         ? InstructionKind.Reassign
         : kind;
       lvalue = instr.value.lvalue.place;
       value = codegenPlace(cx, instr.value.value);
-    } else if (instr.value.kind === "DeclareLocal") {
+    } else if (instr.value.kind === "StoreContext") {
+      lvalue = instr.value.lvalue.place;
+      value = codegenPlace(cx, instr.value.value);
+    } else if (
+      instr.value.kind === "DeclareLocal" ||
+      instr.value.kind === "DeclareContext"
+    ) {
       if (cx.hasDeclared(instr.value.lvalue.place.identifier)) {
         return null;
       }
@@ -507,8 +511,17 @@ function codegenInstructionNullable(
         invariant(value !== null, "Expected a value for reassignment");
         const expr = t.assignmentExpression("=", codegenLValue(lvalue), value);
         if (instr.lvalue !== null) {
-          cx.temp.set(instr.lvalue.identifier.id, expr);
-          return null;
+          if (instr.value.kind !== "StoreContext") {
+            cx.temp.set(instr.lvalue.identifier.id, expr);
+            return null;
+          } else {
+            // Handle chained reassignments for context variables
+            const statement = codegenInstruction(cx, instr, expr);
+            if (statement.type === "EmptyStatement") {
+              return null;
+            }
+            return statement;
+          }
         } else {
           return createExpressionStatement(instr.loc, expr);
         }
@@ -1017,6 +1030,7 @@ function codegenInstructionValue(
     }
     case "Debugger":
     case "DeclareLocal":
+    case "DeclareContext":
     case "Destructure":
     case "StoreLocal":
     case "StoreContext": {

compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -224,6 +224,7 @@ function mayAllocate(value: InstructionValue): boolean {
     }
     case "Await":
     case "DeclareLocal":
+    case "DeclareContext":
     case "StoreLocal":
     case "LoadGlobal":
     case "TypeCastExpression":

compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts

@@ -513,10 +513,7 @@ class PropagationVisitor extends ReactiveFunctionVisitor<Context> {
     value: ReactiveValue,
     lvalue: Place | null
   ): void {
-    if (
-      (value.kind === "LoadLocal" || value.kind === "LoadContext") &&
-      lvalue !== null
-    ) {
+    if (value.kind === "LoadLocal" && lvalue !== null) {
       if (
         value.place.identifier.name !== null &&
         lvalue.identifier.name === null &&
@@ -532,7 +529,7 @@ class PropagationVisitor extends ReactiveFunctionVisitor<Context> {
       } else {
         context.visitProperty(value.object, value.property);
       }
-    } else if (value.kind === "StoreLocal" || value.kind === "StoreContext") {
+    } else if (value.kind === "StoreLocal") {
       context.visitOperand(value.value);
       if (value.lvalue.kind === InstructionKind.Reassign) {
         context.visitReassignment(value.lvalue.place);
@@ -541,11 +538,15 @@ class PropagationVisitor extends ReactiveFunctionVisitor<Context> {
         id,
         scope: context.currentScope,
       });
-    } else if (value.kind === "DeclareLocal") {
+    } else if (value.kind === "DeclareLocal" || value.kind === "DeclareContext") {
       // Some variables may be declared and never initialized. We need
       // to retain (and hoist) these declarations if they are included
       // in a reactive scope. One approach is to simply add all `DeclareLocal`s
       // as scope declarations.
+      
+      // We add context variable declarations here, not at `StoreContext`, since
+      // context Store / Loads are modeled as reads and mutates to the underlying
+      // variable reference (instead of through intermediate / inlined temporaries)
       context.declare(value.lvalue.place.identifier, {
         id,
         scope: context.currentScope,

compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -471,6 +471,19 @@ function computeMemoizationInputs(
         rvalues: [value.place],
       };
     }
+    case "DeclareContext": {
+      const lvalues = [
+        { place: value.lvalue.place, level: MemoizationLevel.Memoized },
+      ];
+      if (lvalue !== null) {
+        lvalues.push({ place: lvalue, level: MemoizationLevel.Unmemoized });
+      }
+      return {
+        lvalues,
+        rvalues: [],
+      };
+    }
+
     case "DeclareLocal": {
       const lvalues = [
         { place: value.lvalue.place, level: MemoizationLevel.Unmemoized },

compiler/forget/src/TypeInference/InferTypes.ts

@@ -202,6 +202,7 @@ function* generateInstructionTypes(
     }
 
     case "DeclareLocal":
+    case "DeclareContext":
     case "Destructure":
     case "NewExpression":
     case "TypeCastExpression":

compiler/forget/src/__tests__/fixtures/compiler/_bug.lambda-reassign-shadowed-primitive.expect.md

@@ -31,7 +31,8 @@ function Component() {
   }
   const x = t0;
 
-  let x_0 = 56;
+  let x_0;
+  x_0 = 56;
   const fn = function () {
     x_0 = 42;
   };

compiler/forget/src/__tests__/fixtures/compiler/capturing-function-alias-computed-load-3.expect.md

@@ -28,7 +28,8 @@ function bar(a, b) {
   if (c_0 || c_1) {
     const x = [a, b];
     y = {};
-    let t = {};
+    let t;
+    t = {};
     (function () {
       y = x[0][1];
       t = x[1][0];

compiler/forget/src/__tests__/fixtures/compiler/chained-assignment-context-variable.expect.md

@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+function Component() {
+  let x,
+    y = (x = {});
+  const foo = () => {
+    x = getObject();
+  };
+  foo();
+  return [y, x];
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component() {
+  const $ = useMemoCache(3);
+  let x;
+  let y;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    y = x = {};
+
+    const foo = () => {
+      x = getObject();
+    };
+    foo();
+    $[0] = x;
+    $[1] = y;
+  } else {
+    x = $[0];
+    y = $[1];
+  }
+  let t0;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = [y, x];
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/chained-assignment-context-variable.js

@@ -0,0 +1,9 @@
+function Component() {
+  let x,
+    y = (x = {});
+  const foo = () => {
+    x = getObject();
+  };
+  foo();
+  return [y, x];
+}

compiler/forget/src/__tests__/fixtures/compiler/declare-reassign-variable-in-closure.expect.md

@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+function Component(p) {
+  let x;
+  const foo = () => {
+    x = {};
+  };
+  foo();
+
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(p) {
+  const $ = useMemoCache(1);
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const foo = () => {
+      x = {};
+    };
+    foo();
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/declare-reassign-variable-in-closure.js

@@ -0,0 +1,9 @@
+function Component(p) {
+  let x;
+  const foo = () => {
+    x = {};
+  };
+  foo();
+
+  return x;
+}