diff --git a/client/modules/Preview/jsPreprocess.js b/client/modules/Preview/jsPreprocess.js index e417285655..2640a58b4d 100644 --- a/client/modules/Preview/jsPreprocess.js +++ b/client/modules/Preview/jsPreprocess.js @@ -39,14 +39,27 @@ function makeVarDecl(varName) { type: 'VariableDeclarator', id: { type: 'Identifier', name: varName }, init: { - type: 'CallExpression', - callee: { - type: 'MemberExpression', - object: { type: 'Identifier', name: 'Date' }, - property: { type: 'Identifier', name: 'now' }, - computed: false - }, - arguments: [] + type: 'ObjectExpression', + properties: [ + { + type: 'Property', + key: { type: 'Identifier', name: 't' }, + value: { + type: 'CallExpression', + callee: { + type: 'MemberExpression', + object: { type: 'Identifier', name: 'Date' }, + property: { type: 'Identifier', name: 'now' }, + computed: false + }, + arguments: [] + }, + kind: 'init', + method: false, + shorthand: false, + computed: false + } + ] } } ] @@ -72,7 +85,12 @@ function makeCheckStatement(varName, line) { }, arguments: [] }, - right: { type: 'Identifier', name: varName } + right: { + type: 'MemberExpression', + object: { type: 'Identifier', name: varName }, + property: { type: 'Identifier', name: 't' }, + computed: false + } }, right: { type: 'Literal', @@ -160,6 +178,8 @@ function collectLoopsToProtect(ast, shaderNames) { walk.ancestor(ast, { ForStatement: visitNode, + ForInStatement: visitNode, + ForOfStatement: visitNode, WhileStatement: visitNode, DoWhileStatement: visitNode }); @@ -179,6 +199,41 @@ function injectProtection(loops) { loop.body = { type: 'BlockStatement', body: [check, loop.body] }; } + walk.ancestor(loop.body, { + AwaitExpression(node, ancestors) { + const isInsideNestedFunction = ancestors.some( + (ancestor) => + ancestor !== loop.body && + (ancestor.type === 'FunctionDeclaration' || + ancestor.type === 'FunctionExpression' || + ancestor.type === 'ArrowFunctionExpression') + ); + + if (!isInsideNestedFunction) { + const originalAwait = { + type: 'AwaitExpression', + argument: node.argument + }; + node.type = 'CallExpression'; + node.callee = { + type: 'MemberExpression', + object: { + type: 'MemberExpression', + object: { type: 'Identifier', name: 'window' }, + property: { type: 'Identifier', name: 'loopProtect' }, + computed: false + }, + property: { type: 'Identifier', name: 'reset' }, + computed: false + }; + node.arguments = [ + { type: 'Identifier', name: varName }, + originalAwait + ]; + } + } + }); + if (parentBlock) { const varDecl = makeVarDecl(varName); const nodeIdx = parentBlock.body.indexOf(loop); diff --git a/client/modules/Preview/jsPreprocess.unit.test.js b/client/modules/Preview/jsPreprocess.unit.test.js index 4fad33c9ae..075b8c308d 100644 --- a/client/modules/Preview/jsPreprocess.unit.test.js +++ b/client/modules/Preview/jsPreprocess.unit.test.js @@ -148,4 +148,51 @@ describe('jsPreprocess', () => { expect(result).toContain('window.loopProtect.hit'); }); }); + + describe('async/await loop protection', () => { + it('wraps await expressions in for loop with window.loopProtect.reset', () => { + const code = ` + async function setup() { + for (let i = 0; i < fns.length; i++) { + imgs.push(await loadImage(fns[i])); + } + } + `; + const result = jsPreprocess(code, ''); + expect(result).toContain('window.loopProtect.hit'); + expect(result).toContain( + 'window.loopProtect.reset(_LP0, await loadImage(fns[i]))' + ); + }); + + it('does not wrap await expressions outside loops', () => { + const code = ` + async function setup() { + const img = await loadImage('a.jpg'); + for (let i = 0; i < 10; i++) {} + } + `; + const result = jsPreprocess(code, ''); + expect(result).toContain("await loadImage('a.jpg')"); + expect(result).not.toContain( + 'window.loopProtect.reset(_LP0, await loadImage' + ); + }); + + it('handles nested loops with await expressions', () => { + const code = ` + async function run() { + for (let i = 0; i < 2; i++) { + for (let j = 0; j < 2; j++) { + await loadData(i, j); + } + } + } + `; + const result = jsPreprocess(code, ''); + expect(result).toContain( + 'window.loopProtect.reset(_LP0, window.loopProtect.reset(_LP1, await loadData(i, j)))' + ); + }); + }); }); diff --git a/client/utils/previewEntry.js b/client/utils/previewEntry.js index 0d369788c3..c71bdbc068 100644 --- a/client/utils/previewEntry.js +++ b/client/utils/previewEntry.js @@ -51,6 +51,12 @@ window.loopProtect = { } return true; + }, + reset: function resetLoopTimer(loopObj, val) { + if (loopObj) { + loopObj.t = Date.now(); + } + return val; } };