#![feature(try_trait)]

use std::marker::PhantomData;
use std::ops::Try;

trait VecExt: Sized {
    type T;

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

        match self.try_map(move |x| Ok::<_, Infallible>(f(x))) {
            Ok(x) => x,
            Err(x) => match x {}
        }
    }

    fn try_map<U, R: Try<Ok = U>, F: FnMut(Self::T) -> R>(self, f: F) -> Result<Vec<U>, R::Error>;
}

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

    fn try_map<U, R: Try<Ok = U>, F: FnMut(Self::T) -> R>(mut self, f: F) -> Result<Vec<U>, R::Error> {
        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.try_into_vec(f)
        } else {
            self.into_iter().map(f).map(R::into_result).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 try_into_vec<R: Try<Ok = U>, F: FnMut(T) -> R>(mut self, mut f: F) -> Result<Vec<U>, R::Error> {
        while self.init_len < self.len {
            unsafe {
                let value = f(self.ptr.read())?;

                (self.ptr as *mut U).write(value);

                self.ptr = self.ptr.add(1);
                self.init_len += 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 {
            Ok(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
            );
        }
    }
}