Extending `impl Trait` to allow multiple return types

language design
1

Summary

Argument-position impl Trait corresponds to a hidden type parameter, which in turn is implemented by monomorphizing the callee function for each set of concrete type arguments.

Return-position impl Trait could be extended to match, by monomorphizing the caller function (or more specifically, the continuation after the call) for each concrete return type.

Motivation

Rust functions that return impl Trait are currently only allowed to return a single concrete type. This can be a papercut—for example, while it can replace some uses of Box<dyn Future>, it can’t handle functions that can return multiple future types.

async fn mitigates this for futures, because it automatically generates an enum-like type to handle its corresponding states. There have also been proposals for something like enum impl Trait, which would automatically wrap all the return types into an enum and derive a delegating implementation of Trait for it.

Monomorphizing caller continuations would enable impl Trait to handle multiple return types via the usual trade-off of code size instead of function pointer-based dispatch (as with dyn Trait) or enum boilerplate (as automated by enum impl Trait).

Aside: some context

There has been quite a bit of discussion in RFC 2071, its associated tracking issue, and the #design channel in the rust-lang Discord server around the syntax of impl Trait/abstract type, as well as how to explain its semantics.

The prevailing explanations in announcements/docs have centered around existential types. The the 1.26 release notes claim "impl Trait is universal in an input position, but existential in an output position." The post "impl Trait is Always Existential" gets closer to the truth, but doesn’t quite capture everything.

More recently, @varkor’s summary of the Discord discussion points out the source of the confusion here: argument-position impl Trait (APIT) and return-position impl Trait (RPIT) can both be expressed as existential types, but only with inconsistent quantifier scopes- thus RPIT’s inability to handle multiple types.

So one solution, which I don’t want to discuss in this thread, is to stop talking about existential types at all and instead explain impl Trait in terms of type inference. @varkor has another summary for this approach, which is basically to treat impl Trait as a placeholder for ML-style let-polymorphism, inferring the “most general type.” This makes the semantics of APIT and RPIT consistent with each other.

Instead, in this post I describe an alternative solution—keep the existential-types interpretation (if not the docs, because “existential” is a terrible way to teach this stuff), and make APIT and RPIT consistent by extending RPIT’s capabilities.

Implementation

View a function call’s return address as an argument to the callee—in this sense, we’re already using a restricted form of continuation-passing style. We can add more return address arguments (i.e. passed continuations). Push them all on the stack like today’s return address, or put them in a vtable and push a pointer to that, etc.

Today, a function returning impl Trait has a single return type that is known to the compiler as it builds the caller (though not exposed to the program). Extend this to a set of N possible return types that can still be known to the compiler, in the same way.

In the caller, for each call to a function returning impl Trait, monomorphize the continuation for each of those N types, and point the extra return address arguments (i.e. passed continuations) to them.

In the callee, each return site still knows its return type statically. Select the return address corresponding to this type, clean up the stack frame, and jump to it.

Example

For example, look at these functions:

fn f() {
    println!("{}", g(false));
    println!("{}", g(true));
}

fn g(x: bool) -> impl Display {
    if x { 42 } else { "hello" }
}

We can lower g to something like this:

struct GVtable {
    i32: extern "rust-continuation" fn(i32) -> !,
    str: extern "rust-continuation" fn(&'static str) -> !,
}

fn _g(r: &'static GVtable, x: bool) -> ! {
    if x {
        become r.i32(42)
    } else {
        become r.str("hello")
    }
}

And we can lower f to something like this:

fn f() {
    _g(F1_GVTABLE, false);

    static F1_GVTABLE: &'static GVtable = &GVtable {
        i32: _f1<i32>,
        str: _f1<&'static str>,
    };
    extern "rust-continuation" fn _f1<T: Display>(t: T) -> ! {
        println!("{}", t);
        _g(F2_GVTABLE, true);
    }

    static F2_GVTABLE: &'static GVTable = &GVtable {
        i32: _f2<i32>,
        str: _f2<str>,
    };
    extern "rust-continuation" fn _f2<T: Display>(t: T) -> ! {
        println!("{}", t);
        return;
    }
}

Notes

The “rust-continuation” calling convention is a bit magical, but then again this syntax is purely for illustration. Its main properties:

  • “rust-continuation” functions must take zero or one arguments, and return !.
  • “rust-continuation” functions share the stack frame of their parent fn item, and thus are only ever called by its callees.
  • While “rust-continuation” functions appear to return !, they can return from their parent fn item.

The lowering of f looks a lot like the lowering of async fns. There can be somewhat of a combinatorial explosion of “states”, depending on the impl Trait functions called and the control flow around them, but this may be the most straightforward way to convince LLVM to generate this code?

