From 60ca5b1f500c0158283f72f9e8a314a0ca92d804 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sun, 15 Mar 2026 13:12:56 +0000
Subject: [PATCH] Use acorn to inject loop guards, removing manual await
 wait(0) from generators

Add `injectLoopGuards(src)` to the flock object, which uses acorn/acorn-walk to
parse generated code and automatically insert `await wait(0)` before the closing
brace of any loop body that doesn't already contain an await expression. This
centralises the loop-guard concern so individual generators no longer need to
embed it inline.

Remove the explicit `await wait(0)` from controls_whileUntil, controls_doWhile,
controls_repeat_ext, and controls_forEach. The more sophisticated
timing/RAF-based guards in controls_for and for_loop are preserved; because
those bodies already contain an AwaitExpression, injectLoopGuards skips them
automatically.

The injection runs in runCode() after AST validation, so it also protects any
hand-written user code that contains unguarded loops.

https://claude.ai/code/session_01Dvex7RparGEMpvAAR2iF1H
---
 flock.js                 | 66 ++++++++++++++++++++++++++++++++++++++--
 generators/generators.js |  5 ---
 2 files changed, 64 insertions(+), 7 deletions(-)

diff --git a/flock.js b/flock.js
index 3d1e531e..a13e2ac9 100644
--- a/flock.js
+++ b/flock.js
@@ -632,9 +632,71 @@ export const flock = {
                         },
                 });
         },
+        injectLoopGuards(src) {
+                const ast = acorn.parse(src, {
+                        ecmaVersion: "latest",
+                        sourceType: "script",
+                        allowAwaitOutsideFunction: true,
+                });
+
+                // Check whether a BlockStatement already contains any await expression
+                // (at any nesting depth). Loops that already yield don't need an extra guard.
+                function bodyHasAwait(node) {
+                        let found = false;
+                        walk.simple(node, {
+                                AwaitExpression() {
+                                        found = true;
+                                },
+                        });
+                        return found;
+                }
+
+                // Collect positions (just before the closing '}') where we'll inject
+                const insertions = [];
+
+                function addGuardIfNeeded(loopBody) {
+                        if (
+                                loopBody?.type === "BlockStatement" &&
+                                !bodyHasAwait(loopBody)
+                        ) {
+                                insertions.push(loopBody.end - 1);
+                        }
+                }
+
+                walk.simple(ast, {
+                        WhileStatement(node) {
+                                addGuardIfNeeded(node.body);
+                        },
+                        DoWhileStatement(node) {
+                                addGuardIfNeeded(node.body);
+                        },
+                        ForStatement(node) {
+                                addGuardIfNeeded(node.body);
+                        },
+                        ForInStatement(node) {
+                                addGuardIfNeeded(node.body);
+                        },
+                        ForOfStatement(node) {
+                                addGuardIfNeeded(node.body);
+                        },
+                });
+
+                // Apply from last to first so earlier positions stay valid
+                insertions.sort((a, b) => b - a);
+
+                let result = src;
+                for (const pos of insertions) {
+                        result =
+                                result.slice(0, pos) +
+                                "await wait(0);\n" +
+                                result.slice(pos);
+                }
+                return result;
+        },
         async runCode(code) {
                 try {
                         flock.validateUserCodeAST(code);
+                        const instrumentedCode = flock.injectLoopGuards(code);
                         await flock.disposeOldScene();
 
                         // --- remove any existing iframe ---
@@ -818,12 +880,12 @@ export const flock = {
                         // Wrap user code to allow top-level await
                         /*const wrapped =
                                 '(async () => {\n"use strict";\n' +
-                                code +
+                                instrumentedCode +
                                 "\n})()\n//# sourceURL=user-code.js";*/
 
                         const wrapped =
                                 '(async function () {\n"use strict";\n' +
-                                code +
+                                instrumentedCode +
                                 "\n}).call(undefined)\n//# sourceURL=user-code.js";
 
                         // Evaluate in SES Compartment
diff --git a/generators/generators.js b/generators/generators.js
index 3ab3da7d..f64d9863 100644
--- a/generators/generators.js
+++ b/generators/generators.js
@@ -3585,7 +3585,6 @@ javascriptGenerator.forBlock["controls_whileUntil"] = function (block) {
                 argument0 +
                 ") {\n" +
                 branch +
-                `\nawait wait(0);\n` +
                 "}\n"
         );
 };
@@ -3602,8 +3601,6 @@ javascriptGenerator.forBlock["controls_doWhile"] = function (block) {
         return `
         do {
                 ${branch}
-
-                await wait(0);
         } while (${condition});\n`;
 };
 
@@ -3651,7 +3648,6 @@ javascriptGenerator.forBlock["controls_repeat_ext"] = function (
                 loopVar +
                 "++) {\n" +
                 branch +
-                "await wait(0);\n" +
                 "}\n";
 
         return code;
@@ -3807,7 +3803,6 @@ javascriptGenerator.forBlock["controls_forEach"] = function (block, generator) {
                 listVar +
                 ") {\n" +
                 branch +
-                "\n  await wait(0);\n" +
                 "}\n";
 
         return code;