Skip to content

Centralize loop guard injection into AST transformation #399

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tracygardner wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Centralize loop guard injection into AST transformation #399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 64 additions & 2 deletions flock.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Comment on lines +642 to +663
Copy link

@coderabbitai coderabbitai bot Mar 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

# First, let's explore the repository structure and locate the file
git ls-files | grep -E "flock\.js|\.js$" | head -20

Repository: flipcomputing/flock

Length of output: 354

🏁 Script executed:

# Read the full context of flock.js around the lines mentioned
wc -l flock.js

Repository: flipcomputing/flock

Length of output: 76

🏁 Script executed:

# Get the context around lines 642-663 and 684-692
sed -n '620,710p' flock.js | cat -n

Repository: flipcomputing/flock

Length of output: 4632

🏁 Script executed:

# Search for walk.simple usage to understand its behavior
rg "walk\." flock.js -A 3 -B 3 | head -100

Repository: flipcomputing/flock

Length of output: 912

🏁 Script executed:

# Look for any test files or documentation about this feature
fd -e test -e spec -e example flock | head -20

Repository: flipcomputing/flock

Length of output: 45

🏁 Script executed:

# Search for continue statements in test context and loop examples
rg "continue" flock.js -B 3 -A 3

Repository: flipcomputing/flock

Length of output: 45

🏁 Script executed:

# Look for async function examples within loops in the codebase
rg "while|for.*{" flock.js -A 10 | grep -E "(while|for|async|await)" | head -30

Repository: flipcomputing/flock

Length of output: 150

🏁 Script executed:

# Let's see the broader context around the problematic lines and the second location
sed -n '630,700p' flock.js | cat -n

Repository: flipcomputing/flock

Length of output: 3397

🏁 Script executed:

# Search for imports to understand which walk library is being used
head -50 flock.js | grep -E "import|require|walk"

Repository: flipcomputing/flock

Length of output: 1256

🏁 Script executed:

# Look for any comments or documentation about loop guards
rg "guard|yield|busy" flock.js -B 2 -A 2 | head -50

Repository: flipcomputing/flock

Length of output: 2029

🏁 Script executed:

# Search for examples of loops with await to understand the pattern
rg "await" flock.js -B 5 -A 5 | head -80

Repository: flipcomputing/flock

Length of output: 3561

🏁 Script executed:

# Check the second location mentioned (684-692)
sed -n '675,700p' flock.js | cat -n

Repository: flipcomputing/flock

Length of output: 1266

🏁 Script executed:

# Let's verify the behavior of acorn-walk's simple function with a quick test
python3 << 'EOF'
# Let's reason through the acorn-walk behavior:
# walk.simple() is a recursive visitor that traverses the entire AST

# Test case 1: continue statement
# Code: while(true) { x++; continue; }
# Guard inserted at loopBody.end - 1 (before closing })
# Result: while(true) { x++; continue; await wait(0); }
# But continue jumps to next iteration, skipping the guard
# TRUE: continue escapes the guard

# Test case 2: single-statement
# Code: while(true) foo();
# loopBody.type would be ExpressionStatement, NOT BlockStatement
# addGuardIfNeeded checks: loopBody?.type === "BlockStatement"
# No guard is added
# TRUE: single-statement bodies not guarded

# Test case 3: nested async
# Code: while(true) { async function f() { await x; } }
# walk.simple traverses entire tree, finds AwaitExpression inside async function
# bodyHasAwait returns true, so NO guard is added
# But the loop itself doesn't yield - only the function yields if called
# If f() is never called, loop blocks UI
# TRUE: nested async awaits incorrectly treated as loop yields

print("All three escape conditions appear to be valid based on code analysis")
EOF

Repository: flipcomputing/flock

Length of output: 132

🏁 Script executed:

# Let's also check if there's a "wait" function defined to understand the guard
rg "function wait|const wait|=.*wait.*=" flock.js | head -10

Repository: flipcomputing/flock

Length of output: 45

🏁 Script executed:

# Check for any PR context or comments about this feature
git log --oneline -20 2>/dev/null | head -10

Repository: flipcomputing/flock

Length of output: 149

