Runtime specialization when types are equal

I would like to write a trait or function that:

  • Uses stable Rust
  • Consumes its arguments (i.e. by value)
  • Is specialized when its arguments are the same type

My motivation is to write a function that maps from one ndarray::ArrayBase to another array of the same shape, but of a possibly different underlying type. If the input and output types are the same then I would like to call ArrayBase::mapv_into() to avoid extra allocations.

To my knowledge, it is not possible do this at compile time in stable rust without some form of specialization. But it looks like I can do dynamic specialization/dispatch at runtime with nominal overhead (playground):

use core::any::TypeId;
use core::mem::transmute;
use core::ptr;

// Some types.
struct A { }
struct B { }

// Do something really efficient if arguments are the same type.
fn same<T>(_t1: T, _t2: T) {
    println!("same");
}

// Do something more generically if arguments could be of different types.
fn diff<T,U>(_t: T, _u: U) {
    println!("different");
}

// My attempt at runtime specialization.  Is this `unsafe` kosher?
fn any<T: 'static, U: 'static>(t: T, u: U) {
    // Dispatch to the specialized `same()` if arguments are the same type.
    if TypeId::of::<T>() == TypeId::of::<U>() {
        // Transmute from U to T.
        let t1 = t;
        let t2 = unsafe {
            let u_ptr: *const U = &u;
            let t_ptr = transmute::<*const U, *const T>(u_ptr);
            let t2 = ptr::read(t_ptr);
            core::mem::forget(u); // EDIT forgot to forget() in original post
            t2
        };
        same(t1,t2)
    }
    // Dispatch to the default `diff()` otherwise.
    else {
        diff(t,u)
    }
}

fn main() {
    any(A{}, A{}); // same
    any(A{}, B{}); // different
}

My question is:

  • Is the above use of unsafe safe?
  • Is there a better way to do this without unsafe?

I think I could do this with safe code using Any::downcast_ref() if my function took its arguments by reference (i.e. fn any<T: Any, U: Any>(&dyn T, &dyn U)), but I cannot seem to figure out how to take the arguments by value using safe code.

Note: this is unnecessarily unsafe. For a better, more idiomatic, 100% safe solution, see @2e71828 's post below.


Original Answer:

The transmute on pointers is unnecessary, you can cast between pointers directly, which is actually safer (slightly harder to use incorrectly, as some properties are checked by the compiler). It is not recommended to ever use transmute, in general. All of its use cases have better alternatives, and it's incredibly easy to misuse.

Here's a slightly more idiomatic (IMO) solution which explicitly asks whether the cast can be performed: Playground.

fn spec_same<T, U>(x: T, y: U)
    where
        T: 'static,
        U: 'static,
{
    if let Some(ptr) = (&y as &dyn std::any::Any).downcast_ref::<T>() {
        let z: T = unsafe { std::ptr::read(ptr) };
        std::mem::forget(y);
        same(x, z);
    } else {
        different(x, y);
    }
}

The trick to do this in safe code is Option::take(), which lets you transform &mut Option<T> into T:

    if TypeId::of::<T>() == TypeId::of::<U>() {
        // Transmute from U to T.
        let t1 = t;
        let mut u_opt = Some(u);
        let t2 = (&mut u_opt as &mut dyn std::any::Any)
            .downcast_mut::<Option<T>>()
            .unwrap()
            .take()
            .unwrap();
        same(t1, t2)
    }
10 Likes

Your version is better, I agree, but I don't agree about the transmute(). That is, I do agree about pointer don't need it, but you should not transmute using pointer. transmute_copy() is better in this case:

fn any<T: 'static, U: 'static>(t: T, u: U) {
    if TypeId::of::<T>() == TypeId::of::<U>() {
        let t1 = t;
        let t2 = unsafe { mem::transmute_copy(&u) };
        mem::forget(u);
        same(t1, t2)
    } else {
        diff(t, u)
    }
}

In fact, the pointers trick is exactly what transmute_copy() does (plus a check for alignment).

@H2CO3, this is indeed better. Thank you for helping me think around my transmute "crutch." I think I probably should have included a mem::forget() in my initial attempt.

@2e71828, I think moving the type into an option and then take()ing it back out of the mutable reference after downcast_mut() may be a little bit less efficient than the unsafe route. But it certainly has the advantage of using only safe code!

@chrefr, mem::transmute_copy(&u) works and is less verbose than my initial attempt, but I am worried it is creating a copy of u in memory instead of moving it?

It creates a copy only if the type is Copy. But I should mem::forget() it. Updated the example.

Why don't you use totally safe code which produces strictly identical machine code?

fn any<T: 'static, U: 'static>(a: T, b: U) {
    use std::any::Any;

    match (Box::new(a) as Box<dyn Any>).downcast() {
        Ok(a) => same::<U>(*a, b),
        Err(a) => {
            let a: T = *a.downcast().unwrap();
            diff::<T, U>(a, b)
        }
    }
}
14 Likes

