How to modify all items of a crate in a compiler plugin?

Question

I'm trying to build a syntax extension that expands an attribute into calls. Before:

#[flame]
fn flamed() {
    ..
}

After:

fn flamed() {
    flame::start_guard("flamed");
    ..
}

This already works. However, I'd also like it to work if I have the #[flame] attribute at the crate level (like #![flame]). Is this possible and if so, how?

Answer 1

@huon's comment

Did you try catching ItemKind::Mod in github.com/llogiq/flamer/blob/… and iterating over its contents (recursively)?

was spot on – I just added a commit that handles mod and trait items by walking them. I'll also probably add code to walk functions to handle inner items and fns.

The code looks like this:

fn flame_item(i: &Item) -> Item {
    let base = i.clone();
    Item {
        node: match i.node {
            ItemKind::Mod(ref m) =>
                ItemKind::Mod(flame_mod(m)),
            ItemKind::Trait(unsafety, ref generic, ref bounds, ref tis) =>
                ItemKind::Trait(unsafety,
                                generic.clone(),
                                bounds.clone(),
                                flame_items(tis)),
        .. // other item types as usual: items, traitimpls, implitems

            _ => return base
        },
        ..base
    }
}

fn flame_mod(m: &Mod) -> Mod {
    Mod {
        inner: m.inner,
        items: m.items.iter().map(|i| P(flame_item(i))).collect()
    }
}

fn flame_items(items: &[TraitItem]) -> Vec<TraitItem> {
    items.iter().map(flame_trait_item).collect()
}