使用 Rust 开发一个 SWC 插件

为了编译效率,你将编译工具从 babel 切换到了 swc,于是你开始研究 swc 插件怎么写了。

swc 相关开发包

js 侧相关的 npm packages:

rust 侧相关的 crates:

swc 插件最终会被编译为 wasm 来运行

初始化项目

本文可配合官网文档的插件部分介绍一并食用,先看官网文档或是先看本文皆可。

从官网文档得知,插件并不向后兼容,因此我们在编写插件时需要保证 swc_core 版本和 js 侧使用的 @swc/core 版本对应,否则可能出现异常报错。

编写本文时,官网文档里列举的最新版是 ”v1.3.81(Unpublished)“,但当我们去 github release 页面看一眼,发现 @swc/core 都发版到 v1.3.99(截止本文编写时)了,可知官网文档并非同步更新的,好在每次发布中标题底部都附带了 swc_core 的版本,我们本次将使用目前最新的 release 版本进行开发:

rust: swc_core@0.86.29
js:   @swc/core@1.3.99

现在根据文档的 Getting stared 指引,安装 cli 并初始化插件开发模板:

cargo install swc_cli
swc plugin new --target-type wasm32-wasi swc-plugin-extract-fn

使用 vscode 打开项目,打开 lib.rs 并等待 rust-analyzer 执行完后,会发现开门一个红色报错 error: no rules expected the token `r#"console.log("transform");"#` ,经过笔者探究,从此 issue 得知不久前该项目将代码测试方式从内联重构成了外部文件的形式,但是并未更新插件模板的代码。那如何使用内联方式呢?只需将 lib.rs 里的 test!(...) 代码改成下述即可:

#[test]
fn demo_test() {
    swc_core::ecma::transforms::testing::test_transform::test_transform(
        Default::default(),
        |_| as_folder(TransformVisitor),
        r#"console.log("transform");"#,
        r#"console.log("transform");"#,
        false,
    );
}

然后点击函数上方的 Run test 或是手动在命令行内执行 cargo test 来执行用例,确保一切运行正常。

了解 swc 插件结构

现在我们简单了解一下这份初始代码。

首先是定义了一个没有任何属性的 struct(这种 struct 被称作 unit struct),通过为其实现不同的 trait 而赋予它不同的能力,swc 为我们提供了 Visit, VisitMut, Fold 等多种 traits,他们都实现了设计模式中访问者模式(Visitor Pattern),这些 trait 的主要区别在于访问函数的签名有所不同,比如 VisitMut 和 Fold:

pub struct TransformVisitor;
impl VisitMut for TransformVisitor {
    // VisitMut 提供了 ast 节点的 mutable 引用,我们可以直接修改它来改变 ast
    fn visit_mut_array_lit(&mut self, n: &mut ArrayLit) {}
}

pub struct TransformVisitorFold;
impl Fold for TransformVisitorFold {
    // Fold 提供了 ast 节点的 immutable 引用,我们需要返回一个新的修改后的 ast 节点
    fn fold_array_lit(&mut self, n: ArrayLit) -> ArrayLit {}
}

当我们在实现某个 visit_xxx 函数时,会发现 vscode 的代码提示需要较长时间才会展示,这可能是因为每个 trait 下的函数个数太多了(数百个),不过我们可以直接查看对应 trait 的文档,里面已经展示了所有的函数,查阅更加方便快捷些。

回到代码,看到接着定义了一个使用了 plugin_transform 修饰的函数,函数里的内容便是最终会被执行的部分,所以我们在这里编写具体的转换逻辑,最终返回一个转换后的 ast。

插件实现

了解了上述一些基础信息后,现在来看看如何实现我们的插件吧。本次我们将实现”移除某个指定函数的包裹“的转换插件,这里可能不太理解意思,我们看一个示例效果,其中的”指定函数“ 假设为 withLog:

// 转换前
withLog(addWithoutLog);
withMonitor(addWithoutLog);

// 转换后
addWithoutLog(); // withLog 被移除了,直接调用参数
withMonitor(addWithoutLog); // 非 withLog 包裹的则不变