Thanks for demonstrating this safe implementation using Box. The optimizer is quite impressive! Here it optimizes away the entire any() function and all of the heap allocations therein, so there's no performance penalty. However, without compiling in release mode the version using Box allocates on the heap whereas the unsafe version does not. I guess I would have to weigh the benefit of avoiding unsafe against the risk of relying on the optimizer to do the right thing.

Although there probably aren't too many cases of runtime specialization on embedded rust, reliance on Box could pose a problem there even though the actual heap allocations would be optimized away.

i'm surprised you cannot use downcast on references.
It would make sense to use &dyn Any instead of Box<dyn Any>

Although you can use &dyn Any::is

i'm surprised you cannot use downcast on references.
It would make sense to use &dyn Any instead of Box<dyn Any>

Yes, but to consume the arguments by value, I think you would then have to use shenanigans like Option and std::mem::take() like in this post to move out of the reference. The Box solution is less verbose... as long as the optimizer does its thing.

yeah, but using box for no reason is disgusting even when optimizer works.
I would just wrap it all into necessary unsafe code and stop to worry about safety

You can downcast on references. But you can't take an owned value from references. &dyn Any -> &T is what <dyn Any>::downcast_ref() does. But to get the owned T you need to consume the dyn Any itself. Since it's DST you can't use it without indirection. The standard way to make owned indirection type is Box<T>.

1 Like

I think in this case you should definitely choose @2e71828 's solution over mine. Of all the solutions here, I think it's the most elegant, idiomatic, safest, and most reliably efficient (although the one by Hyeonu demonstrates that the Box can be optimized away too).

Basically, if you can do something in safe code, there's generally no good reason to drop down to unsafe. That should be a last-resort escape hatch which isn't used lightly. I only provided an unsafe implementation because I did not realize this was possible in safe code, but thinking about it, the use of Option::take is actually a common trick for transforming mutable references into values.

