diff --git a/compiler/forget/src/HIR/FlattenReactiveLoops.ts b/compiler/forget/src/HIR/FlattenReactiveLoops.ts
new file mode 100644
index 0000000000..d07cb8a168
--- /dev/null
+++ b/compiler/forget/src/HIR/FlattenReactiveLoops.ts
@@ -0,0 +1,144 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { assertExhaustive } from "../Common/utils";
+import { ReactiveBasicBlock, ReactiveBlock, ReactiveFunction } from "./HIR";
+
+/**
+ * Given a reactive function, flattens any scopes contained within a loop construct.
+ * We won't initially support memoization within loops though this is possible in the future.
+ */
+export function flattenReactiveLoops(fn: ReactiveFunction): void {
+ visit(fn.body, false);
+}
+
+function visit(block: ReactiveBasicBlock, shouldFlatten: boolean): void {
+ let i = 0;
+ while (i < block.length) {
+ const item = block[i]!;
+ switch (item.kind) {
+ case "block": {
+ if (shouldFlatten) {
+ const successors = block.splice(i + 1);
+ block.pop(); // remove the current element
+ flatten(item, block);
+ i = block.length;
+ block.push(...successors);
+ } else {
+ visit(item.instructions, false);
+ i++;
+ }
+ break;
+ }
+ case "instruction": {
+ i++;
+ break;
+ }
+ case "terminal": {
+ const terminal = item.terminal;
+ switch (terminal.kind) {
+ case "break":
+ case "continue":
+ case "return":
+ case "throw": {
+ break;
+ }
+ case "for": {
+ visit(terminal.loop, true);
+ break;
+ }
+ case "while": {
+ visit(terminal.loop, true);
+ break;
+ }
+ case "if": {
+ visit(terminal.consequent, shouldFlatten);
+ if (terminal.alternate !== null) {
+ visit(terminal.alternate, shouldFlatten);
+ }
+ break;
+ }
+ case "switch": {
+ for (const case_ of terminal.cases) {
+ if (case_.block !== undefined) {
+ visit(case_.block, shouldFlatten);
+ }
+ }
+ break;
+ }
+ default: {
+ assertExhaustive(
+ terminal,
+ `Unexpected terminal kind '${(terminal as any).kind}'`
+ );
+ }
+ }
+ i++;
+ break;
+ }
+ default: {
+ assertExhaustive(item, `Unexpected item`);
+ }
+ }
+ }
+}
+
+function flatten(scope: ReactiveBlock, block: ReactiveBasicBlock): void {
+ for (const item of scope.instructions) {
+ switch (item.kind) {
+ case "block": {
+ flatten(item, block);
+ break;
+ }
+ case "terminal": {
+ const terminal = item.terminal;
+ switch (terminal.kind) {
+ case "break":
+ case "continue":
+ case "return":
+ case "throw": {
+ break;
+ }
+ case "for": {
+ visit(terminal.loop, true);
+ break;
+ }
+ case "while": {
+ visit(terminal.loop, true);
+ break;
+ }
+ case "if": {
+ visit(terminal.consequent, true);
+ if (terminal.alternate !== null) {
+ visit(terminal.alternate, true);
+ }
+ break;
+ }
+ case "switch": {
+ for (const case_ of terminal.cases) {
+ if (case_.block !== undefined) {
+ visit(case_.block, true);
+ }
+ }
+ break;
+ }
+ default: {
+ assertExhaustive(
+ terminal,
+ `Unexpected terminal kind '${(terminal as any).kind}'`
+ );
+ }
+ }
+ block.push(item);
+ break;
+ }
+ default: {
+ block.push(item);
+ }
+ }
+ }
+}
diff --git a/compiler/forget/src/__tests__/fixtures/hir/component.expect.md b/compiler/forget/src/__tests__/fixtures/hir/component.expect.md
index 0596926afc..51445ee584 100644
--- a/compiler/forget/src/__tests__/fixtures/hir/component.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/hir/component.expect.md
@@ -157,9 +157,7 @@ function Component(
}
[22] Call mutate seen$33_@0.add(mutate item$40_@0:TPrimitive)
[23] Const mutate $49:TPrimitive = "div"
- scope @3 [24:25] deps=[read $49:TPrimitive] {
- [24] Const mutate $50_@3 = JSX {freeze item$40_@0:TPrimitive}
- }
+ [24] Const mutate $50_@3 = JSX {freeze item$40_@0:TPrimitive}
[25] Call mutate renderedItems$32_@0.push(read $50_@3)
[26] Const mutate $55:TPrimitive = Binary read renderedItems$32_@0.length >= read max$35_@2:TPrimitive
if (read $55:TPrimitive) {
diff --git a/compiler/forget/src/__tests__/hir-test.ts b/compiler/forget/src/__tests__/hir-test.ts
index 81f15ffb14..86f2802c75 100644
--- a/compiler/forget/src/__tests__/hir-test.ts
+++ b/compiler/forget/src/__tests__/hir-test.ts
@@ -15,6 +15,7 @@ import invariant from "invariant";
import path from "path";
import prettier from "prettier";
import { buildReactiveFunction } from "../HIR/BuildReactiveFunction";
+import { flattenReactiveLoops } from "../HIR/FlattenReactiveLoops";
import { toggleLogging } from "../HIR/logger";
import run from "../HIR/Pipeline";
import { printFunction } from "../HIR/PrintHIR";
@@ -148,6 +149,7 @@ function transform(text: string, file: string): Array {
});
const reactiveFunction = buildReactiveFunction(ir);
+ flattenReactiveLoops(reactiveFunction);
const scopes = printReactiveFunction(reactiveFunction);
const textHIR = printFunction(ir);