From 894eadfa67752666a2d75ea39b6cf477fa9dfbfc Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Tue, 31 Jan 2023 13:39:39 -0800
Subject: [PATCH] [valueblocks] For.update is a proper value block

---
 compiler/forget/src/HIR/HIR.ts                        |  2 +-
 .../src/ReactiveScopes/BuildReactiveFunction.ts       | 11 ++++-------
 .../src/ReactiveScopes/CodegenReactiveFunction.ts     |  2 +-
 .../src/ReactiveScopes/PrintReactiveFunction.ts       |  2 +-
 .../src/ReactiveScopes/PropagateScopeDependencies.ts  |  2 +-
 compiler/forget/src/ReactiveScopes/visitors.ts        |  4 ++--
 6 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/compiler/forget/src/HIR/HIR.ts b/compiler/forget/src/HIR/HIR.ts
index bb89824934..4d34567b79 100644
--- a/compiler/forget/src/HIR/HIR.ts
+++ b/compiler/forget/src/HIR/HIR.ts
@@ -173,7 +173,7 @@ export type ReactiveForTerminal = {
   kind: "for";
   init: ReactiveValueBlock;
   test: ReactiveValue;
-  update: ReactiveValueBlock;
+  update: ReactiveValue;
   loop: ReactiveBlock;
   id: InstructionId;
 };
diff --git a/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts b/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts
index bc7874fc1a..a2f5c1bb2f 100644
--- a/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts
+++ b/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts
@@ -320,13 +320,10 @@ class Driver {
           terminal.loc
         ).value;
 
-        const updateBlock = this.cx.ir.blocks.get(terminal.update)!;
-        const updateTerminal = updateBlock.terminal;
-        invariant(
-          updateTerminal.kind === "goto",
-          "Expected for loop update block to end in a goto"
-        );
-        const updateValue = this.visitValueBlock(blockValue, updateBlock);
+        const updateValue = this.visitValueBlockNew(
+          terminal.update,
+          terminal.loc
+        ).value;
 
         let loopBody: ReactiveBlock;
         if (loopId) {
diff --git a/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts
index e21da55c82..a16bccbae7 100644
--- a/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+++ b/compiler/forget/src/ReactiveScopes/CodegenReactiveFunction.ts
@@ -266,7 +266,7 @@ function codegenTerminal(
       return t.forStatement(
         codegenForInit(cx, terminal.init),
         codegenInstructionValue(cx, terminal.test),
-        codegenValueBlock(cx, terminal.update),
+        codegenInstructionValue(cx, terminal.update),
         codegenBlock(cx, terminal.loop)
       );
     }
diff --git a/compiler/forget/src/ReactiveScopes/PrintReactiveFunction.ts b/compiler/forget/src/ReactiveScopes/PrintReactiveFunction.ts
index 686c683252..596203b78f 100644
--- a/compiler/forget/src/ReactiveScopes/PrintReactiveFunction.ts
+++ b/compiler/forget/src/ReactiveScopes/PrintReactiveFunction.ts
@@ -235,7 +235,7 @@ function printTerminal(writer: Writer, terminal: ReactiveTerminal): void {
       writer.writeLine(";");
       printReactiveValue(writer, terminal.test);
       writer.writeLine(";");
-      printValueBlock(writer, terminal.update);
+      printReactiveValue(writer, terminal.update);
       writer.writeLine(") {");
       printReactiveInstructions(writer, terminal.loop);
       writer.writeLine("}");
diff --git a/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts b/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts
index 7a42a47029..c6acd3de7d 100644
--- a/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts
+++ b/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts
@@ -232,7 +232,7 @@ function visit(context: Context, block: ReactiveBlock): void {
           case "for": {
             visitValueBlock(context, terminal.init);
             visitReactiveValue(context, terminal.test);
-            visitValueBlock(context, terminal.update);
+            visitReactiveValue(context, terminal.update);
             visit(context, terminal.loop);
             break;
           }
diff --git a/compiler/forget/src/ReactiveScopes/visitors.ts b/compiler/forget/src/ReactiveScopes/visitors.ts
index 3543e8953f..518dac5a45 100644
--- a/compiler/forget/src/ReactiveScopes/visitors.ts
+++ b/compiler/forget/src/ReactiveScopes/visitors.ts
@@ -104,7 +104,7 @@ export class ReactiveFunctionVisitor<TState = void> {
       case "for": {
         this.visitValueBlock(terminal.init, state);
         this.visitValue(terminal.id, terminal.test, state);
-        this.visitValueBlock(terminal.update, state);
+        this.visitValue(terminal.id, terminal.update, state);
         this.visitBlock(terminal.loop, state);
         break;
       }
@@ -336,7 +336,7 @@ export function eachTerminalBlock(
       visitValueBlock(terminal.init);
       // TODO
       // visitValueBlock(terminal.test);
-      visitValueBlock(terminal.update);
+      // visitValueBlock(terminal.update);
       visitBlock(terminal.loop);
       break;
     }