use std::{
    alloc::{self, Layout},
    marker::PhantomData,
    ptr::NonNull,
    ptr::{read, write},
};

struct A<'a, X> {
    ptr: NonNull<X>,
    len: usize,
    own: bool,
    _pd: PhantomData<&'a X>,
}

impl<'a, X> A<'a, X> {
    pub fn new(len: usize) -> A<'a, X> {
        let lay = Layout::array::<X>(len).unwrap();
        let ptr = unsafe { alloc::alloc(lay) };
        let ptr = NonNull::new(ptr as *mut X).unwrap();
        A {
            ptr,
            len,
            own: true,
            _pd: PhantomData,
        }
    }

    pub fn set(&mut self, i: usize, x: X) {
        unsafe {
            write(self.ptr.as_ptr().add(i), x);
        }
    }

    pub fn get(&self, i: usize) -> X {
        unsafe { read(self.ptr.as_ptr().add(i)) }
    }

    pub fn as_ref<'b>(&'b self) -> A<'a, X>
    where
        'b: 'a,
    {
        A {
            ptr: self.ptr,
            len: self.len,
            own: false,
            _pd: PhantomData,
        }
    }
}

impl<'a, X> Drop for A<'a, X> {
    fn drop(&mut self) {
        if self.own {
            unsafe {
                alloc::dealloc(
                    self.ptr.as_ptr() as *mut u8,
                    Layout::array::<X>(self.len).unwrap(),
                );
            }
        }
    }
}

fn main() {
    let mut a: A<i32> = A::new(3);
    a.set(0, 1);
    a.set(1, 2);
    a.set(2, 3);

    for i in 0..3 {
        println!("{i}:\t{}", a.get(i));
    }

    let mut b: A<i32> = a.as_ref();
    b.set(0, 1);
    b.set(1, 2);
    b.set(2, 3);

    for i in 0..3 {
        println!("{i}\t{}", b.get(i));
    }
}
