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

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

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

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

プロジェクト作成

今回はRustc GuideにあるStupid Statsというサンプルを作ってみます。これはマクロ展開前のASTを処理し、 「println! が使われている回数」と「グローバル関数の引数の個数のヒストグラム」を分析します。

toolstateとrustc component historyを参考に、RLSなど主要ツールが揃っている中でできるだけ最新のnightlyを選びます。今回は2019-05-22を使います。

$ 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バージョンが異なる可能性に注意が必要)

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 と同様の引数をとるので、自分自身を引数として渡してみます。

$ 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 を直接呼びます。

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 が使えます。