traverse 是遍历 AST,并且遍历的过程中支持 visitor 的调用,在 visitor 里实现对 AST 的增删改。
我们这一节的目的是实现这样的 api:
traverse(ast, {
Identifier(node) {
node.name = 'b';
}
});
path 放到下一节实现。
思路分析
AST 的遍历就是树的遍历,树的遍历就深度优先、广度优先两种方式,而这里明显是深度优先遍历。
深度优先遍历要递归的遍历节点的子节点,那么我们怎么知道对象的属性是可以遍历的子节点呢?
可以维护一份数据来保存不同 AST 的什么属性是可以遍历的,然后在遍历不同节点的时候从中查找应该继续遍历什么属性。这样就实现了深度优先遍历。
在遍历的过程中可以根据类型调用不同的 visitor,然后传入当前节点。
代码实现
首先,我们维护这样一份数据:不同的 AST 有哪些可以遍历的属性。
比如 Program 要遍历 body 属性,VariableDeclarator 要遍历 id、init 属性等:
const astDefinationsMap = new Map();
astDefinationsMap.set('Program', {
visitor: ['body']
});
astDefinationsMap.set('VariableDeclaration', {
visitor: ['declarations']
});
astDefinationsMap.set('VariableDeclarator', {
visitor: ['id', 'init']
});
astDefinationsMap.set('Identifier', {});
astDefinationsMap.set('NumericLiteral', {});
astDefinationsMap.set('FunctionDeclaration', {
visitor: ['id', 'params', 'body']
});
astDefinationsMap.set('BlockStatement', {
visitor: ['body']
});
astDefinationsMap.set('ReturnStatement', {
visitor: ['argument']
});
astDefinationsMap.set('BinaryExpression', {
visitor: ['left', 'right']
});
astDefinationsMap.set('ExpressionStatement', {
visitor: ['expression']
});
astDefinationsMap.set('CallExpression', {
visitor: ['callee', 'arguments']
});
然后实现递归的遍历,遍历到不同节点时,取出不同节点要遍历的属性,然后递归遍历。如果是数组的话,每个元素都是这样处理:
function traverse(node, visitors) {
const defination = astDefinationsMap.get(node.type);
if (defination.visitor) {
defination.visitor.forEach(key => {
const prop = node[key];
if (Array.isArray(prop)) { // 如果该属性是数组
prop.forEach(childNode => {
traverse(childNode, visitors);
})
} else {
traverse(prop, visitors);
}
})
}
}
实现了遍历,当然还要在遍历时支持不同节点的 visitor 回调函数:
visitor 支持 enter 和 exit 阶段,也就是进入节点调用 enter 回调函数,之后遍历子节点,之后再调用 exit 回调函数。
那么分别在遍历前后调用就可以,默认如果没有指定哪个阶段就在 enter 阶段调用。
function traverse(node, visitors) {
const defination = astDefinationsMap.get(node.type);
let visitorFuncs = visitors[node.type] || {};
if(typeof visitorFuncs === 'function') {
visitorFuncs = {
enter: visitorFuncs
}
}
visitorFuncs.enter && visitorFuncs.enter(node);
if (defination.visitor) {
defination.visitor.forEach(key => {
const prop = node[key];
if (Array.isArray(prop)) { // 如果该属性是数组
prop.forEach(childNode => {
traverse(childNode, visitors);
})
} else {
traverse(prop, visitors);
}
})
}
visitorFuncs.exit && visitorFuncs.exit(node);
}
这样,我们就实现了 AST 的遍历和 enter、exit 阶段的 visitor 调用。
为什么要分 enter 和 exit 两个阶段呢?
因为 enter 阶段在遍历子节点之前,那么修改之后就可以立刻遍历子节点,而 exit 是在遍历结束之后了,所以不会继续遍历子节点。如果 enter 阶段修改了 AST 但是不想遍历新生成的子节点,可以用 path.skip 跳过遍历。
可以这样来遍历和修改 AST:
traverse(ast, {
Identifier(node) {
node.name = 'b';
}
});
总结
traverse 就是 AST 的遍历,而树的遍历就深度优先和广度优先两种,这里是深度优先,我们维护了一份什么 AST 遍历什么属性的数据,然后遍历的时候就可以知道如何遍历每一个节点。
遍历的时候调用 visitor 的回调函数,分为 enter 和 exit 阶段来调用,默认是 enter 阶段。
(代码在这里,建议 git clone 下来通过 node 跑一下)