use std::marker::PhantomData;

trait VecExt: Sized {
    type T;

    fn map<U, F: FnMut(Self::T) -> U>(self, f: F) -> Vec<U>;
}

impl<T> VecExt for Vec<T> {
    type T = T;

    fn map<U, F: FnMut(Self::T) -> U>(mut self, f: F) -> Vec<U> {
        use std::alloc::Layout;

        if Layout::new::<T>() == Layout::new::<U>() {
            let len = self.len();
            let cap = self.capacity();
            let ptr = self.as_mut_ptr();

            std::mem::forget(self);

            let iter = MapIter {
                start: ptr as *mut U,
                ptr,
                init_len: 0,
                len,
                cap,
                drop: PhantomData
            };

            iter.into_vec(f)
        } else {
            self.into_iter().map(f).collect()
        }
    }
}

struct MapIter<T, U> {
    start: *mut U,
    ptr: *mut T,
    init_len: usize,
    len: usize,
    cap: usize,
    drop: PhantomData<(T, U)>
}

impl<T, U> MapIter<T, U> {
    fn into_vec<F: FnMut(T) -> U>(mut self, mut f: F) -> Vec<U> {
        for i in 0..self.len {
            unsafe {
                self.init_len = i;

                let value = f(self.ptr.read());
                (self.ptr as *mut U).write(value);

                self.ptr = self.ptr.add(1);
            }
        }

        let ptr = self.start;
        let len = self.len;
        let cap = self.cap;

        // we don't want to free the memory
        // which is what dropping this `MapIter` will do
        std::mem::forget(self);

        unsafe {
            Vec::from_raw_parts(ptr, len, cap)
        }
    }
}

impl<T, U> Drop for MapIter<T, U> {
    fn drop(&mut self) {
        unsafe {
            // assert!(std::thread::panicking());

            // offset by 1 because self.ptr is pointing to
            // memory that was just read from, dropping that
            // would lead to a double free
            std::ptr::drop_in_place(
                std::slice::from_raw_parts_mut(
                    self.ptr.add(1),
                    self.len - self.init_len - 1,
                )
            );

            Vec::from_raw_parts(
                self.start,
                self.init_len,
                self.cap
            );
        }
    }
}