Implementing Index trait for a custom ref type?

help
1

I'm trying to write a struct-of-arrays data structure where a slice is actually a collection of multiple pointers with a shared length. I'd like to be able to index this into a custom reference type that is also multiple pointers instead of a single &[T]. However, I'm having trouble implementing the Index trait (so I can use the [] operator) because Index assumes you always want to return a reference. Here's a min repro of what I'm trying to do with a SOA slice (the equivalent of [T]) and custom reference type (the equivalent of &[T] and &mut [T] but as distinct data structures). Note that the code around pointer arithmetic is just illustrative here to get something that compiles -- don't judge it too harshly!

use std::ops::Range;
use std::ptr::NonNull;
use std::marker::PhantomData;

pub struct SoaSlice3<A, B, C> {
    len: usize,
    ptr_a: NonNull<A>,
    ptr_b: NonNull<B>,
    ptr_c: NonNull<C>,
}

pub struct SoaSliceRef3<'a, A, B, C> {
    slice: SoaSlice3<A, B, C>,
    lifetime: PhantomData<&'a SoaSlice3<A, B, C>>,
}

pub struct SoaSliceRefMut3<'a, A, B, C> {
    slice: SoaSlice3<A, B, C>,
    lifetime: PhantomData<&'a mut SoaSlice3<A, B, C>>,
}

pub unsafe trait IndexSoa<T> {
    type Output<'a> where T: 'a;
    type OutputMut<'a> where T: 'a;

    fn index<'a>(self, slice: &'a T) -> Self::Output<'a>;
    fn index_mut<'a>(self, slice: &'a mut T) -> Self::OutputMut<'a>;

    unsafe fn get_unchecked<'a>(self, slice: &'a T) -> Self::Output<'a>;
    unsafe fn get_unchecked_mut<'a>(self, slice: &'a mut T) -> Self::OutputMut<'a>;
}

unsafe impl<A, B, C> IndexSoa<SoaSlice3<A, B, C>> for Range<usize> {
    type Output<'a> = SoaSliceRef3<'a, A, B, C> where A: 'a, B: 'a, C: 'a;
    type OutputMut<'a> = SoaSliceRefMut3<'a, A, B, C> where A: 'a, B: 'a, C: 'a;

    fn index<'a>(
        self,
        slice: &'a SoaSlice3<A, B, C>,
    ) -> Self::Output<'a> {
        assert!(self.end >= self.start);
        assert!(self.end <= slice.len);
        unsafe { self.get_unchecked(slice) }
    }

    fn index_mut<'a>(
        self,
        slice: &'a mut SoaSlice3<A, B, C>,
    ) -> Self::OutputMut<'a> {
        assert!(self.end >= self.start);
        assert!(self.end <= slice.len);
        unsafe { self.get_unchecked_mut(slice) }
    }

    unsafe fn get_unchecked<'a>(
        self,
        slice: &'a SoaSlice3<A, B, C>,
    ) -> Self::Output<'a> {
        SoaSliceRef3 {
            slice: SoaSlice3 {
                len: self.end - self.start,
                ptr_a: NonNull::new(slice.ptr_a.as_ptr().add(self.start)).unwrap(),
                ptr_b: NonNull::new(slice.ptr_b.as_ptr().add(self.start)).unwrap(),
                ptr_c: NonNull::new(slice.ptr_c.as_ptr().add(self.start)).unwrap(),
            },
            lifetime: PhantomData,
        }
    }

    unsafe fn get_unchecked_mut<'a>(
        self,
        slice: &'a mut SoaSlice3<A, B, C>,
    ) -> Self::OutputMut<'a> {
        SoaSliceRefMut3 {
            slice: SoaSlice3 {
                len: self.end - self.start,
                ptr_a: NonNull::new(slice.ptr_a.as_ptr().add(self.start)).unwrap(),
                ptr_b: NonNull::new(slice.ptr_b.as_ptr().add(self.start)).unwrap(),
                ptr_c: NonNull::new(slice.ptr_c.as_ptr().add(self.start)).unwrap(),
            },
            lifetime: PhantomData,
        }
    }
}

With this though, I'm at a loss for how to write Index, since Output needs a lifetime, and index expects to return a reference.

impl<A, B, C> Index<Range<usize>> for SoaSlice3<A, B, C> {
    type Output = /* ??? */    
}

// or...

impl<'a, A, B, C> Index<Range<usize>> for SoaSlice3<A, B, C> {
    type Output = SoaSliceRef3<'a, A, B, C>;

    fn index(&self, index: Range<usize>) -> /* ??? */
}

Any guidance on how I could get a nice [] operation for my custom SOA slice here? Is it even possible somehow?

2

Insofar as I am aware, this is impossible. Index only works for returning a reference to something that already exists within the value being indexed.

It's less ideal, but you'll have to stick to a custom method.

4 Likes