AST解web控制流平坦化-编程知识

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')