Todo for early return within reactive scopes

Adds a new compiler pass that will eventually actually handle early returns within reactive scopes. For now it just detects them and throws a Todo error.
2025-11-01 09:12:30 +00:00 · 2023-12-20 13:52:37 -08:00
parent f504eaa16e
commit 4c68da2e60
23 changed files with 492 additions and 333 deletions
@@ -49,6 +49,7 @@ import {
  mergeOverlappingReactiveScopes,
  mergeReactiveScopesThatInvalidateTogether,
  promoteUsedTemporaries,
+  propagateEarlyReturns,
  propagateScopeDependencies,
  pruneAllReactiveScopes,
  pruneHoistedContexts,
@@ -283,7 +284,7 @@ function* runWithEnvironment(
  pruneNonEscapingScopes(reactiveFunction);
  yield log({
    kind: "reactive",
-    name: "PruneNonEscapingDependencies",
+    name: "PruneNonEscapingScopes",
    value: reactiveFunction,
  });

@@ -294,6 +295,13 @@ function* runWithEnvironment(
    value: reactiveFunction,
  });

+  propagateEarlyReturns(reactiveFunction);
+  yield log({
+    kind: "reactive",
+    name: "PropagateEarlyReturns",
+    value: reactiveFunction,
+  });
+
  pruneUnusedScopes(reactiveFunction);
  yield log({
    kind: "reactive",
@@ -0,0 +1,120 @@
+/*
+ * 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 { visitReactiveFunction } from ".";
+import { CompilerError } from "..";
+import {
+  ReactiveFunction,
+  ReactiveScopeBlock,
+  ReactiveTerminalStatement,
+} from "../HIR";
+import { ReactiveFunctionVisitor } from "./visitors";
+
+/**
+ * TODO: Actualy propagate early return information, for now we throw a Todo bailout.
+ *
+ * This pass ensures that reactive blocks honor the control flow behavior of the
+ * original code including early return semantics. Specifically, if a reactive
+ * scope early returned during the previous execution and the inputs to that block
+ * have not changed, then the code should early return (with the same value) again.
+ *
+ * Example:
+ *
+ * ```javascript
+ * let x = [];
+ * if (props.cond) {
+ *   x.push(12);
+ *   return x;
+ * } else {
+ *   return foo();
+ * }
+ * ```
+ *
+ * Imagine that this code is called twice in a row with props.cond = true. Both
+ * times it should return the same object (===), an array `[12]`.
+ *
+ * The compilation strategy is as follows. For each top-level reactive scope
+ * that contains (transitively) an early return:
+ *
+ * - Label the scope
+ * - Synthesize a new temporary, eg `t0`, and set it as a declaration of the scope.
+ *   This will represent the possibly-unset return value for that scope.
+ * - Make the first instruction of the scope a reassignment of that temporary,
+ *   assigning a sentinel value (can reuse the same symbol as we use for cache slots).
+ *   This assignment ensures that if we don't take an early return, that the value
+ *   is the sentinel.
+ * - Replace all `return` statements with:
+ *   - An assignment of the temporary with the value being returned.
+ *   - An assignment of the temporary into a cache slot, so it can be retrieved in the
+ *     scope's "else" branch.
+ *   - A `break` to the reactive scope's label.
+ * - Finally, add code _after_ the reactive scope that checks the temporary. If
+ *   it equals the sentinel value do nothing; else return its value.
+ *
+ * For the above example that looks roughly like:
+ *
+ * ```javascript
+ * let t0; // temporary for early return;
+ * bb1: if (props.cond !== $[0]) {
+ *   // reset the temporary
+ *   t0 = Symbol.for('react.forget');
+ *   // original code
+ *   let x = [];
+ *   if (props.cond) {
+ *     x.push(12);
+ *     // replace the early return w assignment and break
+ *     t0 = x;
+ *     $[2] = t0;
+ *     break bb1
+ *   } else {
+ *     let t1;
+ *     if ($[1] === Symbol.for('react.forget')) {
+ *       t1 = foo();
+ *       $[1] = t1;
+ *     } else {
+ *       t1 = $[1];
+ *     }
+ *     // Replace early return w assignment and break;
+ *     t0 = t1;
+ *     $[2] = t0;
+ *     break bb1;
+ *   }
+ * } else {
+ *   t0 = $[2];
+ * }
+ * if (t0 !== Symbol.for('react.forget')) {
+ *   return t0;
+ * }
+ * ```
+ */
+export function propagateEarlyReturns(fn: ReactiveFunction): void {
+  visitReactiveFunction(fn, new Visitor(), false);
+}
+
+class Visitor extends ReactiveFunctionVisitor<boolean> {
+  override visitScope(
+    scopeBlock: ReactiveScopeBlock,
+    _withinReactiveScope: boolean
+  ): void {
+    this.traverseScope(scopeBlock, true);
+  }
+
+  override visitTerminal(
+    stmt: ReactiveTerminalStatement,
+    withinReactiveScope: boolean
+  ): void {
+    if (withinReactiveScope && stmt.terminal.kind === "return") {
+      CompilerError.throwTodo({
+        reason: `Support early return within a reactive scope`,
+        loc: stmt.terminal.value.loc,
+        description: null,
+        suggestions: null,
+      });
+    }
+    this.traverseTerminal(stmt, withinReactiveScope);
+  }
+}
@@ -872,19 +872,35 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
  Set<IdentifierId>
 > {
  override transformScope(
-    scope: ReactiveScopeBlock,
+    scopeBlock: ReactiveScopeBlock,
    state: Set<IdentifierId>
  ): Transformed<ReactiveStatement> {
-    this.visitScope(scope, state);
+    this.visitScope(scopeBlock, state);
+
+    /**
+     * Scopes may initially appear "empty" because the value being memoized
+     * is early-returned from within the scope. For now we intentionaly keep
+     * these scopes, and let them get pruned later by PruneUnusedScopes
+     * _after_ handling the early-return case in PropagateEarlyReturns.
+     */
+    if (
+      scopeBlock.scope.declarations.size === 0 &&
+      scopeBlock.scope.reassignments.size === 0
+    ) {
+      return { kind: "keep" };
+    }
+
    const hasMemoizedOutput =
-      Array.from(scope.scope.declarations.keys()).some((id) => state.has(id)) ||
-      Array.from(scope.scope.reassignments).some((identifier) =>
+      Array.from(scopeBlock.scope.declarations.keys()).some((id) =>
+        state.has(id)
+      ) ||
+      Array.from(scopeBlock.scope.reassignments).some((identifier) =>
        state.has(identifier.id)
      );
    if (hasMemoizedOutput) {
      return { kind: "keep" };
    } else {
-      return { kind: "replace-many", value: scope.instructions };
+      return { kind: "replace-many", value: scopeBlock.instructions };
    }
  }
 }
@@ -23,6 +23,7 @@ export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes
 export { mergeReactiveScopesThatInvalidateTogether } from "./MergeReactiveScopesThatInvalidateTogether";
 export { printReactiveFunction } from "./PrintReactiveFunction";
 export { promoteUsedTemporaries } from "./PromoteUsedTemporaries";
+export { propagateEarlyReturns } from "./PropagateEarlyReturns";
 export { propagateScopeDependencies } from "./PropagateScopeDependencies";
 export { pruneAllReactiveScopes } from "./PruneAllReactiveScopes";
 export { pruneHoistedContexts } from "./PruneHoistedContexts";
@@ -1,158 +0,0 @@
-
-## Input
-
-```javascript
-/**
- * props.b does *not* influence `a`
- */
-function ComponentA(props) {
-  const a_DEBUG = [];
-  a_DEBUG.push(props.a);
-  if (props.b) {
-    return null;
-  }
-  a_DEBUG.push(props.d);
-  return a_DEBUG;
-}
-
-/**
- * props.b *does* influence `a`
- */
-function ComponentB(props) {
-  const a = [];
-  a.push(props.a);
-  if (props.b) {
-    a.push(props.c);
-  }
-  a.push(props.d);
-  return a;
-}
-
-/**
- * props.b *does* influence `a`, but only in a way that is never observable
- */
-function ComponentC(props) {
-  const a = [];
-  a.push(props.a);
-  if (props.b) {
-    a.push(props.c);
-    return null;
-  }
-  a.push(props.d);
-  return a;
-}
-
-/**
- * props.b *does* influence `a`
- */
-function ComponentD(props) {
-  const a = [];
-  a.push(props.a);
-  if (props.b) {
-    a.push(props.c);
-    return a;
-  }
-  a.push(props.d);
-  return a;
-}
-
-```
-
-## Code
-
-```javascript
-import { unstable_useMemoCache as useMemoCache } from "react";
-/**
- * props.b does *not* influence `a`
- */
-function ComponentA(props) {
-  const $ = useMemoCache(4);
-  let a_DEBUG;
-  if ($[0] !== props.a || $[1] !== props.b || $[2] !== props.d) {
-    a_DEBUG = [];
-    a_DEBUG.push(props.a);
-    if (props.b) {
-      return null;
-    }
-
-    a_DEBUG.push(props.d);
-    $[0] = props.a;
-    $[1] = props.b;
-    $[2] = props.d;
-    $[3] = a_DEBUG;
-  } else {
-    a_DEBUG = $[3];
-  }
-  return a_DEBUG;
-}
-
-/**
- * props.b *does* influence `a`
- */
-function ComponentB(props) {
-  const $ = useMemoCache(2);
-  let a;
-  if ($[0] !== props) {
-    a = [];
-    a.push(props.a);
-    if (props.b) {
-      a.push(props.c);
-    }
-
-    a.push(props.d);
-    $[0] = props;
-    $[1] = a;
-  } else {
-    a = $[1];
-  }
-  return a;
-}
-
-/**
- * props.b *does* influence `a`, but only in a way that is never observable
- */
-function ComponentC(props) {
-  const $ = useMemoCache(2);
-  let a;
-  if ($[0] !== props) {
-    a = [];
-    a.push(props.a);
-    if (props.b) {
-      a.push(props.c);
-      return null;
-    }
-
-    a.push(props.d);
-    $[0] = props;
-    $[1] = a;
-  } else {
-    a = $[1];
-  }
-  return a;
-}
-
-/**
- * props.b *does* influence `a`
- */
-function ComponentD(props) {
-  const $ = useMemoCache(2);
-  let a;
-  if ($[0] !== props) {
-    a = [];
-    a.push(props.a);
-    if (props.b) {
-      a.push(props.c);
-      return a;
-    }
-
-    a.push(props.d);
-    $[0] = props;
-    $[1] = a;
-  } else {
-    a = $[1];
-  }
-  return a;
-}
-
-```
-      
@@ -0,0 +1,68 @@
+
+## Input
+
+```javascript
+/**
+ * props.b does *not* influence `a`
+ */
+function ComponentA(props) {
+  const a_DEBUG = [];
+  a_DEBUG.push(props.a);
+  if (props.b) {
+    return null;
+  }
+  a_DEBUG.push(props.d);
+  return a_DEBUG;
+}
+
+/**
+ * props.b *does* influence `a`
+ */
+function ComponentB(props) {
+  const a = [];
+  a.push(props.a);
+  if (props.b) {
+    a.push(props.c);
+  }
+  a.push(props.d);
+  return a;
+}
+
+/**
+ * props.b *does* influence `a`, but only in a way that is never observable
+ */
+function ComponentC(props) {
+  const a = [];
+  a.push(props.a);
+  if (props.b) {
+    a.push(props.c);
+    return null;
+  }
+  a.push(props.d);
+  return a;
+}
+
+/**
+ * props.b *does* influence `a`
+ */
+function ComponentD(props) {
+  const a = [];
+  a.push(props.a);
+  if (props.b) {
+    a.push(props.c);
+    return a;
+  }
+  a.push(props.d);
+  return a;
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] Todo: Support early return within a reactive scope (8:8)
+```
+          
+      
@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = [];
+  if (props.cond) {
+    x.push(props.a);
+    if (props.b) {
+      const y = [props.b];
+      x.push(y);
+      // oops no memo!
+      return x;
+    }
+    // oops no memo!
+    return x;
+  } else {
+    return foo();
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true, a: 42, b: 3.14 }],
+};
+
+```
+
+
+## Error
+
+```
+[ReactForget] Todo: Support early return within a reactive scope (9:9)
+```
+          
+      
@@ -0,0 +1,21 @@
+function Component(props) {
+  let x = [];
+  if (props.cond) {
+    x.push(props.a);
+    if (props.b) {
+      const y = [props.b];
+      x.push(y);
+      // oops no memo!
+      return x;
+    }
+    // oops no memo!
+    return x;
+  } else {
+    return foo();
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true, a: 42, b: 3.14 }],
+};
@@ -0,0 +1,30 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = [];
+  if (props.cond) {
+    x.push(props.a);
+    // oops no memo!
+    return x;
+  } else {
+    return foo();
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true, a: 42 }],
+};
+
+```
+
+
+## Error
+
+```
+[ReactForget] Todo: Support early return within a reactive scope (6:6)
+```
+          
+      
@@ -0,0 +1,15 @@
+function Component(props) {
+  let x = [];
+  if (props.cond) {
+    x.push(props.a);
+    // oops no memo!
+    return x;
+  } else {
+    return foo();
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true, a: 42 }],
+};
@@ -0,0 +1,35 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = [];
+  let y = null;
+  if (props.cond) {
+    x.push(props.a);
+    // oops no memo!
+    return x;
+  } else {
+    y = foo();
+    if (props.b) {
+      return;
+    }
+  }
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true, a: 42 }],
+};
+
+```
+
+
+## Error
+
+```
+[ReactForget] Todo: Support early return within a reactive scope (7:7)
+```
+          
+      
@@ -0,0 +1,20 @@
+function Component(props) {
+  let x = [];
+  let y = null;
+  if (props.cond) {
+    x.push(props.a);
+    // oops no memo!
+    return x;
+  } else {
+    y = foo();
+    if (props.b) {
+      return;
+    }
+  }
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ cond: true, a: 42 }],
+};
@@ -0,0 +1,33 @@
+
+## Input
+
+```javascript
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  try {
+    const y = [];
+    y.push(props.y);
+    throwInput(y);
+  } catch (e) {
+    e.push(props.e);
+    return e;
+  }
+  return null;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ y: "foo", e: "bar" }],
+};
+
+```
+
+
+## Error
+
+```
+[ReactForget] Todo: Support early return within a reactive scope (10:10)
+```
+          
+      
@@ -24,30 +24,11 @@ export const FIXTURE_ENTRYPOINT = {

 ```

