Different return values dependent on value of generic constant

user_8239463 0
Viewed 268

Let's consider a function like this:

fn test<const N: usize>() -> [f64; N] {
    if N == 1 {
        [0.0_f64; 1]
    } else if N == 2 {
        [1.0_f64; 2]
    } else {
        panic!()
    }
}

My understanding is that the compiler would evaluate the value of N at compile time. If this is the case, the if statement could also be evaluated at compile time, thus the right type should be returned since [0.0_f64; 1] is only returned if N == 1 and [1.0_f64; 2] is only returned if N == 2.

Now, when i try to compile this code, the compiler fails, basically telling me that the dimensions of the returned arrays are wrong since they do not explicitly have N as length.

I do realize, that i could implement this specific example as

fn test<const N: usize>() -> [f64; N] {
    match N {
        1 => { [0.0_f64; N] },
        2 => { [1.0_f64; N] },
        _ => { panic!("Invalid value {}", N) },
    }
}

But that does not work in my actual code, since that uses different functions with fixed array sizes for the different branches.

Is there a way to do this at all? Maybe using something like the #![cfg] makro?

To clarify my why my problem does not work, let's write this out:

fn some_fct() -> [f64; 1] {
    [0.0_f64; 1]
}
fn some_other_fct() -> [f64; 2] {
    [1.0_f64; 2]
}

fn test<const N: usize>() -> [f64; N] {
    match N {
        1 => some_fct(),
        2 => some_other_fct(),
        _ => {
            panic!("Invalid value {}", N)
        }
    }
}

And I cannot really write some_fct() and some_other_fct() to return with generic sizes due to other restrictions in the program structure.

rust generics
2 Answers

You can do that with a generic trait:

trait Test<const N: usize> {
    fn test() -> [f64; N];
}

Then you implement it for a zero sized type:

struct T;

impl Test<1> for T {
    fn test() -> [f64; 1] {
        return [0.0_f64; 1];
    }
}

impl Test<2> for T {
    fn test() -> [f64; 2] {
        return [1.0_f64; 2];
    }
}

The drawback is that calling it is a bit cumbersome:

fn main() {
    dbg!(<T as Test<1>>::test());
    dbg!(<T as Test<2>>::test());
}

But as @eggyal comments below, you can add a generic function with a well-written bound to get your required syntax:

fn test<const N: usize>() -> [f64; N]
where
    T: Test<N>
{
    T::test()
}
fn main() {
    dbg!(test::<1>());
    dbg!(test::<2>());
}

Now, you don't have the behavior of "panic! when a wrong N is used". Consider that a feature instead of a limitation: if you use a wrong N your code will fail to compile instead of panic at runtime.

If you really want the panic!() behavior you could get it using the unstable feature of #![feature(specialization)], just adding default to this impl:

impl<const N: usize> Test<N> for T {
    default fn test() -> [f64; N] {
        panic!();
    }
}

But that feature is explicitly marked as incomplete, so I would not count on it, yet.

user_8658740

Here is a solution that is not particularly clever, but is easy to understand and resembles the original:

fn test<const N: usize>() -> [f64; N] {
    match N {
        1 => some_fct().as_slice().try_into().unwrap(),
        2 => some_other_fct().as_slice().try_into().unwrap(),
        _ => {
            panic!("Invalid value {}", N)
        }
    }
}

Although the code looks like it checks array sizes at run time, godbolt shows that rustc/LLVM is able to reason that [f64; N].as_slice().try_into() always succeeds in coercing the array-turned-slice to [f64; N]. Generated code for test<1> and test<2> thus contains no checks or panic, and test<N> for N>2 just panics unconditionally due to the panic in the catch-all match arm.

user_16008980
Related