现在大家应该知道我们的插件要做什么了,那么如同之前的文章 使用 jscodeshift 重构你的代码 里所提到的,编写插件前我们需要列出所有需要处理的写法,因此先写测试用例是个比较好的方式,在 rust 中编写单元测试非常方便,使用 #[test] 修饰即可:

// cfg 条件宏,这里表示里面的代码仅在测试场景下使用到
#[cfg(test)]
mod test {
    use super::*;
    use swc_core::ecma::transforms::testing::test_transform;

    #[test]
    fn test_fn_variable() {
        test_transform(
            Default::default(),
            |_| as_folder(TransformVisitor),
            r#"withLog(originalFn);"#,
            r#"originalFn()"#,
            false,
        );
    }

    #[test]
    fn test_inline_fn() {
        test_transform(
            Default::default(),
            |_| as_folder(TransformVisitor),
            r#"withLog(() => {
                    custom_logic();
                });"#,
            r#"(() => {
                    custom_logic();
                })();"#,
            false,
        );
    }

    #[test]
    fn test_fn_without_change() {
        test_transform(
            Default::default(),
            |_| as_folder(TransformVisitor),
            r#"originalFn();"#,
            r#"originalFn();"#,
            false,
        );
    }
}

上面用例写完后,发现代码颇为冗余,来让我们使用 rust 中强大的自定义宏做一下封装吧,下面是变更后的内容:

#[cfg(test)]
mod test {
    use super::*;
    use swc_core::ecma::transforms::testing::test_transform;

    // 新建一个叫做 test_case 的自定义宏
    macro_rules! test_case {
        ($case_name: ident, $input:expr, $output:expr) => {
            #[test]
            fn $case_name() {
                test_transform(
                    Default::default(),
                    |_| {
                        as_folder(TransformVisitor::new(Config {
                            fn_name: "withLog".into(),
                        }))
                    },
                    $input,
                    $output,
                    false,
                );
            }
        };
    }

    test_case!(
        test_fn_variable,
        r#"withLog(originalFn);"#,
        r#"originalFn()"#
    );

    test_case!(
        test_inline_fn,
        r#"withLog(() => {
            custom_logic();
        });"#,
        r#"(() => {
            custom_logic();
        })();"#
    );

    test_case!(
        test_fn_without_change,
        r#"originalFn();"#,
        r#"originalFn();"#
    );
}

现在测试用例写起来和看起来都感觉清爽多了,宏真的超棒呀!

插件参数处理

回到最初的需求,我们需要移除某个特定函数,但是目前来看并没有参数处理的逻辑,那么 swc 插件如何处理自定义参数呢?

这个官方文档还真没提及,好在 swc 生态良好,现成的开源插件如此之多,看看他们怎么写不就知道了。笔者在官方提供的插件仓库中找到了传参写法,从中可以看到我们可以获取到一个 json 字符串,需要手动反序列化为 rust struct,那么便涉及到 json 的处理,需要安装 serde 和 serde_json 两个 crate:

cargo add serde serde_json

那么现在我们来为插件添加参数,使其接收一个待移除的函数名称字符串:

use serde::Deserialize;
use swc_core::{
    ecma::{
        ast::*,
        atoms::JsWord,
        visit::VisitMutWith,
        visit::{as_folder, FoldWith, VisitMut},
    },
    plugin::{plugin_transform, proxies::TransformPluginProgramMetadata},
};

// 该 struct 定义插件接收的参数内容
#[derive(Deserialize)]
// 该属性用以兼顾命名规范:js 中多以驼峰命名,rust 多以小写+下划线命名
#[serde(rename_all = "camelCase")]
struct Config {
    // JsWord 可理解为一种优化后的字符串,大家可自行 google "copy on write" 做深入了解
    fn_name: JsWord,
}

// 由于插件需要接收参数,所以不再使用 unit struct,这里为其添加一个配置属性
struct TransformVisitor {
    config: Config,
}

impl TransformVisitor {
    // 并新建一个构造函数
    fn new(config: Config) -> Self {
        Self { config }
    }
}

impl VisitMut for TransformVisitor {
    // 具体逻辑待编写...
}