-## Code

-```javascript
-const { throwInput } = require("shared-runtime");
-
-function Component(props) {
-  const x = [];
-  try {
-    throwInput(x);
-  } catch (t22) {
-    const e = t22;
-
-    e.push(null);
-    return e;
-  }
-  return x;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{}],
-};
+## Error

 ```
-      
-### Eval output
-(kind: ok) [null]
+[ReactForget] Todo: Support early return within a reactive scope (11:11)
+```
+          
+      
@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+const { shallowCopy, throwInput } = require("shared-runtime");
+
+// @debug
+function Component(props) {
+  let x = [];
+  try {
+    const y = shallowCopy({});
+    if (y == null) {
+      return;
+    }
+    x.push(throwInput(y));
+  } catch {
+    return null;
+  }
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+
+## Error
+
+```
+[ReactForget] Todo: Support early return within a reactive scope
+```
+          
+      
@@ -1,52 +0,0 @@
-
-## Input
-
-```javascript
-const { throwInput } = require("shared-runtime");
-
-function Component(props) {
-  try {
-    const y = [];
-    y.push(props.y);
-    throwInput(y);
-  } catch (e) {
-    e.push(props.e);
-    return e;
-  }
-  return null;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{ y: "foo", e: "bar" }],
-};
-
-```
-
-## Code
-
-```javascript
-const { throwInput } = require("shared-runtime");
-
-function Component(props) {
-  try {
-    const y = [];
-    y.push(props.y);
-    throwInput(y);
-  } catch (t25) {
-    const e = t25;
-    e.push(props.e);
-    return e;
-  }
-  return null;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{ y: "foo", e: "bar" }],
-};
-
-```
-      
-### Eval output
-(kind: ok) ["foo","bar"]
@@ -1,66 +0,0 @@
-
-## Input
-
-```javascript
-const { shallowCopy, throwInput } = require("shared-runtime");
-
-// @debug
-function Component(props) {
-  let x = [];
-  try {
-    const y = shallowCopy({});
-    if (y == null) {
-      return;
-    }
-    x.push(throwInput(y));
-  } catch {
-    return null;
-  }
-  return x;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{}],
-};
-
-```
-
-## Code
-
-```javascript
-import { unstable_useMemoCache as useMemoCache } from "react";
-const { shallowCopy, throwInput } = require("shared-runtime");
-
-// @debug
-function Component(props) {
-  const $ = useMemoCache(1);
-  let x;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    x = [];
-    try {
-      const y = shallowCopy({});
-      if (y == null) {
-        return;
-      }
-
-      x.push(throwInput(y));
-    } catch {
-      return null;
-    }
-    $[0] = x;
-  } else {
-    x = $[0];
-  }
-  return x;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{}],
-};
-
-```
-      
-### Eval output
-(kind: ok) null
@@ -2,7 +2,7 @@
 ## Input

 ```javascript
