Jon Haddad
Jon Haddad

Reputation: 766

Macro that generates a function with arguments determined by the macro

Is it possible to write a macro that generates a function where the number of arguments to this function to be a determined by the macro? For instance, I'd like to write something to make using prepared statements in the Cassandra driver easier.

let prepared = prepare!(session, "insert into blah (id, name, reading ) values (?, ?, ?)", int, string, float);
let stmt = prepared(1, "test".to_string(), 3.1);
session.execute(stmt);

prepare! would need to generate something like (unwrap only here for brevity):

fn some_func(arg1, arg2, arg3) -> Statement {
    let mut statement = Statement::new("insert into blah (id, name, reading ) values (?, ?, ?)", 3);
    statement.bind_int(0, arg1).unwrap()
        .bind_string(1, arg2).unwrap()
        .bind_float(2, arg3).unwrap()
}

Upvotes: 3

Views: 1936

Answers (1)

DK.
DK.

Reputation: 59155

Two hard things in Rust macros: counting and unique identifers. You have both. Then again, I'm the one writing the answer, so I suppose it's my problem now. At least you didn't ask about parsing the string (which is outright impossible without compiler plugins).

Another impossible thing is mapping types to different methods. You just can't. Instead, I'm going to assume the existence of a helper trait that does this mapping.

Also, Rust doesn't have int, string, or float. I assume you mean i32, String, and f32.

Finally, the way you've written the invocation and expansion don't really gel. I don't see why session is involved; it's not used in the expansion. So I'm going to take the liberty of just pretending you don't need it; if you do, you'll have to hack it back in.

So, with that, here's what I came up with.

// Some dummy types so the following will type-check.

struct Statement;

impl Statement {
    fn new(stmt: &str, args: usize) -> Self { Statement }
    fn bind_int(self, pos: usize, value: i32) -> Result<Self, ()> { Ok(self) }
    fn bind_float(self, pos: usize, value: f32) -> Result<Self, ()> { Ok(self) }
    fn bind_string(self, pos: usize, value: String) -> Result<Self, ()> { Ok(self) }
}

struct Session;

impl Session {
    fn execute(&self, stmt: Statement) {}
}

// The supporting `BindArgument` trait.

trait BindArgument {
    fn bind(stmt: Statement, pos: usize, value: Self) -> Statement;
}

impl BindArgument for i32 {
    fn bind(stmt: Statement, pos: usize, value: Self) -> Statement {
        stmt.bind_int(pos, value).unwrap()
    }
}

impl BindArgument for f32 {
    fn bind(stmt: Statement, pos: usize, value: Self) -> Statement {
        stmt.bind_float(pos, value).unwrap()
    }
}

impl BindArgument for String {
    fn bind(stmt: Statement, pos: usize, value: Self) -> Statement {
        stmt.bind_string(pos, value).unwrap()
    }
}

// The macro itself.

macro_rules! prepare {
    // These three are taken straight from
    // https://danielkeep.github.io/tlborm/book/
    (@as_expr $e:expr) => {$e};

    (@count_tts $($tts:tt)*) => {
        <[()]>::len(&[$(prepare!(@replace_tt $tts ())),*])
    };

    (@replace_tt $_tt:tt $e:expr) => {$e};

    // This is how we bind *one* argument.

    (@bind_arg $stmt:expr, $args:expr, $pos:tt, $t:ty) => {
        prepare!(@as_expr <$t as BindArgument>::bind($stmt, $pos, $args.$pos))
    };

    // This is how we bind *N* arguments.  Note that because you can't do
    // arithmetic in macros, we have to spell out every supported integer.
    // This could *maybe* be factored down with some more work, but that
    // can be homework.  ;)

    (@bind_args $stmt:expr, $args:expr, 0, $next:ty, $($tys:ty,)*) => {
        prepare!(@bind_args prepare!(@bind_arg $stmt, $args, 0, $next), $args, 1, $($tys,)*)
    };

    (@bind_args $stmt:expr, $args:expr, 1, $next:ty, $($tys:ty,)*) => {
        prepare!(@bind_args prepare!(@bind_arg $stmt, $args, 1, $next), $args, 2, $($tys,)*)
    };

    (@bind_args $stmt:expr, $args:expr, 2, $next:ty, $($tys:ty,)*) => {
        prepare!(@bind_args prepare!(@bind_arg $stmt, $args, 2, $next), $args, 3, $($tys,)*)
    };

    (@bind_args $stmt:expr, $_args:expr, $_pos:tt,) => {
        $stmt
    };

    // Finally, the entry point of the macro.

    ($stmt:expr, $($tys:ty),* $(,)*) => {
        {
            // I cheated: rather than face the horror of trying to *also* do
            // unique identifiers, I just shoved the arguments into a tuple, so
            // that I could just re-use the position.
            fn prepared_statement(args: ($($tys,)*)) -> Statement {
                let statement = Statement::new(
                    $stmt,
                    prepare!(@count_tts $(($tys))*));
                prepare!(@bind_args statement, args, 0, $($tys,)*)
            }
            prepared_statement
        }
    };
}

fn main() {
    let session = Session;
    let prepared = prepare!(
        r#"insert into blah (id, name, reading ) values (?, ?, ?)"#,
        i32, String, f32);
    // Don't use .to_string() for &str -> String; it's horribly inefficient.
    let stmt = prepared((1, "test".to_owned(), 3.1));
    session.execute(stmt);
}

And here's what the main function expands to, to give you a frame of reference:

fn main() {
    let session = Session;
    let prepared = {
        fn prepared_statement(args: (i32, String, f32)) -> Statement {
            let statement = Statement::new(
                r#"insert into blah (id, name, reading ) values (?, ?, ?)"#,
                <[()]>::len(&[(), (), ()]));
            <f32 as BindArgument>::bind(
                <String as BindArgument>::bind(
                    <i32 as BindArgument>::bind(
                        statement, 0, args.0),
                    1, args.1),
                2, args.2)
        }
        prepared_statement
    };
    // Don't use .to_string() for &str -> String; it's horribly inefficient.
    let stmt = prepared((1, "test".to_owned(), 3.1));
    session.execute(stmt);
}

Upvotes: 9

Related Questions