rustmacros

How to correctly capture writer in a macro that wraps multiple calls to `write!`


I have a macro that wraps multiple calls to write!. For a toy example, suppose we have this macro:

use std::fmt::Write;

macro_rules! write_twice {
    ($writer:expr, $($fmt:tt)*) => {{
        (|| {
            for _ in 0..2 {
                if let err @ ::core::result::Result::Err(_) = write!($writer, $($fmt)*) {
                    return err;
                }
            }
            Ok(())
        })()
    }};
}

fn main() {
    let v = vec![1, 2];
    let mut buf = String::new();
    write_twice!(&mut buf, "{v:?}").unwrap();
    println!("{buf:?}");  // [1, 2][1, 2]
}

Now, I would like to capture the writer correctly, but I am struggling to make it work in all cases. Namely:

  1. If an expression is passed, it should only be evaluated once, so that something like write_twice(fs::File::create(path)?, args...) works without the second write re-truncating the file.
  2. The writer should still be usable after the macro call, i.e. $writer should not be moved.
  3. If an instance of write!(w, args...) compiles, then it should also compile when replaced with write_twice!(w, args...).

Now, the above version of write_twice! fails #1, as write!($writer, $($fmt)*) will be evaluated twice. We could replace it with this:

macro_rules! write_twice {
    ($writer:expr, $($fmt:tt)*) => {{
        (|| {
            let mut writer = $writer;  // evaluate only once
            for _ in 0..2 {
                if let err @ ::core::result::Result::Err(_) = write!(writer, $($fmt)*) {
                    return err;
                }
            }
            Ok(())
        })()
    }};
}

But this violates #2 because a &mut T is not Copy:

use std::fmt;

struct Double(String);

impl fmt::Display for Double {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "first -> ")?;
        write_twice!(f, "{}", self.0)?;
        write!(f, " <- second")  // borrow of moved value: `f`
    }
}
error[E0382]: borrow of moved value: `f`
  --> src/main.rs:23:10
   |
5  |         (|| {
   |          -- value moved into closure here
...
20 |     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
   |                   - move occurs because `f` has type `&mut Formatter<'_>`, which does not implement the `Copy` trait
21 |         write!(f, "first -> ")?;
22 |         write_twice!(f, "{}", self.0)?;
   |                      - variable moved due to use in closure
23 |         write!(f, " <- second")
   |                ^ value borrowed here after move

So then I tried replacing that line with let writer = &mut $writer; to avoid moving the &mut T, but now we get a new compiler error, violating #3 (although the compiler’s suggestion does fix the issue):

error[E0596]: cannot borrow `f` as mutable, as it is not declared as mutable
  --> src/main.rs:6:17
   |
6  |             let writer = &mut $writer;
   |                          ^^^^^^^^^^^^ cannot borrow as mutable
...
22 |         write_twice!(f, "{}", self.0)?;
   |         ----------------------------- in this macro invocation
   |
   = note: this error originates in the macro `write_twice` (in Nightly builds, run with -Z macro-backtrace for more info)
help: consider changing this to be mutable
   |
20 |     fn fmt(&self, mut f: &mut fmt::Formatter<'_>) -> fmt::Result {
   |                   +++

And if we instead try let mut writer = &mut *$writer, we can no longer call write_twice!(string, ...) because we're now trying to write to a &mut str instead of a &mut String.

So, what is the proper way to capture a writer in a macro that writes to it multiple times?


Solution

  • use std::fmt::{self, Write};
    
    trait WithMut {
        fn with_mut<'a, R>(&'a mut self, f: impl FnOnce(&'a mut Self) -> R) -> R;
    }
    
    impl<T> WithMut for T {
        fn with_mut<'a, R>(&'a mut self, mut f: impl FnOnce(&'a mut Self) -> R) -> R {
            f(self)
        }
    }
    
    macro_rules! write_twice {
        ($writer:expr, $($fmt:tt)*) => {{
            use $crate::WithMut; // make sure the trait is in scope
            $writer.with_mut(|mut writer| {
                for _ in 0..2 {
                    if let err @ ::core::result::Result::Err(_) = write!(writer, $($fmt)*) {
                        return err;
                    }
                }
                Ok(())
            })
        }};
    }
    
    struct Double(String);
    
    impl fmt::Display for Double {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "first -> ")?;
            write_twice!(f, "{}", self.0)?;
            write!(f, " <- second")
        }
    }
    
    fn main() {
        let v = vec![1, 2];
        let mut buf = String::new();
        write_twice!(&mut buf, "{v:?}").unwrap();
        println!("{buf:?}");  // [1, 2][1, 2]
        
        println!("{}", Double("hello".into())); // first -> hellohello <- second
    }
    

    This solution relies on an auxiliary trait, which I've called WithMut and provides a with_mut method. Since the write! macro relies on autoref to call write_fmt, we need to do the same in our macro. The purpose of WithMut::with_mut is to force $writer to be coerced into a mutable reference. The writer is passed back as &mut Self as an argument to the closure. This ensures that $writer is only evaluated once while allowing both values and mutable references to be passed in without moving.