use core::slice;
use std::{
    ops::{Deref, DerefMut},
    sync::Arc,
    thread,
};

fn main() {
    let v = vec![0u32; 256];
    let (handle, refs) = scatter(v, 64);

    refs.into_iter()
        .map(|mut r| thread::spawn(move || r.iter_mut().for_each(|v| *v += 1)))
        .collect::<Vec<_>>()
        .into_iter()
        .for_each(|h| h.join().unwrap());

    let Ok(v) = handle.try_gather() else {
        println!("fail to gather");
        return;
    };

    println!("v_sum = {}", v.into_iter().sum::<u32>());
}

struct GatherHandle<T>(Arc<GatherHanleInner<T>>);

struct GatherHanleInner<T> {
    ptr: *mut T,
    length: usize,
    capacity: usize,
    alive: bool,
}
impl<T> GatherHandle<T> {
    fn try_gather(self) -> Result<Vec<T>, Self> {
        let mut inner = match Arc::try_unwrap(self.0) {
            Ok(v) => v,
            Err(v) => return Err(Self(v)),
        };

        // Safety:
        // disarming check in Drop
        inner.alive = false;
        // Safety:
        // Ok because each ScatteredRef<T> has GatherHandle<T>,
        // so if we have GatherHandleInner<T> there are no ScatteredRef<T>'s left
        Ok(unsafe { Vec::from_raw_parts(inner.ptr, inner.length, inner.capacity) })
    }
}
impl<T> Drop for GatherHanleInner<T> {
    fn drop(&mut self) {
        if self.alive {
            drop(unsafe { Vec::from_raw_parts(self.ptr, self.length, self.capacity) });
        }
    }
}
struct ScatteredRef<T> {
    ptr: *mut T,
    length: usize,

    _h: GatherHandle<T>,
}
// Safety: Ok because Arc::unwrap() and Arc::drop() do aquire/release semantics
unsafe impl<T: Send> Send for ScatteredRef<T> {}

impl<T> Deref for ScatteredRef<T> {
    type Target = [T];
    fn deref(&self) -> &Self::Target {
        // Safety: Ok because Self<T> is basically &mut[T] for current chunk
        unsafe { slice::from_raw_parts(self.ptr, self.length) }
    }
}
impl<T> DerefMut for ScatteredRef<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // Safety: See Self::deref()
        unsafe { slice::from_raw_parts_mut(self.ptr, self.length) }
    }
}

/// Scatter Vec<T> into `count` size owned mutable slices
/// and gather it back after all slices (ScatterRef<T>) are dropped
fn scatter<T>(v: Vec<T>, count: usize) -> (GatherHandle<T>, Vec<ScatteredRef<T>>) {
    assert!(v.len() > count);
    assert!(v.len() % count == 0);

    let refs_count = v.len() / count;

    let (ptr, length, capacity) = Vec::into_raw_parts(v);
    let handle = GatherHandle(Arc::new(GatherHanleInner {
        ptr,
        length,
        capacity,
        alive: true,
    }));

    let refs = (0..refs_count)
        .into_iter()
        .map(|id| unsafe {
            ScatteredRef {
                ptr: ptr.add(count * id),
                length: count,
                _h: GatherHandle(Arc::clone(&handle.0)),
            }
        })
        .collect();

    (handle, refs)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn error_when_refs_are_alive() {
        let v = vec![0u32; 256];
        let (handle, refs) = scatter(v, 64);

        assert!(handle.try_gather().is_err());
        drop(refs);
    }
    #[test]
    fn ok_when_refs_are_dead() {
        let v = vec![0u32; 256];
        let (handle, refs) = scatter(v, 64);

        drop(refs);
        assert!(handle.try_gather().is_ok());
    }

    #[test]
    fn error_when_single_ref_is_alive() {
        let v = vec![0u32; 256];
        let (handle, mut refs) = scatter(v, 64);

        let first = refs.remove(0);
        drop(refs);

        assert!(handle.try_gather().is_err());
        drop(first);
    }
    #[test]
    fn miri_ok_when_dropping_nahdle_but_refs_are_alive() {
        let v = vec![0u32; 256];
        let (handle, mut refs) = scatter(v, 64);

        drop(handle);
        let _first = refs.remove(0);
        drop(refs);
    }
}