Skip to content

Commit cc3ef43

Browse files
Nixxx19Nixxx19
committed
address review feedback on loop protection
1 parent 78eee42 commit cc3ef43

4 files changed

Lines changed: 22157 additions & 26245 deletions

File tree

‎client/modules/Preview/jsPreprocess.js‎

Lines changed: 152 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,16 @@
11
import * as acorn from 'acorn';
22
import * as walk from 'acorn-walk';
3+
import escodegen from 'escodegen';
34

45
const LOOP_TIMEOUT_MS = 100;
56

67
function isShaderCall(node) {
78
const { callee } = node;
89
const isBuildShader =
910
callee.type === 'Identifier' && /^build\w*Shader$/.test(callee.name);
10-
const isBaseShaderModify =
11-
callee.type === 'MemberExpression' &&
12-
callee.property.name === 'modify' &&
13-
callee.object.type === 'CallExpression' &&
14-
callee.object.callee.type === 'Identifier' &&
15-
/^base\w*Shader$/.test(callee.object.callee.name);
16-
return isBuildShader || isBaseShaderModify;
11+
const isModifyCall =
12+
callee.type === 'MemberExpression' && callee.property.name === 'modify';
13+
return isBuildShader || isModifyCall;
1714
}
1815

1916
function collectShaderFunctionNames(ast) {
@@ -77,56 +74,165 @@ function collectLoopsToProtect(ast, shaderNames) {
7774
return loops;
7875
}
7976

80-
export function jsPreprocess(jsText) {
81-
if (/\/\/\s*noprotect/.test(jsText)) {
82-
return jsText;
77+
function makeVarDecl(varName) {
78+
return {
79+
type: 'VariableDeclaration',
80+
kind: 'var',
81+
declarations: [
82+
{
83+
type: 'VariableDeclarator',
84+
id: { type: 'Identifier', name: varName },
85+
init: {
86+
type: 'CallExpression',
87+
callee: {
88+
type: 'MemberExpression',
89+
object: { type: 'Identifier', name: 'Date' },
90+
property: { type: 'Identifier', name: 'now' },
91+
computed: false
92+
},
93+
arguments: []
94+
}
95+
}
96+
]
97+
};
98+
}
99+
100+
function makeCheckStatement(varName, line) {
101+
return {
102+
type: 'IfStatement',
103+
test: {
104+
type: 'BinaryExpression',
105+
operator: '>',
106+
left: {
107+
type: 'BinaryExpression',
108+
operator: '-',
109+
left: {
110+
type: 'CallExpression',
111+
callee: {
112+
type: 'MemberExpression',
113+
object: { type: 'Identifier', name: 'Date' },
114+
property: { type: 'Identifier', name: 'now' },
115+
computed: false
116+
},
117+
arguments: []
118+
},
119+
right: { type: 'Identifier', name: varName }
120+
},
121+
right: {
122+
type: 'Literal',
123+
value: LOOP_TIMEOUT_MS,
124+
raw: String(LOOP_TIMEOUT_MS)
125+
}
126+
},
127+
consequent: {
128+
type: 'BlockStatement',
129+
body: [
130+
{
131+
type: 'ExpressionStatement',
132+
expression: {
133+
type: 'CallExpression',
134+
callee: {
135+
type: 'MemberExpression',
136+
object: {
137+
type: 'MemberExpression',
138+
object: { type: 'Identifier', name: 'window' },
139+
property: { type: 'Identifier', name: 'loopProtect' },
140+
computed: false
141+
},
142+
property: { type: 'Identifier', name: 'hit' },
143+
computed: false
144+
},
145+
arguments: [{ type: 'Literal', value: line, raw: String(line) }]
146+
}
147+
},
148+
{ type: 'BreakStatement', label: null }
149+
]
150+
},
151+
alternate: null
152+
};
153+
}
154+
155+
function injectProtection(loop, idx) {
156+
const varName = `_LP${idx}`;
157+
const { line } = loop.loc.start;
158+
const check = makeCheckStatement(varName, line);
159+
160+
if (loop.body.type === 'BlockStatement') {
161+
loop.body.body.unshift(check);
162+
} else {
163+
loop.body = {
164+
type: 'BlockStatement',
165+
body: [check, loop.body]
166+
};
83167
}
84168

85-
let ast;
169+
return makeVarDecl(varName);
170+
}
171+
172+
function insertVarDeclsIntoBlock(blockBody, loopsWithVarDecls) {
173+
loopsWithVarDecls.forEach(({ loop, varDecl }) => {
174+
const idx = blockBody.indexOf(loop);
175+
if (idx !== -1) {
176+
blockBody.splice(idx, 0, varDecl);
177+
}
178+
});
179+
}
180+
181+
function injectVarDeclsIntoAst(ast, loops) {
182+
const loopToVarDecl = new Map();
183+
loops.forEach((loop, idx) => {
184+
loopToVarDecl.set(loop, injectProtection(loop, idx));
185+
});
186+
187+
walk.simple(ast, {
188+
BlockStatement(node) {
189+
const loopsWithVarDecls = node.body
190+
.filter((child) => loopToVarDecl.has(child))
191+
.map((loop) => ({ loop, varDecl: loopToVarDecl.get(loop) }));
192+
if (loopsWithVarDecls.length > 0) {
193+
insertVarDeclsIntoBlock(node.body, loopsWithVarDecls);
194+
loopsWithVarDecls.forEach(({ loop }) => loopToVarDecl.delete(loop));
195+
}
196+
},
197+
Program(node) {
198+
const loopsWithVarDecls = node.body
199+
.filter((child) => loopToVarDecl.has(child))
200+
.map((loop) => ({ loop, varDecl: loopToVarDecl.get(loop) }));
201+
if (loopsWithVarDecls.length > 0) {
202+
insertVarDeclsIntoBlock(node.body, loopsWithVarDecls);
203+
loopsWithVarDecls.forEach(({ loop }) => loopToVarDecl.delete(loop));
204+
}
205+
}
206+
});
207+
}
208+
209+
function parseJs(jsText) {
210+
const options = { ecmaVersion: 'latest', locations: true };
86211
try {
87-
ast = acorn.parse(jsText, {
88-
ecmaVersion: 'latest',
89-
sourceType: 'script',
90-
locations: true
91-
});
212+
return acorn.parse(jsText, { ...options, sourceType: 'script' });
92213
} catch (e) {
214+
try {
215+
return acorn.parse(jsText, { ...options, sourceType: 'module' });
216+
} catch (e2) {
217+
return null;
218+
}
219+
}
220+
}
221+
222+
export function jsPreprocess(jsText) {
223+
if (/\/\/\s*noprotect/.test(jsText)) {
93224
return jsText;
94225
}
95226

227+
const ast = parseJs(jsText);
228+
if (!ast) return jsText;
229+
96230
const shaderNames = collectShaderFunctionNames(ast);
97231
const loops = collectLoopsToProtect(ast, shaderNames);
98232

99233
if (loops.length === 0) return jsText;
100234

101-
const insertions = [];
102-
103-
loops.forEach((loop) => {
104-
const { line } = loop.loc.start;
105-
const varName = `_LP${loop.start}`;
106-
107-
const beforeCode = `var ${varName} = Date.now(); `;
108-
109-
const checkCode =
110-
`if (Date.now() - ${varName} > ${LOOP_TIMEOUT_MS}) ` +
111-
`{ window.loopProtect.hit(${line}); break; } `;
112-
113-
const { body } = loop;
114-
if (body.type === 'BlockStatement') {
115-
insertions.push({ pos: loop.start, code: beforeCode });
116-
insertions.push({ pos: body.start + 1, code: checkCode });
117-
} else {
118-
insertions.push({ pos: loop.start, code: beforeCode });
119-
insertions.push({ pos: body.start, code: `{ ${checkCode}` });
120-
insertions.push({ pos: body.end, code: ` }` });
121-
}
122-
});
123-
124-
insertions.sort((a, b) => b.pos - a.pos);
125-
126-
let result = jsText;
127-
insertions.forEach(({ pos, code }) => {
128-
result = result.slice(0, pos) + code + result.slice(pos);
129-
});
235+
injectVarDeclsIntoAst(ast, loops);
130236

131-
return result;
237+
return escodegen.generate(ast);
132238
}

‎client/modules/Preview/jsPreprocess.unit.test.js‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,15 @@ describe('jsPreprocess', () => {
3232
const code = 'this is not valid javascript @@@@';
3333
expect(jsPreprocess(code)).toBe(code);
3434
});
35+
36+
it('handles module scripts with import statements', () => {
37+
const code = `
38+
import something from 'somewhere';
39+
for (let i = 0; i < 10; i++) {}
40+
`;
41+
const result = jsPreprocess(code);
42+
expect(result).toContain('window.loopProtect.hit');
43+
});
3544
});
3645

3746
describe('shader strings', () => {
@@ -112,6 +121,17 @@ describe('jsPreprocess', () => {
112121
expect(result).not.toContain('window.loopProtect.hit');
113122
});
114123

124+
it('skips loop protection for function passed to any .modify() call', () => {
125+
const code = `
126+
blur = someCustomShader().modify(doBlur);
127+
function doBlur() {
128+
for (let i = 0; i < 20; i++) {}
129+
}
130+
`;
131+
const result = jsPreprocess(code);
132+
expect(result).not.toContain('window.loopProtect.hit');
133+
});
134+
115135
it('still protects regular loops outside shader functions', () => {
116136
const code = `
117137
blur = buildFilterShader(doBlur);

0 commit comments

Comments
 (0)