AST解web控制流平坦化

此代码可以解决大部分 while if else 控制流平坦化
原理：
- 先将 if 语句转为 switch 语句，再将 switch 分支合并，最后删除已合并的分支（具体看代码）
实现效果图
首先安装依赖：

npm install @babel/parser
npm install @babel/generator
npm install @babel/traverse        
npm install @babel/types

代码：

/*
* 控制流平坦化 if语句转 switch* */
function del_code(name, consequent) {
    // 删除合并分支后多余的 赋值和break代码
    let assignment_bool, break_bool;
    for (let i = consequent.length - 1; i >= 0; i--) {
        if (consequent[i].type === "BreakStatement") {
            if (break_bool) {
                consequent.splice(i, 1);
            } else {
                break_bool = true;
            }
        } else if (consequent[i].type === "ExpressionStatement" && consequent[i].expression.type === "AssignmentExpression" && consequent[i].expression.left.name === name) {
            if (assignment_bool) {
                consequent.splice(i, 1);
            } else {
                assignment_bool = true;
            }
        }
    }
}

function merge_branch(name, key, cases_dict) {
    // 用于递归合并 switch 分支
    let {consequent} = cases_dict[key];
    let value = -1;
    for (let i in consequent) {
        if (consequent[i].type === "ExpressionStatement" && consequent[i].expression.type === "AssignmentExpression" && consequent[i].expression.left.name === name && consequent[i].expression.right.type === "NumericLiteral") {
            value = consequent[i].expression.right.value;
            break;
        }
    }

    if (value !== -1 && cases_dict.hasOwnProperty(value)) {
        del_cases_dict[value] = 1;
        return consequent.concat(merge_branch(name, value, cases_dict));    // 继续下一分支的合并
    }
    return consequent;
}

const fs = require('fs');
const {parse} = require("@babel/parser");
const traverse = require("@babel/traverse").default;
const generator = require("@babel/generator").default;
const types = require("@babel/types");

const js_code = fs.readFileSync("./test.js", 'utf8');
const ast_code = parse(js_code);

switch_cases_dict = {};
break_node = types.breakStatement();

traverse(ast_code, {
    'IfStatement': {
        enter(path) {
            var name = path.node.test.left.name;
            if (path.node.test.operator === "===") {  // 如果判断符号是 ===
                if (switch_cases_dict[name] === undefined) {
                    switch_cases_dict[name] = [];
                }
                path.node.consequent.body.push(break_node);
                switch_cases_dict[name].push(types.switchCase(path.node.test.right, path.node.consequent.body));

                if (path.node.alternate.type === 'BlockStatement') {
                    path.node.alternate.body.push(break_node);
                    let num = path.node.test.right.value + 1;
                    switch_cases_dict[name].push(types.switchCase(
                        types.numericLiteral(num),
                        path.node.alternate.body,
                    ));
                }
            }
        },
        exit(path) {
            var name = path.node.test.left.name;
            if (path.parentPath.parentPath.type === "WhileStatement" && switch_cases_dict[name].length !== 0) {
                console.log(name, "if 已替换 switch");
                path.replaceWith(types.switchStatement(
                    discriminant = types.identifier(name),
                    cases = switch_cases_dict[name]
                ));
            }
        }
    },
    'SwitchStatement': {
        enter(path) {
            del_cases_dict = {}; // 待删除的 分支语句
            let cases_dict = {};
            let cases_list = path.node.cases;
            let {name} = path.node.discriminant;
            if (switch_cases_dict.hasOwnProperty(name)) {
                console.log(name, "switch 分支合并");
                for (let i in cases_list) {
                    cases_dict[cases_list[i].test.value] = cases_list[i]
                }

                for (let key in cases_dict) {   // 合并分支并删除多余代码
                    cases_dict[key].consequent = merge_branch(name, key, cases_dict);
                    del_code(name, cases_dict[key].consequent);
                }
                for (let key in del_cases_dict) {
                    delete cases_dict[key]; // 删除多余分支
                }
                path.node.cases = Object.values(cases_dict);
            }
        },
    }
})

// console.log(generator(ast_code).code)
fs.writeFileSync("./demo.js", generator(ast_code).code, 'utf8')