rust - How to write a custom attribute that injects code into a function

Question

I've gotten as far as having the custom attribute invoked:

#[plugin_registrar]
pub fn registrar(reg: &mut rustc::plugin::Registry) {
  use syntax::parse::token::intern;
  use syntax::ext::base;

  // Register the `#[dummy]` attribute.
  reg.register_syntax_extension(intern("dummy"),
  base::ItemDecorator(dummy_expand));
}

// Decorator for `dummy` attribute
pub fn dummy_expand(context: &mut ext::base::ExtCtxt, span: codemap::Span, meta_item: Gc<ast::MetaItem>, item: Gc<ast::Item>, push: |Gc<ast::Item>|) {
  match item.node {
    ast::ItemFn(decl, ref style, ref abi, ref generics, block) => {
      trace!("{}", decl);
      // ...? Add something here.
    }
    _ => {
      context.span_err(span, "dummy is only permissiable on functions");
    }
  }
}

Invoked via:

#![feature(phase)]

#[phase(plugin)]
extern crate dummy_ext;

#[test]
#[dummy]
fn hello() {
  println!("Put something above this...");
}

...and I've seen a few examples around of things that use quote_expr!( ... ) to do this, but I don't really understand them.

Let's say I want to add this statement (or is it expression?) to the top of any function tagged #[dummy]:

println!("dummy");

How do I achieve that?

Answer 1

There's two tasks here:

• creating the AST you wish to insert
• transforming the AST of some function (e.g. inserting another piece)

Notes:

• when I say "item" in this answer, I specifically meant the item AST node, e.g. fn, struct, impl.
• when doing anything with macros, rustc --pretty expanded foo.rs is your best friend (works best on smallest examples, e.g. avoiding #[deriving] and println!, unless you're trying to debug those specifically).

AST creation

There's 3 basic ways to create chunks of AST from scratch:

• manually writing out the structs & enums,
• using the methods of AstBuilder to abbreviate that, and
• using quotation to avoid that altogether.

In this case, we can use quoting, so I won't waste time on the others. The quote macros take an ExtCtxt ("extension context") and an expression or item etc. and create an AST value that represents that item, e.g.

let x: Gc<ast::Expr> = quote_expr!(cx, 1 + 2);

creates an Expr_ with value ExprBinary, that contains two ExprLits (for the 1 and 2 literals).

Hence, to create the desired expression, quote_expr!(cx, println!("dummy")) should work. Quotation is more powerful than just this: you can use $ it to splice a variable storing AST into an expression, e.g., if we have the x as above, then

let y = quote_expr!(cx, if $x > 0 { println!("dummy") });

will create an AST reprsenting if 1 + 2 > 0 { println!("dummy") }.

This is all very unstable, and the macros are feature gated. A full "working" example:

#![feature(quote)]
#![crate_type = "dylib"]

extern crate syntax;

use syntax::ext::base::ExtCtxt;
use syntax::ast;

use std::gc::Gc;

fn basic_print(cx: &mut ExtCtxt) -> Gc<ast::Expr> {
    quote_expr!(cx, println!("dummy"))
}

fn quoted_print(cx: &mut ExtCtxt) -> Gc<ast::Expr> {
    let p = basic_print(cx);
    quote_expr!(cx, if true { $p })
}

As of 2014-08-29, the list of quoting macros is: quote_tokens, quote_expr, quote_ty, quote_method, quote_item, quote_pat, quote_arm, quote_stmt. (Each essentially creates the similarly-named type in syntax::ast.)

(Be warned: they are implemented in a very hacky way at the moment, by just stringifying their argument and reparsing, so it's relatively easy to encounter confusing behaviour.)

AST transformation

We now know how to make isolated chunks of AST, but how can we feed them back into the main code?

Well, the exact method depends on what you are trying to do. There's a variety of different types of syntax extensions.

• If you just wanted to expand to some expression in place (like println!), NormalTT is correct,
• if you want to create new items based on an existing one, without modifying anything, use ItemDecorator (e.g. #[deriving] creates some impl blocks based on the struct and enum items to which it is attached)
• if you want to take an item and actually change it, use ItemModifier

Thus, in this case, we want an ItemModifier, so that we can change #[dummy] fn foo() { ... } into #[dummy] fn foo() { println!("dummy"); .... }. Let's declare a function with the right signature:

fn dummy_expand(cx: &mut ExtCtxt, sp: Span, _: Gc<ast::MetaItem>, item: Gc<ast::Item>) -> Gc<Item>

This is registered with

reg.register_syntax_extension(intern("dummy"), base::ItemModifier(dummy_expand));

We've got the boilerplate set-up, we just need to write the implementation. There's two approaches. We could just add the println! to the start of the function's contents, or we could change the contents from foo(); bar(); ... to println!("dummy"); { foo(); bar(); ... } by just creating two new expressions.

As you found, an ItemFn can be matched with

ast::ItemFn(decl, ref style, ref abi, ref generics, block)

where block is the actual contents. The second approach I mention above is easiest, just

let new_contents = quote_expr!(cx, 
    println!("dummy");
    $block
);

and then to preserve the old information, we'll construct a new ItemFn and wrap it back up with the right method on AstBuilder. In total:

#![feature(quote, plugin_registrar)]
#![crate_type = "dylib"]

// general boilerplate
extern crate syntax;
extern crate rustc;

use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, ItemModifier};
// NB. this is important or the method calls don't work
use syntax::ext::build::AstBuilder;
use syntax::parse::token;

use std::gc::Gc;

#[plugin_registrar]
pub fn registrar(reg: &mut rustc::plugin::Registry) {
  // Register the `#[dummy]` attribute.
  reg.register_syntax_extension(token::intern("dummy"),
                                ItemModifier(dummy_expand));
}

fn dummy_expand(cx: &mut ExtCtxt, sp: Span, _: Gc<ast::MetaItem>, 
                item: Gc<ast::Item>) -> Gc<ast::Item> {
    match item.node {
        ast::ItemFn(decl, ref style, ref abi, ref generics, block) => {
            let new_contents = quote_expr!(&mut *cx,
                println!("dummy");
                $block
            );
            let new_item_ = ast::ItemFn(decl, style.clone(), 
                                        abi.clone(), generics.clone(),
                                        // AstBuilder to create block from expr
                                        cx.block_expr(new_contents));
            // copying info from old to new
            cx.item(item.span, item.ident, item.attrs.clone(), new_item_)
        }
        _ => {
            cx.span_err(sp, "dummy is only permissible on functions");
            item
        }
    }
}