How to transform JavaScript AST to create infinite loop guards?

Question

I'm using Esprima, estraverse, and escodegen to transform JavaScript code to add infinite loop guards inspired by Codepen

But I have a problem in adding code after the loop.

If I have code like this:

for (let x = 0; i < 100; ++i) {
   console.log('x');
}

it should transform it to:

for (let x = 0; i < 100; ++i) {
    if (koduj.should_break_loop(1)) {
        break;
    }
    console.log('x');
}
koduj.exit_loop(1);

But what I have is:

for (let x = 0; i < 100; ++i) {
    if (koduj.should_break_loop(1)) {
        break;
    }
    console.log('x');
    koduj.exit_loop(1);
}

I'm not sure how to search the AST to find all the loops and modify the code that contains a loop instead of just a loop.

This is the main code:

const patch_loop_body = with_loop_guard();
const loop_types = ['ForStatement', 'ForOfStatement', 'ForInStatement', 'DoWhileStatement', 'WhileStatement'];
function guard_loops(input) {
    loop_count = 0;
    const ast = esprima.parseScript(input);

    estraverse.traverse(ast, {
        enter: function (node, parent) {
            if (loop_types.includes(node.type)) {
                node.body = patch_loop_body(node.body);
            }
        }
    });

    return escodegen.generate(ast);
}

Here is my whole code for the AST transformer.

Answer 1

I've come up with a solution:

const patch_loop_body = with_loop_guard();
const loop_types = ['ForStatement', 'ForOfStatement', 'ForInStatement', 'DoWhileStatement', 'WhileStatement'];

function patch_body(body) {
    return body.flatMap(function(node) {
        if (loop_types.includes(node.type)) {
            node.body = patch_loop_body(node.body);
            return [node, expression_statement(koduj_call('exit_loop', literal(loop_count)))]
        }
        return node;
    });
}

function guard_loops(input) {
    loop_count = 0;
    const ast = esprima.parseScript(input);

    estraverse.traverse(ast, {
        enter: function (node, parent) {
            if (node.type === 'Program') {
                node.body = patch_body(node.body);
            } else if (node?.body?.type === 'BlockStatement') {
                node.body.body = patch_body(node.body.body);
            }
        }
    });

    return escodegen.generate(ast);
}

Working code