The use of that Option will literally compile to merely checking a flag (whether it's None or Some). Compared to performing complicated numerical computations on an ndarray, it likely won't even be a measurable difference in speed at all. And if the wrapped type recursively contains a guaranteed non-zero type (e.g. Vec, a reference, a NonNull pointer or integer, etc.), then the compiler won't even store a separate flag for it, it will just use the otherwise uninhabited zero value for denoting None.

1 Like

FWIW, the whole post here is observing a missing feature of Rust / its stdlib, but a feature which can be implemented using third-party code (granted, in a less ergonomic fashion): &move / &own references.

  • Indeed, we need indirection to perform the dyn Any type erasure, but we need an owning pointer to do it. But nothing requires the owning pointer to be one pointing to the heap: we can have owning pointers that point to local storage (usually the stack)!

And there is indeed a crate –disclaimer: of mine– which features such owning pointers / owning references:

With it, one can do the following:

asciicast

Which is basically the same as @Hyeonu's code, but not relying on the optimizer to elide the heap-allocation: there is no heap allocation anywhere in the code, making it suitable for #[no_std] environments.

  • Aside: the current slightly less pretty StackBoxDynAny<'_> type instead of StackBox<'_, dyn Any> ought to be "fixed" soon: the difference stems from unsized coercions not being automatically available on stable Rust, and thus requiring some manual hand-rolling. With my approach, I had to define a different type (but there is an option for a clever type alias + helper trait usage to make this invisible to the downstream user), but there is a PR on the way which uses a more general stable polyfill of unsize coercions.
1 Like

This post is generating a lot more discussion than I expected for, I think, precisely this reason. I'd be curious if anyone knows why Any has downcast_ref() and downcast_mut() but not downcast() (except for Boxed variables).

There are two ends of the philosophical spectrum here, from @Hyeonu's view that we should avoid unsafe whenever possible to @DoumanAsh's view that unsafe causes needless worry (forgive me if I am misrepresenting your views).

Personally I prefer something like the following, which uses unsafe to accomplish the downcast very succinctly and in a way that is easy to reason about. I don't like using contortions involving Option or Box because they reduce readability and obfuscate the intent of the code, even if they sidestep unsafe.

I have to use unsafe all the time for ffi and splitting borrows, so it's not exactly a dark art. I just like to solicit feedback about whether the code is sound if I'm trying something I'm not used to.

use core::any::TypeId;
use core::mem::ManuallyDrop;
use core::ptr;

// Attempt to cast a value of type T to a value of type U.
// Returns Ok(u) if T and U are the same type.
// Returns Err(t) otherwise.
fn downcast<T: 'static, U: 'static>(t: T) -> Result<U, T> {
    if TypeId::of::<T>() == TypeId::of::<U>() {
        Ok(unsafe {
            ptr::read(&*ManuallyDrop::new(t) as *const T as *const U)
        })
    } else {
        Err(t)
    }
}

Note that this line of reasoning can become dangerous / is a slippery slope: it's not that far from "using iterators reduce readability compared to manually performing unchecked indexing". That being said, I agree that the current situation is suboptimal. FWIW, here is how I'd expect my snippet to be rewritten should we have &move references:

//      upcast
//     vvvvvvvv
match (&move t).downcast::<U>() {
    Ok(&move t_u) => same(t, u),
    Err(any) => diff(*any.downcast::<T>().unwrap(), u),
}

No unsafe needed, and upcasting-then-downcasting is, imho, not that contrived when wanting to "trans-cast".

I do agree that writing a cast by-value helper is the right abstraction here, and based on these &move refs, it could easily be written without unsafe:

fn cast<T : 'static, U: 'static> (t: T)
  -> Result<U, T>
{
    if TypeId::of::<T>() == TypeId::of::<U>() {
        //     upcast
        //    vvvvvvvv
        Ok( *(&move t).downcast::<U>().unwrap() )
    } else {
        Err( t )
    }
}
  • With stackbox:
    fn cast<T : 'static, U: 'static> (t: T)
      -> Result<U, T>
    {
        if TypeId::of::<T>() == TypeId::of::<U>() {
            //              upcast
            //  vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
            Ok( stackbox!(t).into_dyn::<…Any>().downcast::<U>().unwrap().into_inner() )
        } else {
            Err( t )
        }
    }
    

Whether a function or an abstraction has unsafe in it does have a big impact: it makes that code susceptible to UB-inducing bugs, e.g., when maintained by other people in the future. Refactoring becomes harder, and it's yet another part of code that needs auditing. It's always about the UB surface on your code base.


That being said, since, indeed, this is a limitation of the language / stdlib (if it doesn't feature &move references nor specialization, then maybe it should feature this cast method), I do agree that your cast() function, which neatly contains the unsafe within easily auditable boundaries, is not a bad workaround for the moment being :+1: (you could use transmute_copy rather than a ptr-casted-.read(), but that's just a tiny nit).

1 Like

I think calling this a "missing feature" is quite a leap. People here pointed out two efficient and safe solutions. For the same reason, I don't agree that using Option or Box is a "contortion". They are useful types, and they are very fundamental and versatile. That's why they are in the standard library. They exist in order to be used a lot.

In my opinion, it's vastly more useful for Rust programmers to learn some idioms and tricks so that they can also invent their own ones when needed, rather than adding a new language feature for every single not-so-common problem. That's not going to be maintainable in the long run.

Furthermore, this by-value downcasting could just be yet another function in the standard library, with one of the safe implementations provided above, which kind of renders the entire discussion around language features moot.

There is a missing feature here:

  1. When considering trait objects with owned receivers, Any::downcast being one example, calling a dyn FnOnce being another, and so on for an arbitrary number of traits (traits that the stdlib may not be aware of),

    • The stdlib will thus not be able to offer helper functions for all of them, and shouldn't even try (that being said, the cast() function of this thread does look like an interesting candidate :slight_smile:);
  2. if one does not want or cant use heap-allocations (no Box available, or the optimizer does not manage to elide the Boxing in a hot loop),

  3. then, currently, one either needs to use unsafe, or "replace ownership with Option::take()-like patterns", both of which detect less misusage at compile-time. When misused, the latter causes panics, whereas the former causes UB.

The debatable part is actually whether the missing feature / the feature to add to solve this issue is:

  • &move references, which:

    • have clear well-defined semantics (the stackbox crate couldn't have been written without them);

    • fit the trinity / troika / trifecta of &Self, &mut Self, Self

  • or with unsized_locals, which, at the end of the day, are sugar for &move references so that the simpel cases Just Work™ more conveniently, but this also means that people can't have the degree of control (e.g., lifetimes) that StackBoxes / &move refs offer (e.g., exact location of the actual storage of the value, well-defined ABI).

While I am a bit biased towards the former, I do admit that the latter would already fulfill the vast majority of the needs that &move references solve

Well, I'm not sure what you mean.
My view is that we should use unsafe version if it is more straightforward.
I shouldn't need Option or Box for something so simple as downcasting or checking type equality at runtime.
And if Rust has no straightforward safe code, then we should just use unsafe or come up with safe alternative that is straight-forward and performs what we need rather than being work-around for lack of better way to express ourself in safe code

P.s. most definitely we MUST NOT use Box just because we're lazy to write unsafe code (even if compiler can optimize because you can never trust compiler optimization)

No, that's not the usual Rust way. The Rust way is only using unsafe when you absolutely have to. In this case, since safe versions are basically as efficient as the unsafe versions, and using an optional or a box is not a significant complication, there's no good reason to use unsafe.

You are completely misrepresenting and accusing people here. Avoiding unsafe is not done out of sheer laziness. We want to avoid unsafety in order to make our programs more obviously correct. If we simply slapped an unsafe on the code every time we didn't get a pre-baked solution to a problem, there would be no point in using Rust at all, we might as well write everything in C++ or something. I think this would actually count as laziness: not thinking hard enough about a safe solution.

I don't even understand what's the point of a statement like this. Optimize only when you have to, and if your code is correct, look at the generated assembly
if you are really worried about what it's doing.

8 Likes