This technique should still work for trait methods that don’t return associated types- impl Trait was already not going to be object-safe. It does, however, make it less plausible that we might see a dyn Trait<typeof(method)::Output = T> extension, since that would no longer really make sense.

On the other hand, this doubles down on the need for abstract type/RFC 2071, and really calls for a syntax other than type Foo = impl Bar, especially for associated types, which would need to remain as single types. Such a feature would really be something entirely unrelated to impl Trait.

Unanswered questions

How possible is it to convince LLVM to generate this sort of code? Maybe it’s a pipe dream; maybe it’s plausible but a ton of work; maybe it’s easier than I expect. Even if this is only plausible as a far-future extension, maybe we like it enough that we try to stay forward-compatible with it, and make syntactic decisions accordingly.

Does this implementation style overcome the hurdle faced by enum impl Trait? I like that it folds all the “micro-side” runtime cost and implementation complexity back into a dynamic jump that already exists—function return—rather than inventing a wholly new dispatch mechanism, while simultaneously providing the functionality that people expect there.

What do people think of the teachability side of this? Personally I like that it makes impl Trait capable of essentially the same things as dyn Trait, bringing impl/dyn closer to a straightforward “code size”/“run time” knob. It also decouples impl Trait from the abstract type-like functionality that it was created for, which ought to alleviate some of the angst around argument-position impl Trait.

13 Likes
2

For the most part this seems like a great idea, but there’s one small detail that’s bothering me… I really don’t like how magical this makes RPIT, in the sense that if I see

fn foo(..) -> impl Tr { .. }

I now need to look at the definition to find out whether impl Trait is just a compiler fiction for a type the callee doesn’t want to (or can’t) tell me about, or if the caller is implicitly having to assemble a vtable (mind, I think that calling it a vtable is probably not-entirely-correct, because this shouldn’t, if I read the implementation correctly, have the same performance impact a classical vtable would). I feel more ok with async fn doing this because, well, that’s kind of what you’re asking for by using compiler-generated futures, anyways.

Basically, I just want it so that if you want a RPIT that returns an anonymous enum, you have to write enum impl Trait or something like that, so I don’t need to do type inference in my had on the contents of the function to guess whether it’s one or several return types. Alternatively, some kind of attribute that enforces a “thin” RPIT, e.g.

// strawman
fn foo(x: bool) -> #[thin] impl Display { 
    if x { 
        42 
    } else { 
        "hello" // ERROR: expected i32 but found
                // &str instead
    } 
}

impl Trait was stabilized sufficiently recently that this shouldn’t introduce churn? I like dedicated syntax for returning an anonymous enum though (even though, beneath the hood, it’s not really an enum but it’s still not a pure compiler fiction, either).

2 Likes
3

So basically, this would completely tangle the control flow once a multi-type impl Trait is used within a function? That doesn’t sound good. (This sort of increased complexity does tend to introduce compiler bugs, like match ergonomics did. 3 soundness holes and counting.)

It also seems to require turning every such function into CPS, which is a huge pain too, given that Rust allows quite sophisticated procedural-style control flow. I’m not pleased to imagine what such an impl Trait inside two levels of nested loops breaking and continueing at funny locations would look like.

By the way, you didn’t explain (or I didn’t notice the explanation) as to why this is superior to the enum-based approach performance-wise. It still has some runtime overhead. (I didn’t expect it not to have any, because there is a runtime decision to be made, but the description is unclear about that.)

4 Likes
4

@mcy

Presumably the tool to enforce a single return type would be (or perhaps, if we feel like we need it, sugar for) abstract type. If your function happens to have only a single return type, but you don’t feel the need to enforce that, there would be no difference in the generated code.

But either way, you’re right that the performance impact is different from a “classical vtable.” It’s much less—function return is already always an easily-predicted indirect branch. So this doesn’t add any new indirection, just gives a different return address to different return points in the callee.

@H2CO3

No, absolutely not. The control flow doesn’t get “tangled” the way it does in an async fn, no code runs in a different order than it already would. It doesn’t require an actual CPS transform or real become support, either.

The difference is that post-call, there can be multiple copies of the code in the binary, just as with any other instance of generics/monomorphization. You could alternatively explain it as monomorphizing the entire caller for each possible set of return types, and returning “into the middle” of the appropriate copy.

The enum-based approach requires a branch on every method call on the return value. This approach instead, to quote the post, “folds all the “micro-side” runtime cost and implementation complexity back into a dynamic jump that already exists—function return.” In comparison to a normal function call, the usual “push the return address and jump” is replaced with “push the vptr and jump,” and the usual “pop the stack and indirect branch” is replaced with “index the vtable and indirect branch.” The actual runtime decision of which vtable entry to use is folded entirely into the callee’s existing control flow.

