rustc_driverとrustc_interfaceを使ってみる (2019年5月版)

Rustコンパイラと連携するプログラムを書くには主に5つの方法があります。

- Rustのソースコードを単に生成する。 (`lalrpop` など)
- deriveマクロ (`serde_derive` など)
- 手続きマクロ (`runtime` など)
- コンパイラプラグイン ※nightly限定
- コンパイラの一部を使った独立ツール ※nightly限定

ここでは最後の「コンパイラの一部を使った独立ツール」を作る方法を考えます。以前はRustのコンパイラフェーズを呼び出す部分は [`rustc_driver::driver`](https://github.com/rust-lang/rust/blob/1.34.2/src/librustc_driver/driver.rs) にハードコードされていたため、複雑な処理をする場合はdriverのコードの大部分をコピーする必要がありました。しかし、2019年3月末にマージされた[#56732](https://github.com/rust-lang/rust/pull/56732)により、コンパイラフェーズの呼び出し部分が [`rustc_interface`](https://github.com/rust-lang/rust/blob/1.35.0/src/librustc_interface/lib.rs) として分離されたため、ツールの作成が比較的容易になるはずです。

## プロジェクト作成

今回はRustc Guideにある[Stupid Stats](https://rust-lang.github.io/rustc-guide/appendix/stupid-stats.html)というサンプルを作ってみます。これはマクロ展開前のASTを処理し、 「`println!` が使われている回数」と「グローバル関数の引数の個数のヒストグラム」を分析します。

[toolstate](https://rust-lang-nursery.github.io/rust-toolstate/)と[rustc component history](https://rust-lang.github.io/rustup-components-history/)を参考に、RLSなど主要ツールが揃っている中でできるだけ最新のnightlyを選びます。今回は2019-05-22を使います。

```console
$ cargo new --bin stupid-stats
$ cd stupid-stats
$ echo nightly-2019-05-22 > rust-toolchain
$ rustup component add rustfmt clippy rls rust-analysis rust-src
```

## `rustc_driver` を使った実装

Stupid Statsのページに記載のものを実装すればOKです。ただし2019-05-26時点のrustc-guideの内容とrustcの実際の構造には以下のような違いがありました。

- `CompilerCalls` ではなく `Callbacks` になっています。
- 差し込めるコールバックは減っていました。ここでは `after_parsing` を使います。
- コールバックはboolを返します。falseを返すと処理を止められるようなので、ここでは `after_parsing` で必要な処理をしてから `false` を返します。
- `rustc_driver::run_compiler` に第3引数と第4引数があります。 `None` でうまくいきそうなので `None` を渡しておきます。
- `ast::Item_` は `ast::ItemKind` になっていました。
- `ast::Mac_` はenumではなくなっていました。

コンパイラ自身のrustdocは https://doc.rust-lang.org/nightly/nightly-rustc/rustc_driver/index.html で公開されているので、必要ならこれを参照しながらプログラムを書くとよいでしょう。 (ただし、使っているnightlyのバージョンとドキュメントの元になったnightlyバージョンが異なる可能性に注意が必要)

```rust:main.rs
// Rustコンパイラ自身を使うので必要
#![feature(rustc_private)]

extern crate rustc_driver;
extern crate rustc_interface;
extern crate syntax;

use rustc_driver::Callbacks;
use rustc_interface::interface;
use syntax::{ast, visit};

fn main() {
    let args: Vec<_> = std::env::args().collect();
    rustc_driver::run_compiler(&args, &mut StupidCallbacks, None, None).unwrap();
}

struct StupidCallbacks;

impl Callbacks for StupidCallbacks {
    fn after_parsing(&mut self, compiler: &interface::Compiler) -> bool {
        // 既にパースした結果のキャッシュを取り出す
        let cratename = compiler.crate_name().unwrap().peek().clone();
        let krate = compiler.parse().unwrap().peek();

        // ASTをトラバースする
        let mut visitor = StupidVisitor::new();
        visit::walk_crate(&mut visitor, &krate);

        // 結果を出力
        println!("In crate: {},\n", cratename);
        println!("Found {} uses of `println!`;", visitor.println_count);

        let (common, common_percent, four_percent) = visitor.compute_arg_stats();
        println!(
            "The most common number of arguments is {} ({:.0}% of all functions);",
            common, common_percent
        );
        println!(
            "{:.0}% of functions have four or more arguments.",
            four_percent
        );

        // falseを投げると以降の処理を中止できる
        false
    }
}

struct StupidVisitor {
    println_count: usize,
    arg_counts: Vec<usize>,
}

impl StupidVisitor {
    fn new() -> Self {
        Self {
            println_count: 0,
            arg_counts: Vec::new(),
        }
    }

    fn increment_args(&mut self, args: usize) {
        if self.arg_counts.len() <= args {
            self.arg_counts.resize(args + 1, 0);
        }
        self.arg_counts[args] += 1;
    }

    fn compute_arg_stats(&self) -> (usize, f64, f64) {
        let iter = || self.arg_counts.iter().copied();
        let sum = iter().sum::<usize>();
        let sum4 = iter().skip(4).sum::<usize>();
        let max = iter().max().unwrap_or(0);
        let max_idx = iter().position(|x| x == max).unwrap_or(0);
        (
            max_idx,
            max as f64 / sum as f64 * 100.0,
            sum4 as f64 / sum as f64 * 100.0,
        )
    }
}

// visit::Visitor を実装することでASTのトラバースができる
impl<'v> visit::Visitor<'v> for StupidVisitor {
    fn visit_item(&mut self, i: &'v ast::Item) {
        match i.node {
            ast::ItemKind::Fn(ref decl, _, _, _) => {
                self.increment_args(decl.inputs.len());
            }
            _ => {}
        }

        // 続けて、再帰的にスキャンしたいときは `walk_*` を呼ぶ
        visit::walk_item(self, i)
    }

    fn visit_mac(&mut self, mac: &'v ast::Mac) {
        if mac.node.path.to_string() == "println" {
            self.println_count += 1;
        }

        // 続けて、再帰的にスキャンしたいときは `walk_*` を呼ぶ
        visit::walk_mac(self, mac)
    }
}
```

## 実行してみる

`driver` を使って作ったツールは `rustc` と同様の引数をとるので、自分自身を引数として渡してみます。

```console
$ cargo run src/main.rs
   Compiling stupid-stats v0.1.0 (/home/qnighy/workdir/stupid-stats)
    Finished dev [unoptimized + debuginfo] target(s) in 6.24s
     Running `target/debug/stupid-stats src/main.rs`
In crate: main,

Found 4 uses of `println!`;
The most common number of arguments is 0 (100% of all functions);
0% of functions have four or more arguments.
```

## `rustc_interface` を使った実装

`rustc_driver` はあくまで `rustc` コマンドの仕組みがベースにあり、そこにフックを差し込む形で利用するものでした。より主体的に `rustc` 内部のコンパイラフェーズを利用したい場合は `rustc_interface` を直接呼びます。

```rust:main.rs
// Rustコンパイラ自身を使うので必要
#![feature(rustc_private)]

extern crate rustc;
extern crate rustc_interface;
extern crate syntax;

use rustc::session::config::Input;
use rustc::session::DiagnosticOutput;
use rustc_interface::interface;
use syntax::{ast, visit};

fn main() {
    // コマンドライン引数を自分で解析する
    // rustc_driverにコマンドライン引数用のヘルパーもあるので、
    // rustcコマンドに近づけるならそれを使う手もある
    let input: std::path::PathBuf = std::env::args().nth(1).unwrap().into();

    // 今回はパースするだけなのであまりオプションは必要ない。
    // ほぼデフォルトで組み立てる
    let config = interface::Config {
        opts: Default::default(),
        crate_cfg: Default::default(),
        input: Input::File(input.clone()),
        input_path: Some(input.clone()),
        output_file: None,
        output_dir: None,
        file_loader: None,
        diagnostic_output: DiagnosticOutput::Default,
        stderr: None,
        crate_name: None,
        lint_caps: Default::default(),
    };

    // コンパイラのためのリソース (スレッドプールやアロケータ) を起動する
    interface::run_compiler(config, |compiler| {
        // クレート名の取得とパースをする
        let cratename = compiler.crate_name().unwrap().peek().clone();
        let krate = compiler.parse().unwrap().peek();

        // ASTをトラバースする
        let mut visitor = StupidVisitor::new();
        visit::walk_crate(&mut visitor, &krate);

        // 結果を出力
        println!("In crate: {},\n", cratename);
        println!("Found {} uses of `println!`;", visitor.println_count);

        let (common, common_percent, four_percent) = visitor.compute_arg_stats();
        println!(
            "The most common number of arguments is {} ({:.0}% of all functions);",
            common, common_percent
        );
        println!(
            "{:.0}% of functions have four or more arguments.",
            four_percent
        );
    });
}

struct StupidVisitor {
    println_count: usize,
    arg_counts: Vec<usize>,
}

impl StupidVisitor {
    fn new() -> Self {
        Self {
            println_count: 0,
            arg_counts: Vec::new(),
        }
    }

    fn increment_args(&mut self, args: usize) {
        if self.arg_counts.len() <= args {
            self.arg_counts.resize(args + 1, 0);
        }
        self.arg_counts[args] += 1;
    }

    fn compute_arg_stats(&self) -> (usize, f64, f64) {
        let iter = || self.arg_counts.iter().copied();
        let sum = iter().sum::<usize>();
        let sum4 = iter().skip(4).sum::<usize>();
        let max = iter().max().unwrap_or(0);
        let max_idx = iter().position(|x| x == max).unwrap_or(0);
        (
            max_idx,
            max as f64 / sum as f64 * 100.0,
            sum4 as f64 / sum as f64 * 100.0,
        )
    }
}

// visit::Visitor を実装することでASTのトラバースができる
impl<'v> visit::Visitor<'v> for StupidVisitor {
    fn visit_item(&mut self, i: &'v ast::Item) {
        match i.node {
            ast::ItemKind::Fn(ref decl, _, _, _) => {
                self.increment_args(decl.inputs.len());
            }
            _ => {}
        }

        // 続けて、再帰的にスキャンしたいときは `walk_*` を呼ぶ
        visit::walk_item(self, i)
    }

    fn visit_mac(&mut self, mac: &'v ast::Mac) {
        if mac.node.path.to_string() == "println" {
            self.println_count += 1;
        }

        // 続けて、再帰的にスキャンしたいときは `walk_*` を呼ぶ
        visit::walk_mac(self, mac)
    }
}
```

実行結果は元のものと同じです。

## まとめ

- `rustc` コマンドをベースに、途中にフックを挟む形で独自コマンドを作るには、 `rustc_driver::run_compiler` が使えます。
- より独自のフローでrustcの一部を利用するには、 `rustc_interface::run_compiler` が使えます。