Gist & Cookbook ...

小目标:每天能够提交一行代码 ...

编写babel插件,让ES6中的for/of改成并行

吴亮's Avatar 2018-02-02

  1. 1. 背景
  2. 2. 预习
  3. 3. Let’s Go
    1. 3.1. 目标
    2. 3.2. 开始开发插件了哦
    3. 3.3. 如何测试?
  4. 4. 进一步完善
    1. 4.1. 我不是所有for/of代码块都想转换
    2. 4.2. 将continue转换成return
  5. 5. DEMO代码

背景

大家都知道,ES6中的for/of是串行执行的,如果其中有await那就比较影响性能了,因此在eslint都不推荐在for/of中使用await: no-await-in-loop

但我实在又觉得for/of比较方便,又不想像eslint中推荐的那般人为了改成Promise.all,我始终觉得这个应该是编译器搞的事情,开发者不应该关心。因此我就想到我们借用babel的插件机制,自己写一个插件来将for/of编译成Promise.all形式

预习

在进行babel的插件开发时,需要先学习一下babel如何开发插件,建议看一下以下两个文档:

Let’s Go

目标

首先,我们先定好我们所期望的代码是怎么样的?下面这段是我们的测试代码:

1
2
3
4
5
6
7
8
var a = [1, 2, 3];
async function start() {
console.log("before ...");
for (let i of a) {
console.log(i);
}
console.log("after ...");
}

我们希望把其中的for/of改成Promise.all形式,那就是:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
var a = [1, 2, 3];
async function start() {
console.log("before ...");
var _executor = async function (i) {
console.log(i);
};
var tasks = [];
for (let i of a) {
tasks.push(_executor(i));
}
await Promise.all(tasks);

console.log("after ...");
}

这样看上去不错,但我想看上去更简洁一点,类似async似的,单独剥离成一个函数,因此我最终想达到的代码是:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
async function parallel(iterator, handler) {
let tasks = iterator.map((...args) => {
return handler.apply(undefined, args);
});
await Promise.all(tasks);
}

var a = [1, 2, 3];
async function start() {
console.log("before ...");
var _executor = async function (i) {
console.log(i);
};
await parallel(a, _executor);

console.log("after ...");
}

开始开发插件了哦

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import _ from 'lodash';
import {parse as babylonParse} from "babylon";

const PARSE_OPTIONS = {
allowImportExportEverywhere: true,
allowReturnOutsideFunction: true,
allowSuperOutsideMethod: true,
sourceType: 'module',
plugins: [
'estree',
'jsx',
'flow',
'doExpressions',
'objectRestSpread',
'decorators',
'classProperties',
'exportExtensions',
'asyncGenerators',
'functionBind',
'functionSent',
'dynamicImport',
'templateInvalidEscapes'
]
};

// 因为需要创建一个parallel函数,为了防止名字重复,因此取了一个时间戳+随机数作为函数名
const RANDOM = `${_.now()}_${_.parseInt(Math.random() * 100000)}`;