5 Likes
5

Ahha! Thanks for the explanation. So, instead of illustrating it with functions, I think it would be better described in terms of basic blocks then. The notion of functions in Rust has so much more to them than just control flow that I think that’s what confused me. Basically, what you propose as a vtable and functions is more like a computed goto.

3 Likes
6

It’s fascinating to me that we could do this. However, I very much don’t want us to do this, precisely because impl Trait is currently guaranteed to represent a specific concrete type. I look forward to being able to use that unspecified concrete type in other places, such as structs, and this would prevent doing so.

10 Likes
7

I have various concerns about the implementation, and I don’t think we can backwards-compatibly change existing impl Trait return types to this meaning 1 (and even besides that I think a feature with such huge non-local codegen effects should be opt-in), and I agree with what @josh just said.

However, even besides that and perhaps more fundamentally, I am skeptical of the motivation: when would such monomorphization actually be more desirable than what can be done today? The OP states the desire to avoid branches internal to the returned type, but since each branch eliminated by this feature duplicates the entire tail of the surrounding function(s), it results in a code size explosion that’s roughly exponential in the number of eliminated branches. Thus, eliminating more than a very small handful of such branches at once would be absolutely catastrophic for code size, which seems to limit its applicability to niche cases (few branches to be eliminated, small continuation, and the still-considerable code size increase from that is acceptable). Typical futures code, for example, doesn’t seem like it would benefit.

And of course, one can always write the generic continuation out explicitly (and in the process avoid all the hard questions about automatic generation of such continuations).

1 Consider for example let mut x = returns_impl_trait(1); x = returns_impl_trait(2);

8 Likes
8

The enum version is a lot more powerful because it gives a concrete type that can be used with an impl Trait existencial type. It also as @hanna-kruppe pointed out allows you to overwrite a mutable value returned by your function using an other call to the function. Maybe your implementation can be an optimisation in cases where it would not generate to much code (the possibilities can blow up when using multiple functions that all have multiple possible return values).

1 Like
9

