#!/usr/bin/env node

var recast = require("recast");
var Syntax = recast.Syntax;
var Visitor = recast.Visitor;
var b = recast.types.builders;

require("recast").run(function(ast, callback) {
    callback(new Orthodontist().visit(ast));
});

var Orthodontist = Visitor.extend({
    visitIfStatement: function(stmt) {
        stmt.consequent = fix(stmt.consequent);
        var alt = stmt.alternate;
        if (alt && alt.type !== Syntax.IfStatement)
            stmt.alternate = fix(alt);
        this.genericVisit(stmt);
    },

    visitWhileStatement: visitLoop,
    visitForStatement: visitLoop,
    visitForInStatement: visitLoop,
});

function visitLoop(loop) {
    loop.body = fix(loop.body);
    this.genericVisit(loop);
}

function fix(clause) {
    if (clause) {
        if (clause.type !== Syntax.BlockStatement)
            clause = b.blockStatement([clause]);
    }
    return clause;
}