export default function ({types}) {
return {
visitor: {

// 为了把并行的函数放到文件的最头部
Program(nodePath, state) {
// 这块我偷懒,原本应该操作语法树来生成的,我直接让babel解析一段代码,然后把它的语法树结果拿来用就可以了
// 这个其实也是为了减轻插件开发工作量的一个好办法
const PARALLEL_FUNCTION = `
async function parallel_${RANDOM}(iterator, handler) {
let tasks = iterator.map((...args) => {
return handler.apply(undefined, args);
});
await Promise.all(tasks);
};
`;
let ast = babylonParse(PARALLEL_FUNCTION, PARSE_OPTIONS);
let executeDeclaration = _.first(ast.program.body);
nodePath.scope.block.body.unshift(executeDeclaration);
},

// 遇到for/of语句
ForOfStatement(nodePath, state) {
let { left, right, body} = nodePath.node;

let parentPath = nodePath.parentPath;
let parentScope = parentPath.scope;

// 我这边又偷懒,直接还是借用babel生成的语法树来生成一个: var executor = async function() {} 这般的语句
let executorName = parentScope.generateUidIdentifier('executor');
let codeTemplate = `
var ${executorName.name} = async function(${left.declarations[0].id.name}) {};
`;
let ast = babylonParse(codeTemplate, PARSE_OPTIONS);
let executeDeclaration = _.first(ast.program.body);
executeDeclaration.declarations[0].init.body = body;

// 生成一个: await parallel(a, executor);的语句
let taskExpression = types.awaitExpression(types.callExpression(
types.Identifier(`parallel_${RANDOM}`),
[
right,
executorName
]
));

// 把生成的语句加到这个节点中
nodePath.insertAfter(types.expressionStatement(taskExpression));
nodePath.insertAfter(executeDeclaration);
// 删除原本的for/of语句
nodePath.remove();
}
}
};
};

如何测试?

测试的话,其实很简单,写一个简单的测试文件就可以:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import plugin from "../src/transformer";
const babel = require("babel-core");

const rawCode = `
var a = [1, 2, 3];
async function start() {
console.log("before ...");
// parallel
for (let i of a) {
console.log(i);
}
console.log("after ...");
}
`;
const {code} = babel.transform(rawCode, {
plugins: [
plugin
]
});
console.log(code);

进一步完善

我不是所有for/of代码块都想转换

因为使用Promise.all转换后,所有for/of都被转换了,但这样带来一个问题就是如果我的for/of原本是保序的,就坑了。

因此,我想指定哪些for/of进行这个转换。

解决办法也很简单,我们可以再for/of上加一个注释,比如: // parallel,包含这个注释的,我们才转换:

1
2
3
4
// parallel
for (let i of a) {
console.log(i);
}

要实现这个功能,我们只要再转换先取for/of的注释来判断一下即可

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
ForOfStatement(nodePath, state) {
let { left, right, body} = nodePath.node;

// 判断一下for/of前面是否有 // parallel 这般的注释
let leadingComments = nodePath.node.leadingComments;
let comment = null;
if (!_.isArray(leadingComments) && !_.some(leadingComments, (leadingComment) => {
if (!_.isObject(leadingComment)) {
return false;
}
let comment1 = _.toLower(_.trim(leadingComment.value));
if (comment1 !== "parallel") {
return false;
}
comment = comment1;
return true
})) {
return;
}

let parentPath = nodePath.parentPath;
let parentScope = parentPath.scope;

....
}

将continue转换成return

如果原本的for/of代码块中有continue,那么需要转换成return

做法也很简单:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
// 引入一个babel-traverse,用于遍历语法树
import babelTraverse from "babel-traverse";

....

let executorName = parentScope.generateUidIdentifier('executor');
let codeTemplate = `
var ${executorName.name} = async function(${left.declarations[0].id.name}) {};
`;
let ast = babylonParse(codeTemplate, PARSE_OPTIONS);
let executeDeclaration = _.first(ast.program.body);
executeDeclaration.declarations[0].init.body = body;

// 替換try/catch中的continue
babelTraverse(ast, {
enter(nodePath) {
if (!nodePath.isContinueStatement()) {
return;
}
nodePath.replaceWith(types.returnStatement(null));
}
});

let taskExpression = types.awaitExpression(types.callExpression(
types.Identifier(`parallel_${RANDOM}`),
[
right,
executorName
]
));

注意:这个方法有一个问题,在于查找continue语句时,应该只要找一层continue语句就可以了,但我却遍历了一把,结果把所有continue语句都替换了,这块还需要优化。

DEMO代码

本示例的测试代码在sample-codes/for-of-to-parallel

本文作者 : 吴亮
本文使用 署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0) 协议
本文链接 : https://www.wuliang.me/for-of-to-parallel/

本文最后更新于 天前,文中所描述的信息可能已发生改变