I’d prefer the variant where something like an enum keyword is required. Reasons:

  • I feel that single-type return values are the common case, at least for me.
  • Implicit multi-value requires more care. For example, impl fmt::Display might allow a lot of unexpected values as a result. Every trait implemented for () would also require a lot of care, as every ; becomes a hazard.
  • Using an externally declared impl Trait would be a lot more work for the common, safer, and more restrictive variant.
  • Using an externally declared impl Trait would move the information out of the function signature, while an extra keyword like enum makes it clear in the return signature where this happens, which makes mixing these a lot more ergonomic (e.g. -> Result<impl AsRef<str>, enum impl Error>.

In general, since the enum behavior requires additional thinking about the code, it should stand out in the signature.

3 Likes
10

I feel that this is conflating two different questions:

  1. Design: Assuming we want an enum impl Trait feature someday, should it require some kind of explicit opt-in / visible piece of syntax, or be the default behavior of impl Trait?

  2. Implementation: Should enum impl Trait be implemented literally as an enum or as this fancy CPS magic?

I think we should keep these separate, and in particular we need to figure out #1 before we spend any significant amount of effort on #2.

My knee-jerk reaction is of course that I want opt-in for this. It’s explicit. Explicit makes me feel warm and fuzzy.

My nagging-feeling reaction is that we already decided using impl Trait for both sort-of-not-really existentials and sort-of-not-really universals was a net win because the typical user both doesn’t care and doesn’t need to care about that distinction (unlike, say, the ownership rules). Whether impl Trait is a single concrete type or not seems like it’d be a very similar “why should I care?” situation to a lot of Rustaceans, especially those who tend not to participate in these threads. Of course, that decision remains unusually controversial, and our “trust budget” is running low enough as it is, so at the very least we need to wait and see if it pans out in practice before pursuing extensions like this with the same reasoning.

So… I’m torn.

1 Like
11

Not true for me. I need multiple return values when I manipulate those Future values in tokio, as every branch will use different logic and so have different concret type. My current solution depends highly on the Either type.

12

@hanna-kruppe's backcompat hazard only scrapes the surface:

// This compiles today; as I certainly hope it would!

fn foo(n: u32) -> impl Debug { ... }

fn main() {
    let xs = (0..5).map(foo).collect::<Vec<_>>();
    println!("{:?}", xs);
}

so to me, the simple answer to

is that making this the default is indefensible. (unless, well... I guess it could be an edition-related change)

13

can you desugar into a closure with TCO?

14

If we compile an impl Trait function that returns multiple types that implement Trait to a function that returns an enum that implements Trait by delegating the implementation to the enum variants it would be backwards compatible. The function still returns a value that implements Trait. However we could run into some problems for some traits when delegating.

15

Say we have a trait Trait that looks as follows:

trait Trait {
    type Type;

    fn foo(self) -> Self::Type;
    
    fn bar();
}

and a function

fn func(x: bool) -> impl Trait<Type = i32> {
    if x {
        TypeA::new()
    } else {
        TypeB::new()
    }
}

then we compile it to

fn func(x: bool) -> Hidden {
    if x {
        Hidden::HiddenTypeA(TypeA::new())
    } else {
        Hidden::HiddenTypeB(TypeB::new())
    }
}

enum Hidden {
    HiddenTypeA(TypeA),
    HiddenTypeB(TypeB),
}

impl Trait for Hidden {
    type Type = i32;

    fn foo(self) -> Self::Type {
        match self {
            Hidden::HiddenTypeA(a) => a.foo(),
            Hidden::HiddenTypeB(b) => b.foo(),
        }
    }

    fn bar() {
        match &self {
            &Hidden::HiddenTypeA(_) => TypeA::bar(),
            &Hidden::HiddenTypeB(_) => TypeB::bar(),
        }
    }
}

I don’t know if this would be possible for every trait. The only case I see a problem with is traits with a const value. That is a niche case and I think it would be alright to disallow impl Trait with multiple types when Trait contains a value. All cases with impl Trait would still work because we still allow all functions that return a single type and functions that return multiple types really only return one. For that reason I personally we don’t need the enum impl Trait because it does not matter for people that call your function.

16

There's a lot of prior discussion of the autogenerated sum type idea here: Auto-generated sum types · Issue #2414 · rust-lang/rfcs · GitHub (along with a bit of prior discussion across other earlier RFCs).

There's many things that will cause issues with generating the delegation e.g. fn foo(self, other: Self), but hopefully that can all be handled by a generalized delegation framework that an autogenerated sum type could just reuse.

It doesn't matter for people calling your function, but it does matter for writing the function itself, there are different tradeoffs between an auto-generated sum type and a dynamically dispatched type that you should be aware of.

3 Likes
17

I would have :+1:'ed the -> enum impl Trait syntax if there was no APIT, but given APIT it feels asymmetric if this isn’t allowed:

fn output() -> enum impl Trait { ... }

fn inpit(a: enum impl Trait) { ... }

Perhaps it should just be a function attribute (since the CPS transform is applied to the whole function)

#[allow(multiple_return_types)]
fn output() -> impl Trait { ... }

fn input(a: impl Trait) { ... }
5 Likes
18

I dropped by the paint store:

  • enum impl Trait APIT is reserved as sugar for the Any mechanism (gross, but the only use of APIT enum impl I can think of).
  • An attribute, as you suggest. I like #[polymorphic_return], even though it’s a heinous lie because no monomorphizaton happens. See also #[virtual_return].
  • On that note, virtual is reserved. fn foo() -> virtual Trait?
  • Drop impl altogether and say enum Trait. This is probably confusing with how items are declared though, especially since this type syntax (from C) suggests that Trait is an enum.
  • Move enum to the start of the function: enum fn foo() -> impl Trait. Is there something interesting we can do with a return type that isn’t an impl Trait?
  • Since we’re using continuations, -> become Trait? Kinda silly, and maybe we want become T in type position for something else.
  • Accept that we can’t have symmetry because covariant (return) position is special.

Relatedly; how many of the problems you’d want to solve with enum impl Trait can be solved by having anonymous coproduct types? We have anonymous product types (tuples) so I hallucinate something like

fn foo(x: (i32 + &'static str)) {
    match x {
        0(i) => println!("int: {}", i),
        1(s) => println!("str: {}", s),
    }
}
// desugar
enum __magic {
    v0(i32), v1(&'static str),
}
fn foo(x: __magic) {
    match x {
        v0(i) => println!("int: {}", i),
        v1(s) => println!("str: {}", s),
    }
}

(Yes, I’m aware that enum impl Trait is not returning an anonymous enum but doing fancy become continuation.)

19

Without considering enum impl which is totally fine (the function still returns one specific type), how would you typecheck this with your « multiple return types » setting:

fn bar<T>(x: T, y: T) { }

fn main() {
    bar(g(false), g(true));
}

From a typeck point of view, would the return type of the g function still considered as an enum?

Edit: I see that @hanna-kruppe actually raised a very similar concern

3 Likes
20

I assume T will tyck as g::return (whatever that is) which is known to be an enum impl Trait. The monomorphization of bar for that value of T will include the necessary become glue, with main simply passing bar::GVtable to g.