Safe and sound `try_map` that works on rust stable

try_map.rs

use std::mem::{self, MaybeUninit};

/// This is a stable alternative implementation of array_try_map
///
/// TODO: remove this when array_try_map is stabilized
///
/// This follows the example of MaybeUninit usage
/// https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#initializing-an-array-element-by-element
pub fn try_map<T, R, E, const N: usize, F>(source: [T; N], mut cb: F) -> Result<[R; N], E>
where
    F: FnMut(T) -> Result<R, E>,
    R: Sized,
{
    struct Guard<T, const N: usize> {
        dst: [MaybeUninit<T>; N],
        initialized: usize,
    }

    impl<T, const N: usize> Drop for Guard<T, N> {
        fn drop(&mut self) {
            debug_assert!(self.initialized <= N);

            for i in 0..self.initialized {
                unsafe {
                    self.dst[i].assume_init_drop();
                }
            }
        }
    }

    let dst: [MaybeUninit<R>; N] = unsafe { MaybeUninit::uninit().assume_init() };
    let mut guard: Guard<R, N> = Guard {
        dst,
        initialized: 0,
    };

    for (i, s) in source.into_iter().enumerate() {
        match cb(s) {
            Ok(elem) => {
                guard.dst[i].write(elem);
                guard.initialized += 1;
            }
            Err(err) => {
                return Err(err);
            }
        }
    }

    guard.initialized = 0;
    Ok(unsafe { mem::transmute_copy::<[MaybeUninit<R>; N], [R; N]>(&guard.dst) })
}

#[cfg(test)]
mod test {
    use std::{panic, rc::Rc};

    use crate::util::try_map;

    #[test]
    fn drop_on_err() {
        let rc = Rc::new(());

        let f = |i| match i {
            0 => Ok(rc.clone()),
            _ => Err(()),
        };

        {
            let x = [0, 0, 0, 0];
            let _res = try_map(x, f);
            assert_eq!(Rc::strong_count(&rc), 5);
        }

        {
            let x = [0, 0, 0, 1];
            let _res = try_map(x, f);
            assert_eq!(Rc::strong_count(&rc), 1);
        }

        {
            let x = [0, 1, 0, 1];
            let _res = try_map(x, f);
            assert_eq!(Rc::strong_count(&rc), 1);
        }
    }

    #[test]
    fn drop_on_panic() {
        let rc = Rc::new(());

        let f = |i| match i {
            0 => Ok::<_, ()>(rc.clone()),
            _ => panic!("expected panic"),
        };

        {
            let x = [0, 0, 0, 0];
            let _res = try_map(x, f);
            assert_eq!(Rc::strong_count(&rc), 5);
        }

        {
            let x = [0, 0, 0, 1];
            let _res = panic::catch_unwind(move || try_map(x, f));
            assert_eq!(Rc::strong_count(&rc), 1);
        }
    }
}