Reduce stack footprint by using `core::mem::transmute_copy`

[maxdexh]

The stack footprint of several functions can be improved by using core::mem::transmute_copy on rust versions starting at 1.83 (using e.g. the rustversion crate).

For example, const_transmute can be changed to

pub const unsafe fn const_transmute<A, B>(a: A) -> B {
    core::mem::transmute_copy(&core::mem::ManuallyDrop::new(a))
}

With this, assume_init can also be changed to just defer to const_transmute (which is a valid way of performing MaybeUninit::assume_init)

impl<T, N: ArrayLength> GenericArray<T, N> {
    pub const unsafe fn assume_init(array: GenericArray<MaybeUninit<T>, N>) -> Self {
        const_transmute(array)
    }
}

This is more stack-friendly than even the current implementation of IntrusiveArrayBuilder::array_assume_init, though that one can also be improved further by using transmute_copy(&array) directly, since GenericArray<MaybeUninit<T>, N> is known to not have any drop glue.

[novacrazy]

Using MaybeUninit::assume_init invokes intrinsics::assert_inhabited::<T>() internally, which I presume can be important. transmute_copy just does a pointer read inside, too, so it should be equivalent.

How are you testing the stack size here? I'd like to look into that.

[maxdexh]

Using MaybeUninit::assume_init invokes intrinsics::assert_inhabited::<T>() internally, which I presume can be important. transmute_copy just does a pointer read inside, too, so it should be equivalent.

That's just a UB check for never-like types, e.g. the assertion fails for core::marker::Infallible.

How are you testing the stack size here? I'd like to look into that.

Sry, I must have deleted that part of the message on accident. Basically I poked at the example from #108 to find out when a stack overflow occurs.

[maxdexh]

The part about the implementation of transmute_copy is a good point. The current implementation goes through assume_init though, which requires an extra copy through MaybeUninit<A> (at least I assume that's why I got different results for the maximum length).

[novacrazy]

Sry, I must have deleted that part of the message on accident. Basically I poked at the example from #108 to find out when a stack overflow occurs.

Cool, I should play around with that issue again, too. Thanks for all the ideas lately!

[novacrazy]

Using something like this:

#[inline(never)]
#[unsafe(no_mangle)]
pub fn test_stack_size() -> Array {
    Array::generate(|i| i as u64 + 1)
}

// like black_box but doesn't return, and in optimized builds forces it to do something
#[inline(never)]
pub fn consume(a: Array) {
    println!(
        "Wrapping sum: {}",
        a.iter().fold(0u64, |acc, &x| acc.wrapping_add(x))
    );
}

#[inline(never)]
#[unsafe(no_mangle)]
pub fn overflow_stack() {
    consume(test_stack_size());
}

fn main() {
    println!("Size of Array: {} bytes", core::mem::size_of::<Array>());

    overflow_stack();
}

in an examples/ binary, run with cargo rustc --release --example stack_test -- -C opt-level=0 --emit asm -C target-cpu=native, I can look at the asm and search for .seh_stackalloc values to see what is allocated.

I was able to minimize codegen and stack allocations significantly by merging IntrusiveArrayConsumer::finish and assume_array_init like so:

#[rustversion::attr(since(1.71), const)]
#[inline(always)]
pub unsafe fn finish_assume_init(self) -> GenericArray<T, N> {
    debug_assert!(self.is_full());
    // Manual implementation of transmute_copy with better optimization
    let out = ptr::read(self.array as *const _ as *const GenericArray<T, N>);
    mem::forget(self);
    out
}

The asm also reveals that ::default()/test_stack_size still allocates 3x as much as needed after this (was 6x before), so I'm still looking into ways to improve it.

[novacrazy]

Was able to shrink stack usage even further, down to 1x, with only a single memcpy.

Using MaybeUninit::<GenericArray<T, N>>::uninit() instead of GenericArray::uninit(), then casting a mutable reference of that to &mut GenericArray<MaybeUninit<T>, N> to pass into IntrusiveArrayBuilder::new, then after finish(), do a:

ptr::read(&array as *const MaybeUninit<GenericArray<T, N>> as *const GenericArray<T, N>)

to "assume_init", results in very clean and efficient codegen. Miri is totally fine with it, too, so okay then.

[novacrazy]

I've committed my solution to master if you want to try it out. Nothing will ever be perfect in debug builds, but it should help significantly.

[maxdexh]

Haven't had time to test this yet but it looks good, that's basically what I did in my testing.

Have you actually tested compiling on 1.74-1.82? The implementation of const_transmute looks to me like it shouldn't compile there due to the generics potentially having interior mutability, so I would be a bit confused if this works.

Edit: Yea it doesn't work, you need to use version 1.83 as your cutoff

error[E0658]: cannot borrow here, since the borrowed element may contain interior mutability
    --> src/lib.rs:1320:29
     |
1320 |         mem::transmute_copy(&ManuallyDrop::new(a))
     |                             ^^^^^^^^^^^^^^^^^^^^^
     |
     = note: see issue #80384 <https://github.com/rust-lang/rust/issues/80384> for more information

[novacrazy]

Good catch, I was just going off when transmute_copy was stabilized as const. Put out 1.3.5 with the fix and yanked 1.3.4.