Ensure valid mutable ranges for all scopes; fix ranges for context vars

Our current validation fails to detect some invalid cases of mutable ranges — 
namely, ranges that are fully or partially uninitialized, with start or 
start+end still set to zero. 

This PR fixes these cases, starting by validating that the ranges for _all_ 
reactive scopes are valid: start >= 1, and end <= (last instr id + 1). This 
exposed the invalid cases, which are also fixed here: 

* During AnalyzeFunctions, we need to reset identifier ranges and scopes when 
exiting an inner function. This has to happen *after* the effects have been 
translated to the function deps/context operands, in order for 
InferRefenceEffects to continue working on the outer function. Previously I did 
this at the start of InferReactiveScopeVariables, but that's insufficient bc the 
incorrect ranges could also influence InferMutableRanges. AnalyzeFunctions is 
the point at which we compute the ranges for identifiers in the inner function, 
so it's the most ideal place to clean those up so they don't influence the outer 
function. 

* In InferMutableLifetimes, we need to ensure that context variable identifiers 
end up with a mutable range starting where they are declared, and ending with 
their last assignment. We now track declarations and extend their mutable range 
to account for each reassignment.

compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts

@@ -633,10 +633,14 @@ function isMutable(range: MutableRange): boolean {
 }
 
 function printMutableRange(identifier: Identifier): string {
-  const range =
-    identifier.scope !== null
-      ? identifier.scope.range
-      : identifier.mutableRange;
+  const range = identifier.mutableRange;
+  const scopeRange = identifier.scope?.range;
+  if (
+    scopeRange != null &&
+    (scopeRange.start !== range.start || scopeRange.end !== range.end)
+  ) {
+    return `[${range.start}:${range.end}] scope=[${scopeRange.start}:${scopeRange.end}]`;
+  }
   return isMutable(range) ? `[${range.start}:${range.end}]` : "";
 }
 

compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts

