Use esprima's syntax delegate to remove use of escodegen. fixes #93 fixes #73

src/utils.js

@@ -42,52 +42,55 @@
 		}
 	}
 
-	// Generate 2 ASTs for the code to be inserted in loops for infinite run protection.
-	// The `myVar` variable names would be changed later for every insertion.
-	function getLoopProtectorBlocks() {
-		var ast1 = esprima.parse('var myvar = Date.now();');
-		var ast2 = esprima.parse('while(a){if (Date.now() - a787897 > 1000) { window.top.previewException(new Error("Infinite loop")); break;}}');
-		return {
-			before: ast1.body[0],
-			inside: ast2.body[0].body.body[0]
-		}
-	}
-
 	/**
-	 * Add timed limit on the loops found in the passed AST body
-	 * @param {ASTBody}	Body of an AST generated by esprima or any ES compliant AST
+	 * Adds timed limit on the loops found in the passed code.
+	 * Contributed by Ariya Hidayat!
+	 * @param code {string}	Code to be protected from infinite loops.
 	 */
-	function addInfiniteLoopProtection(astBody) {
-		if (!astBody) { return; }
-		if (!Array.isArray(astBody)) {
-			addInfiniteLoopProtection(astBody.body);
-			return;
-		}
-		var el, randomVariableName, insertionBLocks;
+	function addInfiniteLoopProtection(code) {
+		var loopId = 1;
+		var patches = [];
+		var varPrefix = '_wmloopvar';
+		var varStr = 'var %d = Date.now();\n'
+		var checkStr = '\nif (Date.now() - %d > 1000) { window.top.previewException(new Error("Infinite loop")); break;}\n'
 
-		for (var i = astBody.length; i--;) {
-			el = astBody[i];
-			if (el && el.type === 'ForStatement' || el.type === 'WhileStatement' || el.type === 'DoWhileStatement') {
-				randomVariableName = '_' + generateRandomId(3);
-				insertionBLocks = getLoopProtectorBlocks();
-				insertionBLocks.before.declarations[0].id.name = insertionBLocks.inside.test.left.right.name = randomVariableName;
-				// Insert time variable assignment
-				astBody.splice(i, 0, insertionBLocks.before);
-				// If the loop's body is a single statement, then convert it into a block statement
-				// so that we can insert our conditional break inside it.
-				if (!Array.isArray(el.body)) {
-					el.body = {
-						body: [ el.body ],
-						type: 'BlockStatement'
-					};
+		esprima.parse(code, { tolerant: true, range: true }, function (node) {
+			switch (node.type) {
+			case 'DoWhileStatement':
+			case 'ForStatement':
+			case 'ForInStatement':
+			case 'ForOfStatement':
+			case 'WhileStatement':
+				var start = 1 + node.body.range[0];
+				var end = node.body.range[1];
+				var prolog = checkStr.replace('%d', varPrefix + loopId);
+				var epilog = '';
+
+				if (node.body.type !== 'BlockStatement') {
+					// `while(1) doThat()` becomes `while(1) {doThat()}`
+					prolog = '{' + prolog;
+					epilog = '}';
+					--start;
 				}
-				// Insert IfStatement
-				el.body.body.unshift(insertionBLocks.inside);
+
+				patches.push({ pos: start, str: prolog });
+				patches.push({ pos: end, str: epilog });
+				patches.push({ pos: node.range[0], str: varStr.replace('%d', varPrefix + loopId) });
+				++loopId;
+				break;
+
+			default:
+				break;
 			}
-			if (el.body) {
-				addInfiniteLoopProtection(el.body);
-			}
-		}
+		});
+
+		/* eslint-disable no-param-reassign */
+		patches.sort(function (a, b) { return b.pos - a.pos }).forEach(function (patch) {
+			code = code.slice(0, patch.pos) + patch.str + code.slice(patch.pos);
+		});
+
+		/* eslint-disable no-param-reassign */
+		return code;
 	}
 
 	function getHumanDate(timestamp) {