Dynamic dispatch on display trait

Hello
So, I have a array of structs I want to print, which are the result of some computation.
But depending on some condition I want to use a special display implementation that compares them to the original data. I would like to reuse the printing code, so it's not duplicated.

I've made it work with a custom enum wrapper so I do the dispatch to the display implementations manually. But I was curious to see if there was a more dynamic way to do it that didn't require an add-hoc wrapper.

I've tried using closures but I'm hitting lifetime issues and I don't know if this is even doable. Here's a simplified version of the code, if anyone has comments or suggestions :slight_smile: :

use std::fmt::{self, Display};
use rand::random;

// Some data to display. Not Copy.
#[derive(Debug)]
struct A;

// A regular display implementation.
impl Display for A {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

fn main() {
    let x = A;
    let xs = vec![A, A, A, A];
    print_xs(&x, &xs);
}

fn print_xs(x: &A, xs: &[A]) {
    if xs.is_empty() { return }
    
    // I can't get the dynamic display to work with closures.
    let special_display = &mut SpecialDisplay(&x, &xs[0]);
    let display: &dyn Fn(&A) -> &dyn Display = if need_special_display() {
        &|x: &A| {
            special_display.1 = x;
            special_display
        }
    } else {
        &|x: &A| x
    };

    // some printing code below
    for x in xs {
        println!("{}", display(x));
    }
}

fn need_special_display() -> bool{
    random()
}

struct SpecialDisplay<'a>(&'a A, &'a A);

// A special display implementation.
impl Display for SpecialDisplay<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}-{:?}", self.0, self.1)
    }
}

That won't work. There are two tricky parts here. First, if you're returning a reference (like &dyn Display in your example), you can't construct the value inside the closure and return a reference to it. The referenced object must outlive the reference, so it must be stored somewhere while you're using the reference. That's not specific to trait objects, that's just how references work in general.

Second, there is a similar issue with the display variable. You can't construct it out of references to temporaries. The borrowed closures returned from both if branches must live long enough. For example, they could be stored in local variables. Here's an example when it would work:

fn print_xs(x: &A, xs: &[A]) {
    if xs.is_empty() { return }
    
    let special_display = |t: &A| println!("{}", SpecialDisplay(t, x));
    
    let display: &dyn Fn(&A) = if need_special_display() {
        &special_display
    } else {
        &|t| println!("{}", t)
        
    };
    
    // some printing code below
    for x in xs {
        display(x);
    }
}

Note that the first variant of the closure is stored as a local variable. The second closure doesn't have any context, so it lives forever.

When you have a problem with dynamic return value, it's usually better to refactor the code so that the return type is generic. For example:

fn print_xs_with_display<'a, F, D>(xs: &'a [A], mut display: F)
where
    F: FnMut(&'a A) -> D + 'a,
    D: Display,
{
    for x in xs {
        println!("{}", display(x));
    }
}

fn print_xs(x: &A, xs: &[A]) {
    if need_special_display() {
        print_xs_with_display(xs, |t| SpecialDisplay(t, x));
    } else {
        print_xs_with_display(xs, |t| t);
    }
}
1 Like

First, if you're returning a reference (like &dyn Display in your example), you can't construct the value inside the closure and return a reference to it. The referenced object must outlive the reference, so it must be stored somewhere while you're using the reference. That's not specific to trait objects, that's just how references work in general.

Yes I understand that. I don't see that I have stored anything on the closures' stacks though.

Second, there is a similar issue with the display variable. You can't construct it out of references to temporaries. The borrowed closures returned from both if branches must live long enough.

I'm not sure what you're saying here. Why couldn't I make a reference to a temporary value? As long as the reference does not outlive the value that's fine. The closures are on the stack while in use and deallocated at the end of the function's scope, which is fine because I'm not returning them or holding any reference to them. Also I think the borrows are necessary to cast the closures to dyn Fn types. Well that or boxing them.

When you have a problem with dynamic return value, it's usually better to refactor the code so that the return type is generic.

Yes that's good advice but I already know how to do that. I would like to know if it was possible to do it the way I showed (or something equivalent), as an exercise. The error message of the compiler makes me think that it's more a matter of not being able to name all the lifetimes and explaining their relationships to compiler. (expected &dyn std::fmt::Display found &dyn std::fmt::Display)

One strategy when seeing lifetimes complaints near something with a lot of lifetimes, is to break things down to avoid confounding factors. This helps you comprehend the compiler's actually inferred lifetime relationships much more quickly. Help the compiler help you. I can already see one issue off the bat, for example: you're assigning to an &Fn a closure that mutates through a captured &mut. It should be &FnMut, but where's that error?

Starting point

let special_display = &mut SpecialDisplay(&x, &xs[0]);
let display: &dyn Fn(&A) -> &dyn Display = if need_special_display() {
    &|x: &A| {
        special_display.1 = x; // <--- Remove for now.
        special_display
    }
} else {
    &|x: &A| x
};

(playground)

Let's first remove the mutation in the closure to make progress, and perhaps tease apart the various lifetimes in &dyn Fn(&A) -> &dyn Display by making an explicit type that explicitly specifies which lifetimes are related.

type Foo<'a, 'b> = &'b dyn Fn(&'a A) -> &'a dyn Display;