-// @validateNoSetStateInRender
+// @validateNoSetStateInRender @enableAssumeHooksFollowRulesOfReact
 function Component(props) {
  const logEvent = useLogging(props.appId);
  const [currentStep, setCurrentStep] = useState(0);
@@ -33,47 +33,62 @@ function Component(props) {
 ## Code

 ```javascript
-import { unstable_useMemoCache as useMemoCache } from "react"; // @validateNoSetStateInRender
+import { unstable_useMemoCache as useMemoCache } from "react"; // @validateNoSetStateInRender @enableAssumeHooksFollowRulesOfReact
 function Component(props) {
-  const $ = useMemoCache(3);
+  const $ = useMemoCache(7);
  const logEvent = useLogging(props.appId);
  const [currentStep, setCurrentStep] = useState(0);
-
-  const onSubmit = (errorEvent) => {
-    logEvent(errorEvent);
-    setCurrentStep(1);
-  };
+  let t0;
+  if ($[0] !== logEvent) {
+    t0 = (errorEvent) => {
+      logEvent(errorEvent);
+      setCurrentStep(1);
+    };
+    $[0] = logEvent;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const onSubmit = t0;
  switch (currentStep) {
    case 0: {
-      let t0;
-      if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-        t0 = <OtherComponent data={{ foo: "bar" }} />;
-        $[0] = t0;
+      let t1;
+      if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+        t1 = <OtherComponent data={{ foo: "bar" }} />;
+        $[2] = t1;
      } else {
-        t0 = $[0];
+        t1 = $[2];
      }
-      return t0;
+      return t1;
    }
    case 1: {
-      let t1;
-      if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-        t1 = { foo: "joe" };
-        $[1] = t1;
+      let t2;
+      if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
+        t2 = { foo: "joe" };
+        $[3] = t2;
      } else {
-        t1 = $[1];
+        t2 = $[3];
      }
-      return <OtherComponent data={t1} onSubmit={onSubmit} />;
+      let t3;
+      if ($[4] !== onSubmit) {
+        t3 = <OtherComponent data={t2} onSubmit={onSubmit} />;
+        $[4] = onSubmit;
+        $[5] = t3;
+      } else {
+        t3 = $[5];
+      }
+      return t3;
    }
    default: {
      logEvent("Invalid step");
-      let t2;
-      if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-        t2 = <OtherComponent data={null} />;
-        $[2] = t2;
+      let t4;
+      if ($[6] === Symbol.for("react.memo_cache_sentinel")) {
+        t4 = <OtherComponent data={null} />;
+        $[6] = t4;
      } else {
-        t2 = $[2];
+        t4 = $[6];
      }
-      return t2;
+      return t4;
    }
  }
 }
@@ -1,4 +1,4 @@
-// @validateNoSetStateInRender
+// @validateNoSetStateInRender @enableAssumeHooksFollowRulesOfReact
 function Component(props) {
  const logEvent = useLogging(props.appId);
  const [currentStep, setCurrentStep] = useState(0);