Fix bug with partially memoized destructuring

compiler/forget/src/CompilerPipeline.ts

@@ -32,6 +32,7 @@ import {
   buildReactiveBlocks,
   buildReactiveFunction,
   codegenReactiveFunction,
+  extractScopeDeclarationsFromDestructuring,
   flattenReactiveLoops,
   flattenScopesWithHooks,
   inferReactiveScopeVariables,
@@ -223,6 +224,13 @@ export function* run(
     value: reactiveFunction,
   });
 
+  extractScopeDeclarationsFromDestructuring(reactiveFunction);
+  yield log({
+    kind: "reactive",
+    name: "ExtractScopeDeclarationsFromDestructuring",
+    value: reactiveFunction,
+  });
+
   renameVariables(reactiveFunction);
   yield log({
     kind: "reactive",

compiler/forget/src/HIR/HIR.ts

@@ -585,12 +585,7 @@ export type InstructionValue =
       value: Place;
       loc: SourceLocation;
     }
-  | {
-      kind: "Destructure";
-      lvalue: LValuePattern;
-      value: Place;
-      loc: SourceLocation;
-    }
+  | Destructure
   | {
       kind: "Primitive";
       value: number | boolean | string | null | undefined;
@@ -756,6 +751,14 @@ export type FunctionExpression = {
   expr: t.ArrowFunctionExpression | t.FunctionExpression;
   loc: SourceLocation;
 };
+
+export type Destructure = {
+  kind: "Destructure";
+  lvalue: LValuePattern;
+  value: Place;
+  loc: SourceLocation;
+};
+
 /**
  * A place where data may be read from / written to:
  * - a variable (identifier)

compiler/forget/src/ReactiveScopes/ExtractScopeDeclarationsFromDestructuring.ts

@@ -0,0 +1,184 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {
+  Destructure,
+  Environment,
+  IdentifierId,
+  InstructionKind,
+  Place,
+  ReactiveBlock,
+  ReactiveFunction,
+  ReactiveInstruction,
+  ReactiveScopeBlock,
+} from "../HIR";
+import { eachPatternOperand, mapPatternOperands } from "../HIR/visitors";
+import { ReactiveFunctionTransform, visitReactiveFunction } from "./visitors";
+
+/**
+ * Destructuring statements may sometimes define some variables which are declared by the scope,
+ * and others that are only used locally within the scope, for example:
+ *
+ * ```
+ * const {x, ...rest} = value;
+ * return rest;
+ * ```
+ *
+ * Here the scope structure turns into:
+ *
+ * ```
+ * let c_0 = $[0] !== value;
+ * let rest;
+ * if (c_0) {
+ *   // OOPS! we want to reassign `rest` here, but
+ *   // `x` isn't declared anywhere!
+ *   {x, ...rest} = value;
+ *   $[0] = value;
+ *   $[1] = rest;
+ * } else {
+ *   rest = $[1];
+ * }
+ * return rest;
+ * ```
+ *
+ * Note that because `rest` is declared by the scope, we can't redeclare it in the
+ * destructuring statement. But we have to declare `x`!
+ *
+ * This pass finds destructuring instructions that contain mixed values such as this,
+ * and rewrites them to ensure that any scope variable assignments are extracted first
+ * to a temporary and reassigned in a separate instruction. For example, the output
+ * for the above would be along the lines of:
+ *
+ * ```
+ * let c_0 = $[0] !== value;
+ * let rest;
+ * if (c_0) {
+ *   const {x, ...t0} = value; <-- replace `rest` with a temporary
+ *   rest = t0; // <-- and create a separate instruction to assign that to `rest`
+ *   $[0] = value;
+ *   $[1] = rest;
+ * } else {
+ *   rest = $[1];
+ * }
+ * return rest;
+ * ```
+ *
+ */
+export function extractScopeDeclarationsFromDestructuring(
+  fn: ReactiveFunction
+): void {
+  const state = new State(fn.env);
+  visitReactiveFunction(fn, new Visitor(), state);
+}
+
+class State {
+  env: Environment;
+  declared: Set<IdentifierId> = new Set();
+
+  constructor(env: Environment) {
+    this.env = env;
+  }
+}
+
+class Visitor extends ReactiveFunctionTransform<State> {
+  override visitScope(scope: ReactiveScopeBlock, state: State): void {
+    for (const [, declaration] of scope.scope.declarations) {
+      state.declared.add(declaration.identifier.id);
+    }
+    this.traverseScope(scope, state);
+  }
+
+  override visitBlock(block: ReactiveBlock, state: State): void {
+    // Traverse first to transform inner items
+    this.traverseBlock(block, state);
+
+    // Then transform any mixed destructuring instructions
+    let nextBlock: ReactiveBlock | null = null;
+    for (let i = 0; i < block.length; i++) {
+      const instr = block[i];
+      if (
+        instr.kind === "instruction" &&
+        instr.instruction.value.kind === "Destructure"
+      ) {
+        const transformed = transformDestructuring(
+          state,
+          instr.instruction,
+          instr.instruction.value
+        );
+        if (transformed) {
+          nextBlock ??= block.slice(0, i);
+          transformed.forEach((instruction) => {
+            nextBlock?.push({
+              kind: "instruction",
+              instruction,
+            });
+          });
+          continue;
+        }
+      } else if (nextBlock !== null) {
+        nextBlock.push(instr);
+      }
+    }
+    if (nextBlock !== null) {
+      block.length = 0;
+      block.push(...nextBlock);
+    }
+  }
+}
+
+function transformDestructuring(
+  state: State,
+  instr: ReactiveInstruction,
+  destructure: Destructure
+): null | Array<ReactiveInstruction> {
+  let reassigned: Set<IdentifierId> = new Set();
+  let hasDeclaration = false;
+  for (const place of eachPatternOperand(destructure.lvalue.pattern)) {
+    const isDeclared = state.declared.has(place.identifier.id);
+    if (isDeclared) {
+      reassigned.add(place.identifier.id);
+    }
+    hasDeclaration ||= !isDeclared;
+  }
+  if (reassigned.size === 0 || !hasDeclaration) {
+    return null;
+  }
+  // Else it's a mix, replace the reassigned items in the destructuring with temporary
+  // variables and emit separate assignment statements for them
+  const instructions: Array<ReactiveInstruction> = [];
+  const renamed: Map<Place, Place> = new Map();
+  mapPatternOperands(destructure.lvalue.pattern, (place) => {
+    if (!reassigned.has(place.identifier.id)) {
+      return place;
+    }
+    const tempId = state.env.nextIdentifierId;
+    const temporary = {
+      ...place,
+      identifier: { ...place.identifier, id: tempId, name: `t${tempId}` },
+    };
+    renamed.set(place, temporary);
+    return temporary;
+  });
+  instructions.push(instr);
+  for (const [original, temporary] of renamed) {
+    instructions.push({
+      id: instr.id,
+      lvalue: null,
+      value: {
+        kind: "StoreLocal",
+        lvalue: {
+          kind: InstructionKind.Reassign,
+          place: original,
+        },
+        value: temporary,
+        loc: destructure.loc,
+      },
+      loc: instr.loc,
+    });
+  }
+  return instructions;
+}

compiler/forget/src/ReactiveScopes/index.ts

@@ -9,6 +9,7 @@ export { alignReactiveScopesToBlockScopes } from "./AlignReactiveScopesToBlockSc
 export { buildReactiveBlocks } from "./BuildReactiveBlocks";
 export { buildReactiveFunction } from "./BuildReactiveFunction";
 export { codegenReactiveFunction } from "./CodegenReactiveFunction";
+export { extractScopeDeclarationsFromDestructuring } from "./ExtractScopeDeclarationsFromDestructuring";
 export { flattenReactiveLoops } from "./FlattenReactiveLoops";
 export { flattenScopesWithHooks } from "./FlattenScopesWithHooks";
 export { inferReactiveScopeVariables } from "./InferReactiveScopeVariables";

compiler/forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.expect.md

@@ -0,0 +1,101 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const post = useFragment(graphql`...`, props.post);
+  const allUrls = [];
+  // `media` and `urls` are exported from the scope that will wrap this code,
+  // but `comments` is not (it doesn't need to be memoized, bc the callback
+  // only checks `comments.length`)
+  // because of the scope, the let declaration for media and urls are lifted
+  // out of the scope, and the destructure statement ends up turning into
+  // a reassignment, instead of a const declaration. this means we try to
+  // reassign `comments` when there's no declaration for it.
+  const { media = null, comments = [], urls = [] } = post;
+  const onClick = (e) => {
+    if (!comments.length) {
+      return;
+    }
+    log(comments.length);
+  };
+  allUrls.push(...urls);
+  return <Media media={media} onClick={onClick} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(12);
+  const post = useFragment(graphql`...`, props.post);
+  const c_0 = $[0] !== post;
+  let media;
+  let onClick;
+  if (c_0) {
+    const allUrls = [];
+
+    const { media: t0, comments: t2, urls: t81 } = post;
+    const c_3 = $[3] !== t0;
+    let t1;
+    if (c_3) {
+      t1 = t0 === undefined ? null : t0;
+      $[3] = t0;
+      $[4] = t1;
+    } else {
+      t1 = $[4];
+    }
+    media = t1;
+    const c_5 = $[5] !== t2;
+    let t3;
+    if (c_5) {
+      t3 = t2 === undefined ? [] : t2;
+      $[5] = t2;
+      $[6] = t3;
+    } else {
+      t3 = $[6];
+    }
+    const comments = t3;
+    const urls = t81 === undefined ? [] : t81;
+    const c_7 = $[7] !== comments.length;
+    let t4;
+    if (c_7) {
+      t4 = (e) => {
+        if (!comments.length) {
+          return;
+        }
+        log(comments.length);
+      };
+      $[7] = comments.length;
+      $[8] = t4;
+    } else {
+      t4 = $[8];
+    }
+    onClick = t4;
+    allUrls.push(...urls);
+    $[0] = post;
+    $[1] = media;
+    $[2] = onClick;
+  } else {
+    media = $[1];
+    onClick = $[2];
+  }
+  const c_9 = $[9] !== media;
+  const c_10 = $[10] !== onClick;
+  let t5;
+  if (c_9 || c_10) {
+    t5 = <Media media={media} onClick={onClick} />;
+    $[9] = media;
+    $[10] = onClick;
+    $[11] = t5;
+  } else {
+    t5 = $[11];
+  }
+  return t5;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.js

@@ -1,7 +1,3 @@
-
-## Input
-
-```javascript
 function Component(props) {
   const post = useFragment(graphql`...`, props.post);
   const allUrls = [];
@@ -12,7 +8,7 @@ function Component(props) {
   // out of the scope, and the destructure statement ends up turning into
   // a reassignment, instead of a const declaration. this means we try to
   // reassign `comments` when there's no declaration for it.
-  const { media, comments, urls } = post;
+  const { media = null, comments = [], urls = [] } = post;
   const onClick = (e) => {
     if (!comments.length) {
       return;
@@ -22,14 +18,3 @@ function Component(props) {
   allUrls.push(...urls);
   return <Media media={media} onClick={onClick} />;
 }
-
-```
-
-
-## Error
-
-```
-[ReactForget] Invariant: Encountered a destructuring operation where some identifiers are already declared (reassignments) but others are not (declarations) (11:11)
-```
-          
-      

compiler/forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md

@@ -0,0 +1,81 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const post = useFragment(graphql`...`, props.post);
+  const allUrls = [];
+  // `media` and `urls` are exported from the scope that will wrap this code,
+  // but `comments` is not (it doesn't need to be memoized, bc the callback
+  // only checks `comments.length`)
+  // because of the scope, the let declaration for media and urls are lifted
+  // out of the scope, and the destructure statement ends up turning into
+  // a reassignment, instead of a const declaration. this means we try to
+  // reassign `comments` when there's no declaration for it.
+  const { media, comments, urls } = post;
+  const onClick = (e) => {
+    if (!comments.length) {
+      return;
+    }
+    log(comments.length);
+  };
+  allUrls.push(...urls);
+  return <Media media={media} onClick={onClick} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(8);
+  const post = useFragment(graphql`...`, props.post);
+  const c_0 = $[0] !== post;
+  let media;
+  let onClick;
+  if (c_0) {
+    const allUrls = [];
+
+    const { media: t83, comments, urls } = post;
+    media = t83;
+    const c_3 = $[3] !== comments.length;
+    let t0;
+    if (c_3) {
+      t0 = (e) => {
+        if (!comments.length) {
+          return;
+        }
+        log(comments.length);
+      };
+      $[3] = comments.length;
+      $[4] = t0;
+    } else {
+      t0 = $[4];
+    }
+    onClick = t0;
+    allUrls.push(...urls);
+    $[0] = post;
+    $[1] = media;
+    $[2] = onClick;
+  } else {
+    media = $[1];
+    onClick = $[2];
+  }
+  const c_5 = $[5] !== media;
+  const c_6 = $[6] !== onClick;
+  let t1;
+  if (c_5 || c_6) {
+    t1 = <Media media={media} onClick={onClick} />;
+    $[5] = media;
+    $[6] = onClick;
+    $[7] = t1;
+  } else {
+    t1 = $[7];
+  }
+  return t1;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/error.escape-analysis-destructured-rest-element.expect.md

@@ -1,22 +0,0 @@
-
-## Input
-
-```javascript
-function Component(props) {
-  // b is an object, must be memoized even though the input is not memoized
-  const { a, ...b } = props.a;
-  // d is an array, mut be memoized even though the input is not memoized
-  const [c, ...d] = props.c;
-  return <div b={b} d={d}></div>;
-}
-
-```
-
-
-## Error
-
-```
-[ReactForget] Invariant: Encountered a destructuring operation where some identifiers are already declared (reassignments) but others are not (declarations) (3:3)
-```
-          
-      

compiler/forget/src/__tests__/fixtures/compiler/error.unused-object-element-with-rest.expect.md

@@ -1,20 +0,0 @@
-
-## Input
-
-```javascript
-function Foo(props) {
-  // can't remove `unused` since it affects which properties are copied into `rest`
-  const { unused, ...rest } = props.a;
-  return rest;
-}
-
-```
-
-
-## Error
-
-```
-[ReactForget] Invariant: Encountered a destructuring operation where some identifiers are already declared (reassignments) but others are not (declarations) (3:3)
-```
-          
-      

compiler/forget/src/__tests__/fixtures/compiler/escape-analysis-destructured-rest-element.expect.md

@@ -0,0 +1,56 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  // b is an object, must be memoized even though the input is not memoized
+  const { a, ...b } = props.a;
+  // d is an array, mut be memoized even though the input is not memoized
+  const [c, ...d] = props.c;
+  return <div b={b} d={d}></div>;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(7);
+  const c_0 = $[0] !== props.a;
+  let b;
+  if (c_0) {
+    const { a, ...t30 } = props.a;
+    b = t30;
+    $[0] = props.a;
+    $[1] = b;
+  } else {
+    b = $[1];
+  }
+  const c_2 = $[2] !== props.c;
+  let d;
+  if (c_2) {
+    const [c, ...t31] = props.c;
+    d = t31;
+    $[2] = props.c;
+    $[3] = d;
+  } else {
+    d = $[3];
+  }
+  const c_4 = $[4] !== b;
+  const c_5 = $[5] !== d;
+  let t0;
+  if (c_4 || c_5) {
+    t0 = <div b={b} d={d} />;
+    $[4] = b;
+    $[5] = d;
+    $[6] = t0;
+  } else {
+    t0 = $[6];
+  }
+  return t0;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/unused-object-element-with-rest.expect.md

@@ -0,0 +1,33 @@
+
+## Input
+
+```javascript
+function Foo(props) {
+  // can't remove `unused` since it affects which properties are copied into `rest`
+  const { unused, ...rest } = props.a;
+  return rest;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Foo(props) {
+  const $ = useMemoCache(2);
+  const c_0 = $[0] !== props.a;
+  let rest;
+  if (c_0) {
+    const { unused, ...t16 } = props.a;
+    rest = t16;
+    $[0] = props.a;
+    $[1] = rest;
+  } else {
+    rest = $[1];
+  }
+  return rest;
+}
+
+```
+