Loop guards can be bypassed in three ways, reintroducing UI-blocking loops.

The current implementation has three bypass mechanisms that defeat the guard:

  1. continue statements skip the tail-inserted guard — Since the guard is injected just before the closing }, a continue jumps to the next iteration without executing it.

  2. Single-statement loop bodies are never guarded — The code checks loopBody?.type === "BlockStatement", so loops like while (true) foo(); are never instrumented.

  3. bodyHasAwait() treats awaits in nested async functions as loop yields — Since walk.simple() traverses the entire AST, an async function() { await x; } inside the loop body causes the guard to be skipped, even though the loop itself doesn't yield if the function is never called.

All three reintroduce the "busy loop blocks the UI" failure mode. Normalizing eligible loop bodies to blocks and injecting the guard as the first statement of each iteration would close these gaps.

Also applies to: 684-692

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@flock.js` around lines 642 - 663, The loop-guardging logic must (1) normalize
single-statement loop bodies into BlockStatement so constructs like while(true)
foo(); get guarded, (2) inject the guard as the first statement of the block
(not just before the closing brace) so a continue still executes the guard on
the next iteration, and (3) make bodyHasAwait ignore awaits inside nested
function declarations/expressions/arrow functions so an inner async function
containing await does not cause skipping; to implement: update addGuardIfNeeded
to wrap non-BlockStatement loop bodies in a BlockStatement (preserving the
original node) and record insertion positions at the block start (or better,
record the AST block node to insert the guard as its first statement instead of
using end-1), and change bodyHasAwait to use walk.simple but override/skipping
handlers for FunctionDeclaration/FunctionExpression/ArrowFunctionExpression (and
similar) so traversal does not descend into nested functions when detecting
AwaitExpression.

}

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" +
Comment on lines +689 to +691
Copy link

@coderabbitai coderabbitai bot Mar 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

git ls-files | head -20

Repository: flipcomputing/flock

Length of output: 437

🏁 Script executed:

fd -t f "flock.js" | head -5

Repository: flipcomputing/flock

Length of output: 71

🏁 Script executed:

wc -l flock.js

Repository: flipcomputing/flock

Length of output: 76

🏁 Script executed:

sed -n '680,700p' flock.js

Repository: flipcomputing/flock

Length of output: 875

🏁 Script executed:

rg "wait" flock.js -A 2 -B 2 | head -50

Repository: flipcomputing/flock

Length of output: 2176

🏁 Script executed:

rg "createWhitelist\|wait" flock.js -B 3 -A 5 | head -100

Repository: flipcomputing/flock

Length of output: 45

🏁 Script executed:

rg "function wait\|const wait\|=.*wait" flock.js | head -20

Repository: flipcomputing/flock

Length of output: 45

🏁 Script executed:

rg "wait" flock.js -B 5 -A 5 | grep -A 10 -B 10 "await wait"

Repository: flipcomputing/flock

Length of output: 756

🏁 Script executed:

rg "function createWhitelist|createWhitelist" flock.js -B 2 -A 20

Repository: flipcomputing/flock

Length of output: 2343

🏁 Script executed:

sed -n '900,1100p' flock.js | grep -A 50 "createWhitelist"

Repository: flipcomputing/flock

Length of output: 2747

🏁 Script executed:

rg "return.*\{" flock.js -A 30 | grep -B 10 -A 30 "wait"

Repository: flipcomputing/flock

Length of output: 2170

🏁 Script executed:

sed -n '900,1100p' flock.js

Repository: flipcomputing/flock

Length of output: 11739

🏁 Script executed:

sed -n '1000,1200p' flock.js

Repository: flipcomputing/flock

Length of output: 11072

🏁 Script executed:

rg "validateUserCodeAST" flock.js -B 3 -A 30

Repository: flipcomputing/flock

Length of output: 2926

🏁 Script executed:

sed -n '500,700p' flock.js | grep -A 100 "REJECT_IDENTIFIERS"

Repository: flipcomputing/flock

Length of output: 7628

🏁 Script executed:

rg "REJECT_IDENTIFIERS" flock.js -A 100 | head -150

Repository: flipcomputing/flock

Length of output: 6659

Use a hygienic helper name for the injected guard.

The injected code calls bare wait, which is exposed as a public API in the whitelist. User code can shadow this identifier (const wait = 0, function wait() {}), causing the guard to malfunction at runtime. Route this through a reserved name like __flockLoopGuardWait__ that is both exposed from createWhitelist() and rejected in validateUserCodeAST() so user code cannot redefine it.

Suggested fix 
-                                "await wait(0);\n" +
+                                "await __flockLoopGuardWait__(0);\n" +

Then expose __flockLoopGuardWait__ from createWhitelist() and add it to the REJECT_IDENTIFIERS set in validateUserCodeAST().

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@flock.js` around lines 689 - 691, The injected guard currently calls the bare
identifier wait which can be shadowed by user code; change the injected call to
a reserved helper name like __flockLoopGuardWait__ (replace the "await wait(0);"
injection with "await __flockLoopGuardWait__(0);"), export that identifier from
createWhitelist() so it is available to the runtime, and add
__flockLoopGuardWait__ to the REJECT_IDENTIFIERS set checked by
validateUserCodeAST() so user code cannot define or shadow it.