let special_display = &mut SpecialDisplay(&x, &xs[0]);
let display: Foo<'a, 'b> = if need_special_display() {
    &|_x: &A| { special_display }
} else {
    &|x: &A| x
};

(playground)
The error now says something more useful! Specifically &|_x: &A| { special_display } itself is a temporary that will be freed and it suggests to extend its lifetime with a let binding. Let's bind both closures to increase their lifetimes.

type Foo<'a, 'b> = &'b dyn Fn(&'a A) -> &'a dyn Display;

let special_display = &mut SpecialDisplay(&x, &xs[0]);
let foo: Foo<'_, '_> = &|x: &A| {
    special_display
};
let bar: Foo<'_, '_> = &|x: &A| x;
let display: Foo<'_, '_> = if need_special_display() { foo } else { bar };

(playground)
aaaand, that works! Let's add the mutation back in with:

let foo: Foo<'_, '_> = &|x: &A| {
    special_display.1 = x;
    special_display
};

(playground)
And that gives

error[E0594]: cannot assign to `special_display.1`, as `Fn` closures cannot mutate their captured variables

Now we're talkin'! Now change Foo to an FnMut and mutably borrow the closure (as required to invoke the closure):

type Foo<'a, 'b> = &'b mut dyn FnMut(&'a A) -> &'a dyn Display;

and update the assigned closures so that they are borrowed mutably as well.
(playground)

error: captured variable cannot escape `FnMut` closure body
  --> src/main.rs:30:9
   |
28 |     let foo: Foo<'_, '_> = &mut |x: &A| {
   |                                       - inferred to be a `FnMut` closure
29 |         special_display.1 = x;
30 |         special_display
   |         ^^^^^^^^^^^^^^^ returns a reference to a captured variable which escapes the closure body
   |
   = note: `FnMut` closures only have access to their captured variables while they are executing...
   = note: ...therefore, they cannot allow references to captured variables to escape

Oh...right...cause a closure can be run multiple times and could return the same mutable reference twice, thereby violating rust's invariants.

Personally, I find handling the numerous subtle consequences of the uniqueness property of &mut very difficult to work with sometimes, especially when nested. Part of the problem is that the output should borrow from the closure as well; maybe someone can make it work with an appropriate definition for Foo (I tried a few). But let's revert to the working non-mut version using shared references and simply introduce a Cell, which has no additional space or computational cost at the expense of a more restrictive API.

let special_display = &SpecialDisplay(&x, Cell::new(&xs[0]));
let foo: Foo<'_, '_> = &|x: &A| {
    special_display.1.set(x);
    special_display
};

// ...

struct SpecialDisplay<'a>(&'a A, Cell<&'a A>);

// A special display implementation.
impl Display for SpecialDisplay<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}-{:?}", self.0, self.1.get())
    }
}

(playground)

And this works. Does that work for you?

2 Likes

Here's the relevant temporary:

        &|x: &A| {
            special_display.1 = x;
            special_display
        }

The closure is a temporary value. As specified here, it will be dropped at the end of the block. One exception from this rule is constant promotion (I believe that does in fact apply to closures without context), and the second exception is temporary lifetime extension that will make the temporary value live until the end of the block if a reference to it is stored in a let binding. Neither of these cases applies here. This closure will be dropped immediately, so you can't store a reference to it.

That's true.

1 Like

Thank you for the detailed answer!

Yeah, I tried breaking up the lifetimes like you did, but I suppose I failed. I try not to use them if I can avoid it, they make reading the code pretty terrible I think.

I hadn't foreseen the issue with not being able to return the reference from the closure because it still holds onto the mutable ref. I should have seen it was moot from the get-go (at least without some internal mutability).

the second exception is temporary lifetime extension that will make the temporary value live until the end of the block if a reference to it is stored in a let binding.

I felt that putting the let in front of the if expression should have been enough to extend the lifetimes of the closures. I was wrong obviously. But what's even the point of being able to return them if they don't live, though? Having to write the lets before manually feels hackish.

Anyway thanks for clarifying, that was helpful !

It probably did—to the end of the if let statement. But if the same expression is assigned to a variable, it stays alive until the end of the variable's scope or until it's no longer needed, whichever comes first. That's potentially many computations later.

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.