#[plugin_transform]
pub fn process_transform(program: Program, metadata: TransformPluginProgramMetadata) -> Program {
    // 获取插件参数
    let config = serde_json::from_str::<Config>(
        &metadata
            .get_transform_plugin_config()
            .expect("Failed to get plugin config"),
    )
    .expect("Invalid config");

    // 操作 ast 并返回新的 ast
    program.fold_with(&mut as_folder(TransformVisitor::new(config)))
}

访问函数实现

现在则可以实现具体的访问函数了,由于我们需要处理的是函数调用部分,故将使用到涉及 CallExpression 的访问函数:

impl VisitMut for TransformVisitor {

    fn visit_mut_call_expr(&mut self, n: &mut CallExpr) {
        // 需要手动调用该函数确保子节点可以被处理到
        n.visit_mut_children_with(self);

        match &n.callee {
            Callee::Expr(expr) => match &**expr {
                Expr::Ident(ident) => {
                    // 获取到函数调用的函数名称
                    let val = ident.sym.as_str();
                    // 如何等于插件配置里设置的名称,则进一步做处理
                    if val == self.config.fn_name.as_str() {
                        // 处理一些边界场景,比如如果函数参数没有或存在多个时,提醒但不做处理
                        if n.args.len() == 0 || n.args.len() > 1 {
                            // 这里是 swc 中的报错、告警等场景的处理方式
                            HANDLER.with(|handler| {
                                handler
                                    .struct_span_warn(
                                        n.span,
                                        format!(
                                        "Arguments of {val} should not be zero or more than one"
                                    )
                                        .as_str(),
                                    )
                                    .emit();
                            });
                            return;
                        }

                        // 修改本次访问的节点,变成第一个参数作为函数进行调用
                        let arg = n.args.first().unwrap();
                        *n = CallExpr {
                            span: DUMMY_SP,
                            callee: Callee::Expr(arg.expr.clone()),
                            args: vec![],
                            type_args: None,
                        };
                    }
                }
                _ => {}
            },
            _ => {}
        }
    }
}

上述代码若对于接触 rust 不多的读者,可能发现看起来有点头疼,但本文核心并非介绍 rust,我们唯有不断学习和主动探索才能毫无畏惧。

现在我们执行测试用例,all passed!

在 JS 侧使用插件

插件就算编写完了,那么如何在 js 侧使用呢?

正常流程是将其发布为 npm 包,在项目初始化时发现已经有了一个 package.json,我们按照正常的 npm 发包流程走便可以了。当然在发布前需要本地验证下,这里为了简单,我们直接在同一个目录下进行验证吧:

  1. 手动编译插件为 wasm 文件:
cargo build-wasi --release
  1. 安装 js 侧所需依赖:
npm i @swc/cli @swc/core -D
  1. 新建一个配置文件 .swcrc:
{
  "jsc": {
    "experimental": {
      "plugins": [
        [
          // 由于我们没有发成 npm 包,此处不能直接写成包名 swc-plugin-extract-fn,我们可以传入一个相对路径
          "./",
          {
            "fnName": "withLog"
          }
        ]
      ]
    }
  }
}
  1. 新建一个测试文件 demo.js:
withLog(originalFn);
  1. 现在在命令行内执行转换:
npx swc ./demo.js

如果正常输出了转换后的结果,则表明成功了。但是,笔者略有不幸,它报错了:

thread '<unnamed>' panicked at library/std/src/sys/wasi/mod.rs:117:37:
random_get failure: Errno { code: 0, name: "SUCCESS", message: "No error occurred. System call completed successfully." }
...

看着让人一脸懵的报错我挠了挠秃秃的脑袋。报错信息里也没有打印我们代码里的错误提示,难道是那里写法不兼容?尝试安装了官方仓库的插件来执行转换,发现是正常的,于是和官方仓库的文件逐个对比,最终发现是 Cargo.toml 里的一个配置所致:

[profile.release]
# 将此处配置移除或是改为 false 重新编译便可
lto = true

修改上述后重新编译,成功!swc 插件开发便至此结束。

后记

由于官方文档并非面面俱到,我们很多时候需要需要去参考代码,下述是笔者在编写插件期间所搜集的部分信息,可以方便更好地掌握 swc 插件开发: