Writing a plugin to instrument code

17 May 2016

This assumes some nightly Rust 1.10.0 2016-04-28

A few weeks ago, I stumbled over flame, a library to produce beautiful SVG flame graphs by instrumenting the code (which

Redditor /u/kibwen thankfully reminded me that I wasn’t doing this complex topic justice. Let’s just say that sampling and instrumentation suffer from different sets of problems, so it’s good to have both in our toolbox).

However, adding instrumentation calls to every method is a bit tedious, and I thought it’d be useful to have a #[flame] annotation that would be used to instrument all methods found within its scope. This requires a syntax extension plugin, which is pretty foreign territory to me though I have some experience with lints, which are another type of plugin after all.

Unfortunately, the documentation is pretty bare at the moment; the book shows how to create roman numerals, but that’s a bang-macro, while I wanted an annotation. Digging into the source brought up rustc_plugin::registry::Registry, which lets us register our plugin.

So, let’s take the plunge, shall we?

Like with macros, we need to activate the features plugin_registrar and rustc_private, so this is nightly only. Since we want to modify the items in place rather than adding new ones, we’ll need a MultiItemModifier, which is a trait that is already implemented for functions with the right signature. Cool.

So let’s register an identity function first, here comes flamer.rs:

#![feature(plugin_registrar, rustc_private)]

extern crate rustc_plugin;
extern crate syntax;

use rustc_plugin::registry::Registry;
use syntax::ast::MetaItem;
use syntax::codemap::Span;
use syntax::ext::base::{Annotatable, ExtCtxt, MultiItemModifier, SyntaxExtension};
use syntax::parse::token;

pub fn expand(cx: &mut ExtCtxt, span: Span, mi: &MetaItem, a: Annotatable) -> 
        Annotatable {
    a // identity function for now
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(token::intern("flame"),
        SyntaxExtension::MultiModifier(
            Box::new(expand) as Box<MultiItemModifier + 'static>));
}

This at least compiles with rustc --crate-type=dylib. Let’s try our new plugin using a flamertest.rs:

#[feature(plugin)]
#[plugin(flamer)]

#[flame]
fn main() {
    println!("It does nothing");
}

Compile with rustc -L . flamertest.rs and lo and behold: It does nothing.

That’s admittedly somewhat underwhelming. But it’s not too hard to let it do something useful; as the syntax::ast classes have a lot of functionality we can use. Of note are the following items:

• syntax::ptr::P is both a type and a method to create that type (as in P(_))
• syntax::codemap::DUMMY_SP is a dummy span. the dummy_spanned(_) function in the same module adds the dummy span to a node to create a Spanned
• syntax::ast::DUMMY_NODE_ID is the NodeId of choice when expanding ASTs
• syntax::parse::token::intern(&str) gets us a Name wherever we need one
• InternedString::new_from_name(Name) from the same module creates an InternedString for whatever Name we have (e.g. for string literals)
• manishearth.github.io/rust-internals-docs has fairly up-to-date compiler docs, including libsyntax, which is handy when dealing with AST types

Also note that the plugin only gets called when there are actual annotations, so our plugin has to walk all items if we intend to inherit annotations to inner items.

The full code is available on github.