Search code examples
macrosrustrust-macros

Is there a way to count with macros?


I want to create a macro that prints "Hello" a specified number of times. It's used like:

many_greetings!(3);  // expands to three `println!("Hello");` statements

The naive way to create that macro is:

macro_rules! many_greetings {
    ($times:expr) => {{
        println!("Hello");
        many_greetings!($times - 1);
    }};
    (0) => ();
}

However, this doesn't work because the compiler does not evaluate expressions; $times - 1 isn't calculated, but fed as a new expression into the macro.


Solution

  • While the ordinary macro system does not enable you to repeat the macro expansion many times, there is no problem with using a for loop in the macro:

    macro_rules! many_greetings {
        ($times:expr) => {{
            for _ in 0..$times {
                println!("Hello");
            }
        }};
    }
    

    If you really need to repeat the macro, you have to look into procedural macros/compiler plugins (which as of 1.4 are unstable, and a bit harder to write).

    Edit: There are probably better ways of implementing this, but I've spent long enough on this for today, so here goes. repeat!, a macro that actually duplicates a block of code a number of times:

    main.rs

    #![feature(plugin)]
    #![plugin(repeat)]
    
    fn main() {
        let mut n = 0;
        repeat!{ 4 {
            println!("hello {}", n);
            n += 1;
        }};
    }
    

    lib.rs

    #![feature(plugin_registrar, rustc_private)]
    
    extern crate syntax;
    extern crate rustc;
    
    use syntax::codemap::Span;
    use syntax::ast::TokenTree;
    use syntax::ext::base::{ExtCtxt, MacResult, MacEager, DummyResult};
    use rustc::plugin::Registry;
    use syntax::util::small_vector::SmallVector;
    use syntax::ast::Lit_;
    use std::error::Error;
    
    fn expand_repeat(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
        let mut parser = cx.new_parser_from_tts(tts);
        let times = match parser.parse_lit() {
            Ok(lit) => match lit.node {
                Lit_::LitInt(n, _) => n,
                _ => {
                    cx.span_err(lit.span, "Expected literal integer");
                    return DummyResult::any(sp);
                }
            },
            Err(e) => {
                cx.span_err(sp, e.description());
                return DummyResult::any(sp);
            }
        };
        let res = parser.parse_block();
    
        match res {
            Ok(block) => {
                let mut stmts = SmallVector::many(block.stmts.clone());
                for _ in 1..times {
                    let rep_stmts = SmallVector::many(block.stmts.clone());
                    stmts.push_all(rep_stmts);
                }
                MacEager::stmts(stmts)
            }
            Err(e) => {
                cx.span_err(sp, e.description());
                DummyResult::any(sp)
            }
        }
    }
    
    #[plugin_registrar]
    pub fn plugin_registrar(reg: &mut Registry) {
        reg.register_macro("repeat", expand_repeat);
    }
    

    added to Cargo.toml

    [lib]
    name = "repeat"
    plugin = true
    

    Note that if we really don't want to do looping, but expanding at compile-time, we have to do things like requiring literal numbers. After all, we are not able to evaluate variables and function calls that reference other parts of the program at compile time.