@@ -32,8 +32,9 @@ export function* eachInstructionValueLValue(
   value: ReactiveValue
 ): Iterable<Place> {
   switch (value.kind) {
-    case "DeclareLocal":
     case "DeclareContext":
+    case "StoreContext":
+    case "DeclareLocal":
     case "StoreLocal": {
       yield value.lvalue.place;
       break;

compiler/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts

@@ -16,6 +16,7 @@ import {
   ReactiveScopeDependency,
   isRefValueType,
   isUseRefType,
+  makeInstructionId,
 } from "../HIR";
 import { deadCodeElimination } from "../Optimization";
 import { inferReactiveScopeVariables } from "../ReactiveScopes";
@@ -126,7 +127,6 @@ function infer(
     ) {
       mutations.set(operand.identifier.name.value, operand.effect);
     }
-    operand.identifier.mutableRange.end = operand.identifier.mutableRange.start;
   }
 
   for (const dep of loweredFunc.dependencies) {
@@ -178,6 +178,12 @@ function infer(
       loweredFunc.dependencies.push(place);
     }
   }
+
+  for (const operand of loweredFunc.func.context) {
+    operand.identifier.mutableRange.start = makeInstructionId(0);
+    operand.identifier.mutableRange.end = makeInstructionId(0);
+    operand.identifier.scope = null;
+  }
 }
 
 function isMutatedOrReassigned(id: Identifier): boolean {

compiler/packages/babel-plugin-react-forget/src/Inference/InferMutableLifetimes.ts

@@ -8,7 +8,9 @@
 import {
   Effect,
   HIRFunction,
+  Identifier,
   InstructionId,
+  InstructionKind,
   makeInstructionId,
   Place,
 } from "../HIR/HIR";
@@ -99,6 +101,15 @@ export function inferMutableLifetimes(
   func: HIRFunction,
   inferMutableRangeForStores: boolean
 ): void {
+  /*
+   * Context variables only appear to mutate where they are assigned, but we need
+   * to force their range to start at their declaration. Track the declaring instruction
+   * id so that the ranges can be extended if/when they are reassigned
+   */
+  const contextVariableDeclarationInstructions = new Map<
+    Identifier,
+    InstructionId
+  >();
   for (const [_, block] of func.body.blocks) {
     for (const phi of block.phis) {
       for (const [_, operand] of phi.operands) {
@@ -146,6 +157,35 @@ export function inferMutableLifetimes(
       for (const operand of eachInstructionOperand(instr)) {
         inferPlace(operand, instr.id, inferMutableRangeForStores);
       }
+
+      if (
+        instr.value.kind === "DeclareContext" ||
+        (instr.value.kind === "StoreContext" &&
+          instr.value.lvalue.kind !== InstructionKind.Reassign)
+      ) {
+        // Save declarations of context variables
+        contextVariableDeclarationInstructions.set(
+          instr.value.lvalue.place.identifier,
+          instr.id
+        );
+      } else if (instr.value.kind === "StoreContext") {
+        /*
+         * Else this is a reassignment, extend the range from the declaration (if present).
+         * Note that declarations may not be present for context variables that are reassigned
+         * within a function expression before (or without) a read of the same variable
+         */
+        const declaration = contextVariableDeclarationInstructions.get(
+          instr.value.lvalue.place.identifier
+        );
+        if (declaration != null) {
+          const range = instr.value.lvalue.place.identifier.mutableRange;
+          if (range.start === 0) {
+            range.start = declaration;
+          } else {
+            range.start = makeInstructionId(Math.min(range.start, declaration));
+          }
+        }
+      }
     }
     for (const operand of eachTerminalOperand(block.terminal)) {
       inferPlace(operand, block.terminal.id, inferMutableRangeForStores);

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -5,21 +5,21 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import { CompilerError } from "..";
 import { Environment } from "../HIR";
 import {
+  GeneratedSource,
   HIRFunction,
   Identifier,
   Instruction,
-  makeInstructionId,
   Place,
   ReactiveScope,
+  makeInstructionId,
 } from "../HIR/HIR";
 import {
   doesPatternContainSpreadElement,
-  eachInstructionLValue,
   eachInstructionOperand,
   eachPatternOperand,
-  eachTerminalOperand,
 } from "../HIR/visitors";
 import DisjointSet from "../Utils/DisjointSet";
 import { assertExhaustive } from "../Utils/utils";
@@ -81,27 +81,6 @@ import { assertExhaustive } from "../Utils/utils";
  * ```
  */
 export function inferReactiveScopeVariables(fn: HIRFunction): void {
-  // First reset any scopes that may have been created from inner functions
-  for (const [, block] of fn.body.blocks) {
-    for (const phi of block.phis) {
-      phi.id.scope = null;
-      for (const [, operand] of phi.operands) {
-        operand.scope = null;
-      }
-    }
-    for (const instr of block.instructions) {
-      for (const lvalue of eachInstructionLValue(instr)) {
-        lvalue.identifier.scope = null;
-      }
-      for (const operand of eachInstructionOperand(instr)) {
-        operand.identifier.scope = null;
-      }
-    }
-    for (const operand of eachTerminalOperand(block.terminal)) {
-      operand.identifier.scope = null;
-    }
-  }
-
   /*
    * Represents the set of reactive scopes as disjoint sets of identifiers
    * that mutate together.
@@ -141,7 +120,40 @@ export function inferReactiveScopeVariables(fn: HIRFunction): void {
       );
     }
     identifier.scope = scope;
+    identifier.mutableRange = scope.range;
   });
+
+  let maxInstruction = 0;
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      maxInstruction = makeInstructionId(Math.max(maxInstruction, instr.id));
+    }
+    maxInstruction = makeInstructionId(
+      Math.max(maxInstruction, block.terminal.id)
+    );
+  }
+
+  /*
+   * Validate that all scopes have properly intialized, valid mutable ranges
+   * within the span of instructions for this function, ie from 1 to 1 past
+   * the last instruction id.
+   */
+  for (const [, scope] of scopes) {
+    if (
+      scope.range.start === 0 ||
+      scope.range.end === 0 ||
+      maxInstruction === 0 ||
+      scope.range.end > maxInstruction + 1
+    ) {
+      CompilerError.invariant(false, {
+        reason: `Invalid mutable range for scope`,
+        loc: GeneratedSource,
+        description: `Scope @${scope.id} has range [${scope.range.start}:${
+          scope.range.end
+        }] but the valid range is [1:${maxInstruction + 1}]`,
+      });
+    }
+  }
 }
 
 // Is the operand mutable at this given instruction

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructure-array-declaration-to-context-var.expect.md

@@ -27,32 +27,31 @@ import { unstable_useMemoCache as useMemoCache } from "react";
 import { identity } from "shared-runtime";
 
 function Component(props) {
-  const $ = useMemoCache(5);
-  const [t0] = props.value;
+  const $ = useMemoCache(4);
   let x;
-  if ($[0] !== t0 || $[1] !== props.value) {
+  if ($[0] !== props.value) {
+    const [t0] = props.value;
     x = t0;
     const foo = () => {
       x = identity(props.value[0]);
     };
 
     foo();
-    $[0] = t0;
-    $[1] = props.value;
-    $[2] = x;
+    $[0] = props.value;
+    $[1] = x;
   } else {
-    x = $[2];
+    x = $[1];
   }
-  const t1 = x;
-  let t2;
-  if ($[3] !== t1) {
-    t2 = { x: t1 };
+  const t0 = x;
+  let t1;
+  if ($[2] !== t0) {
+    t1 = { x: t0 };
+    $[2] = t0;
     $[3] = t1;
-    $[4] = t2;
   } else {
-    t2 = $[4];
+    t1 = $[3];
   }
-  return t2;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructure-object-declaration-to-context-var.expect.md

@@ -27,32 +27,31 @@ import { unstable_useMemoCache as useMemoCache } from "react";
 import { identity } from "shared-runtime";
 
 function Component(props) {
-  const $ = useMemoCache(5);
-  const { x: t0 } = props;
+  const $ = useMemoCache(4);
   let x;
-  if ($[0] !== t0 || $[1] !== props.x) {
+  if ($[0] !== props) {
+    const { x: t0 } = props;
     x = t0;
     const foo = () => {
       x = identity(props.x);
     };
 
     foo();
-    $[0] = t0;
-    $[1] = props.x;
-    $[2] = x;
+    $[0] = props;
+    $[1] = x;
   } else {
-    x = $[2];
+    x = $[1];
   }
-  const t1 = x;
-  let t2;
-  if ($[3] !== t1) {
-    t2 = { x: t1 };
+  const t0 = x;
+  let t1;
+  if ($[2] !== t0) {
+    t1 = { x: t0 };
+    $[2] = t0;
     $[3] = t1;
-    $[4] = t2;
   } else {
-    t2 = $[4];
+    t1 = $[3];
   }
-  return t2;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hoisting-nested-const-declaration-2.expect.md

@@ -27,7 +27,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function hoisting(cond) {
-  const $ = useMemoCache(3);
+  const $ = useMemoCache(2);
   let items;
   if ($[0] !== cond) {
     items = [];
@@ -35,14 +35,8 @@ function hoisting(cond) {
       const foo = () => {
         items.push(bar());
       };
-      let t0;
-      if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-        t0 = () => true;
-        $[2] = t0;
-      } else {
-        t0 = $[2];
-      }
-      const bar = t0;
+
+      const bar = () => true;
       foo();
     }
     $[0] = cond;