result.slice(pos);
Comment on lines +666 to +692
Copy link

@coderabbitai coderabbitai bot Mar 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🌐 Web query:

In JavaScript, is await syntactically valid inside a non-async function body when the surrounding program is executed from an async IIFE?

💡 Result:

No. await is not syntactically valid inside the body of a non-async function, even if that function is called from (or nested inside) an async IIFE—the code will fail at parse time with a syntax error. await is only allowed inside async functions (and async generators), or at the top level of an ES module (“top-level await”). [1] [2]

Only inject guards inside async-capable scopes.

This code instruments loops anywhere in the AST, producing invalid JavaScript. For example, function tick() { while (true) {} } becomes syntactically invalid after transformation because await is not legal inside non-async function bodies—the code fails at parse time. The wrapper async IIFE does not change the syntax rules for nested functions, methods, or generators. Track async ancestry before inserting guards to ensure guards are only injected where await is syntactically valid.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@flock.js` around lines 666 - 692, The loop instrumentation currently inserts
"await wait(0)" everywhere (in walk.simple handlers for WhileStatement,
DoWhileStatement, ForStatement, ForInStatement, ForOfStatement), which yields
invalid syntax inside non-async scopes; update the AST traversal to track
async-capable ancestry (push/pop a flag when entering FunctionDeclaration,
FunctionExpression, ArrowFunctionExpression, MethodDefinition, and
GeneratorFunction nodes and consider top-level/module await rules) and only call
addGuardIfNeeded when the current ancestry indicates await is legal (i.e.,
inside an async function or allowed top-level); keep the insertion logic
(insertions sort + splice into result) unchanged but ensure you reference the
existing addGuardIfNeeded and the walk/simple handlers so guards are only added
where await is syntactically valid.

}
return result;
},
async runCode(code) {
try {
flock.validateUserCodeAST(code);
const instrumentedCode = flock.injectLoopGuards(code);
await flock.disposeOldScene();

// --- remove any existing iframe ---
Expand Down Expand Up @@ -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
Expand Down
5 changes: 0 additions & 5 deletions generators/generators.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3585,7 +3585,6 @@ javascriptGenerator.forBlock["controls_whileUntil"] = function (block) {
argument0 +
") {\n" +
branch +
`\nawait wait(0);\n` +
"}\n"
);
};
Expand All @@ -3602,8 +3601,6 @@ javascriptGenerator.forBlock["controls_doWhile"] = function (block) {
return `
do {
${branch}

await wait(0);
} while (${condition});\n`;
};

Expand Down Expand Up @@ -3651,7 +3648,6 @@ javascriptGenerator.forBlock["controls_repeat_ext"] = function (
loopVar +
"++) {\n" +
branch +
"await wait(0);\n" +
"}\n";

return code;
Expand Down Expand Up @@ -3807,7 +3803,6 @@ javascriptGenerator.forBlock["controls_forEach"] = function (block, generator) {
listVar +
") {\n" +
branch +
"\n await wait(0);\n" +
"}\n";

return code;
Expand Down
Loading